sglang_v0.5.2/pytorch_2.8.0/third_party/XNNPACK/test/transpose-reshape.cc

71 lines
3.0 KiB
C++

// Copyright 2023 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 <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <numeric>
#include <vector>
#include <gtest/gtest.h>
#include "xnnpack.h"
#include "xnnpack/node-type.h"
#include "xnnpack/subgraph.h"
TEST(TransposeTestF32, 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=*/3, /*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_NODE_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_NODE_ID);
uint32_t output_id = XNN_INVALID_NODE_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_NODE_ID);
std::vector<size_t> perm{2, 1, 0};
ASSERT_EQ(xnn_status_success, xnn_define_static_transpose(subgraph, perm.size(), perm.data(), 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_static_transpose);
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[perm[0]]);
ASSERT_EQ(output_shape->dim[1], dims[perm[1]]);
ASSERT_EQ(output_shape->dim[2], dims[perm[2]]);
ASSERT_EQ(runtime->values[node->outputs[0]].size, num_input_elements * sizeof(float));
}