sglang_v0.5.2/pytorch_2.8.0/third_party/XNNPACK/test/unary.cc

268 lines
10 KiB
C++

// Copyright 2020 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.
#include <sys/types.h>
#include <algorithm>
#include <array>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <functional>
#include <numeric>
#include <random>
#include <sstream>
#include <string>
#include <tuple>
#include <vector>
#include <gtest/gtest.h>
#include "xnnpack.h"
#include "xnnpack/datatype.h"
#include "xnnpack/buffer.h"
#include "xnnpack/log.h"
#include "xnnpack/operator-utils.h"
#include "xnnpack/subgraph.h"
#include "replicable_random_device.h"
#include "unary-ops.h"
struct Param {
using UnaryT = std::tuple<xnn_unary_operator, xnn_datatype>;
using ConvertT = std::tuple<xnn_unary_operator, xnn_datatype, xnn_datatype>;
explicit Param(UnaryT p)
: unary_operator(std::get<0>(p)),
input_datatype(std::get<1>(p)),
output_datatype(std::get<1>(p)) {}
explicit Param(ConvertT p)
: unary_operator(std::get<0>(p)),
input_datatype(std::get<1>(p)),
output_datatype(std::get<2>(p)) {}
std::string Name() const {
std::stringstream sstr;
sstr << xnn_unary_operator_to_string(unary_operator) << "_"
<< xnn_datatype_to_string(input_datatype);
if (input_datatype != output_datatype) {
sstr << "_" << xnn_datatype_to_string(output_datatype);
}
std::string s = sstr.str();
// Test names must be alphanumeric with no spaces
std::replace(s.begin(), s.end(), ' ', '_');
std::replace(s.begin(), s.end(), '(', '_');
std::replace(s.begin(), s.end(), ')', '_');
return s;
}
xnn_unary_operator unary_operator;
xnn_datatype input_datatype;
xnn_datatype output_datatype;
};
class UnaryTest : public testing::TestWithParam<Param> {
public:
xnnpack::ReplicableRandomDevice rng_;
};
TEST_P(UnaryTest, matches_operator_api) {
const xnn_unary_operator unary_operator = GetParam().unary_operator;
const xnn_datatype input_datatype = GetParam().input_datatype;
const xnn_datatype output_datatype = GetParam().output_datatype;
const size_t sizeof_input = xnn_datatype_size_bytes(input_datatype);
const size_t sizeof_output = xnn_datatype_size_bytes(output_datatype);
ASSERT_EQ(xnn_status_success, xnn_initialize(/*allocator=*/nullptr));
std::uniform_int_distribution<> rank_dist(0, XNN_MAX_TENSOR_DIMS);
std::uniform_int_distribution<> dim_dist(1, 10);
std::vector<size_t> dims(rank_dist(rng_));
std::generate(dims.begin(), dims.end(), [&]() { return dim_dist(rng_); });
size_t size =
std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<size_t>());
size_t channels = dims.empty() ? 1 : dims.back();
size_t batch_size = size / channels;
xnnpack::Buffer<char> input(size * sizeof_input + XNN_EXTRA_BYTES);
xnnpack::fill_uniform_random_bits(input.data(), input.size(), rng_);
xnnpack::Buffer<char> subgraph_output(size * sizeof_output);
xnnpack::Buffer<char> operator_output(size * sizeof_output);
const UnaryOpInfo* op_info = GetUnaryOpInfo(unary_operator);
xnn_unary_params params = op_info->DefaultParams();
const xnn_quantization_params input_quantization =
op_info->InputQuantizationParams(input_datatype);
const xnn_quantization_params output_quantization =
op_info->OutputQuantizationParams(output_datatype);
// Call operator API.
const xnn_status status = xnn_run_unary_elementwise_nc(
unary_operator, input_datatype, output_datatype, &params,
&input_quantization, &output_quantization, /*flags=*/0, batch_size,
channels, channels, channels, /*thread_pool=*/nullptr, input.data(),
operator_output.data());
if (status == xnn_status_unsupported_parameter) {
GTEST_SKIP();
}
ASSERT_EQ(xnn_status_success, status);
// Call subgraph API.
xnn_subgraph_t subgraph = nullptr;
ASSERT_EQ(xnn_status_success, xnn_create_subgraph(/*external_value_ids=*/2, /*flags=*/0, &subgraph));
std::unique_ptr<xnn_subgraph, decltype(&xnn_delete_subgraph)> auto_subgraph(subgraph, xnn_delete_subgraph);
uint32_t input_id = XNN_INVALID_VALUE_ID;
if (xnn_datatype_is_quantized(input_datatype)) {
ASSERT_EQ(xnn_status_success,
xnn_define_quantized_tensor_value(
subgraph, input_datatype, input_quantization.zero_point,
input_quantization.scale, dims.size(), dims.data(), nullptr,
/*external_id=*/0,
/*flags=*/XNN_VALUE_FLAG_EXTERNAL_INPUT, &input_id));
} else {
ASSERT_EQ(xnn_status_success,
xnn_define_tensor_value(subgraph, input_datatype, dims.size(),
dims.data(), nullptr, /*external_id=*/0,
/*flags=*/XNN_VALUE_FLAG_EXTERNAL_INPUT,
&input_id));
}
ASSERT_NE(input_id, XNN_INVALID_VALUE_ID);
uint32_t output_id = XNN_INVALID_VALUE_ID;
if (xnn_datatype_is_quantized(output_datatype)) {
ASSERT_EQ(xnn_status_success,
xnn_define_quantized_tensor_value(
subgraph, output_datatype, output_quantization.zero_point,
output_quantization.scale, dims.size(), dims.data(), nullptr,
/*external_id=*/1,
/*flags=*/XNN_VALUE_FLAG_EXTERNAL_OUTPUT, &output_id));
} else {
ASSERT_EQ(xnn_status_success,
xnn_define_tensor_value(subgraph, output_datatype, dims.size(),
dims.data(), nullptr, /*external_id=*/1,
/*flags=*/XNN_VALUE_FLAG_EXTERNAL_OUTPUT,
&output_id));
}
ASSERT_NE(output_id, XNN_INVALID_VALUE_ID);
xnn_runtime_t runtime = nullptr;
ASSERT_EQ(xnn_status_success,
xnn_define_unary(subgraph, unary_operator, &params, input_id,
output_id, /*flags=*/0));
ASSERT_EQ(xnn_status_success, xnn_create_runtime_v3(subgraph, nullptr, nullptr, /*flags=*/0, &runtime));
ASSERT_NE(nullptr, runtime);
std::unique_ptr<xnn_runtime, decltype(&xnn_delete_runtime)> auto_runtime(runtime, xnn_delete_runtime);
std::array<xnn_external_value, 2> external = {
xnn_external_value{input_id, input.data()}, xnn_external_value{output_id, subgraph_output.data()}};
ASSERT_EQ(xnn_status_success, xnn_setup_runtime(runtime, external.size(), external.data()));
ASSERT_EQ(xnn_status_success, xnn_invoke_runtime(runtime));
ASSERT_EQ(subgraph_output, operator_output);
}
xnn_unary_operator all_unary_ops[] = {
xnn_unary_clamp,
xnn_unary_abs,
xnn_unary_bankers_rounding,
xnn_unary_ceiling,
xnn_unary_elu,
xnn_unary_exp,
xnn_unary_floor,
xnn_unary_gelu,
xnn_unary_hardswish,
xnn_unary_leaky_relu,
xnn_unary_log,
xnn_unary_negate,
xnn_unary_sigmoid,
xnn_unary_square,
xnn_unary_square_root,
xnn_unary_reciprocal_square_root,
xnn_unary_tanh,
xnn_unary_cube_root,
xnn_unary_cosine,
xnn_unary_sine,
xnn_unary_count_leading_zeros,
xnn_unary_bitwise_not,
xnn_unary_popcount,
xnn_unary_sign,
};
const xnn_datatype all_datatypes[] = {
xnn_datatype_quint8,
xnn_datatype_qint8,
#ifndef XNN_EXCLUDE_F16_TESTS
xnn_datatype_fp16,
#endif
xnn_datatype_bf16,
xnn_datatype_fp32,
xnn_datatype_int32,
};
INSTANTIATE_TEST_SUITE_P(
UnaryTest, UnaryTest,
testing::ConvertGenerator<Param::UnaryT>(testing::Combine(
testing::ValuesIn(all_unary_ops), testing::ValuesIn(all_datatypes))),
[](const auto& info) { return info.param.Name(); });
INSTANTIATE_TEST_SUITE_P(
ConvertTest, UnaryTest,
testing::ConvertGenerator<Param::ConvertT>(testing::Combine(
testing::Values(xnn_unary_convert), testing::ValuesIn(all_datatypes),
testing::ValuesIn(all_datatypes))),
[](const auto& info) { return info.param.Name(); });
TEST(AbsTest, reshape) {
ASSERT_EQ(xnn_status_success, xnn_initialize(/*allocator=*/nullptr));
xnn_subgraph_t subgraph = nullptr;
ASSERT_EQ(xnn_status_success, xnn_create_subgraph(/*external_value_ids=*/2, /*flags=*/0, &subgraph));
std::unique_ptr<xnn_subgraph, decltype(&xnn_delete_subgraph)> auto_subgraph(subgraph, xnn_delete_subgraph);
std::vector<size_t> dims{2, 3, 4};
uint32_t input_id = XNN_INVALID_VALUE_ID;
ASSERT_EQ(
xnn_status_success, xnn_define_tensor_value(
subgraph, xnn_datatype_fp32, dims.size(), dims.data(), nullptr, 0,
/*flags=*/XNN_VALUE_FLAG_EXTERNAL_INPUT, &input_id));
ASSERT_NE(input_id, XNN_INVALID_VALUE_ID);
uint32_t output_id = XNN_INVALID_VALUE_ID;
ASSERT_EQ(
xnn_status_success, xnn_define_tensor_value(
subgraph, xnn_datatype_fp32, dims.size(), dims.data(), nullptr, 1,
/*flags=*/XNN_VALUE_FLAG_EXTERNAL_OUTPUT, &output_id));
ASSERT_NE(output_id, XNN_INVALID_VALUE_ID);
ASSERT_EQ(xnn_status_success, xnn_define_unary(subgraph, xnn_unary_abs, /*params=*/nullptr, input_id, output_id, /*flags=*/0));
ASSERT_EQ(subgraph->num_nodes, 1);
struct xnn_node* node = &subgraph->nodes[0];
ASSERT_EQ(node->type, xnn_node_type_unary_elementwise);
ASSERT_EQ(node->unary_operator, xnn_unary_abs);
ASSERT_EQ(node->num_inputs, 1);
ASSERT_EQ(node->inputs[0], input_id);
ASSERT_EQ(node->num_outputs, 1);
ASSERT_EQ(node->outputs[0], output_id);
ASSERT_EQ(node->flags, 0);
xnn_runtime_t runtime = nullptr;
ASSERT_EQ(xnn_status_success, xnn_create_runtime_v3(subgraph, nullptr, nullptr, /*flags=*/0, &runtime));
ASSERT_NE(nullptr, runtime);
std::unique_ptr<xnn_runtime, decltype(&xnn_delete_runtime)> auto_runtime(runtime, xnn_delete_runtime);
ASSERT_EQ(node->reshape(&runtime->opdata[0], subgraph->values, subgraph->num_values, /*threadpool=*/nullptr), xnn_status_success);
dims[0] = 7;
ASSERT_EQ(xnn_status_success, xnn_reshape_external_value(runtime, 0, dims.size(), dims.data()));
ASSERT_EQ(node->reshape(&runtime->opdata[0], runtime->values, runtime->num_values, /*threadpool=*/nullptr), xnn_status_reallocation_required);
const xnn_shape* output_shape = &runtime->values[node->outputs[0]].shape;
const size_t num_input_elements = std::accumulate(dims.cbegin(), dims.cend(), size_t{1}, std::multiplies<size_t>());
ASSERT_EQ(output_shape->dim[0], dims[0]);
ASSERT_EQ(runtime->values[node->outputs[0]].size, num_input_elements * sizeof(float));
}