已合并
add aicpu:tile_with_axis&topkpqdistance #4174
sujunwei3创建于 7月21日
add aicpu:tile_with_axis&topkpqdistance #4174
已合并
共 33 个文件变更+4568-0
| @@ -0,0 +1,46 @@ | |||
| 1 | +# TileWithAxis | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 产品 | 是否支持 | | ||
| 6 | +| :----------------------------------------------------------- | :------: | | ||
| 7 | +| Ascend 950PR/Ascend 950DT | √ | | ||
| 8 | +| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ | | ||
| 9 | +| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ | | ||
| 10 | +| Atlas 200I/500 A2 推理产品 | √ | | ||
| 11 | +| Atlas 推理系列产品 | √ | | ||
| 12 | +| Atlas 训练系列产品 | √ | | ||
| 13 | + | ||
| 14 | +## 功能说明 | ||
| 15 | + | ||
| 16 | +- 算子功能：沿指定维度复制输入Tensor数据，扩展输出Tensor。 | ||
| 17 | + | ||
| 18 | +- 计算示例： | ||
| 19 | + - 输入 x = [[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]，shape = (2, 3, 2) | ||
| 20 | + - axis = 1，tiles = 2 | ||
| 21 | + - 输出 y = [[[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12], [7, 8], [9, 10], [11, 12]]]，shape = (2, 6, 2) | ||
| 22 | + | ||
| 23 | +- 计算公式： | ||
| 24 | + 输出 y 的 shape 与输入 x 相同，仅 axis 维度变为 `x.shape[axis] * tiles`。 | ||
| 25 | + | ||
| 26 | +## 参数说明 | ||
| 27 | + | ||
| 28 | +| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 | | ||
| 29 | +|--------|---------------|------|----------|----------| | ||
| 30 | +| x | 输入 | 输入张量 | FLOAT16、FLOAT、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64 | ND | | ||
| 31 | +| axis | 属性（可选） | 指定复制的维度，默认为1 | INT | - | | ||
| 32 | +| tiles | 属性（必选） | 复制次数，必须大于0 | INT | - | | ||
| 33 | +| y | 输出 | 输出张量，与x具有相同的数据类型和格式 | 与x一致 | ND | | ||
| 34 | + | ||
| 35 | +## 约束说明 | ||
| 36 | + | ||
| 37 | +- axis 必须在输入Tensor的维度范围内（支持负数索引，如 axis=-1 表示最后一维）。 | ||
| 38 | +- tiles 必须大于0。 | ||
| 39 | +- 输出Tensor在 axis 维度的大小必须等于输入Tensor在 axis 维度的大小乘以 tiles。 | ||
| 40 | +- 输入Tensor维度不超过8维。 | ||
| 41 | + | ||
| 42 | +## 调用说明 | ||
| 43 | + | ||
| 44 | +| 调用方式 | 样例代码 | 说明 | | ||
| 45 | +|------------|----------|------| | ||
| 46 | +| 图模式调用 | [test_geir_tile_with_axis](./examples/test_geir_tile_with_axis.cpp) | 通过[算子IR](./op_graph/tile_with_axis_proto.h)构图方式调用TileWithAxis算子。 | | ||
| @@ -0,0 +1,328 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +/*! | ||
| 12 | + * \file test_geir_tile_with_axis.cpp | ||
| 13 | + * \brief GE IR test for TileWithAxis operator | ||
| 14 | + * | ||
| 15 | + * TileWithAxis: input x[2,3,2], axis=1, tiles=2 -> output y[2,6,2] | ||
| 16 | + * 使用 DT_INT64 数据类型，AICore 不支持该类型，自动走 AICPU 算子路径。 | ||
| 17 | + */ | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + | ||
| 35 | + | ||
| 36 | + | ||
| 37 | + | ||
| 38 | + | ||
| 39 | + | ||
| 40 | + | ||
| 41 | + | ||
| 42 | + | ||
| 43 | +using namespace ge; | ||
| 44 | +using std::map; | ||
| 45 | +using std::string; | ||
| 46 | +using std::vector; | ||
| 47 | + | ||
| 48 | +string GetTime() | ||
| 49 | +{ | ||
| 50 | + time_t timep; | ||
| 51 | + time(&timep); | ||
| 52 | + char tmp[64]; | ||
| 53 | + strftime(tmp, sizeof(tmp), "%Y-%m-%d %H:%M:%S,000", localtime(&timep)); | ||
| 54 | + return tmp; | ||
| 55 | +} | ||
| 56 | + | ||
| 57 | +uint32_t GetDataTypeSize(DataType dt) | ||
| 58 | +{ | ||
| 59 | + switch (dt) { | ||
| 60 | + case ge::DT_BOOL: | ||
| 61 | + return 1U; | ||
| 62 | + case ge::DT_INT8: | ||
| 63 | + case ge::DT_UINT8: | ||
| 64 | + return 1U; | ||
| 65 | + case ge::DT_FLOAT16: | ||
| 66 | + case ge::DT_INT16: | ||
| 67 | + case ge::DT_UINT16: | ||
| 68 | + return 2U; | ||
| 69 | + case ge::DT_FLOAT: | ||
| 70 | + case ge::DT_INT32: | ||
| 71 | + case ge::DT_UINT32: | ||
| 72 | + return 4U; | ||
| 73 | + case ge::DT_INT64: | ||
| 74 | + case ge::DT_UINT64: | ||
| 75 | + return 8U; | ||
| 76 | + default: | ||
| 77 | + return 0U; | ||
| 78 | + } | ||
| 79 | +} | ||
| 80 | + | ||
| 81 | +int32_t GenData(vector<int64_t> shapes, Tensor& tensor, TensorDesc& desc, DataType dt, double value) | ||
| 82 | +{ | ||
| 83 | + desc.SetRealDimCnt(shapes.size()); | ||
| 84 | + size_t size = 1; | ||
| 85 | + for (size_t i = 0; i < shapes.size(); i++) { | ||
| 86 | + size *= static_cast<size_t>(shapes[i]); | ||
| 87 | + } | ||
| 88 | + uint32_t type_size = GetDataTypeSize(dt); | ||
| 89 | + if (type_size == 0U) { | ||
| 90 | + printf("%s - ERROR - [XIR]: GenData: unsupported data type %d\n", GetTime().c_str(), dt); | ||
| 91 | + return FAILED; | ||
| 92 | + } | ||
| 93 | + uint32_t data_len = static_cast<uint32_t>(size * type_size); | ||
| 94 | + uint8_t* buf = new (std::nothrow) uint8_t[data_len]; | ||
| 95 | + if (buf == nullptr) { | ||
| 96 | + return FAILED; | ||
| 97 | + } | ||
| 98 | + | ||
| 99 | + switch (dt) { | ||
| 100 | + case DT_FLOAT: { | ||
| 101 | + float* p = reinterpret_cast<float*>(buf); | ||
| 102 | + for (size_t i = 0; i < size; i++) { | ||
| 103 | + p[i] = static_cast<float>(value + static_cast<double>(i)); | ||
| 104 | + } | ||
| 105 | + break; | ||
| 106 | + } | ||
| 107 | + case DT_INT64: { | ||
| 108 | + int64_t* p = reinterpret_cast<int64_t*>(buf); | ||
| 109 | + for (size_t i = 0; i < size; i++) { | ||
| 110 | + p[i] = static_cast<int64_t>(value + static_cast<double>(i)); | ||
| 111 | + } | ||
| 112 | + break; | ||
| 113 | + } | ||
| 114 | + case DT_INT32: { | ||
| 115 | + int32_t* p = reinterpret_cast<int32_t*>(buf); | ||
| 116 | + for (size_t i = 0; i < size; i++) { | ||
| 117 | + p[i] = static_cast<int32_t>(value + static_cast<double>(i)); | ||
| 118 | + } | ||
| 119 | + break; | ||
| 120 | + } | ||
| 121 | + case DT_INT16: { | ||
| 122 | + int16_t* p = reinterpret_cast<int16_t*>(buf); | ||
| 123 | + for (size_t i = 0; i < size; i++) { | ||
| 124 | + p[i] = static_cast<int16_t>(value + static_cast<double>(i)); | ||
| 125 | + } | ||
| 126 | + break; | ||
| 127 | + } | ||
| 128 | + case DT_INT8: { | ||
| 129 | + int8_t* p = reinterpret_cast<int8_t*>(buf); | ||
| 130 | + for (size_t i = 0; i < size; i++) { | ||
| 131 | + p[i] = static_cast<int8_t>(value + static_cast<double>(i)); | ||
| 132 | + } | ||
| 133 | + break; | ||
| 134 | + } | ||
| 135 | + case DT_UINT64: { | ||
| 136 | + uint64_t* p = reinterpret_cast<uint64_t*>(buf); | ||
| 137 | + for (size_t i = 0; i < size; i++) { | ||
| 138 | + p[i] = static_cast<uint64_t>(value + static_cast<double>(i)); | ||
| 139 | + } | ||
| 140 | + break; | ||
| 141 | + } | ||
| 142 | + case DT_UINT32: { | ||
| 143 | + uint32_t* p = reinterpret_cast<uint32_t*>(buf); | ||
| 144 | + for (size_t i = 0; i < size; i++) { | ||
| 145 | + p[i] = static_cast<uint32_t>(value + static_cast<double>(i)); | ||
| 146 | + } | ||
| 147 | + break; | ||
| 148 | + } | ||
| 149 | + case DT_UINT16: { | ||
| 150 | + uint16_t* p = reinterpret_cast<uint16_t*>(buf); | ||
| 151 | + for (size_t i = 0; i < size; i++) { | ||
| 152 | + p[i] = static_cast<uint16_t>(value + static_cast<double>(i)); | ||
| 153 | + } | ||
| 154 | + break; | ||
| 155 | + } | ||
| 156 | + case DT_UINT8: { | ||
| 157 | + uint8_t* p = buf; | ||
| 158 | + for (size_t i = 0; i < size; i++) { | ||
| 159 | + p[i] = static_cast<uint8_t>(value + static_cast<double>(i)); | ||
| 160 | + } | ||
| 161 | + break; | ||
| 162 | + } | ||
| 163 | + default: | ||
| 164 | + printf("%s - ERROR - [XIR]: GenData: unsupported data type %d\n", GetTime().c_str(), dt); | ||
| 165 | + delete[] buf; | ||
| 166 | + return FAILED; | ||
| 167 | + } | ||
| 168 | + tensor = Tensor(desc, buf, data_len); | ||
| 169 | + delete[] buf; | ||
| 170 | + return SUCCESS; | ||
| 171 | +} | ||
| 172 | + | ||
| 173 | +int32_t WriteDataToFile(string bin_file, uint64_t data_size, uint8_t* inputData) | ||
| 174 | +{ | ||
| 175 | + FILE* fp = fopen(bin_file.c_str(), "wb"); | ||
| 176 | + if (fp == nullptr) { | ||
| 177 | + return FAILED; | ||
| 178 | + } | ||
| 179 | + size_t written = fwrite(inputData, 1, data_size, fp); | ||
| 180 | + fclose(fp); | ||
| 181 | + if (written != data_size) { | ||
| 182 | + return FAILED; | ||
| 183 | + } | ||
| 184 | + return SUCCESS; | ||
| 185 | +} | ||
| 186 | + | ||
| 187 | +void ProcessInputData(vector<Tensor>& input) | ||
| 188 | +{ | ||
| 189 | + for (size_t i = 0; i < input.size(); i++) { | ||
| 190 | + string input_file = "./tc_ge_irrun_test_0008_npu_input_" + std::to_string(i) + ".bin"; | ||
| 191 | + uint8_t* input_data_i = input[i].GetData(); | ||
| 192 | + int64_t input_shape = input[i].GetTensorDesc().GetShape().GetShapeSize(); | ||
| 193 | + uint32_t type_size = GetDataTypeSize(input[i].GetTensorDesc().GetDataType()); | ||
| 194 | + if (type_size == 0U) { | ||
| 195 | + printf("ERROR: input %zu has unsupported dtype\n", i); | ||
| 196 | + continue; | ||
| 197 | + } | ||
| 198 | + uint32_t data_size = static_cast<uint32_t>(input_shape * type_size); | ||
| 199 | + WriteDataToFile(input_file.c_str(), data_size, input_data_i); | ||
| 200 | + } | ||
| 201 | +} | ||
| 202 | + | ||
| 203 | +void ProcessOutputData(vector<Tensor>& output) | ||
| 204 | +{ | ||
| 205 | + for (size_t i = 0; i < output.size(); i++) { | ||
| 206 | + string output_file = "./tc_ge_irrun_test_0008_npu_output_" + std::to_string(i) + ".bin"; | ||
| 207 | + uint8_t* output_data_i = output[i].GetData(); | ||
| 208 | + int64_t output_shape = output[i].GetTensorDesc().GetShape().GetShapeSize(); | ||
| 209 | + uint32_t type_size = GetDataTypeSize(output[i].GetTensorDesc().GetDataType()); | ||
| 210 | + if (type_size == 0U) { | ||
| 211 | + printf("ERROR: output %zu has unsupported dtype\n", i); | ||
| 212 | + continue; | ||
| 213 | + } | ||
| 214 | + uint32_t data_size = static_cast<uint32_t>(output_shape * type_size); | ||
| 215 | + WriteDataToFile(output_file.c_str(), data_size, output_data_i); | ||
| 216 | + } | ||
| 217 | +} | ||
| 218 | + | ||
| 219 | +int CreateOppInGraph(DataType inDtype, vector<Tensor>& input, vector<Operator>& inputs, vector<Operator>& outputs, | ||
| 220 | + Graph& graph) | ||
| 221 | +{ | ||
| 222 | + Status ret = SUCCESS; | ||
| 223 | + auto node = op::TileWithAxis("tile_with_axis_1"); | ||
| 224 | + | ||
| 225 | + vector<int64_t> xShape = {2, 3, 2}; | ||
| 226 | + TensorDesc xDesc(ge::Shape(xShape), FORMAT_ND, inDtype); | ||
| 227 | + xDesc.SetPlacement(ge::kPlacementHost); | ||
| 228 | + xDesc.SetFormat(FORMAT_ND); | ||
| 229 | + | ||
| 230 | + Tensor xTensor; | ||
| 231 | + ret = GenData(xShape, xTensor, xDesc, inDtype, 1.0); | ||
| 232 | + if (ret != SUCCESS) { | ||
| 233 | + printf("%s - ERROR - [XIR]: Generate input data failed\n", GetTime().c_str()); | ||
| 234 | + return FAILED; | ||
| 235 | + } | ||
| 236 | + input.push_back(xTensor); | ||
| 237 | + | ||
| 238 | + auto xData = op::Data("x_data").set_attr_index(0); | ||
| 239 | + xData.update_input_desc_x(xDesc); | ||
| 240 | + xData.update_output_desc_y(xDesc); | ||
| 241 | + | ||
| 242 | + node.set_input_x(xData); | ||
| 243 | + node.SetAttr("axis", 1); | ||
| 244 | + node.SetAttr("tiles", 2); | ||
| 245 | + | ||
| 246 | + vector<int64_t> yShape = {2, 6, 2}; | ||
| 247 | + TensorDesc yDesc(ge::Shape(yShape), FORMAT_ND, inDtype); | ||
| 248 | + node.update_output_desc_y(yDesc); | ||
| 249 | + | ||
| 250 | + graph.AddOp(xData); | ||
| 251 | + graph.AddOp(node); | ||
| 252 | + | ||
| 253 | + inputs.push_back(xData); | ||
| 254 | + outputs.push_back(node); | ||
| 255 | + | ||
| 256 | + return SUCCESS; | ||
| 257 | +} | ||
| 258 | + | ||
| 259 | +int main(int argc, char* argv[]) | ||
| 260 | +{ | ||
| 261 | + const char* graph_name = "tc_ge_irrun_test"; | ||
| 262 | + Graph graph(graph_name); | ||
| 263 | + vector<Tensor> input; | ||
| 264 | + | ||
| 265 | + printf("%s - INFO - [XIR]: Start to initialize ge\n", GetTime().c_str()); | ||
| 266 | + std::map<AscendString, AscendString> global_options = {{"ge.exec.deviceId", "0"}, {"ge.graphRunMode", "1"}}; | ||
| 267 | + Status ret = ge::GEInitialize(global_options); | ||
| 268 | + if (ret != SUCCESS) { | ||
| 269 | + printf("%s - ERROR - [XIR]: Initialize ge failed\n", GetTime().c_str()); | ||
| 270 | + return FAILED; | ||
| 271 | + } | ||
| 272 | + printf("%s - INFO - [XIR]: Initialize ge success\n", GetTime().c_str()); | ||
| 273 | + | ||
| 274 | + vector<Operator> inputs{}; | ||
| 275 | + vector<Operator> outputs{}; | ||
| 276 | + | ||
| 277 | + // DT_INT64: AICore def 不支持，自动走 AICPU 算子路径 | ||
| 278 | + DataType inDtype = DT_INT64; | ||
| 279 | + | ||
| 280 | + ret = CreateOppInGraph(inDtype, input, inputs, outputs, graph); | ||
| 281 | + if (ret != SUCCESS) { | ||
| 282 | + printf("%s - ERROR - [XIR]: Create graph failed\n", GetTime().c_str()); | ||
| 283 | + return FAILED; | ||
| 284 | + } | ||
| 285 | + | ||
| 286 | + if (!inputs.empty() && !outputs.empty()) { | ||
| 287 | + graph.SetInputs(inputs).SetOutputs(outputs); | ||
| 288 | + } | ||
| 289 | + | ||
| 290 | + std::map<AscendString, AscendString> build_options = {}; | ||
| 291 | + ge::Session* session = new Session(build_options); | ||
| 292 | + if (session == nullptr) { | ||
| 293 | + printf("%s - ERROR - [XIR]: Create session failed\n", GetTime().c_str()); | ||
| 294 | + return FAILED; | ||
| 295 | + } | ||
| 296 | + | ||
| 297 | + uint32_t graph_id = 0; | ||
| 298 | + std::map<AscendString, AscendString> graph_options = {}; | ||
| 299 | + ret = session->AddGraph(graph_id, graph, graph_options); | ||
| 300 | + if (ret != SUCCESS) { | ||
| 301 | + printf("%s - ERROR - [XIR]: Add graph failed\n", GetTime().c_str()); | ||
| 302 | + delete session; | ||
| 303 | + GEFinalize(); | ||
| 304 | + return FAILED; | ||
| 305 | + } | ||
| 306 | + | ||
| 307 | + vector<Tensor> output; | ||
| 308 | + ret = session->RunGraph(graph_id, input, output); | ||
| 309 | + if (ret != SUCCESS) { | ||
| 310 | + printf("%s - ERROR - [XIR]: Run graph failed\n", GetTime().c_str()); | ||
| 311 | + delete session; | ||
| 312 | + GEFinalize(); | ||
| 313 | + return FAILED; | ||
| 314 | + } | ||
| 315 | + printf("%s - INFO - [XIR]: Run graph success\n", GetTime().c_str()); | ||
| 316 | + | ||
| 317 | + ProcessInputData(input); | ||
| 318 | + ProcessOutputData(output); | ||
| 319 | + | ||
| 320 | + delete session; | ||
| 321 | + ret = ge::GEFinalize(); | ||
| 322 | + if (ret != SUCCESS) { | ||
| 323 | + printf("%s - ERROR - [XIR]: Finalize failed\n", GetTime().c_str()); | ||
| 324 | + return FAILED; | ||
| 325 | + } | ||
| 326 | + printf("%s - INFO - [XIR]: Finalize ir graph session success\n", GetTime().c_str()); | ||
| 327 | + return SUCCESS; | ||
| 328 | +} | ||
| @@ -0,0 +1,54 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +namespace ge { | ||
| 17 | +/** | ||
| 18 | +* @brief Extends the input with copies of data along a specified dimension. For example: \n | ||
| 19 | +* (1) If x = [[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]], with shape (2, 3, 2); \n | ||
| 20 | +* (2) axis = 1; \n | ||
| 21 | +* (3) tiles = 2; \n | ||
| 22 | +* (4) Then, y = [[[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6]], [[7, 8], | ||
| 23 | +* [9, 10], [11, 12], [7, 8], [9, 10], [11, 12]]], | ||
| 24 | +* with shape (2, 6, 2). | ||
| 25 | + | ||
| 26 | +* @par Inputs: | ||
| 27 | +* One input: | ||
| 28 | +* x: A Tensor with any format. Must be one of the following types: | ||
| 29 | +* bfloat16, float16, float32, int8, int16, int32, int64, uint8, uint16, uint32, uint64 . \n | ||
| 30 | + | ||
| 31 | +* @par Attributes: | ||
| 32 | +* @li axis: An optional int, specifying the axis to tile. Defaults to 1. | ||
| 33 | +* @li tiles: A required int, specifying the number of copies (tiles) to output . \n | ||
| 34 | + | ||
| 35 | +* @par Outputs: | ||
| 36 | +* y: A Tensor with the same type and format of x. \n | ||
| 37 | + | ||
| 38 | +* @attention Constraints: | ||
| 39 | +* @li "axis" must be within the rank of the input tensor. | ||
| 40 | +* @li "tiles" must be greater than 1. | ||
| 41 | +* @par Third-party framework compatibility | ||
| 42 | +* Compatible with the Caffe operator Tile. | ||
| 43 | +*/ | ||
| 44 | +REG_OP(TileWithAxis) | ||
| 45 | + .INPUT(x, TensorType({DT_BF16, DT_FLOAT16, DT_FLOAT, DT_INT64, DT_INT32, DT_INT16, DT_INT8, DT_UINT64, DT_UINT32, | ||
| 46 | + DT_UINT16, DT_UINT8})) | ||
 | |||
| 47 | + .OUTPUT(y, TensorType({DT_BF16, DT_FLOAT16, DT_FLOAT, DT_INT64, DT_INT32, DT_INT16, DT_INT8, DT_UINT64, DT_UINT32, | ||
| 48 | + DT_UINT16, DT_UINT8})) | ||
| 49 | + .ATTR(axis, Int, 1) | ||
| 50 | + .REQUIRED_ATTR(tiles, Int) | ||
| 51 | + .OP_END_FACTORY_REG(TileWithAxis) | ||
| 52 | +} // namespace ge | ||
| 53 | + | ||
| 54 | + | ||
| @@ -0,0 +1,196 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +namespace { | ||
| 19 | +const char* const kTileWithAxis = "TileWithAxis"; | ||
| 20 | +} | ||
| 21 | +namespace aicpu { | ||
| 22 | +template <typename T, int32_t OPTION, int32_t DIMS> | ||
| 23 | +uint32_t TileWithAxisCpuKernel::TileComputeByAxis(const CpuKernelContext& ctx) | ||
| 24 | +{ | ||
| 25 | + Tensor* input = ctx.Input(kFirstInputIndex); | ||
| 26 | + Tensor* output = ctx.Output(kFirstOutputIndex); | ||
| 27 | + | ||
| 28 | + int64_t axis = 1; | ||
| 29 | + AttrValue* axis_attr = ctx.GetAttr("axis"); | ||
| 30 | + if (axis_attr != nullptr) { | ||
| 31 | + axis = axis_attr->GetInt(); | ||
| 32 | + axis = (axis < 0) ? (axis + DIMS) : axis; | ||
| 33 | + } | ||
| 34 | + | ||
| 35 | + // tiles == 1, output is identical to input | ||
| 36 | + int64_t tiles = ctx.GetAttr("tiles")->GetInt(); | ||
| 37 | + if (tiles == 1) { | ||
| 38 | + T* input0 = reinterpret_cast<T*>(input->GetData()); | ||
| 39 | + T* output0 = reinterpret_cast<T*>(output->GetData()); | ||
| 40 | + KERNEL_CHECK_FALSE((memcpy_s(output0, output->GetDataSize(), input0, output->GetDataSize()) == EOK), | ||
| 41 | + KERNEL_STATUS_INNER_ERROR, "TileWithAxis memcpy failed, dst len is %ld, src size is %ld.", | ||
| 42 | + output->GetDataSize(), output->GetDataSize()); | ||
| 43 | + | ||
| 44 | + return KERNEL_STATUS_OK; | ||
| 45 | + } | ||
| 46 | + | ||
| 47 | + // Get input and output shapes | ||
| 48 | + std::vector<int64_t> input_shape = input->GetTensorShape()->GetDimSizes(); | ||
| 49 | + std::vector<int64_t> output_shape = output->GetTensorShape()->GetDimSizes(); | ||
| 50 | + // Reshape and broadcast output | ||
| 51 | + Eigen::DSizes<Eigen::DenseIndex, DIMS> in_reshape; | ||
| 52 | + Eigen::DSizes<Eigen::DenseIndex, DIMS> out_reshape; | ||
| 53 | + Eigen::array<Eigen::DenseIndex, DIMS> bcast; | ||
| 54 | + for (int32_t i = 0; i < DIMS; i++) { | ||
| 55 | + in_reshape[(DIMS - i) - 1] = input_shape[i]; | ||
| 56 | + out_reshape[(DIMS - i) - 1] = output_shape[i]; | ||
| 57 | + bcast[i] = (i == ((DIMS - axis) - 1)) ? tiles : 1; | ||
| 58 | + } | ||
| 59 | + | ||
| 60 | + Eigen::TensorMap<Eigen::Tensor<T, 1>, OPTION> input0(static_cast<T*>(input->GetData()), | ||
| 61 | + input->GetTensorShape()->NumElements()); | ||
| 62 | + Eigen::TensorMap<Eigen::Tensor<T, 1>, OPTION> output0(static_cast<T*>(output->GetData()), | ||
| 63 | + output->GetTensorShape()->NumElements()); | ||
| 64 | + output0.reshape(out_reshape) = input0.reshape(in_reshape).broadcast(bcast); | ||
| 65 | + | ||
| 66 | + return KERNEL_STATUS_OK; | ||
| 67 | +} | ||
| 68 | + | ||
| 69 | +template <typename T, int32_t OPTION> | ||
| 70 | +uint32_t TileWithAxisCpuKernel::TileComputeInDims(const CpuKernelContext& ctx) | ||
| 71 | +{ | ||
| 72 | + int32_t dims = ctx.Output(kFirstOutputIndex)->GetTensorShape()->GetDims(); | ||
| 73 | + switch (dims) { | ||
| 74 | + case 0: { | ||
| 75 | + T* input0 = reinterpret_cast<T*>(ctx.Input(kFirstInputIndex)->GetData()); | ||
| 76 | + T* output0 = reinterpret_cast<T*>(ctx.Output(kFirstOutputIndex)->GetData()); | ||
| 77 | + KERNEL_CHECK_NULLPTR(input0, KERNEL_STATUS_PARAM_INVALID, "input data is null."); | ||
| 78 | + KERNEL_CHECK_NULLPTR(output0, KERNEL_STATUS_PARAM_INVALID, "output data is null."); | ||
| 79 | + *output0 = *input0; | ||
| 80 | + return KERNEL_STATUS_OK; | ||
| 81 | + } | ||
| 82 | + case 1: | ||
| 83 | + return TileComputeByAxis<T, OPTION, 1>(ctx); | ||
| 84 | + case 2: | ||
| 85 | + return TileComputeByAxis<T, OPTION, 2>(ctx); | ||
| 86 | + case 3: | ||
| 87 | + return TileComputeByAxis<T, OPTION, 3>(ctx); | ||
| 88 | + case 4: | ||
| 89 | + return TileComputeByAxis<T, OPTION, 4>(ctx); | ||
| 90 | + case 5: | ||
| 91 | + return TileComputeByAxis<T, OPTION, 5>(ctx); | ||
| 92 | + case 6: | ||
| 93 | + return TileComputeByAxis<T, OPTION, 6>(ctx); | ||
| 94 | + case 7: | ||
| 95 | + return TileComputeByAxis<T, OPTION, 7>(ctx); | ||
| 96 | + case 8: | ||
| 97 | + return TileComputeByAxis<T, OPTION, 8>(ctx); | ||
| 98 | + default: | ||
| 99 | + KERNEL_LOG_ERROR("[%s] Rank of output should less than 8 but get [%d].", ctx.GetOpType().c_str(), dims); | ||
| 100 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 101 | + } | ||
| 102 | +} | ||
| 103 | + | ||
| 104 | +template <typename T> | ||
| 105 | +uint32_t TileWithAxisCpuKernel::TileCompute(const CpuKernelContext& ctx) | ||
| 106 | +{ | ||
| 107 | + bool flag = AddrAlignedCheck(ctx.Input(kFirstInputIndex)->GetData()); | ||
| 108 | + if (flag) { | ||
| 109 | + return TileComputeInDims<T, Eigen::Aligned>(ctx); | ||
| 110 | + } else { | ||
| 111 | + return TileComputeInDims<T, Eigen::Unaligned>(ctx); | ||
| 112 | + } | ||
| 113 | +} | ||
| 114 | + | ||
| 115 | +uint32_t TileWithAxisCpuKernel::TileParaCheck(const CpuKernelContext& ctx) const | ||
| 116 | +{ | ||
| 117 | + Tensor* input = ctx.Input(kFirstInputIndex); | ||
| 118 | + Tensor* output = ctx.Output(kFirstOutputIndex); | ||
| 119 | + int64_t axis = 1; | ||
| 120 | + | ||
| 121 | + // Check axis is within input dimensions range (optional parameter) | ||
| 122 | + AttrValue* axis_attr = ctx.GetAttr("axis"); | ||
| 123 | + if (axis_attr != nullptr) { | ||
| 124 | + axis = axis_attr->GetInt(); | ||
| 125 | + auto dims = input->GetTensorShape()->GetDims(); | ||
| 126 | + axis = (axis < 0) ? (axis + dims) : axis; | ||
| 127 | + if (axis < 0 || axis >= dims) { | ||
| 128 | + KERNEL_LOG_ERROR("TileWithAxis axis[%ld] is invalid.", axis); | ||
| 129 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 130 | + } | ||
| 131 | + } | ||
| 132 | + | ||
| 133 | + // tiles must be greater than 0 (required parameter) | ||
| 134 | + AttrValue* tiles_attr = ctx.GetAttr("tiles"); | ||
| 135 | + if (tiles_attr == nullptr) { | ||
| 136 | + KERNEL_LOG_ERROR("TileWithAxis tiles is null."); | ||
| 137 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 138 | + } | ||
| 139 | + int64_t tiles = tiles_attr->GetInt(); | ||
| 140 | + if (tiles <= 0) { | ||
| 141 | + KERNEL_LOG_ERROR("TileWithAxis tiles[%ld] is invalid.", tiles); | ||
| 142 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 143 | + } | ||
| 144 | + | ||
| 145 | + // Check output shape on axis equals input shape on axis * tiles | ||
| 146 | + std::vector<int64_t> shape_input = input->GetTensorShape()->GetDimSizes(); | ||
| 147 | + std::vector<int64_t> shape_output = output->GetTensorShape()->GetDimSizes(); | ||
| 148 | + if (axis < 0 || axis >= static_cast<int64_t>(shape_input.size()) || | ||
| 149 | + axis >= static_cast<int64_t>(shape_output.size())) { | ||
| 150 | + KERNEL_LOG_ERROR("TileWithAxis axis[%ld] is out of range (input_dims=%zu, output_dims=%zu).", axis, | ||
| 151 | + shape_input.size(), shape_output.size()); | ||
| 152 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 153 | + } | ||
| 154 | + if (shape_output[axis] != shape_input[axis] * tiles) { | ||
| 155 | + KERNEL_LOG_ERROR("TileWithAxis output_shape[%ld] is invalid.", axis); | ||
| 156 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 157 | + } | ||
| 158 | + | ||
| 159 | + return KERNEL_STATUS_OK; | ||
| 160 | +} | ||
| 161 | + | ||
| 162 | +uint32_t TileWithAxisCpuKernel::Compute(CpuKernelContext& ctx) | ||
| 163 | +{ | ||
| 164 | + KERNEL_HANDLE_ERROR(NormalCheck(ctx, 1, 1), "TileWithAxis NormalCheck fail."); | ||
| 165 | + KERNEL_HANDLE_ERROR(TileParaCheck(ctx), "TileWithAxis TileParaCheck fail."); | ||
| 166 | + | ||
| 167 | + auto data_type = static_cast<DataType>(ctx.Input(kFirstInputIndex)->GetDataType()); | ||
| 168 | + switch (data_type) { | ||
| 169 | + case DT_FLOAT16: | ||
| 170 | + return TileCompute<Eigen::half>(ctx); | ||
| 171 | + case DT_FLOAT: | ||
| 172 | + return TileCompute<float>(ctx); | ||
| 173 | + case DT_INT64: | ||
| 174 | + return TileCompute<int64_t>(ctx); | ||
| 175 | + case DT_INT32: | ||
| 176 | + return TileCompute<int32_t>(ctx); | ||
| 177 | + case DT_INT16: | ||
| 178 | + return TileCompute<int16_t>(ctx); | ||
| 179 | + case DT_INT8: | ||
| 180 | + return TileCompute<int8_t>(ctx); | ||
| 181 | + case DT_UINT64: | ||
| 182 | + return TileCompute<uint64_t>(ctx); | ||
| 183 | + case DT_UINT32: | ||
| 184 | + return TileCompute<uint32_t>(ctx); | ||
| 185 | + case DT_UINT16: | ||
| 186 | + return TileCompute<uint16_t>(ctx); | ||
| 187 | + case DT_UINT8: | ||
| 188 | + return TileCompute<uint8_t>(ctx); | ||
| 189 | + default: | ||
| 190 | + KERNEL_LOG_ERROR("TileWithAxis dtype is invalid."); | ||
| 191 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 192 | + } | ||
| 193 | +} | ||
| 194 | + | ||
| 195 | +REGISTER_CPU_KERNEL(kTileWithAxis, TileWithAxisCpuKernel); | ||
| 196 | +} // namespace aicpu | ||
| @@ -0,0 +1,36 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +namespace aicpu { | ||
| 17 | +class TileWithAxisCpuKernel : public CpuKernel { | ||
| 18 | +public: | ||
| 19 | + TileWithAxisCpuKernel() = default; | ||
| 20 | + ~TileWithAxisCpuKernel() = default; | ||
| 21 | + uint32_t Compute(CpuKernelContext& ctx) override; | ||
| 22 | + | ||
| 23 | +private: | ||
| 24 | + uint32_t TileParaCheck(const CpuKernelContext& ctx) const; | ||
| 25 | + | ||
| 26 | + template <typename T, int32_t OPTION, int32_t DIMS> | ||
| 27 | + uint32_t TileComputeByAxis(const CpuKernelContext& ctx); | ||
| 28 | + | ||
| 29 | + template <typename T, int32_t OPTION> | ||
| 30 | + uint32_t TileComputeInDims(const CpuKernelContext& ctx); | ||
| 31 | + | ||
| 32 | + template <typename T> | ||
| 33 | + uint32_t TileCompute(const CpuKernelContext& ctx); | ||
| 34 | +}; | ||
| 35 | +} // namespace aicpu | ||
| 36 | + | ||
| @@ -0,0 +1,32 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +namespace ops { | ||
| 15 | +class TileWithAxis : public OpDef { | ||
| 16 | +public: | ||
| 17 | + explicit TileWithAxis(const char* name) : OpDef(name) | ||
| 18 | + { | ||
| 19 | + this->Input("x").DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64, | ||
| 20 | + ge::DT_UINT8, ge::DT_UINT16, ge::DT_UINT32, ge::DT_UINT64}); | ||
| 21 | + this->Output("y").DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64, | ||
| 22 | + ge::DT_UINT8, ge::DT_UINT16, ge::DT_UINT32, ge::DT_UINT64}); | ||
| 23 | + this->Attr("axis").AttrType(OPTIONAL).Int(); | ||
| 24 | + this->Attr("tiles").AttrType(REQUIRED).Int(); | ||
| 25 | + | ||
| 26 | + ApplyMathAicpuDefaultCfg(*this); | ||
| 27 | + this->AICPU().ExtendCfgInfo(OP_INFO_OPS_FLAG.c_str(), OPEN_OPS_FLAG.c_str()); | ||
| 28 | + } | ||
| 29 | +}; | ||
| 30 | + | ||
| 31 | +OP_ADD(TileWithAxis); | ||
| 32 | +} // namespace ops | ||
| @@ -0,0 +1,393 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +using namespace std; | ||
| 28 | +using namespace aicpu; | ||
| 29 | + | ||
| 30 | +class TEST_TILEWITHAXIS_UT : public testing::Test {}; | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + auto node_def = CpuKernelUtils::CreateNodeDef(); \ | ||
| 34 | + NodeDefBuilder(node_def.get(), "TileWithAxis", "TileWithAxis") \ | ||
| 35 | + .Input({"x", data_types[0], shapes[0], datas[0]}) \ | ||
| 36 | + .Output({"y", data_types[1], shapes[1], datas[1]}) \ | ||
| 37 | + .Attr("axis", axis) \ | ||
| 38 | + .Attr("tiles", tiles); | ||
| 39 | + | ||
| 40 | +// Helper: compute expected output for TileWithAxis using modular indexing | ||
| 41 | +template <typename T> | ||
| 42 | +void ComputeTileExpected(const T* input, const vector<int64_t>& in_shape, T* expected, const vector<int64_t>& out_shape, | ||
| 43 | + int64_t axis) | ||
| 44 | +{ | ||
| 45 | + int64_t ndim = static_cast<int64_t>(in_shape.size()); | ||
| 46 | + if (axis < 0) { | ||
| 47 | + axis += ndim; | ||
| 48 | + } | ||
| 49 | + | ||
| 50 | + int64_t out_total = 1; | ||
| 51 | + for (int64_t i = 0; i < ndim; ++i) { | ||
| 52 | + out_total *= out_shape[i]; | ||
| 53 | + } | ||
| 54 | + | ||
| 55 | + for (int64_t out_idx = 0; out_idx < out_total; ++out_idx) { | ||
| 56 | + int64_t tmp = out_idx; | ||
| 57 | + vector<int64_t> out_indices(ndim); | ||
| 58 | + for (int64_t i = ndim - 1; i >= 0; --i) { | ||
| 59 | + out_indices[i] = tmp % out_shape[i]; | ||
| 60 | + tmp /= out_shape[i]; | ||
| 61 | + } | ||
| 62 | + | ||
| 63 | + vector<int64_t> in_indices = out_indices; | ||
| 64 | + in_indices[axis] = out_indices[axis] % in_shape[axis]; | ||
| 65 | + | ||
| 66 | + int64_t in_idx = 0; | ||
| 67 | + int64_t stride = 1; | ||
| 68 | + for (int64_t i = ndim - 1; i >= 0; --i) { | ||
| 69 | + in_idx += in_indices[i] * stride; | ||
| 70 | + stride *= in_shape[i]; | ||
| 71 | + } | ||
| 72 | + | ||
| 73 | + expected[out_idx] = input[in_idx]; | ||
| 74 | + } | ||
| 75 | +} | ||
| 76 | + | ||
| 77 | +// Test 1: Float 3D, axis=1, tiles=2 | ||
| 78 | +TEST_F(TEST_TILEWITHAXIS_UT, Float3D_Axis1_Tiles2) | ||
| 79 | +{ | ||
| 80 | + vector<DataType> data_types = {DT_FLOAT, DT_FLOAT}; | ||
| 81 | + vector<int64_t> in_shape = {2, 3, 2}; | ||
| 82 | + vector<int64_t> out_shape = {2, 6, 2}; | ||
| 83 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 84 | + | ||
| 85 | + float input[12] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f}; | ||
| 86 | + float output[24] = {0}; | ||
| 87 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 88 | + | ||
| 89 | + CREATE_NODEDEF(shapes, data_types, datas, 1, 2); | ||
| 90 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 91 | + | ||
| 92 | + float expect[24] = {0}; | ||
| 93 | + ComputeTileExpected<float>(input, in_shape, expect, out_shape, 1); | ||
| 94 | + CompareResult<float>(output, expect, 24); | ||
| 95 | +} | ||
| 96 | + | ||
| 97 | +// Test 2: Int32 3D, axis=1, tiles=2 | ||
| 98 | +TEST_F(TEST_TILEWITHAXIS_UT, Int32_3D_Axis1_Tiles2) | ||
| 99 | +{ | ||
| 100 | + vector<DataType> data_types = {DT_INT32, DT_INT32}; | ||
| 101 | + vector<int64_t> in_shape = {3, 2, 1}; | ||
| 102 | + vector<int64_t> out_shape = {3, 4, 1}; | ||
| 103 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 104 | + | ||
| 105 | + int32_t input[6] = {1, 2, 3, 4, 5, 6}; | ||
| 106 | + int32_t output[12] = {0}; | ||
| 107 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 108 | + | ||
| 109 | + CREATE_NODEDEF(shapes, data_types, datas, 1, 2); | ||
| 110 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 111 | + | ||
| 112 | + int32_t expect[12] = {0}; | ||
| 113 | + ComputeTileExpected<int32_t>(input, in_shape, expect, out_shape, 1); | ||
| 114 | + CompareResult<int32_t>(output, expect, 12); | ||
| 115 | +} | ||
| 116 | + | ||
| 117 | +// Test 3: Int64 4D, axis=1, tiles=2 | ||
| 118 | +TEST_F(TEST_TILEWITHAXIS_UT, Int64_4D_Axis1_Tiles2) | ||
| 119 | +{ | ||
| 120 | + vector<DataType> data_types = {DT_INT64, DT_INT64}; | ||
| 121 | + vector<int64_t> in_shape = {1, 2, 3, 4}; | ||
| 122 | + vector<int64_t> out_shape = {1, 4, 3, 4}; | ||
| 123 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 124 | + | ||
| 125 | + int64_t input[24]; | ||
| 126 | + for (int i = 0; i < 24; ++i) { | ||
| 127 | + input[i] = i + 1; | ||
| 128 | + } | ||
| 129 | + int64_t output[48] = {0}; | ||
| 130 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 131 | + | ||
| 132 | + CREATE_NODEDEF(shapes, data_types, datas, 1, 2); | ||
| 133 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 134 | + | ||
| 135 | + int64_t expect[48] = {0}; | ||
| 136 | + ComputeTileExpected<int64_t>(input, in_shape, expect, out_shape, 1); | ||
| 137 | + CompareResult<int64_t>(output, expect, 48); | ||
| 138 | +} | ||
| 139 | + | ||
| 140 | +// Test 4: Negative axis (-3 on 4D = axis 1) | ||
| 141 | +TEST_F(TEST_TILEWITHAXIS_UT, Int64_4D_NegAxis) | ||
| 142 | +{ | ||
| 143 | + vector<DataType> data_types = {DT_INT64, DT_INT64}; | ||
| 144 | + vector<int64_t> in_shape = {1, 2, 3, 4}; | ||
| 145 | + vector<int64_t> out_shape = {1, 4, 3, 4}; | ||
| 146 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 147 | + | ||
| 148 | + int64_t input[24]; | ||
| 149 | + for (int i = 0; i < 24; ++i) { | ||
| 150 | + input[i] = i + 1; | ||
| 151 | + } | ||
| 152 | + int64_t output[48] = {0}; | ||
| 153 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 154 | + | ||
| 155 | + CREATE_NODEDEF(shapes, data_types, datas, -3, 2); | ||
| 156 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 157 | + | ||
| 158 | + int64_t expect[48] = {0}; | ||
| 159 | + ComputeTileExpected<int64_t>(input, in_shape, expect, out_shape, -3); | ||
| 160 | + CompareResult<int64_t>(output, expect, 48); | ||
| 161 | +} | ||
| 162 | + | ||
| 163 | +// Test 5: tiles=1, identity copy | ||
| 164 | +TEST_F(TEST_TILEWITHAXIS_UT, Int32_Tiles1_Identity) | ||
| 165 | +{ | ||
| 166 | + vector<DataType> data_types = {DT_INT32, DT_INT32}; | ||
| 167 | + vector<int64_t> in_shape = {2, 3}; | ||
| 168 | + vector<int64_t> out_shape = {2, 3}; | ||
| 169 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 170 | + | ||
| 171 | + int32_t input[6] = {1, 2, 3, 4, 5, 6}; | ||
| 172 | + int32_t output[6] = {0}; | ||
| 173 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 174 | + | ||
| 175 | + CREATE_NODEDEF(shapes, data_types, datas, 0, 1); | ||
| 176 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 177 | + | ||
| 178 | + int32_t expect[6] = {0}; | ||
| 179 | + ComputeTileExpected<int32_t>(input, in_shape, expect, out_shape, 0); | ||
| 180 | + CompareResult<int32_t>(output, expect, 6); | ||
| 181 | +} | ||
| 182 | + | ||
| 183 | +// Test 6: Int8 3D, axis=2, tiles=2 | ||
| 184 | +TEST_F(TEST_TILEWITHAXIS_UT, Int8_3D_Axis2_Tiles2) | ||
| 185 | +{ | ||
| 186 | + vector<DataType> data_types = {DT_INT8, DT_INT8}; | ||
| 187 | + vector<int64_t> in_shape = {2, 3, 1}; | ||
| 188 | + vector<int64_t> out_shape = {2, 3, 2}; | ||
| 189 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 190 | + | ||
| 191 | + int8_t input[6] = {1, 2, 3, 4, 5, 6}; | ||
| 192 | + int8_t output[12] = {0}; | ||
| 193 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 194 | + | ||
| 195 | + CREATE_NODEDEF(shapes, data_types, datas, 2, 2); | ||
| 196 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 197 | + | ||
| 198 | + int8_t expect[12] = {0}; | ||
| 199 | + ComputeTileExpected<int8_t>(input, in_shape, expect, out_shape, 2); | ||
| 200 | + CompareResult<int8_t>(output, expect, 12); | ||
| 201 | +} | ||
| 202 | + | ||
| 203 | +// Test 7: Uint8 3D, axis=2, tiles=2 | ||
| 204 | +TEST_F(TEST_TILEWITHAXIS_UT, Uint8_3D_Axis2_Tiles2) | ||
| 205 | +{ | ||
| 206 | + vector<DataType> data_types = {DT_UINT8, DT_UINT8}; | ||
| 207 | + vector<int64_t> in_shape = {2, 3, 1}; | ||
| 208 | + vector<int64_t> out_shape = {2, 3, 2}; | ||
| 209 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 210 | + | ||
| 211 | + uint8_t input[6] = {10, 20, 30, 40, 50, 60}; | ||
| 212 | + uint8_t output[12] = {0}; | ||
| 213 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 214 | + | ||
| 215 | + CREATE_NODEDEF(shapes, data_types, datas, 2, 2); | ||
| 216 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 217 | + | ||
| 218 | + uint8_t expect[12] = {0}; | ||
| 219 | + ComputeTileExpected<uint8_t>(input, in_shape, expect, out_shape, 2); | ||
| 220 | + CompareResult<uint8_t>(output, expect, 12); | ||
| 221 | +} | ||
| 222 | + | ||
| 223 | +// Test 8: Int16 3D, axis=0, tiles=2 | ||
| 224 | +TEST_F(TEST_TILEWITHAXIS_UT, Int16_3D_Axis0_Tiles2) | ||
| 225 | +{ | ||
| 226 | + vector<DataType> data_types = {DT_INT16, DT_INT16}; | ||
| 227 | + vector<int64_t> in_shape = {2, 3, 2}; | ||
| 228 | + vector<int64_t> out_shape = {4, 3, 2}; | ||
| 229 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 230 | + | ||
| 231 | + int16_t input[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; | ||
| 232 | + int16_t output[24] = {0}; | ||
| 233 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 234 | + | ||
| 235 | + CREATE_NODEDEF(shapes, data_types, datas, 0, 2); | ||
| 236 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 237 | + | ||
| 238 | + int16_t expect[24] = {0}; | ||
| 239 | + ComputeTileExpected<int16_t>(input, in_shape, expect, out_shape, 0); | ||
| 240 | + CompareResult<int16_t>(output, expect, 24); | ||
| 241 | +} | ||
| 242 | + | ||
| 243 | +// Test 9: Uint16 3D, axis=0, tiles=2 | ||
| 244 | +TEST_F(TEST_TILEWITHAXIS_UT, Uint16_3D_Axis0_Tiles2) | ||
| 245 | +{ | ||
| 246 | + vector<DataType> data_types = {DT_UINT16, DT_UINT16}; | ||
| 247 | + vector<int64_t> in_shape = {2, 3, 2}; | ||
| 248 | + vector<int64_t> out_shape = {4, 3, 2}; | ||
| 249 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 250 | + | ||
| 251 | + uint16_t input[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; | ||
| 252 | + uint16_t output[24] = {0}; | ||
| 253 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 254 | + | ||
| 255 | + CREATE_NODEDEF(shapes, data_types, datas, 0, 2); | ||
| 256 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 257 | + | ||
| 258 | + uint16_t expect[24] = {0}; | ||
| 259 | + ComputeTileExpected<uint16_t>(input, in_shape, expect, out_shape, 0); | ||
| 260 | + CompareResult<uint16_t>(output, expect, 24); | ||
| 261 | +} | ||
| 262 | + | ||
| 263 | +// Test 10: Uint32 3D, axis=1, tiles=2 | ||
| 264 | +TEST_F(TEST_TILEWITHAXIS_UT, Uint32_3D_Axis1_Tiles2) | ||
| 265 | +{ | ||
| 266 | + vector<DataType> data_types = {DT_UINT32, DT_UINT32}; | ||
| 267 | + vector<int64_t> in_shape = {3, 2, 1}; | ||
| 268 | + vector<int64_t> out_shape = {3, 4, 1}; | ||
| 269 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 270 | + | ||
| 271 | + uint32_t input[6] = {100, 200, 300, 400, 500, 600}; | ||
| 272 | + uint32_t output[12] = {0}; | ||
| 273 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 274 | + | ||
| 275 | + CREATE_NODEDEF(shapes, data_types, datas, 1, 2); | ||
| 276 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 277 | + | ||
| 278 | + uint32_t expect[12] = {0}; | ||
| 279 | + ComputeTileExpected<uint32_t>(input, in_shape, expect, out_shape, 1); | ||
| 280 | + CompareResult<uint32_t>(output, expect, 12); | ||
| 281 | +} | ||
| 282 | + | ||
| 283 | +// Test 11: Uint64 3D, axis=1, tiles=2 | ||
| 284 | +TEST_F(TEST_TILEWITHAXIS_UT, Uint64_3D_Axis1_Tiles2) | ||
| 285 | +{ | ||
| 286 | + vector<DataType> data_types = {DT_UINT64, DT_UINT64}; | ||
| 287 | + vector<int64_t> in_shape = {1, 2, 3}; | ||
| 288 | + vector<int64_t> out_shape = {1, 4, 3}; | ||
| 289 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 290 | + | ||
| 291 | + uint64_t input[6] = {1, 2, 3, 4, 5, 6}; | ||
| 292 | + uint64_t output[12] = {0}; | ||
| 293 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 294 | + | ||
| 295 | + CREATE_NODEDEF(shapes, data_types, datas, 1, 2); | ||
| 296 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 297 | + | ||
| 298 | + uint64_t expect[12] = {0}; | ||
| 299 | + ComputeTileExpected<uint64_t>(input, in_shape, expect, out_shape, 1); | ||
| 300 | + CompareResult<uint64_t>(output, expect, 12); | ||
| 301 | +} | ||
| 302 | + | ||
| 303 | +// Test 12: 1D input, axis=0, tiles=2 | ||
| 304 | +TEST_F(TEST_TILEWITHAXIS_UT, Int32_1D_Axis0_Tiles2) | ||
| 305 | +{ | ||
| 306 | + vector<DataType> data_types = {DT_INT32, DT_INT32}; | ||
| 307 | + vector<int64_t> in_shape = {2}; | ||
| 308 | + vector<int64_t> out_shape = {4}; | ||
| 309 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 310 | + | ||
| 311 | + int32_t input[2] = {42, 99}; | ||
| 312 | + int32_t output[4] = {0}; | ||
| 313 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 314 | + | ||
| 315 | + CREATE_NODEDEF(shapes, data_types, datas, 0, 2); | ||
| 316 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 317 | + | ||
| 318 | + int32_t expect[4] = {0}; | ||
| 319 | + ComputeTileExpected<int32_t>(input, in_shape, expect, out_shape, 0); | ||
| 320 | + CompareResult<int32_t>(output, expect, 4); | ||
| 321 | +} | ||
| 322 | + | ||
| 323 | +// Test 13: 2D input, axis=0, tiles=2 | ||
| 324 | +TEST_F(TEST_TILEWITHAXIS_UT, Int32_2D_Axis0_Tiles2) | ||
| 325 | +{ | ||
| 326 | + vector<DataType> data_types = {DT_INT32, DT_INT32}; | ||
| 327 | + vector<int64_t> in_shape = {1, 2}; | ||
| 328 | + vector<int64_t> out_shape = {2, 2}; | ||
| 329 | + vector<vector<int64_t>> shapes = {in_shape, out_shape}; | ||
| 330 | + | ||
| 331 | + int32_t input[2] = {7, 8}; | ||
| 332 | + int32_t output[4] = {0}; | ||
| 333 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 334 | + | ||
| 335 | + CREATE_NODEDEF(shapes, data_types, datas, 0, 2); | ||
| 336 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 337 | + | ||
| 338 | + int32_t expect[4] = {0}; | ||
| 339 | + ComputeTileExpected<int32_t>(input, in_shape, expect, out_shape, 0); | ||
| 340 | + CompareResult<int32_t>(output, expect, 4); | ||
| 341 | +} | ||
| 342 | + | ||
| 343 | +// Test 14: FAIL - axis out of range (axis=4, 2D input) | ||
| 344 | +TEST_F(TEST_TILEWITHAXIS_UT, Fail_AxisOutOfRange) | ||
| 345 | +{ | ||
| 346 | + vector<DataType> data_types = {DT_INT64, DT_INT64}; | ||
| 347 | + vector<vector<int64_t>> shapes = {{1, 2}, {2, 2}}; | ||
| 348 | + int64_t input[2] = {0}; | ||
| 349 | + int64_t output[4] = {0}; | ||
| 350 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 351 | + | ||
| 352 | + CREATE_NODEDEF(shapes, data_types, datas, 4, 2); | ||
| 353 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 354 | +} | ||
| 355 | + | ||
| 356 | +// Test 15: FAIL - negative axis out of range (axis=-4, 2D input) | ||
| 357 | +TEST_F(TEST_TILEWITHAXIS_UT, Fail_NegAxisOutOfRange) | ||
| 358 | +{ | ||
| 359 | + vector<DataType> data_types = {DT_INT64, DT_INT64}; | ||
| 360 | + vector<vector<int64_t>> shapes = {{1, 2}, {2, 2}}; | ||
| 361 | + int64_t input[2] = {0}; | ||
| 362 | + int64_t output[4] = {0}; | ||
| 363 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 364 | + | ||
| 365 | + CREATE_NODEDEF(shapes, data_types, datas, -4, 2); | ||
| 366 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 367 | +} | ||
| 368 | + | ||
| 369 | +// Test 16: FAIL - negative tiles (tiles=-2) | ||
| 370 | +TEST_F(TEST_TILEWITHAXIS_UT, Fail_NegativeTiles) | ||
| 371 | +{ | ||
| 372 | + vector<DataType> data_types = {DT_INT64, DT_INT64}; | ||
| 373 | + vector<vector<int64_t>> shapes = {{1, 2}, {2, 2}}; | ||
| 374 | + int64_t input[2] = {0}; | ||
| 375 | + int64_t output[4] = {0}; | ||
| 376 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 377 | + | ||
| 378 | + CREATE_NODEDEF(shapes, data_types, datas, 0, -2); | ||
| 379 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 380 | +} | ||
| 381 | + | ||
| 382 | +// Test 17: FAIL - shape mismatch (input[axis]*tiles != output[axis]) | ||
| 383 | +TEST_F(TEST_TILEWITHAXIS_UT, Fail_ShapeMismatch) | ||
| 384 | +{ | ||
| 385 | + vector<DataType> data_types = {DT_INT64, DT_INT64}; | ||
| 386 | + vector<vector<int64_t>> shapes = {{1, 2}, {4, 2}}; | ||
| 387 | + int64_t input[2] = {0}; | ||
| 388 | + int64_t output[8] = {0}; | ||
| 389 | + vector<void*> datas = {(void*)input, (void*)output}; | ||
| 390 | + | ||
| 391 | + CREATE_NODEDEF(shapes, data_types, datas, 1, 2); | ||
| 392 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 393 | +} | ||
| @@ -2066,6 +2066,26 @@ | |||
| 2066 | <td>AI CPU</td> | 2066 | <td>AI CPU</td> |
| 2067 | <td>沿指定维度找出输入张量中最大或最小的k个元素及其索引。</td> | 2067 | <td>沿指定维度找出输入张量中最大或最小的k个元素及其索引。</td> |
| 2068 | </tr> | 2068 | </tr> |
| 2069 | + <tr> | ||
| 2070 | + <td>math</td> | ||
| 2071 | + <td><a href="../../math/top_k_v2_d/README.md">top_k_v2_d</a></td> | ||
| 2072 | + <td>×</td> | ||
| 2073 | + <td>√</td> | ||
| 2074 | + <td>×</td> | ||
| 2075 | + <td>√</td> | ||
| 2076 | + <td>AI CPU</td> | ||
| 2077 | + <td>沿指定维度找出输入张量中最大或最小的k个元素及其索引，兼容TensorFlow TopKV2。</td> | ||
| 2078 | + </tr> | ||
| 2079 | + <tr> | ||
| 2080 | + <td>math</td> | ||
| 2081 | + <td><a href="../../math/top_k_pq_distance/README.md">top_k_pq_distance</a></td> | ||
| 2082 | + <td>×</td> | ||
| 2083 | + <td>×</td> | ||
| 2084 | + <td>×</td> | ||
| 2085 | + <td>√</td> | ||
| 2086 | + <td>AI CPU</td> | ||
| 2087 | + <td>基于分组极值剪枝的TopK距离计算，用于向量检索场景。</td> | ||
| 2088 | + </tr> | ||
| 2069 | <tr> | 2089 | <tr> |
| 2070 | <td>math</td> | 2090 | <td>math</td> |
| 2071 | <td><a href="../../math/trace/README.md">trace</a></td> | 2091 | <td><a href="../../math/trace/README.md">trace</a></td> |
| @@ -2156,6 +2176,16 @@ | |||
| 2156 | <td>AI Core</td> | 2176 | <td>AI Core</td> |
| 2157 | <td>计算x * log(1 + y)，当x == 0时结果为0。支持broadcast。</td> | 2177 | <td>计算x * log(1 + y)，当x == 0时结果为0。支持broadcast。</td> |
| 2158 | </tr> | 2178 | </tr> |
| 2179 | + <tr> | ||
| 2180 | + <td>math</td> | ||
| 2181 | + <td><a href="../../math/zeta/README.md">zeta</a></td> | ||
| 2182 | + <td>×</td> | ||
| 2183 | + <td>×</td> | ||
| 2184 | + <td>×</td> | ||
| 2185 | + <td>√</td> | ||
| 2186 | + <td>AI CPU</td> | ||
| 2187 | + <td>计算Hurwitz zeta函数 ζ(x, q)。</td> | ||
| 2188 | + </tr> | ||
| 2159 | <tr> | 2189 | <tr> |
| 2160 | <td>math</td> | 2190 | <td>math</td> |
| 2161 | <td><a href="../../math/zero_op/README.md">zeros_like</a></td> | 2191 | <td><a href="../../math/zero_op/README.md">zeros_like</a></td> |
| @@ -2826,6 +2856,16 @@ | |||
| 2826 | <td>AI Core</td> | 2856 | <td>AI Core</td> |
| 2827 | <td>将输入tensor的值搬运到输出tensor中。</td> | 2857 | <td>将输入tensor的值搬运到输出tensor中。</td> |
| 2828 | </tr> | 2858 | </tr> |
| 2859 | + <tr> | ||
| 2860 | + <td>conversion</td> | ||
| 2861 | + <td><a href="../../conversion/tile_with_axis/README.md">tile_with_axis</a></td> | ||
| 2862 | + <td>√</td> | ||
| 2863 | + <td>√</td> | ||
| 2864 | + <td>×</td> | ||
| 2865 | + <td>√</td> | ||
| 2866 | + <td>AI Core/AI CPU</td> | ||
| 2867 | + <td>沿指定维度复制输入Tensor数据，扩展输出Tensor。</td> | ||
| 2868 | + </tr> | ||
| 2829 | <tr> | 2869 | <tr> |
| 2830 | <td>conversion</td> | 2870 | <td>conversion</td> |
| 2831 | <td><a href="../../conversion/trans_data/README.md">trans_data</a></td> | 2871 | <td><a href="../../conversion/trans_data/README.md">trans_data</a></td> |
| @@ -0,0 +1,12 @@ | |||
| 1 | +# ---------------------------------------------------------------------------- | ||
| 2 | +# This program is free software, you can redistribute it and/or modify it. | ||
| 3 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 4 | +# This file is a part of the CANN Open Software. | ||
| 5 | +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 6 | +# Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 7 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING | ||
| 8 | +# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See LICENSE in the root of the software repository for the full text of the License. | ||
| 10 | +# ---------------------------------------------------------------------------- | ||
| 11 | + | ||
| 12 | +add_all_modules_sources(OPTYPE top_k_pq_distance ACLNNTYPE aclnn_exclude) | ||
| @@ -0,0 +1,50 @@ | |||
| 1 | +# TopKPQDistance | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 产品 | 是否支持 | | ||
| 6 | +| :----------------------------------------------------------- | :------: | | ||
| 7 | +| Ascend 950PR/Ascend 950DT | √ | | ||
| 8 | +| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ | | ||
| 9 | +| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ | | ||
| 10 | +| Atlas 200I/500 A2 推理产品 | √ | | ||
| 11 | +| Atlas 推理系列产品 | √ | | ||
| 12 | +| Atlas 训练系列产品 | √ | | ||
| 13 | + | ||
| 14 | +## 功能说明 | ||
| 15 | + | ||
| 16 | +- 算子功能：基于分组极值剪枝的TopK距离计算，用于向量检索场景。从多组PQ距离数据中找出最大或最小的k个距离及其对应的ivf和index。 | ||
| 17 | + | ||
| 18 | +- 计算流程： | ||
| 19 | + 1. 从grouped_extreme_distance中选出TopK个极值（分组剪枝） | ||
| 20 | + 2. 从选中的分组中展开计算精确TopK距离 | ||
| 21 | + 3. 输出TopK距离、ivf和index | ||
| 22 | + | ||
| 23 | +## 参数说明 | ||
| 24 | + | ||
| 25 | +| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 | | ||
| 26 | +|--------|---------------|------|----------|----------| | ||
| 27 | +| actual_count | 输入（动态） | 每组实际元素个数 | INT32 | ND | | ||
| 28 | +| pq_distance | 输入（动态） | PQ距离数据 | FLOAT16、FLOAT | ND | | ||
| 29 | +| grouped_extreme_distance | 输入（动态） | 每组极值距离 | FLOAT16、FLOAT | ND | | ||
| 30 | +| pq_ivf | 输入（动态） | IVF索引 | INT32 | ND | | ||
| 31 | +| pq_index | 输入（动态） | PQ索引 | INT32 | ND | | ||
| 32 | +| order | 属性（可选） | 排序方式，"ASC"升序或"DES"降序，默认"ASC" | STRING | - | | ||
| 33 | +| k | 属性（必选） | TopK的k值 | INT | - | | ||
| 34 | +| group_size | 属性（必选） | 分组大小 | INT | - | | ||
| 35 | +| topk_distance | 输出 | TopK距离值 | FLOAT16、FLOAT | ND | | ||
| 36 | +| topk_ivf | 输出 | TopK对应的ivf | INT32 | ND | | ||
| 37 | +| topk_index | 输出 | TopK对应的index | INT32 | ND | | ||
| 38 | + | ||
| 39 | +## 约束说明 | ||
| 40 | + | ||
| 41 | +- k不能为0，且不能大于actual_count的总和。 | ||
| 42 | +- group_size不能为0。 | ||
| 43 | +- actual_count必须能被group_size整除。 | ||
| 44 | +- 支持多组动态输入（data_batch），每组包含5个输入。 | ||
| 45 | + | ||
| 46 | +## 调用说明 | ||
| 47 | + | ||
| 48 | +| 调用方式 | 样例代码 | 说明 | | ||
| 49 | +|------------|----------|------| | ||
| 50 | +| 图模式调用 | [test_geir_top_k_pq_distance](./examples/test_geir_top_k_pq_distance.cpp) | 通过[算子IR](./op_graph/top_k_pq_distance_proto.h)构图方式调用TopKPQDistance算子。 | | ||
| @@ -0,0 +1,287 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +/*! | ||
| 12 | + * \file test_geir_top_k_pq_distance.cpp | ||
| 13 | + * \brief GE IR test for TopKPQDistance operator | ||
| 14 | + */ | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + | ||
| 35 | + | ||
| 36 | + | ||
| 37 | + | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +using namespace ge; | ||
| 41 | +using std::map; | ||
| 42 | +using std::string; | ||
| 43 | +using std::vector; | ||
| 44 | + | ||
| 45 | +string GetTime() | ||
| 46 | +{ | ||
| 47 | + time_t timep; | ||
| 48 | + time(&timep); | ||
| 49 | + char tmp[64]; | ||
| 50 | + strftime(tmp, sizeof(tmp), "%Y-%m-%d %H:%M:%S,000", localtime(&timep)); | ||
| 51 | + return tmp; | ||
| 52 | +} | ||
| 53 | + | ||
| 54 | +uint32_t GetDataTypeSize(DataType dt) | ||
| 55 | +{ | ||
| 56 | + switch (dt) { | ||
| 57 | + case ge::DT_BOOL: | ||
| 58 | + return 1U; | ||
| 59 | + case ge::DT_INT8: | ||
| 60 | + case ge::DT_UINT8: | ||
| 61 | + return 1U; | ||
| 62 | + case ge::DT_FLOAT16: | ||
| 63 | + case ge::DT_INT16: | ||
| 64 | + case ge::DT_UINT16: | ||
| 65 | + return 2U; | ||
| 66 | + case ge::DT_FLOAT: | ||
| 67 | + case ge::DT_INT32: | ||
| 68 | + case ge::DT_UINT32: | ||
| 69 | + return 4U; | ||
| 70 | + case ge::DT_INT64: | ||
| 71 | + case ge::DT_UINT64: | ||
| 72 | + return 8U; | ||
| 73 | + default: | ||
| 74 | + return 0U; | ||
| 75 | + } | ||
| 76 | +} | ||
| 77 | + | ||
| 78 | +int32_t GenData(vector<int64_t> shapes, Tensor& tensor, TensorDesc& desc, DataType dt, double value) | ||
| 79 | +{ | ||
| 80 | + desc.SetRealDimCnt(shapes.size()); | ||
| 81 | + size_t size = 1; | ||
| 82 | + for (size_t i = 0; i < shapes.size(); i++) { | ||
| 83 | + size *= static_cast<size_t>(shapes[i]); | ||
| 84 | + } | ||
| 85 | + uint32_t type_size = GetDataTypeSize(dt); | ||
| 86 | + if (type_size == 0U) { | ||
| 87 | + printf("%s - ERROR - [XIR]: GenData: unsupported data type\n", GetTime().c_str()); | ||
| 88 | + return FAILED; | ||
| 89 | + } | ||
| 90 | + uint32_t data_len = static_cast<uint32_t>(size * type_size); | ||
| 91 | + uint8_t* buf = new (std::nothrow) uint8_t[data_len]; | ||
| 92 | + if (buf == nullptr) { | ||
| 93 | + return FAILED; | ||
| 94 | + } | ||
| 95 | + if (dt == DT_FLOAT) { | ||
| 96 | + float* p = reinterpret_cast<float*>(buf); | ||
| 97 | + for (size_t i = 0; i < size; i++) { | ||
| 98 | + p[i] = static_cast<float>(value + static_cast<double>(i)); | ||
| 99 | + } | ||
| 100 | + } else if (dt == DT_INT32) { | ||
| 101 | + int32_t* p = reinterpret_cast<int32_t*>(buf); | ||
| 102 | + for (size_t i = 0; i < size; i++) { | ||
| 103 | + p[i] = static_cast<int32_t>(value + static_cast<double>(i)); | ||
| 104 | + } | ||
| 105 | + } else { | ||
| 106 | + delete[] buf; | ||
| 107 | + return FAILED; | ||
| 108 | + } | ||
| 109 | + tensor = Tensor(desc, buf, data_len); | ||
| 110 | + delete[] buf; | ||
| 111 | + return SUCCESS; | ||
| 112 | +} | ||
| 113 | + | ||
| 114 | +int32_t WriteDataToFile(string bin_file, uint64_t data_size, uint8_t* inputData) | ||
| 115 | +{ | ||
| 116 | + FILE* fp = fopen(bin_file.c_str(), "wb"); | ||
| 117 | + if (fp == nullptr) { | ||
| 118 | + return FAILED; | ||
| 119 | + } | ||
| 120 | + size_t written = fwrite(inputData, 1, data_size, fp); | ||
| 121 | + fclose(fp); | ||
| 122 | + if (written != data_size) { | ||
| 123 | + return FAILED; | ||
| 124 | + } | ||
| 125 | + return SUCCESS; | ||
| 126 | +} | ||
| 127 | + | ||
| 128 | +void ProcessInputData(vector<Tensor>& input) | ||
| 129 | +{ | ||
| 130 | + for (size_t i = 0; i < input.size(); i++) { | ||
| 131 | + string input_file = "./tc_ge_irrun_test_0008_npu_input_" + std::to_string(i) + ".bin"; | ||
| 132 | + uint8_t* data = input[i].GetData(); | ||
| 133 | + int64_t shape_size = input[i].GetTensorDesc().GetShape().GetShapeSize(); | ||
| 134 | + uint32_t type_size = GetDataTypeSize(input[i].GetTensorDesc().GetDataType()); | ||
| 135 | + if (type_size == 0U) { | ||
| 136 | + printf("ERROR: input %zu has unsupported dtype\n", i); | ||
| 137 | + continue; | ||
| 138 | + } | ||
| 139 | + uint32_t data_size = static_cast<uint32_t>(shape_size * type_size); | ||
| 140 | + WriteDataToFile(input_file.c_str(), data_size, data); | ||
| 141 | + } | ||
| 142 | +} | ||
| 143 | + | ||
| 144 | +void ProcessOutputData(vector<Tensor>& output) | ||
| 145 | +{ | ||
| 146 | + for (size_t i = 0; i < output.size(); i++) { | ||
| 147 | + string output_file = "./tc_ge_irrun_test_0008_npu_output_" + std::to_string(i) + ".bin"; | ||
| 148 | + uint8_t* data = output[i].GetData(); | ||
| 149 | + int64_t shape_size = output[i].GetTensorDesc().GetShape().GetShapeSize(); | ||
| 150 | + uint32_t type_size = GetDataTypeSize(output[i].GetTensorDesc().GetDataType()); | ||
| 151 | + if (type_size == 0U) { | ||
| 152 | + printf("ERROR: output %zu has unsupported dtype\n", i); | ||
| 153 | + continue; | ||
| 154 | + } | ||
| 155 | + uint32_t data_size = static_cast<uint32_t>(shape_size * type_size); | ||
| 156 | + WriteDataToFile(output_file.c_str(), data_size, data); | ||
| 157 | + } | ||
| 158 | +} | ||
| 159 | + | ||
| 160 | + | ||
| 161 | + do { \ | ||
| 162 | + auto dataName = op::Data(nodeName).set_attr_index(attrIndex); \ | ||
| 163 | + TensorDesc dataName##Desc(ge::Shape(inputShape), FORMAT_ND, inputDtype); \ | ||
| 164 | + dataName##Desc.SetPlacement(ge::kPlacementHost); \ | ||
| 165 | + dataName##Desc.SetFormat(FORMAT_ND); \ | ||
| 166 | + Tensor dataName##Tensor; \ | ||
| 167 | + GenData(inputShape, dataName##Tensor, dataName##Desc, inputDtype, genValue); \ | ||
| 168 | + input.push_back(dataName##Tensor); \ | ||
| 169 | + dataName.update_input_desc_x(dataName##Desc); \ | ||
| 170 | + dataName.update_output_desc_y(dataName##Desc); \ | ||
| 171 | + node.UpdateDynamicInputDesc(#portName, 0, dataName##Desc); \ | ||
| 172 | + node.set_dynamic_input_##portName(0, dataName); \ | ||
| 173 | + graph.AddOp(dataName); \ | ||
| 174 | + inputs.push_back(dataName); \ | ||
| 175 | + } while (0) | ||
| 176 | + | ||
| 177 | +int CreateOppInGraph(DataType inDtype, vector<Tensor>& input, vector<Operator>& inputs, vector<Operator>& outputs, | ||
| 178 | + Graph& graph) | ||
| 179 | +{ | ||
| 180 | + constexpr int32_t k = 3; | ||
| 181 | + constexpr int32_t group_size = 2; | ||
| 182 | + constexpr int32_t actual_count = 6; | ||
| 183 | + | ||
| 184 | + auto node = op::TopKPQDistance("topkpqdistance_1"); | ||
| 185 | + node.create_dynamic_input_actual_count(1); | ||
| 186 | + node.create_dynamic_input_pq_distance(1); | ||
| 187 | + node.create_dynamic_input_grouped_extreme_distance(1); | ||
| 188 | + node.create_dynamic_input_pq_ivf(1); | ||
| 189 | + node.create_dynamic_input_pq_index(1); | ||
| 190 | + | ||
| 191 | + vector<int64_t> actualCountShape = {1}; | ||
| 192 | + vector<int64_t> pqDistanceShape = {actual_count}; | ||
| 193 | + vector<int64_t> groupedExtremeShape = {actual_count / group_size}; | ||
| 194 | + vector<int64_t> topkShape = {k}; | ||
| 195 | + | ||
| 196 | + ADD_DYNAMIC_INPUT(actual_count, 0, DT_INT32, actualCountShape, "actual_count_data", acData, | ||
| 197 | + static_cast<double>(actual_count)); | ||
| 198 | + ADD_DYNAMIC_INPUT(pq_distance, 1, inDtype, pqDistanceShape, "pq_distance_data", pqData, 1.0); | ||
| 199 | + ADD_DYNAMIC_INPUT(grouped_extreme_distance, 2, inDtype, groupedExtremeShape, "grouped_extreme_data", gedData, 1.0); | ||
| 200 | + ADD_DYNAMIC_INPUT(pq_ivf, 3, DT_INT32, pqDistanceShape, "pq_ivf_data", ivfData, 1.0); | ||
| 201 | + ADD_DYNAMIC_INPUT(pq_index, 4, DT_INT32, pqDistanceShape, "pq_index_data", idxData, 1.0); | ||
| 202 | + | ||
| 203 | + node.set_attr_order("DES"); | ||
| 204 | + node.set_attr_k(k); | ||
| 205 | + node.set_attr_group_size(group_size); | ||
| 206 | + | ||
| 207 | + TensorDesc topkDistDesc(ge::Shape(topkShape), FORMAT_ND, inDtype); | ||
| 208 | + node.update_output_desc_topk_distance(topkDistDesc); | ||
| 209 | + TensorDesc topkIvfDesc(ge::Shape(topkShape), FORMAT_ND, DT_INT32); | ||
| 210 | + node.update_output_desc_topk_ivf(topkIvfDesc); | ||
| 211 | + TensorDesc topkIdxDesc(ge::Shape(topkShape), FORMAT_ND, DT_INT32); | ||
| 212 | + node.update_output_desc_topk_index(topkIdxDesc); | ||
| 213 | + | ||
| 214 | + graph.AddOp(node); | ||
| 215 | + outputs.push_back(node); | ||
| 216 | + return SUCCESS; | ||
| 217 | +} | ||
| 218 | + | ||
| 219 | +int main(int argc, char* argv[]) | ||
| 220 | +{ | ||
| 221 | + const char* graph_name = "tc_ge_irrun_test"; | ||
| 222 | + Graph graph(graph_name); | ||
| 223 | + vector<Tensor> input; | ||
| 224 | + | ||
| 225 | + printf("%s - INFO - [XIR]: Start to initialize ge\n", GetTime().c_str()); | ||
| 226 | + std::map<AscendString, AscendString> global_options = {{"ge.exec.deviceId", "0"}, {"ge.graphRunMode", "1"}}; | ||
| 227 | + Status ret = ge::GEInitialize(global_options); | ||
| 228 | + if (ret != SUCCESS) { | ||
| 229 | + printf("%s - ERROR - [XIR]: Initialize ge failed\n", GetTime().c_str()); | ||
| 230 | + return FAILED; | ||
| 231 | + } | ||
| 232 | + printf("%s - INFO - [XIR]: Initialize ge success\n", GetTime().c_str()); | ||
| 233 | + | ||
| 234 | + vector<Operator> inputs{}; | ||
| 235 | + vector<Operator> outputs{}; | ||
| 236 | + | ||
| 237 | + DataType inDtype = DT_FLOAT; | ||
| 238 | + | ||
| 239 | + ret = CreateOppInGraph(inDtype, input, inputs, outputs, graph); | ||
| 240 | + if (ret != SUCCESS) { | ||
| 241 | + printf("%s - ERROR - [XIR]: Create graph failed\n", GetTime().c_str()); | ||
| 242 | + return FAILED; | ||
| 243 | + } | ||
| 244 | + | ||
| 245 | + if (!inputs.empty() && !outputs.empty()) { | ||
| 246 | + graph.SetInputs(inputs).SetOutputs(outputs); | ||
| 247 | + } | ||
| 248 | + | ||
| 249 | + std::map<AscendString, AscendString> build_options = {}; | ||
| 250 | + ge::Session* session = new Session(build_options); | ||
| 251 | + if (session == nullptr) { | ||
| 252 | + printf("%s - ERROR - [XIR]: Create session failed\n", GetTime().c_str()); | ||
| 253 | + return FAILED; | ||
| 254 | + } | ||
| 255 | + | ||
| 256 | + uint32_t graph_id = 0; | ||
| 257 | + std::map<AscendString, AscendString> graph_options = {}; | ||
| 258 | + ret = session->AddGraph(graph_id, graph, graph_options); | ||
| 259 | + if (ret != SUCCESS) { | ||
| 260 | + printf("%s - ERROR - [XIR]: Add graph failed\n", GetTime().c_str()); | ||
| 261 | + delete session; | ||
| 262 | + GEFinalize(); | ||
| 263 | + return FAILED; | ||
| 264 | + } | ||
| 265 | + | ||
| 266 | + vector<Tensor> output; | ||
| 267 | + ret = session->RunGraph(graph_id, input, output); | ||
| 268 | + if (ret != SUCCESS) { | ||
| 269 | + printf("%s - ERROR - [XIR]: Run graph failed\n", GetTime().c_str()); | ||
| 270 | + delete session; | ||
| 271 | + GEFinalize(); | ||
| 272 | + return FAILED; | ||
| 273 | + } | ||
| 274 | + printf("%s - INFO - [XIR]: Run graph success\n", GetTime().c_str()); | ||
| 275 | + | ||
| 276 | + ProcessInputData(input); | ||
| 277 | + ProcessOutputData(output); | ||
| 278 | + | ||
| 279 | + delete session; | ||
| 280 | + ret = ge::GEFinalize(); | ||
| 281 | + if (ret != SUCCESS) { | ||
| 282 | + printf("%s - ERROR - [XIR]: Finalize failed\n", GetTime().c_str()); | ||
| 283 | + return FAILED; | ||
| 284 | + } | ||
| 285 | + printf("%s - INFO - [XIR]: Finalize ir graph session success\n", GetTime().c_str()); | ||
| 286 | + return SUCCESS; | ||
| 287 | +} | ||
| @@ -0,0 +1,52 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +namespace ge { | ||
| 17 | +/** | ||
| 18 | + * @brief Finds values and indices of the "k" largest or least elements for the last dimension. \n | ||
| 19 | + * | ||
| 20 | + * @par Inputs: | ||
| 21 | + * Dynamin inputs, including: | ||
| 22 | + * @li actual_count: A Tensor of type int32, the actual number of pq_distance. | ||
| 23 | + * @li pq_distance: A Tensor, Will be updated after calculation. Must be one of the following types: float32, float16. | ||
| 24 | + * @li grouped_extreme_distance: A Tensor, the extremum in each group. Must be one of the following types: float32, | ||
| 25 | + * float16. | ||
| 26 | + * @li pq_index: A Tensor of type int32, index corresponding to pq_distance. | ||
| 27 | + * @li pq_ivf: A Tensor of type int32 , the bucket number corresponding to pq_distance. | ||
| 28 | + * | ||
| 29 | + * @par Attributes: | ||
| 30 | + * @li order: A string, indicates the sorting method of topk_pq_distance. \n | ||
| 31 | + * @li k: Int, k maximum or minimum values. \n | ||
| 32 | + * @li group_size: Int, the group size of the extremum. \n | ||
| 33 | + * | ||
| 34 | + * @par Restrictions: | ||
| 35 | + * Warning: THIS FUNCTION IS EXPERIMENTAL. Please do not use. | ||
| 36 | + */ | ||
| 37 | +REG_OP(TopKPQDistance) | ||
| 38 | + .DYNAMIC_INPUT(actual_count, TensorType({DT_INT32})) | ||
| 39 | + .DYNAMIC_INPUT(pq_distance, TensorType({DT_FLOAT16, DT_FLOAT})) | ||
| 40 | + .DYNAMIC_INPUT(grouped_extreme_distance, TensorType({DT_FLOAT16, DT_FLOAT})) | ||
| 41 | + .DYNAMIC_INPUT(pq_ivf, TensorType({DT_INT32})) | ||
| 42 | + .DYNAMIC_INPUT(pq_index, TensorType({DT_INT32})) | ||
| 43 | + .OUTPUT(topk_distance, TensorType({DT_FLOAT16, DT_FLOAT})) | ||
| 44 | + .OUTPUT(topk_ivf, TensorType({DT_INT32})) | ||
| 45 | + .OUTPUT(topk_index, TensorType({DT_INT32})) | ||
| 46 | + .ATTR(order, String, "ASC") | ||
| 47 | + .REQUIRED_ATTR(k, Int) | ||
| 48 | + .REQUIRED_ATTR(group_size, Int) | ||
| 49 | + .OP_END_FACTORY_REG(TopKPQDistance) | ||
| 50 | +} // namespace ge | ||
| 51 | + | ||
| 52 | + | ||
| @@ -0,0 +1,328 @@ | |||||||||||||
| 1 | +/** | ||||||||||||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||||||||||||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||||||||||||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||||||||||||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||||||||||||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||||||||||||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||||||||||||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||||||||||||
| 9 | + */ | ||||||||||||
| 10 | + | ||||||||||||
| 11 | + | ||||||||||||
| 12 | + | ||||||||||||
| 13 | + | ||||||||||||
| 14 | + | ||||||||||||
| 15 | + | ||||||||||||
| 16 | + | ||||||||||||
| 17 | + | ||||||||||||
| 18 | + | ||||||||||||
| 19 | +namespace { | ||||||||||||
| 20 | +const char* const kTopKPQDistance = "TopKPQDistance"; | ||||||||||||
| 21 | +constexpr uint32_t kParamNum = 5; | ||||||||||||
| 22 | +const uint32_t kOutputNum = 3; | ||||||||||||
| 23 | +} // namespace | ||||||||||||
| 24 | + | ||||||||||||
| 25 | +namespace aicpu { | ||||||||||||
| 26 | +uint32_t TopKPQDistanceCpuKernel::Compute(CpuKernelContext& ctx) | ||||||||||||
| 27 | +{ | ||||||||||||
| 28 | + uint32_t inputs_size = ctx.GetInputsSize(); | ||||||||||||
| 29 | + KERNEL_HANDLE_ERROR(NormalCheck(ctx, inputs_size, kOutputNum), "Check TopKPQDistance params failed."); | ||||||||||||
| 30 | + data_batch_ = inputs_size / kParamNum; | ||||||||||||
| 31 | + Tensor* pq_distance = ctx.Input(data_batch_ * 1); | ||||||||||||
| 32 | + | ||||||||||||
| 33 | + DataType data_type = pq_distance->GetDataType(); | ||||||||||||
| 34 | + uint32_t res = KERNEL_STATUS_OK; | ||||||||||||
| 35 | + switch (data_type) { | ||||||||||||
| 36 | + case DT_FLOAT16: | ||||||||||||
| 37 | + KERNEL_LOG_DEBUG("TopKPQDistance compute DT_FLOAT16"); | ||||||||||||
| 38 | + res = DoCompute<FP16>(ctx); | ||||||||||||
| 39 | + break; | ||||||||||||
| 40 | + case DT_FLOAT: | ||||||||||||
| 41 | + KERNEL_LOG_DEBUG("TopKPQDistance compute DT_FLOAT"); | ||||||||||||
| 42 | + res = DoCompute<float>(ctx); | ||||||||||||
| 43 | + break; | ||||||||||||
| 44 | + default: | ||||||||||||
| 45 | + KERNEL_LOG_ERROR("TopKPQDistance input pq_distance only support type[DT_FLOAT16, " | ||||||||||||
| 46 | + "DT_FLOAT], but got type[%s]", | ||||||||||||
| 47 | + DTypeStr(data_type).c_str()); | ||||||||||||
| 48 | + return KERNEL_STATUS_PARAM_INVALID; | ||||||||||||
| 49 | + } | ||||||||||||
| 50 | + if (res != KERNEL_STATUS_OK) { | ||||||||||||
| 51 | + KERNEL_LOG_ERROR("TopKPQDistance kernel compute failed, KernelStatus is [%d]", res); | ||||||||||||
| 52 | + } | ||||||||||||
| 53 | + return res; | ||||||||||||
| 54 | +} | ||||||||||||
| 55 | + | ||||||||||||
| 56 | +template <typename T> | ||||||||||||
| 57 | +uint32_t TopKPQDistanceCpuKernel::DoCompute(const CpuKernelContext& ctx) | ||||||||||||
| 58 | +{ | ||||||||||||
| 59 | + InputsData<T> input_data; | ||||||||||||
| 60 | + uint32_t ret = GetInputAndCheck(ctx, input_data); | ||||||||||||
| 61 | + if (ret != KERNEL_STATUS_OK) { | ||||||||||||
| 62 | + return ret; | ||||||||||||
| 63 | + } | ||||||||||||
| 64 | + | ||||||||||||
| 65 | + Item<T> grp_extreme_ptr[k_]; | ||||||||||||
| 66 | + Item<T> topk_ptr[k_]; | ||||||||||||
🟠 High Priority
改动建议
  🟠 High Priority 文件 | |||||||||||||
| 67 | + ret = GetGroupedDistanceTopKHeap(grp_extreme_ptr, input_data); | ||||||||||||
| 68 | + if (ret != KERNEL_STATUS_OK) { | ||||||||||||
| 69 | + return ret; | ||||||||||||
| 70 | + } | ||||||||||||
| 71 | + ret = GetDistanceTopKHeap(topk_ptr, grp_extreme_ptr, input_data); | ||||||||||||
| 72 | + if (ret != KERNEL_STATUS_OK) { | ||||||||||||
| 73 | + return ret; | ||||||||||||
| 74 | + } | ||||||||||||
| 75 | + ret = ProcessResult(ctx, input_data, topk_ptr); | ||||||||||||
| 76 | + if (ret != KERNEL_STATUS_OK) { | ||||||||||||
| 77 | + return ret; | ||||||||||||
| 78 | + } | ||||||||||||
| 79 | + return KERNEL_STATUS_OK; | ||||||||||||
| 80 | +} | ||||||||||||
| 81 | + | ||||||||||||
| 82 | +template <typename T> | ||||||||||||
| 83 | +uint32_t TopKPQDistanceCpuKernel::GetInputAndCheck(const CpuKernelContext& ctx, InputsData<T>& input_data) | ||||||||||||
| 84 | +{ | ||||||||||||
| 85 | + auto* order = ctx.GetAttr("order"); | ||||||||||||
| 86 | + KERNEL_CHECK_NULLPTR(order, KERNEL_STATUS_PARAM_INVALID, "Get attr [order] failed."); | ||||||||||||
| 87 | + order_ = order->GetString(); | ||||||||||||
| 88 | + is_min_heap_ = ("ASC" == order_) ? false : true; | ||||||||||||
| 89 | + KERNEL_LOG_DEBUG("TopKPQDistance getInputAndCheck order_[%s], is_min_heap_ is[%d], " | ||||||||||||
| 90 | + "data_batch_ is[%d]", | ||||||||||||
| 91 | + order_.c_str(), is_min_heap_, data_batch_); | ||||||||||||
| 92 | + auto* k = ctx.GetAttr("k"); | ||||||||||||
| 93 | + KERNEL_CHECK_NULLPTR(k, KERNEL_STATUS_PARAM_INVALID, "Get attr [k] failed."); | ||||||||||||
| 94 | + k_ = static_cast<int32_t>(k->GetInt()); | ||||||||||||
| 95 | + KERNEL_CHECK_FALSE((k_ > 0), KERNEL_STATUS_PARAM_INVALID, "k_[%d] should be bigger than zero.", k_); | ||||||||||||
| 96 | + auto* group_size = ctx.GetAttr("group_size"); | ||||||||||||
| 97 | + KERNEL_CHECK_NULLPTR(group_size, KERNEL_STATUS_PARAM_INVALID, "Get attr [group_size] failed."); | ||||||||||||
| 98 | + group_size_ = static_cast<int32_t>(group_size->GetInt()); | ||||||||||||
| 99 | + KERNEL_CHECK_FALSE((group_size_ > 0), KERNEL_STATUS_PARAM_INVALID, "group_size[%d] should be bigger than zero.", | ||||||||||||
| 100 | + group_size_); | ||||||||||||
| 101 | + for (uint32_t i = 0; i < data_batch_; i++) { | ||||||||||||
| 102 | + Tensor* actual_count_tensor = ctx.Input(data_batch_ * 0 + i); | ||||||||||||
| 103 | + KERNEL_CHECK_NULLPTR(actual_count_tensor->GetData(), KERNEL_STATUS_PARAM_INVALID, | ||||||||||||
| 104 | + "actual_count tensor data is null."); | ||||||||||||
| 105 | + int32_t actual_count = *(static_cast<int32_t*>(actual_count_tensor->GetData())); | ||||||||||||
| 106 | + KERNEL_CHECK_FALSE((actual_count >= 0 && actual_count % group_size_ == 0), KERNEL_STATUS_PARAM_INVALID, | ||||||||||||
| 107 | + "actual_count[%d] should be non-negative and an integer multiple of group_size[%d].", | ||||||||||||
| 108 | + actual_count, group_size_); | ||||||||||||
| 109 | + input_data.actual_count += actual_count; | ||||||||||||
| 110 | + | ||||||||||||
| 111 | + Tensor* pq_distance_tensor = ctx.Input(data_batch_ * 1 + i); | ||||||||||||
| 112 | + input_data.pq_distances.Add(static_cast<T*>(pq_distance_tensor->GetData()), actual_count); | ||||||||||||
| 113 | + | ||||||||||||
| 114 | + Tensor* grouped_extreme_distance_tensor = ctx.Input(data_batch_ * 2 + i); | ||||||||||||
| 115 | + input_data.grouped_extreme_distances.Add(static_cast<T*>(grouped_extreme_distance_tensor->GetData()), | ||||||||||||
| 116 | + actual_count / group_size_); | ||||||||||||
| 117 | + | ||||||||||||
| 118 | + Tensor* pq_ivf_tensor = ctx.Input(data_batch_ * 3 + i); | ||||||||||||
| 119 | + input_data.pq_ivfs.Add(static_cast<int32_t*>(pq_ivf_tensor->GetData()), actual_count); | ||||||||||||
| 120 | + | ||||||||||||
| 121 | + Tensor* pq_index_tensor = ctx.Input(data_batch_ * 4 + i); | ||||||||||||
| 122 | + input_data.pq_indexs.Add(static_cast<int32_t*>(pq_index_tensor->GetData()), actual_count); | ||||||||||||
| 123 | + } | ||||||||||||
| 124 | + | ||||||||||||
| 125 | + KERNEL_CHECK_FALSE((input_data.actual_count >= k_), KERNEL_STATUS_PARAM_INVALID, | ||||||||||||
| 126 | + "k_[%d] should not be greater than actual_count[%d].", k_, input_data.actual_count); | ||||||||||||
| 127 | + return KERNEL_STATUS_OK; | ||||||||||||
| 128 | +} | ||||||||||||
| 129 | + | ||||||||||||
| 130 | +template <typename T> | ||||||||||||
| 131 | +uint32_t TopKPQDistanceCpuKernel::ProcessResult(const CpuKernelContext& ctx, const InputsData<T>& input_data, | ||||||||||||
| 132 | + Item<T> topk_ptr[]) | ||||||||||||
| 133 | +{ | ||||||||||||
| 134 | + Tensor* topk_distance = ctx.Output(0); | ||||||||||||
| 135 | + Tensor* topk_ivf = ctx.Output(1); | ||||||||||||
| 136 | + Tensor* topk_index = ctx.Output(2); | ||||||||||||
| 137 | + | ||||||||||||
| 138 | + if (topk_distance->GetDataSize() < static_cast<uint64_t>(sizeof(T) * k_) || | ||||||||||||
| 139 | + topk_ivf->GetDataSize() < static_cast<uint64_t>(sizeof(int32_t) * k_) || | ||||||||||||
| 140 | + topk_index->GetDataSize() < static_cast<uint64_t>(sizeof(int32_t) * k_)) { | ||||||||||||
| 141 | + KERNEL_LOG_ERROR("outputs data size error"); | ||||||||||||
| 142 | + return KERNEL_STATUS_PARAM_INVALID; | ||||||||||||
| 143 | + } | ||||||||||||
| 144 | + | ||||||||||||
| 145 | + T* topk_distance_ptr = static_cast<T*>(topk_distance->GetData()); | ||||||||||||
| 146 | + int32_t* topk_ivf_ptr = static_cast<int32_t*>(topk_ivf->GetData()); | ||||||||||||
| 147 | + int32_t* topk_index_ptr = static_cast<int32_t*>(topk_index->GetData()); | ||||||||||||
| 148 | + | ||||||||||||
| 149 | + for (int32_t i = k_; i > 0; i--) { | ||||||||||||
| 150 | + Item<T> res; | ||||||||||||
| 151 | + PopHeap<T>(topk_ptr, i, &res); | ||||||||||||
| 152 | + int32_t grp = res.grp; | ||||||||||||
| 153 | + int32_t grpi = res.grpi; | ||||||||||||
| 154 | + | ||||||||||||
| 155 | + topk_distance_ptr[i - 1] = res.val; | ||||||||||||
| 156 | + topk_ivf_ptr[i - 1] = input_data.pq_ivfs.Getv(grp, grpi); | ||||||||||||
| 157 | + topk_index_ptr[i - 1] = input_data.pq_indexs.Getv(grp, grpi); | ||||||||||||
| 158 | + } | ||||||||||||
| 159 | + return KERNEL_STATUS_OK; | ||||||||||||
| 160 | +} | ||||||||||||
| 161 | + | ||||||||||||
| 162 | +template <typename T> | ||||||||||||
| 163 | +void TopKPQDistanceCpuKernel::InitTopKHeap(int& cnt, int& cntk, Item<T> topk_ptr[], const Item<T> grp_extreme_ptr[], | ||||||||||||
| 164 | + const InputsData<T>& inputs_data) | ||||||||||||
| 165 | +{ | ||||||||||||
| 166 | + T** ptr = inputs_data.pq_distances.GetPointer(); | ||||||||||||
| 167 | + for (; cntk < k_; cntk++) { | ||||||||||||
| 168 | + int32_t grp = grp_extreme_ptr[cntk].grp; | ||||||||||||
| 169 | + int32_t grpi = grp_extreme_ptr[cntk].grpi * group_size_; | ||||||||||||
| 170 | + T* itemvalptr = ptr[grp] + grpi; | ||||||||||||
| 171 | + for (int32_t index = 0; index < group_size_; index++, cnt++) { | ||||||||||||
| 172 | + if (cnt == k_) { | ||||||||||||
| 173 | + return; | ||||||||||||
| 174 | + } | ||||||||||||
| 175 | + topk_ptr[cnt] = {itemvalptr[index], grp, grpi + index}; | ||||||||||||
| 176 | + } | ||||||||||||
| 177 | + } | ||||||||||||
| 178 | +} | ||||||||||||
| 179 | + | ||||||||||||
| 180 | +template <typename T> | ||||||||||||
| 181 | +uint32_t TopKPQDistanceCpuKernel::GetDistanceTopKHeap(Item<T> topk_ptr[], const Item<T> grp_extreme_ptr[], | ||||||||||||
| 182 | + const InputsData<T>& inputs_data) | ||||||||||||
| 183 | +{ | ||||||||||||
| 184 | + int cnt = 0; | ||||||||||||
| 185 | + int cntk = 0; | ||||||||||||
| 186 | + T** ptr = inputs_data.pq_distances.GetPointer(); | ||||||||||||
| 187 | + InitTopKHeap(cnt, cntk, topk_ptr, grp_extreme_ptr, inputs_data); | ||||||||||||
| 188 | + MakeHeap(topk_ptr, k_); | ||||||||||||
| 189 | + int32_t extreme_size = inputs_data.actual_count / group_size_; | ||||||||||||
| 190 | + int32_t size = std::min(k_, extreme_size); | ||||||||||||
| 191 | + int32_t index = k_ % group_size_; | ||||||||||||
| 192 | + for (; cntk < size; cntk++, index = 0) { | ||||||||||||
| 193 | + if (is_min_heap_) { | ||||||||||||
| 194 | + if (grp_extreme_ptr[cntk].val <= topk_ptr[0].val) { | ||||||||||||
| 195 | + continue; | ||||||||||||
| 196 | + } | ||||||||||||
| 197 | + } else if (grp_extreme_ptr[cntk].val >= topk_ptr[0].val) { | ||||||||||||
| 198 | + continue; | ||||||||||||
| 199 | + } | ||||||||||||
| 200 | + int32_t grp = grp_extreme_ptr[cntk].grp; | ||||||||||||
| 201 | + int32_t grpi = grp_extreme_ptr[cntk].grpi * group_size_; | ||||||||||||
| 202 | + T* itemvalptr = ptr[grp] + grpi; | ||||||||||||
| 203 | + for (; index < group_size_; index++) { | ||||||||||||
| 204 | + T& itemval = itemvalptr[index]; | ||||||||||||
| 205 | + if (is_min_heap_) { | ||||||||||||
| 206 | + if (itemval <= topk_ptr[0].val) { | ||||||||||||
| 207 | + continue; | ||||||||||||
| 208 | + } | ||||||||||||
| 209 | + } else if (itemval >= topk_ptr[0].val) { | ||||||||||||
| 210 | + continue; | ||||||||||||
| 211 | + } | ||||||||||||
| 212 | + topk_ptr[0] = {itemval, grp, grpi + index}; | ||||||||||||
| 213 | + HeapFixdown(topk_ptr, 0, k_); | ||||||||||||
| 214 | + } | ||||||||||||
| 215 | + } | ||||||||||||
| 216 | + KERNEL_LOG_DEBUG("GetDistanceTopKHeap end"); | ||||||||||||
| 217 | + return KERNEL_STATUS_OK; | ||||||||||||
| 218 | +} | ||||||||||||
| 219 | + | ||||||||||||
| 220 | +template <typename T> | ||||||||||||
| 221 | +uint32_t TopKPQDistanceCpuKernel::GetGroupedDistanceTopKHeap(Item<T> grp_extreme_ptr[], const InputsData<T>& input_data) | ||||||||||||
| 222 | +{ | ||||||||||||
| 223 | + T** ptr = input_data.grouped_extreme_distances.GetPointer(); | ||||||||||||
| 224 | + int32_t extreme_size = input_data.actual_count / group_size_; | ||||||||||||
| 225 | + int32_t size = std::min(k_, extreme_size); | ||||||||||||
| 226 | + int32_t grp = 0; | ||||||||||||
| 227 | + int32_t grpi = 0; | ||||||||||||
| 228 | + InitGrpExtreme<T>(grp_extreme_ptr, input_data, grp, grpi); | ||||||||||||
| 229 | + MakeHeap(grp_extreme_ptr, size); | ||||||||||||
| 230 | + int32_t grp_size = static_cast<int32_t>(input_data.grouped_extreme_distances.data_count.size()); | ||||||||||||
| 231 | + for (; grp < grp_size; grpi = 0, grp++) { | ||||||||||||
| 232 | + for (; grpi < input_data.grouped_extreme_distances.data_count[grp]; grpi++) { | ||||||||||||
| 233 | + T temp = ptr[grp][grpi]; | ||||||||||||
| 234 | + if (is_min_heap_) { | ||||||||||||
| 235 | + if (grp_extreme_ptr[0].val > temp) { | ||||||||||||
| 236 | + continue; | ||||||||||||
| 237 | + } | ||||||||||||
| 238 | + } else if (grp_extreme_ptr[0].val < temp) { | ||||||||||||
| 239 | + continue; | ||||||||||||
| 240 | + } | ||||||||||||
| 241 | + grp_extreme_ptr[0] = {temp, grp, grpi}; | ||||||||||||
| 242 | + HeapFixdown(grp_extreme_ptr, 0, size); | ||||||||||||
| 243 | + } | ||||||||||||
| 244 | + } | ||||||||||||
| 245 | + SortHeap(grp_extreme_ptr, size); | ||||||||||||
| 246 | + return KERNEL_STATUS_OK; | ||||||||||||
| 247 | +} | ||||||||||||
| 248 | + | ||||||||||||
| 249 | +template <typename T> | ||||||||||||
| 250 | +void TopKPQDistanceCpuKernel::InitGrpExtreme(Item<T> grp_extreme_ptr[], const InputsData<T>& input_data, int32_t& grp, | ||||||||||||
| 251 | + int32_t& grpi) | ||||||||||||
| 252 | +{ | ||||||||||||
| 253 | + int32_t extreme_size = input_data.actual_count / group_size_; | ||||||||||||
| 254 | + int32_t size = std::min(k_, extreme_size); | ||||||||||||
| 255 | + int32_t n = 0; | ||||||||||||
| 256 | + T** ptr = input_data.grouped_extreme_distances.GetPointer(); | ||||||||||||
| 257 | + int32_t grp_size = static_cast<int32_t>(input_data.grouped_extreme_distances.data_count.size()); | ||||||||||||
| 258 | + for (; grp < grp_size; grpi = 0, grp++) { | ||||||||||||
| 259 | + for (; grpi < input_data.grouped_extreme_distances.data_count[grp]; grpi++, n++) { | ||||||||||||
| 260 | + if (n == size) { | ||||||||||||
| 261 | + return; | ||||||||||||
| 262 | + } | ||||||||||||
| 263 | + grp_extreme_ptr[n] = {ptr[grp][grpi], grp, grpi}; | ||||||||||||
| 264 | + } | ||||||||||||
| 265 | + } | ||||||||||||
| 266 | +} | ||||||||||||
| 267 | + | ||||||||||||
| 268 | +template <typename T> | ||||||||||||
| 269 | +void TopKPQDistanceCpuKernel::MakeHeap(Item<T> arr_ptr[], const int32_t n) | ||||||||||||
| 270 | +{ | ||||||||||||
| 271 | + for (int32_t i = (static_cast<uint32_t>(n) >> 1) - 1; i >= 0; i--) { | ||||||||||||
| 272 | + HeapFixdown(arr_ptr, i, n); | ||||||||||||
| 273 | + } | ||||||||||||
| 274 | +} | ||||||||||||
| 275 | + | ||||||||||||
| 276 | +template <typename T> | ||||||||||||
| 277 | +void TopKPQDistanceCpuKernel::PopHeap(Item<T> arr_ptr[], const int32_t n, Item<T>* res) | ||||||||||||
| 278 | +{ | ||||||||||||
| 279 | + *res = arr_ptr[0]; | ||||||||||||
| 280 | + arr_ptr[0] = arr_ptr[n - 1]; | ||||||||||||
| 281 | + HeapFixdown(arr_ptr, 0, n - 1); | ||||||||||||
| 282 | +} | ||||||||||||
| 283 | + | ||||||||||||
| 284 | +template <typename T> | ||||||||||||
| 285 | +inline void TopKPQDistanceCpuKernel::HeapFixdown(Item<T> a[], const int32_t index, const int32_t n) | ||||||||||||
| 286 | +{ | ||||||||||||
| 287 | + int32_t j = 0; | ||||||||||||
| 288 | + int32_t i = index; | ||||||||||||
| 289 | + Item<T> temp = a[i]; | ||||||||||||
| 290 | + | ||||||||||||
| 291 | + j = (i << 1) + 1; | ||||||||||||
| 292 | + while (j < n) { | ||||||||||||
| 293 | + if (is_min_heap_) { | ||||||||||||
| 294 | + if (j + 1 < n && a[j].val > a[j + 1].val) { | ||||||||||||
| 295 | + j++; | ||||||||||||
| 296 | + } | ||||||||||||
| 297 | + if (a[j].val >= temp.val) { | ||||||||||||
| 298 | + break; | ||||||||||||
| 299 | + } | ||||||||||||
| 300 | + } else { | ||||||||||||
| 301 | + if (j + 1 < n && a[j].val < a[j + 1].val) { | ||||||||||||
| 302 | + j++; | ||||||||||||
| 303 | + } | ||||||||||||
| 304 | + if (a[j].val <= temp.val) { | ||||||||||||
| 305 | + break; | ||||||||||||
| 306 | + } | ||||||||||||
| 307 | + } | ||||||||||||
| 308 | + | ||||||||||||
| 309 | + a[i] = a[j]; | ||||||||||||
| 310 | + i = j; | ||||||||||||
| 311 | + j = (static_cast<uint32_t>(i) << 1) + 1; | ||||||||||||
| 312 | + } | ||||||||||||
| 313 | + a[i] = temp; | ||||||||||||
| 314 | +} | ||||||||||||
| 315 | + | ||||||||||||
| 316 | +template <typename T> | ||||||||||||
| 317 | +void TopKPQDistanceCpuKernel::SortHeap(Item<T> arr_ptr[], const int32_t n) | ||||||||||||
| 318 | +{ | ||||||||||||
| 319 | + Item<T> temp; | ||||||||||||
| 320 | + for (int i = n - 1; i >= 0; i--) { | ||||||||||||
| 321 | + temp = arr_ptr[0]; | ||||||||||||
| 322 | + arr_ptr[0] = arr_ptr[i]; | ||||||||||||
| 323 | + arr_ptr[i] = temp; | ||||||||||||
| 324 | + HeapFixdown(arr_ptr, 0, i); | ||||||||||||
| 325 | + } | ||||||||||||
| 326 | +} | ||||||||||||
| 327 | +REGISTER_CPU_KERNEL(kTopKPQDistance, TopKPQDistanceCpuKernel); | ||||||||||||
| 328 | +} // namespace aicpu | ||||||||||||
| @@ -0,0 +1,112 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +namespace aicpu { | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +class TopKPQDistanceCpuKernel : public CpuKernel { | ||
| 26 | +public: | ||
| 27 | + TopKPQDistanceCpuKernel() = default; | ||
| 28 | + | ||
| 29 | + ~TopKPQDistanceCpuKernel() = default; | ||
| 30 | + | ||
| 31 | + uint32_t Compute(CpuKernelContext& ctx) override; | ||
| 32 | + | ||
| 33 | +private: | ||
| 34 | + template <typename T> | ||
| 35 | + struct Item { | ||
| 36 | + T val; | ||
| 37 | + int32_t grp; | ||
| 38 | + int32_t grpi; | ||
| 39 | + }; | ||
| 40 | + | ||
| 41 | + template <typename T> | ||
| 42 | + class InputGroup { | ||
| 43 | + private: | ||
| 44 | + std::vector<T*> data_; | ||
| 45 | + T** data_ptr_ = nullptr; | ||
| 46 | + | ||
| 47 | + public: | ||
| 48 | + std::vector<int32_t> data_count; | ||
| 49 | + | ||
| 50 | + void Add(T* data, int32_t count) | ||
| 51 | + { | ||
| 52 | + this->data_.push_back(data); | ||
| 53 | + this->data_count.push_back(count); | ||
| 54 | + data_ptr_ = this->data_.data(); | ||
| 55 | + } | ||
| 56 | + | ||
| 57 | + T** GetPointer() const { return data_ptr_; } | ||
| 58 | + | ||
| 59 | + auto Getv(const int32_t grp, const int32_t grpi) const -> T { return data_ptr_[grp][grpi]; } | ||
| 60 | + }; | ||
| 61 | + | ||
| 62 | + template <typename T> | ||
| 63 | + struct InputsData { | ||
| 64 | + int32_t actual_count = 0; | ||
| 65 | + InputGroup<T> pq_distances; | ||
| 66 | + InputGroup<T> grouped_extreme_distances; | ||
| 67 | + InputGroup<int32_t> pq_ivfs; | ||
| 68 | + InputGroup<int32_t> pq_indexs; | ||
| 69 | + }; | ||
| 70 | + | ||
| 71 | + std::string order_; | ||
| 72 | + int32_t k_ = 0; | ||
| 73 | + int32_t group_size_ = 0; | ||
| 74 | + uint32_t data_batch_ = 0; | ||
| 75 | + bool is_min_heap_ = true; | ||
| 76 | + | ||
| 77 | + template <typename T> | ||
| 78 | + uint32_t GetInputAndCheck(const CpuKernelContext& ctx, InputsData<T>& input_data); | ||
| 79 | + | ||
| 80 | + template <typename T> | ||
| 81 | + uint32_t DoCompute(const CpuKernelContext& ctx); | ||
| 82 | + | ||
| 83 | + template <typename T> | ||
| 84 | + uint32_t ProcessResult(const CpuKernelContext& ctx, const InputsData<T>& input_data, Item<T> topk_ptr[]); | ||
| 85 | + | ||
| 86 | + template <typename T> | ||
| 87 | + void InitTopKHeap(int& cnt, int& cntk, Item<T> topk_ptr[], const Item<T> grp_extreme_ptr[], | ||
| 88 | + const InputsData<T>& inputs_data); | ||
| 89 | + | ||
| 90 | + template <typename T> | ||
| 91 | + uint32_t GetDistanceTopKHeap(Item<T> topk_ptr[], const Item<T> grp_extreme_ptr[], const InputsData<T>& inputs_data); | ||
| 92 | + | ||
| 93 | + template <typename T> | ||
| 94 | + uint32_t GetGroupedDistanceTopKHeap(Item<T> grp_extreme_ptr[], const InputsData<T>& input_data); | ||
| 95 | + | ||
| 96 | + template <typename T> | ||
| 97 | + void MakeHeap(Item<T> arr_ptr[], const int32_t n); | ||
| 98 | + | ||
| 99 | + template <typename T> | ||
| 100 | + void PopHeap(Item<T> arr_ptr[], const int32_t n, Item<T>* const res); | ||
| 101 | + | ||
| 102 | + template <typename T> | ||
| 103 | + inline void HeapFixdown(Item<T> a[], const int32_t index, const int32_t n); | ||
| 104 | + | ||
| 105 | + template <typename T> | ||
| 106 | + void SortHeap(Item<T> arr_ptr[], const int32_t n); | ||
| 107 | + | ||
| 108 | + template <typename T> | ||
| 109 | + void InitGrpExtreme(Item<T> grp_extreme_ptr[], const InputsData<T>& input_data, int32_t& grp, int32_t& grpi); | ||
| 110 | +}; | ||
| 111 | +} // namespace aicpu | ||
| 112 | + | ||
| @@ -0,0 +1,38 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +namespace ops { | ||
| 15 | +class TopKPQDistance : public OpDef { | ||
| 16 | +public: | ||
| 17 | + explicit TopKPQDistance(const char* name) : OpDef(name) | ||
| 18 | + { | ||
| 19 | + this->Input("actual_count").ParamType(DYNAMIC).DataType({ge::DT_INT32}); | ||
| 20 | + this->Input("pq_distance").ParamType(DYNAMIC).DataType({ge::DT_FLOAT, ge::DT_FLOAT16}); | ||
| 21 | + this->Input("grouped_extreme_distance").ParamType(DYNAMIC).DataType({ge::DT_FLOAT, ge::DT_FLOAT16}); | ||
| 22 | + this->Input("pq_ivf").ParamType(DYNAMIC).DataType({ge::DT_INT32}); | ||
| 23 | + this->Input("pq_index").ParamType(DYNAMIC).DataType({ge::DT_INT32}); | ||
| 24 | + this->Output("topk_distance").DataType({ge::DT_FLOAT, ge::DT_FLOAT16}); | ||
| 25 | + this->Output("topk_ivf").DataType({ge::DT_INT32}); | ||
| 26 | + this->Output("topk_index").DataType({ge::DT_INT32}); | ||
| 27 | + this->Attr("order").AttrType(OPTIONAL).String("ASC"); | ||
| 28 | + this->Attr("k").AttrType(REQUIRED).Int(); | ||
| 29 | + this->Attr("group_size").AttrType(REQUIRED).Int(); | ||
| 30 | + | ||
| 31 | + ApplyMathAicpuDefaultCfg(*this); | ||
| 32 | + this->AICPU().ExtendCfgInfo(OP_INFO_OPS_FLAG.c_str(), OPEN_OPS_FLAG.c_str()); | ||
| 33 | + this->AICPU().ExtendCfgInfo(OP_INFO_FORMAT_AGNOSTIC.c_str(), TRUE_FORMAT_AGNOSTIC.c_str()); | ||
| 34 | + } | ||
| 35 | +}; | ||
| 36 | + | ||
| 37 | +OP_ADD(TopKPQDistance); | ||
| 38 | +} // namespace ops | ||
| @@ -0,0 +1,359 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | +using namespace std; | ||
| 29 | +using namespace aicpu; | ||
| 30 | + | ||
| 31 | +class TEST_TOPKPQ_DISTANCE_UT : public testing::Test {}; | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + auto node_def = CpuKernelUtils::CreateNodeDef(); \ | ||
| 35 | + NodeDefBuilder(node_def.get(), "TopKPQDistance", "TopKPQDistance") \ | ||
| 36 | + .Input({"actual_count0", data_types[3], shapes[3], datas[0]}) \ | ||
| 37 | + .Input({"actual_count1", data_types[3], shapes[3], datas[1]}) \ | ||
| 38 | + .Input({"pq_distance0", data_types[1], shapes[0], datas[2]}) \ | ||
| 39 | + .Input({"pq_distance1", data_types[1], shapes[0], datas[3]}) \ | ||
| 40 | + .Input({"grouped_extreme_distance0", data_types[1], shapes[2], datas[4]}) \ | ||
| 41 | + .Input({"grouped_extreme_distance1", data_types[1], shapes[2], datas[5]}) \ | ||
| 42 | + .Input({"pq_ivf0", data_types[2], shapes[0], datas[6]}) \ | ||
| 43 | + .Input({"pq_ivf1", data_types[2], shapes[0], datas[7]}) \ | ||
| 44 | + .Input({"pq_index0", data_types[2], shapes[0], datas[8]}) \ | ||
| 45 | + .Input({"pq_index1", data_types[2], shapes[0], datas[9]}) \ | ||
| 46 | + .Output({"topk_distance", data_types[1], shapes[1], datas[10]}) \ | ||
| 47 | + .Output({"topk_ivf", data_types[2], shapes[1], datas[11]}) \ | ||
| 48 | + .Output({"topk_index", data_types[2], shapes[1], datas[12]}) \ | ||
| 49 | + .Attr("order", std::string("DES")) \ | ||
| 50 | + .Attr("k", k) \ | ||
| 51 | + .Attr("group_size", group_size) | ||
| 52 | + | ||
| 53 | + | ||
| 54 | + auto node_def = CpuKernelUtils::CreateNodeDef(); \ | ||
| 55 | + NodeDefBuilder(node_def.get(), "TopKPQDistance", "TopKPQDistance") \ | ||
| 56 | + .Input({"actual_count0", data_types[3], shapes[3], datas[0]}) \ | ||
| 57 | + .Input({"pq_distance0", data_types[1], shapes[0], datas[1]}) \ | ||
| 58 | + .Input({"grouped_extreme_distance0", data_types[1], shapes[2], datas[2]}) \ | ||
| 59 | + .Input({"pq_ivf0", data_types[2], shapes[0], datas[3]}) \ | ||
| 60 | + .Input({"pq_index0", data_types[2], shapes[0], datas[4]}) \ | ||
| 61 | + .Output({"topk_distance", data_types[1], shapes[1], datas[5]}) \ | ||
| 62 | + .Output({"topk_ivf", data_types[2], shapes[1], datas[6]}) \ | ||
| 63 | + .Output({"topk_index", data_types[2], shapes[1], datas[7]}) \ | ||
| 64 | + .Attr("order", std::string("DES")) \ | ||
| 65 | + .Attr("k", k) \ | ||
| 66 | + .Attr("group_size", group_size) | ||
| 67 | + | ||
| 68 | +TEST_F(TEST_TOPKPQ_DISTANCE_UT, DATA_TYPE_FLOAT) | ||
| 69 | +{ | ||
| 70 | + vector<DataType> data_types = {DT_FLOAT16, DT_FLOAT, DT_INT32, DT_INT32}; | ||
| 71 | + std::string order = "DES"; | ||
| 72 | + constexpr int32_t k = 5; | ||
| 73 | + constexpr int32_t group_size = 2; | ||
| 74 | + constexpr int32_t actual_count = 6; | ||
| 75 | + | ||
| 76 | + vector<vector<int64_t>> shapes = {{actual_count}, {k}, {actual_count / group_size}, {}}; | ||
| 77 | + float pq_distance0[actual_count] = {1, 2, 3, 4, 12, 13}; | ||
| 78 | + float grouped_extreme_distance0[actual_count / group_size] = {2, 4, 13}; | ||
| 79 | + int32_t pq_ivf0[actual_count] = {1, 1, 1, 1, 1, 1}; | ||
| 80 | + int32_t pq_index0[actual_count] = {1, 2, 3, 4, 5, 6}; | ||
| 81 | + | ||
| 82 | + float pq_distance1[actual_count] = {5, 1, 4, 2, 3, 2}; | ||
| 83 | + float grouped_extreme_distance1[actual_count / group_size] = {5, 4, 3}; | ||
| 84 | + int32_t pq_ivf1[actual_count] = {1, 1, 1, 1, 1, 1}; | ||
| 85 | + int32_t pq_index1[actual_count] = {7, 8, 9, 10, 11, 12}; | ||
| 86 | + | ||
| 87 | + // output | ||
| 88 | + float topk_distance[k]; | ||
| 89 | + int32_t topk_ivf[k]; | ||
| 90 | + int32_t topk_index[k]; | ||
| 91 | + vector<void*> datas = {(void*)(&actual_count), | ||
| 92 | + (void*)(&actual_count), | ||
| 93 | + (void*)pq_distance0, | ||
| 94 | + (void*)pq_distance1, | ||
| 95 | + (void*)grouped_extreme_distance0, | ||
| 96 | + (void*)grouped_extreme_distance1, | ||
| 97 | + (void*)pq_ivf0, | ||
| 98 | + (void*)pq_ivf1, | ||
| 99 | + (void*)pq_index0, | ||
| 100 | + (void*)pq_index1, | ||
| 101 | + (void*)topk_distance, | ||
| 102 | + (void*)topk_ivf, | ||
| 103 | + (void*)topk_index}; | ||
| 104 | + CREATE_NODEDEF(shapes, data_types, datas, k, group_size); | ||
| 105 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 106 | + | ||
| 107 | + vector<float> expect_vec(pq_distance0, pq_distance0 + actual_count); | ||
| 108 | + expect_vec.insert(expect_vec.end(), pq_distance1, pq_distance1 + actual_count); | ||
| 109 | + sort(expect_vec.begin(), expect_vec.end(), [order](float a, float b) { | ||
| 110 | + if (order == "DES") { | ||
| 111 | + return a > b; | ||
| 112 | + } | ||
| 113 | + return a < b; | ||
| 114 | + }); | ||
| 115 | + for (int i = 0; i < k; i++) { | ||
| 116 | + EXPECT_EQ(expect_vec[i], topk_distance[i]); | ||
| 117 | + } | ||
| 118 | +} | ||
| 119 | + | ||
| 120 | +TEST_F(TEST_TOPKPQ_DISTANCE_UT, DATA_TYPE_FLOAT_5INPUT) | ||
| 121 | +{ | ||
| 122 | + vector<DataType> data_types = {DT_FLOAT16, DT_FLOAT, DT_INT32, DT_INT32}; | ||
| 123 | + std::string order = "DES"; | ||
| 124 | + constexpr int32_t k = 5; | ||
| 125 | + constexpr int32_t group_size = 2; | ||
| 126 | + constexpr int32_t actual_count = 12; | ||
| 127 | + | ||
| 128 | + vector<vector<int64_t>> shapes = {{actual_count}, {k}, {actual_count / group_size}, {}}; | ||
| 129 | + float pq_distance0[actual_count] = {1, 2, 3, 4, 12, 13, 5, 1, 4, 2, 3, 2}; | ||
| 130 | + float grouped_extreme_distance0[actual_count / group_size] = {2, 4, 13, 5, 4, 3}; | ||
| 131 | + int32_t pq_ivf0[actual_count] = {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1}; | ||
| 132 | + int32_t pq_index0[actual_count] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; | ||
| 133 | + | ||
| 134 | + // output | ||
| 135 | + float topk_distance[k]; | ||
| 136 | + int32_t topk_ivf[k]; | ||
| 137 | + int32_t topk_index[k]; | ||
| 138 | + vector<void*> datas = {(void*)(&actual_count), (void*)pq_distance0, (void*)grouped_extreme_distance0, | ||
| 139 | + (void*)pq_ivf0, (void*)pq_index0, (void*)topk_distance, | ||
| 140 | + (void*)topk_ivf, (void*)topk_index}; | ||
| 141 | + CREATE_NODEDEF2(shapes, data_types, datas, k, group_size); | ||
| 142 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 143 | + | ||
| 144 | + vector<float> expect_vec(pq_distance0, pq_distance0 + actual_count); | ||
| 145 | + sort(expect_vec.begin(), expect_vec.end(), [order](float a, float b) { | ||
| 146 | + if (order == "DES") { | ||
| 147 | + return a > b; | ||
| 148 | + } | ||
| 149 | + return a < b; | ||
| 150 | + }); | ||
| 151 | + for (int i = 0; i < k; i++) { | ||
| 152 | + EXPECT_EQ(expect_vec[i], topk_distance[i]); | ||
| 153 | + } | ||
| 154 | +} | ||
| 155 | + | ||
| 156 | +template <typename T> | ||
| 157 | +void fillExtreme(T extremeArr[], T pq_distance[], int actual_count, int group_size, string& order) | ||
| 158 | +{ | ||
| 159 | + bool is_max = order == "DES" ? true : false; | ||
| 160 | + int32_t grp_num = actual_count / group_size; | ||
| 161 | + for (int32_t i = 0; i < grp_num; i++) { | ||
| 162 | + T temp_extreme = pq_distance[i * group_size]; | ||
| 163 | + for (int32_t j = i * group_size; j < (i + 1) * group_size; j++) { | ||
| 164 | + if (is_max) { | ||
| 165 | + if (pq_distance[j] > temp_extreme) { | ||
| 166 | + temp_extreme = pq_distance[j]; | ||
| 167 | + } | ||
| 168 | + } else { | ||
| 169 | + if (pq_distance[j] < temp_extreme) { | ||
| 170 | + temp_extreme = pq_distance[j]; | ||
| 171 | + } | ||
| 172 | + } | ||
| 173 | + } | ||
| 174 | + extremeArr[i] = temp_extreme; | ||
| 175 | + } | ||
| 176 | +} | ||
| 177 | + | ||
| 178 | +TEST_F(TEST_TOPKPQ_DISTANCE_UT, DATA_TYPE_FLOAT_TIME) | ||
| 179 | +{ | ||
| 180 | + vector<DataType> data_types = {DT_FLOAT16, DT_FLOAT, DT_INT32, DT_INT32}; | ||
| 181 | + std::string order = "DES"; | ||
| 182 | + constexpr int32_t k = 1024; | ||
| 183 | + constexpr int32_t group_size = 60; | ||
| 184 | + // 尾部 40 个元素放入全局最大值，验证整除校验是否拦截（不校验则结果错误） | ||
| 185 | + constexpr int32_t actual_count0 = 60060; | ||
| 186 | + constexpr int32_t actual_count1 = 60060; | ||
| 187 | + | ||
| 188 | + vector<float> expect_vec; | ||
| 189 | + | ||
| 190 | + vector<vector<int64_t>> shapes = { | ||
| 191 | + {actual_count0}, {k}, {actual_count0 / group_size}, {}, {actual_count1}, {actual_count1 / group_size}}; | ||
| 192 | + static float pq_distance0[actual_count0]; | ||
| 193 | + static float grouped_extreme_distance0[actual_count0 / group_size]; | ||
| 194 | + static int32_t pq_ivf0[actual_count0]; | ||
| 195 | + static int32_t pq_index0[actual_count0]; | ||
| 196 | + | ||
| 197 | + // 前部 60000 个元素（1000 个完整分组）填小值 [1,200] | ||
| 198 | + for (int32_t i = 0; i < 60000; i++) { | ||
| 199 | + pq_distance0[i] = static_cast<float>((i % 200) + 1); | ||
| 200 | + pq_ivf0[i] = i; | ||
| 201 | + pq_index0[i] = 1; | ||
| 202 | + expect_vec.emplace_back(pq_distance0[i]); | ||
| 203 | + } | ||
| 204 | + // 尾部 60 个元素（第 1001 组完整分组）填全局最大值，验证尾部数据能被正确选入 TopK | ||
| 205 | + for (int32_t i = 60000; i < actual_count0; i++) { | ||
| 206 | + pq_distance0[i] = 999999.0f; | ||
| 207 | + pq_ivf0[i] = i; | ||
| 208 | + pq_index0[i] = 1; | ||
| 209 | + expect_vec.emplace_back(pq_distance0[i]); | ||
| 210 | + } | ||
| 211 | + fillExtreme(grouped_extreme_distance0, pq_distance0, actual_count0, group_size, order); | ||
| 212 | + | ||
| 213 | + static float pq_distance1[actual_count1]; | ||
| 214 | + static float grouped_extreme_distance1[actual_count1 / group_size]; | ||
| 215 | + static int32_t pq_ivf1[actual_count1]; | ||
| 216 | + static int32_t pq_index1[actual_count1]; | ||
| 217 | + | ||
| 218 | + for (int32_t i = 0; i < 60000; i++) { | ||
| 219 | + pq_distance1[i] = static_cast<float>((i % 200) + 1); | ||
| 220 | + pq_ivf1[i] = i; | ||
| 221 | + pq_index1[i] = 1; | ||
| 222 | + expect_vec.emplace_back(pq_distance1[i]); | ||
| 223 | + } | ||
| 224 | + for (int32_t i = 60000; i < actual_count1; i++) { | ||
| 225 | + pq_distance1[i] = 888888.0f; | ||
| 226 | + pq_ivf1[i] = i; | ||
| 227 | + pq_index1[i] = 1; | ||
| 228 | + expect_vec.emplace_back(pq_distance1[i]); | ||
| 229 | + } | ||
| 230 | + fillExtreme(grouped_extreme_distance1, pq_distance1, actual_count1, group_size, order); | ||
| 231 | + | ||
| 232 | + // output | ||
| 233 | + float topk_distance[k]; | ||
| 234 | + int32_t topk_ivf[k]; | ||
| 235 | + int32_t topk_index[k]; | ||
| 236 | + vector<void*> datas = {(void*)(&actual_count0), | ||
| 237 | + (void*)(&actual_count1), | ||
| 238 | + (void*)pq_distance0, | ||
| 239 | + (void*)pq_distance1, | ||
| 240 | + (void*)grouped_extreme_distance0, | ||
| 241 | + (void*)grouped_extreme_distance1, | ||
| 242 | + (void*)pq_ivf0, | ||
| 243 | + (void*)pq_ivf1, | ||
| 244 | + (void*)pq_index0, | ||
| 245 | + (void*)pq_index1, | ||
| 246 | + (void*)topk_distance, | ||
| 247 | + (void*)topk_ivf, | ||
| 248 | + (void*)topk_index}; | ||
| 249 | + CREATE_NODEDEF(shapes, data_types, datas, k, group_size); | ||
| 250 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 251 | + | ||
| 252 | + sort(expect_vec.begin(), expect_vec.end(), [order](float a, float b) { | ||
| 253 | + if (order == "DES") { | ||
| 254 | + return a > b; | ||
| 255 | + } else { | ||
| 256 | + return a < b; | ||
| 257 | + } | ||
| 258 | + }); | ||
| 259 | + | ||
| 260 | + for (int i = 0; i < k; i++) { | ||
| 261 | + EXPECT_EQ(expect_vec[i], topk_distance[i]); | ||
| 262 | + } | ||
| 263 | +} | ||
| 264 | + | ||
| 265 | +TEST_F(TEST_TOPKPQ_DISTANCE_UT, DATA_TYPE_FLOAT_TIME_5INPUT) | ||
| 266 | +{ | ||
| 267 | + vector<DataType> data_types = {DT_FLOAT16, DT_FLOAT, DT_INT32, DT_INT32}; | ||
| 268 | + std::string order = "DES"; | ||
| 269 | + constexpr int32_t k = 1024; | ||
| 270 | + constexpr int32_t group_size = 60; | ||
| 271 | + constexpr int32_t actual_count0 = 63960; | ||
| 272 | + | ||
| 273 | + vector<float> expect_vec; | ||
| 274 | + | ||
| 275 | + vector<vector<int64_t>> shapes = { | ||
| 276 | + {actual_count0}, {k}, {actual_count0 / group_size}, {}, {actual_count0}, {actual_count0 / group_size}}; | ||
| 277 | + static float pq_distance0[actual_count0]; | ||
| 278 | + static float grouped_extreme_distance0[actual_count0 / group_size]; | ||
| 279 | + static int32_t pq_ivf0[actual_count0]; | ||
| 280 | + static int32_t pq_index0[actual_count0]; | ||
| 281 | + | ||
| 282 | + for (int32_t i = 0; i < actual_count0; i++) { | ||
| 283 | + pq_distance0[i] = (rand() % 240000 + 1); | ||
| 284 | + pq_ivf0[i] = i; | ||
| 285 | + pq_index0[i] = 1; | ||
| 286 | + expect_vec.emplace_back(pq_distance0[i]); | ||
| 287 | + } | ||
| 288 | + fillExtreme(grouped_extreme_distance0, pq_distance0, actual_count0, group_size, order); | ||
| 289 | + | ||
| 290 | + // output | ||
| 291 | + float topk_distance[k]; | ||
| 292 | + int32_t topk_ivf[k]; | ||
| 293 | + int32_t topk_index[k]; | ||
| 294 | + vector<void*> datas = {(void*)(&actual_count0), (void*)pq_distance0, (void*)grouped_extreme_distance0, | ||
| 295 | + (void*)pq_ivf0, (void*)pq_index0, (void*)topk_distance, | ||
| 296 | + (void*)topk_ivf, (void*)topk_index}; | ||
| 297 | + CREATE_NODEDEF2(shapes, data_types, datas, k, group_size); | ||
| 298 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 299 | + | ||
| 300 | + sort(expect_vec.begin(), expect_vec.end(), [order](float a, float b) { | ||
| 301 | + if (order == "DES") { | ||
| 302 | + return a > b; | ||
| 303 | + } else { | ||
| 304 | + return a < b; | ||
| 305 | + } | ||
| 306 | + }); | ||
| 307 | + | ||
| 308 | + for (int i = 0; i < k; i++) { | ||
| 309 | + EXPECT_EQ(expect_vec[i], topk_distance[i]); | ||
| 310 | + } | ||
| 311 | +} | ||
| 312 | + | ||
| 313 | +TEST_F(TEST_TOPKPQ_DISTANCE_UT, DATA_TYPE_FLOAT16_TIME_6INPUT) | ||
| 314 | +{ | ||
| 315 | + vector<DataType> data_types = {DT_FLOAT16, DT_FLOAT16, DT_INT32, DT_INT32}; | ||
| 316 | + std::string order = "DES"; | ||
| 317 | + constexpr int32_t k = 1024; | ||
| 318 | + constexpr int32_t group_size = 60; | ||
| 319 | + constexpr int32_t actual_count0 = 63960; | ||
| 320 | + | ||
| 321 | + vector<Eigen::half> expect_vec; | ||
| 322 | + | ||
| 323 | + vector<vector<int64_t>> shapes = { | ||
| 324 | + {actual_count0}, {k}, {actual_count0 / group_size}, {}, {actual_count0}, {actual_count0 / group_size}}; | ||
| 325 | + static Eigen::half pq_distance0[actual_count0]; | ||
| 326 | + static Eigen::half grouped_extreme_distance0[actual_count0 / group_size]; | ||
| 327 | + static int32_t pq_ivf0[actual_count0]; | ||
| 328 | + static int32_t pq_index0[actual_count0]; | ||
| 329 | + | ||
| 330 | + for (int32_t i = 0; i < actual_count0; i++) { | ||
| 331 | + pq_distance0[i] = Eigen::half(rand() % 24000 + 1); | ||
| 332 | + pq_ivf0[i] = i; | ||
| 333 | + pq_index0[i] = 1; | ||
| 334 | + expect_vec.emplace_back(pq_distance0[i]); | ||
| 335 | + } | ||
| 336 | + fillExtreme(grouped_extreme_distance0, pq_distance0, actual_count0, group_size, order); | ||
| 337 | + | ||
| 338 | + // output | ||
| 339 | + Eigen::half topk_distance[k]; | ||
| 340 | + int32_t topk_ivf[k]; | ||
| 341 | + int32_t topk_index[k]; | ||
| 342 | + vector<void*> datas = {(void*)(&actual_count0), (void*)pq_distance0, (void*)grouped_extreme_distance0, | ||
| 343 | + (void*)pq_ivf0, (void*)pq_index0, (void*)topk_distance, | ||
| 344 | + (void*)topk_ivf, (void*)topk_index}; | ||
| 345 | + CREATE_NODEDEF2(shapes, data_types, datas, k, group_size); | ||
| 346 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 347 | + | ||
| 348 | + sort(expect_vec.begin(), expect_vec.end(), [order](Eigen::half a, Eigen::half b) { | ||
| 349 | + if (order == "DES") { | ||
| 350 | + return a > b; | ||
| 351 | + } else { | ||
| 352 | + return a < b; | ||
| 353 | + } | ||
| 354 | + }); | ||
| 355 | + | ||
| 356 | + for (int i = 0; i < k; i++) { | ||
| 357 | + EXPECT_EQ(expect_vec[i], topk_distance[i]); | ||
| 358 | + } | ||
| 359 | +} | ||
| @@ -0,0 +1,12 @@ | |||
| 1 | +# ---------------------------------------------------------------------------- | ||
| 2 | +# This program is free software, you can redistribute it and/or modify it. | ||
| 3 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 4 | +# This file is a part of the CANN Open Software. | ||
| 5 | +# Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 6 | +# Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 7 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING | ||
| 8 | +# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See LICENSE in the root of the software repository for the full text of the License. | ||
| 10 | +# ---------------------------------------------------------------------------- | ||
| 11 | + | ||
| 12 | +add_all_modules_sources(OPTYPE top_k_v2_d ACLNNTYPE aclnn_exclude) | ||
| @@ -0,0 +1,103 @@ | |||
| 1 | +# TopKV2D | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 产品 | 是否支持 | | ||
| 6 | +| :------------------------------------------ | :------: | | ||
| 7 | +| Ascend 950PR/Ascend 950DT | √ | | ||
| 8 | +| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ | | ||
| 9 | +| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ | | ||
| 10 | +| Atlas 200I/500 A2 推理产品 | √ | | ||
| 11 | +| Atlas 推理系列产品 | √ | | ||
| 12 | +| Atlas 训练系列产品 | √ | | ||
| 13 | + | ||
| 14 | +## 功能说明 | ||
| 15 | + | ||
| 16 | +- 算子功能：沿指定维度找出输入张量中最大或最小的k个元素及其索引。与TopK相比，TopKV2D多了一个assist_seq辅助输入，兼容TensorFlow的TopKV2算子。 | ||
| 17 | + | ||
| 18 | +## 参数说明 | ||
| 19 | + | ||
| 20 | +<table style="undefined;table-layout: fixed; width: 1005px"><colgroup> | ||
| 21 | +<col style="width: 140px"> | ||
| 22 | +<col style="width: 140px"> | ||
| 23 | +<col style="width: 180px"> | ||
| 24 | +<col style="width: 213px"> | ||
| 25 | +<col style="width: 100px"> | ||
| 26 | +</colgroup> | ||
| 27 | +<thead> | ||
| 28 | + <tr> | ||
| 29 | + <th>参数名</th> | ||
| 30 | + <th>输入/输出/属性</th> | ||
| 31 | + <th>描述</th> | ||
| 32 | + <th>数据类型</th> | ||
| 33 | + <th>数据格式</th> | ||
| 34 | + </tr></thead> | ||
| 35 | +<tbody> | ||
| 36 | + <tr> | ||
| 37 | + <td>x</td> | ||
| 38 | + <td>输入</td> | ||
| 39 | + <td>输入张量。</td> | ||
| 40 | + <td>FLOAT16、FLOAT、DOUBLE、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64</td> | ||
| 41 | + <td>ND</td> | ||
| 42 | + </tr> | ||
| 43 | + <tr> | ||
| 44 | + <td>k</td> | ||
| 45 | + <td>输入</td> | ||
| 46 | + <td>要取出的元素个数。</td> | ||
| 47 | + <td>INT32</td> | ||
| 48 | + <td>ND</td> | ||
| 49 | + </tr> | ||
| 50 | + <tr> | ||
| 51 | + <td>assist_seq</td> | ||
| 52 | + <td>输入</td> | ||
| 53 | + <td>辅助序列张量。</td> | ||
| 54 | + <td>FLOAT16</td> | ||
| 55 | + <td>ND</td> | ||
| 56 | + </tr> | ||
| 57 | + <tr> | ||
| 58 | + <td>sorted</td> | ||
| 59 | + <td>属性</td> | ||
| 60 | + <td>是否对输出结果排序，默认为true。</td> | ||
| 61 | + <td>Bool</td> | ||
| 62 | + <td>-</td> | ||
| 63 | + </tr> | ||
| 64 | + <tr> | ||
| 65 | + <td>dim</td> | ||
| 66 | + <td>属性</td> | ||
| 67 | + <td>指定沿哪个维度进行操作，默认为-1（最后一维）。</td> | ||
| 68 | + <td>Int</td> | ||
| 69 | + <td>-</td> | ||
| 70 | + </tr> | ||
| 71 | + <tr> | ||
| 72 | + <td>largest</td> | ||
| 73 | + <td>属性</td> | ||
| 74 | + <td>是否取最大值，默认为true。若为false则取最小值。</td> | ||
| 75 | + <td>Bool</td> | ||
| 76 | + <td>-</td> | ||
| 77 | + </tr> | ||
| 78 | + <tr> | ||
| 79 | + <td>values</td> | ||
| 80 | + <td>输出</td> | ||
| 81 | + <td>输出的top-k个元素值。</td> | ||
| 82 | + <td>FLOAT16、FLOAT、DOUBLE、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64</td> | ||
| 83 | + <td>ND</td> | ||
| 84 | + </tr> | ||
| 85 | + <tr> | ||
| 86 | + <td>indices</td> | ||
| 87 | + <td>输出</td> | ||
| 88 | + <td>输出的top-k个元素在输入张量中的索引。</td> | ||
| 89 | + <td>INT32</td> | ||
| 90 | + <td>ND</td> | ||
| 91 | + </tr> | ||
| 92 | +</tbody></table> | ||
| 93 | + | ||
| 94 | +## 约束说明 | ||
| 95 | + | ||
| 96 | +- dim必须在[-x维度数, x维度数)范围内 | ||
| 97 | +- k必须大于等于0且小于等于x在dim维度上的大小 | ||
| 98 | + | ||
| 99 | +## 调用说明 | ||
| 100 | + | ||
| 101 | +| 调用方式 | 样例代码 | 说明 | | ||
| 102 | +| --------- | ------------------------------------------------------------ | ------------------------------------------------------------ | | ||
| 103 | +| 图模式接口 | [test_geir_top_k_v2_d](examples/test_geir_top_k_v2_d.cpp) | 通过[算子IR](op_graph/top_k_v2_d_proto.h)接口方式调用TopKV2D算子。 | | ||
| @@ -0,0 +1,339 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +/*! | ||
| 12 | + * \file test_geir_topkv2d.cpp | ||
| 13 | + * \brief GE IR test for TopKV2D operator (force AICPU by excluding AiCore engine) | ||
| 14 | + */ | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + | ||
| 35 | + | ||
| 36 | + | ||
| 37 | + | ||
| 38 | + | ||
| 39 | + | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +using namespace ge; | ||
| 43 | +using std::map; | ||
| 44 | +using std::string; | ||
| 45 | +using std::vector; | ||
| 46 | + | ||
| 47 | +string GetTime() | ||
| 48 | +{ | ||
| 49 | + time_t timep; | ||
| 50 | + time(&timep); | ||
| 51 | + char tmp[64]; | ||
| 52 | + strftime(tmp, sizeof(tmp), "%Y-%m-%d %H:%M:%S,000", localtime(&timep)); | ||
| 53 | + return tmp; | ||
| 54 | +} | ||
| 55 | + | ||
| 56 | +uint32_t GetDataTypeSize(DataType dt) | ||
| 57 | +{ | ||
| 58 | + switch (dt) { | ||
| 59 | + case ge::DT_BOOL: | ||
| 60 | + return 1U; | ||
| 61 | + case ge::DT_INT8: | ||
| 62 | + case ge::DT_UINT8: | ||
| 63 | + return 1U; | ||
| 64 | + case ge::DT_FLOAT16: | ||
| 65 | + case ge::DT_INT16: | ||
| 66 | + case ge::DT_UINT16: | ||
| 67 | + return 2U; | ||
| 68 | + case ge::DT_FLOAT: | ||
| 69 | + case ge::DT_INT32: | ||
| 70 | + case ge::DT_UINT32: | ||
| 71 | + return 4U; | ||
| 72 | + case ge::DT_INT64: | ||
| 73 | + case ge::DT_UINT64: | ||
| 74 | + return 8U; | ||
| 75 | + default: | ||
| 76 | + return 0U; | ||
| 77 | + } | ||
| 78 | +} | ||
| 79 | + | ||
| 80 | +int32_t GenData(vector<int64_t> shapes, Tensor& tensor, TensorDesc& desc, DataType dtype, double value) | ||
| 81 | +{ | ||
| 82 | + desc.SetRealDimCnt(shapes.size()); | ||
| 83 | + size_t size = 1; | ||
| 84 | + for (size_t i = 0; i < shapes.size(); i++) { | ||
| 85 | + size *= static_cast<size_t>(shapes[i]); | ||
| 86 | + } | ||
| 87 | + uint32_t type_size = GetDataTypeSize(dtype); | ||
| 88 | + if (type_size == 0U) { | ||
| 89 | + printf("%s - ERROR - [XIR]: GenData: unsupported data type %d\n", GetTime().c_str(), static_cast<int>(dtype)); | ||
| 90 | + return FAILED; | ||
| 91 | + } | ||
| 92 | + uint32_t data_len = static_cast<uint32_t>(size * type_size); | ||
| 93 | + uint8_t* buf = new (std::nothrow) uint8_t[data_len]; | ||
| 94 | + if (buf == nullptr) { | ||
| 95 | + printf("%s - ERROR - [XIR]: GenData: allocation failed\n", GetTime().c_str()); | ||
| 96 | + return FAILED; | ||
| 97 | + } | ||
| 98 | + | ||
| 99 | + switch (dtype) { | ||
| 100 | + case ge::DT_FLOAT: { | ||
| 101 | + float* p = reinterpret_cast<float*>(buf); | ||
| 102 | + for (size_t i = 0; i < size; ++i) { | ||
| 103 | + p[i] = static_cast<float>(value + i); | ||
| 104 | + } | ||
| 105 | + break; | ||
| 106 | + } | ||
| 107 | + case ge::DT_INT32: { | ||
| 108 | + int32_t* p = reinterpret_cast<int32_t*>(buf); | ||
| 109 | + for (size_t i = 0; i < size; ++i) { | ||
| 110 | + p[i] = static_cast<int32_t>(value + i); | ||
| 111 | + } | ||
| 112 | + break; | ||
| 113 | + } | ||
| 114 | + case ge::DT_INT64: { | ||
| 115 | + int64_t* p = reinterpret_cast<int64_t*>(buf); | ||
| 116 | + for (size_t i = 0; i < size; ++i) { | ||
| 117 | + p[i] = static_cast<int64_t>(value + i); | ||
| 118 | + } | ||
| 119 | + break; | ||
| 120 | + } | ||
| 121 | + case ge::DT_FLOAT16: { | ||
| 122 | + // Fill as uint16_t bit pattern for float16 (simple sequential values) | ||
| 123 | + uint16_t* p = reinterpret_cast<uint16_t*>(buf); | ||
| 124 | + for (size_t i = 0; i < size; ++i) { | ||
| 125 | + p[i] = static_cast<uint16_t>(i + 1); | ||
| 126 | + } | ||
| 127 | + break; | ||
| 128 | + } | ||
| 129 | + default: { | ||
| 130 | + // Generic byte fill for other types | ||
| 131 | + for (size_t i = 0; i < data_len; ++i) { | ||
| 132 | + buf[i] = static_cast<uint8_t>(static_cast<int64_t>(value + i) & 0xFF); | ||
| 133 | + } | ||
| 134 | + break; | ||
| 135 | + } | ||
| 136 | + } | ||
| 137 | + tensor = Tensor(desc, buf, data_len); | ||
| 138 | + delete[] buf; | ||
| 139 | + return SUCCESS; | ||
| 140 | +} | ||
| 141 | + | ||
| 142 | +int32_t WriteDataToFile(string bin_file, uint64_t data_size, uint8_t* inputData) | ||
| 143 | +{ | ||
| 144 | + FILE* fp = fopen(bin_file.c_str(), "wb"); | ||
| 145 | + if (fp == nullptr) { | ||
| 146 | + return FAILED; | ||
| 147 | + } | ||
| 148 | + fwrite(inputData, sizeof(uint8_t), data_size, fp); | ||
| 149 | + fclose(fp); | ||
| 150 | + return SUCCESS; | ||
| 151 | +} | ||
| 152 | + | ||
| 153 | +void ProcessInputData(vector<Tensor>& input) | ||
| 154 | +{ | ||
| 155 | + for (size_t i = 0; i < input.size(); i++) { | ||
| 156 | + string input_file = "./tc_ge_irrun_test_0008_npu_input_" + std::to_string(i) + ".bin"; | ||
| 157 | + uint8_t* input_data_i = input[i].GetData(); | ||
| 158 | + int64_t input_shape = input[i].GetTensorDesc().GetShape().GetShapeSize(); | ||
| 159 | + uint32_t type_size = GetDataTypeSize(input[i].GetTensorDesc().GetDataType()); | ||
| 160 | + if (type_size == 0U) { | ||
| 161 | + printf("ERROR: input %zu has unsupported dtype\n", i); | ||
| 162 | + continue; | ||
| 163 | + } | ||
| 164 | + uint32_t data_size = static_cast<uint32_t>(input_shape * type_size); | ||
| 165 | + WriteDataToFile(input_file.c_str(), data_size, input_data_i); | ||
| 166 | + } | ||
| 167 | +} | ||
| 168 | + | ||
| 169 | +void ProcessOutputData(vector<Tensor>& output) | ||
| 170 | +{ | ||
| 171 | + for (size_t i = 0; i < output.size(); i++) { | ||
| 172 | + string output_file = "./tc_ge_irrun_test_0008_npu_output_" + std::to_string(i) + ".bin"; | ||
| 173 | + uint8_t* output_data_i = output[i].GetData(); | ||
| 174 | + int64_t output_shape = output[i].GetTensorDesc().GetShape().GetShapeSize(); | ||
| 175 | + uint32_t type_size = GetDataTypeSize(output[i].GetTensorDesc().GetDataType()); | ||
| 176 | + if (type_size == 0U) { | ||
| 177 | + printf("ERROR: output %zu has unsupported dtype\n", i); | ||
| 178 | + continue; | ||
| 179 | + } | ||
| 180 | + uint32_t data_size = static_cast<uint32_t>(output_shape * type_size); | ||
| 181 | + WriteDataToFile(output_file.c_str(), data_size, output_data_i); | ||
| 182 | + } | ||
| 183 | +} | ||
| 184 | + | ||
| 185 | +int CreateOppInGraph(DataType inDtype, vector<Tensor>& input, vector<Operator>& inputs, vector<Operator>& outputs, | ||
| 186 | + Graph& graph) | ||
| 187 | +{ | ||
| 188 | + vector<int64_t> x_shape = {4, 6}; | ||
| 189 | + vector<int64_t> k_shape = {1}; | ||
| 190 | + vector<int64_t> assist_shape = {4, 6}; | ||
| 191 | + vector<int64_t> values_shape = {4, 3}; | ||
| 192 | + vector<int64_t> indices_shape = {4, 3}; | ||
| 193 | + | ||
| 194 | + TensorDesc x_desc = TensorDesc(ge::Shape(x_shape), FORMAT_ND, inDtype); | ||
| 195 | + x_desc.SetPlacement(ge::kPlacementHost); | ||
| 196 | + x_desc.SetFormat(FORMAT_ND); | ||
| 197 | + | ||
| 198 | + TensorDesc k_desc = TensorDesc(ge::Shape(k_shape), FORMAT_ND, DT_INT32); | ||
| 199 | + k_desc.SetPlacement(ge::kPlacementHost); | ||
| 200 | + k_desc.SetFormat(FORMAT_ND); | ||
| 201 | + | ||
| 202 | + TensorDesc assist_desc = TensorDesc(ge::Shape(assist_shape), FORMAT_ND, DT_FLOAT16); | ||
| 203 | + assist_desc.SetPlacement(ge::kPlacementHost); | ||
| 204 | + assist_desc.SetFormat(FORMAT_ND); | ||
| 205 | + | ||
| 206 | + Tensor x_tensor; | ||
| 207 | + if (GenData(x_shape, x_tensor, x_desc, inDtype, 1.0) != SUCCESS) { | ||
| 208 | + return FAILED; | ||
| 209 | + } | ||
| 210 | + input.push_back(x_tensor); | ||
| 211 | + | ||
| 212 | + Tensor k_tensor; | ||
| 213 | + int32_t k_val = 3; | ||
| 214 | + TensorDesc k_tensor_desc = TensorDesc(ge::Shape(k_shape), FORMAT_ND, DT_INT32); | ||
| 215 | + k_tensor = Tensor(k_tensor_desc, reinterpret_cast<uint8_t*>(&k_val), sizeof(int32_t)); | ||
| 216 | + input.push_back(k_tensor); | ||
| 217 | + | ||
| 218 | + Tensor assist_tensor; | ||
| 219 | + if (GenData(assist_shape, assist_tensor, assist_desc, DT_FLOAT16, 0.0) != SUCCESS) { | ||
| 220 | + return FAILED; | ||
| 221 | + } | ||
| 222 | + input.push_back(assist_tensor); | ||
| 223 | + | ||
| 224 | + auto x_data = op::Data("x_data").set_attr_index(0); | ||
| 225 | + x_data.update_input_desc_x(x_desc); | ||
| 226 | + x_data.update_output_desc_y(x_desc); | ||
| 227 | + | ||
| 228 | + auto k_data = op::Data("k_data").set_attr_index(1); | ||
| 229 | + k_data.update_input_desc_x(k_desc); | ||
| 230 | + k_data.update_output_desc_y(k_desc); | ||
| 231 | + | ||
| 232 | + auto assist_data = op::Data("assist_data").set_attr_index(2); | ||
| 233 | + assist_data.update_input_desc_x(assist_desc); | ||
| 234 | + assist_data.update_output_desc_y(assist_desc); | ||
| 235 | + | ||
| 236 | + auto topkv2d_op = op::TopKV2D("topkv2d_op"); | ||
| 237 | + topkv2d_op.set_input_x(x_data); | ||
| 238 | + topkv2d_op.set_input_k(k_data); | ||
| 239 | + topkv2d_op.set_input_assist_seq(assist_data); | ||
| 240 | + topkv2d_op.set_attr_dim(-1); | ||
| 241 | + topkv2d_op.set_attr_largest(true); | ||
| 242 | + topkv2d_op.set_attr_sorted(true); | ||
| 243 | + | ||
| 244 | + TensorDesc values_desc = TensorDesc(ge::Shape(values_shape), FORMAT_ND, inDtype); | ||
| 245 | + TensorDesc indices_desc = TensorDesc(ge::Shape(indices_shape), FORMAT_ND, DT_INT32); | ||
| 246 | + topkv2d_op.update_output_desc_values(values_desc); | ||
| 247 | + topkv2d_op.update_output_desc_indices(indices_desc); | ||
| 248 | + | ||
| 249 | + graph.AddOp(x_data); | ||
| 250 | + graph.AddOp(k_data); | ||
| 251 | + graph.AddOp(assist_data); | ||
| 252 | + graph.AddOp(topkv2d_op); | ||
| 253 | + | ||
| 254 | + inputs.push_back(x_data); | ||
| 255 | + inputs.push_back(k_data); | ||
| 256 | + inputs.push_back(assist_data); | ||
| 257 | + outputs.push_back(topkv2d_op); | ||
| 258 | + | ||
| 259 | + return SUCCESS; | ||
| 260 | +} | ||
| 261 | + | ||
| 262 | +int main(int argc, char* argv[]) | ||
| 263 | +{ | ||
| 264 | + const char* graph_name = "tc_ge_irrun_test"; | ||
| 265 | + Graph graph(graph_name); | ||
| 266 | + vector<Tensor> input; | ||
| 267 | + | ||
| 268 | + printf("%s - INFO - [XIR]: Start to initialize ge using ge global options\n", GetTime().c_str()); | ||
| 269 | + std::map<AscendString, AscendString> global_options = {{"ge.exec.deviceId", "0"}, {"ge.graphRunMode", "1"}}; | ||
| 270 | + Status ret = ge::GEInitialize(global_options); | ||
| 271 | + if (ret != SUCCESS) { | ||
| 272 | + printf("%s - INFO - [XIR]: Initialize ge using ge global options failed\n", GetTime().c_str()); | ||
| 273 | + return FAILED; | ||
| 274 | + } | ||
| 275 | + printf("%s - INFO - [XIR]: Initialize ge using ge global options success\n", GetTime().c_str()); | ||
| 276 | + | ||
| 277 | + vector<Operator> inputs{}; | ||
| 278 | + vector<Operator> outputs{}; | ||
| 279 | + | ||
| 280 | + DataType inDtype = DT_FLOAT; | ||
| 281 | + | ||
| 282 | + ret = CreateOppInGraph(inDtype, input, inputs, outputs, graph); | ||
| 283 | + if (ret != SUCCESS) { | ||
| 284 | + printf("%s - ERROR - [XIR]: Create graph failed\n", GetTime().c_str()); | ||
| 285 | + return FAILED; | ||
| 286 | + } | ||
| 287 | + | ||
| 288 | + if (!inputs.empty() && !outputs.empty()) { | ||
| 289 | + graph.SetInputs(inputs).SetOutputs(outputs); | ||
| 290 | + } | ||
| 291 | + | ||
| 292 | + std::map<AscendString, AscendString> build_options = {}; | ||
| 293 | + printf("%s - INFO - [XIR]: Start to create ir session\n", GetTime().c_str()); | ||
| 294 | + ge::Session* session = new Session(build_options); | ||
| 295 | + | ||
| 296 | + if (session == nullptr) { | ||
| 297 | + printf("%s - ERROR - [XIR]: Create ir session failed\n", GetTime().c_str()); | ||
| 298 | + GEFinalize(); | ||
| 299 | + return FAILED; | ||
| 300 | + } | ||
| 301 | + printf("%s - INFO - [XIR]: Create ir session success\n", GetTime().c_str()); | ||
| 302 | + | ||
| 303 | + uint32_t graph_id = 0; | ||
| 304 | + | ||
| 305 | + std::map<AscendString, AscendString> graph_options = {{"ge.exec.exclude_engines", "AiCore"}}; | ||
| 306 | + | ||
| 307 | + printf("%s - INFO - [XIR]: Add graph with exclude AiCore engine (force AICPU)\n", GetTime().c_str()); | ||
| 308 | + ret = session->AddGraph(graph_id, graph, graph_options); | ||
| 309 | + if (ret != SUCCESS) { | ||
| 310 | + printf("%s - ERROR - [XIR]: Add graph failed\n", GetTime().c_str()); | ||
| 311 | + delete session; | ||
| 312 | + GEFinalize(); | ||
| 313 | + return FAILED; | ||
| 314 | + } | ||
| 315 | + | ||
| 316 | + printf("%s - INFO - [XIR]: Start to run graph\n", GetTime().c_str()); | ||
| 317 | + vector<Tensor> output; | ||
| 318 | + ret = session->RunGraph(graph_id, input, output); | ||
| 319 | + if (ret != SUCCESS) { | ||
| 320 | + printf("%s - ERROR - [XIR]: Run graph failed\n", GetTime().c_str()); | ||
| 321 | + delete session; | ||
| 322 | + GEFinalize(); | ||
| 323 | + return FAILED; | ||
| 324 | + } | ||
| 325 | + printf("%s - INFO - [XIR]: Run graph success\n", GetTime().c_str()); | ||
| 326 | + | ||
| 327 | + ProcessInputData(input); | ||
| 328 | + ProcessOutputData(output); | ||
| 329 | + | ||
| 330 | + printf("%s - INFO - [XIR]: Start to finalize ir graph session\n", GetTime().c_str()); | ||
| 331 | + delete session; | ||
| 332 | + ret = ge::GEFinalize(); | ||
| 333 | + if (ret != SUCCESS) { | ||
| 334 | + printf("%s - INFO - [XIR]: Finalize ir graph session failed\n", GetTime().c_str()); | ||
| 335 | + return FAILED; | ||
| 336 | + } | ||
| 337 | + printf("%s - INFO - [XIR]: Finalize ir graph session success\n", GetTime().c_str()); | ||
| 338 | + return SUCCESS; | ||
| 339 | +} | ||
| @@ -0,0 +1,44 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +namespace ge { | ||
| 17 | +/** | ||
| 18 | + *@brief Finds the k largest or smallest values and indices along a dimension. | ||
| 19 | + *@par Inputs: | ||
| 20 | + * Three inputs, including: | ||
| 21 | + *@li x: A Tensor of type float16/float/double/int8/int16/int32/int64/uint8/uint16/uint32/uint64. \n | ||
| 22 | + *@li k: A Tensor of type int32, specifying the number of top elements. | ||
| 23 | + *@li assist_seq: A Tensor of type float16, assisting in sequence computation. \n | ||
| 24 | + *@par Attributes: | ||
| 25 | + *@li sorted: An optional bool, specifying whether to sort the output. Default: true. | ||
| 26 | + *@li dim: An optional int, specifying the dimension along which to perform topk. Default: -1. | ||
| 27 | + *@li largest: An optional bool, specifying whether to select largest or smallest values. Default: true. | ||
| 28 | + *@par Outputs: | ||
| 29 | + *@li values: A Tensor of same type as x, containing the k largest/smallest values. | ||
| 30 | + *@li indices: A Tensor of type int32, containing the indices of the k largest/smallest values. | ||
| 31 | + */ | ||
| 32 | +REG_OP(TopKV2D) | ||
| 33 | + .INPUT(x, TensorType::RealNumberType()) | ||
| 34 | + .INPUT(k, TensorType({DT_INT32})) | ||
| 35 | + .INPUT(assist_seq, TensorType({DT_FLOAT16})) | ||
| 36 | + .OUTPUT(values, TensorType::RealNumberType()) | ||
| 37 | + .OUTPUT(indices, TensorType({DT_INT32})) | ||
| 38 | + .ATTR(sorted, Bool, true) | ||
| 39 | + .ATTR(dim, Int, -1) | ||
| 40 | + .ATTR(largest, Bool, true) | ||
| 41 | + .OP_END_FACTORY_REG(TopKV2D) | ||
| 42 | +} // namespace ge | ||
| 43 | + | ||
| 44 | + | ||
| @@ -0,0 +1,82 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +/*! | ||
| 12 | + * \file topkv2d_infershape.cpp | ||
| 13 | + * \brief InferShape for TopKV2D operator: output shape = input shape with dim-th axis replaced by k | ||
| 14 | + */ | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +static constexpr int OUTPUT_VALUES_INDEX = 0; | ||
| 20 | +static constexpr int OUTPUT_INDICES_INDEX = 1; | ||
| 21 | +static constexpr int INPUT_X_INDEX = 0; | ||
| 22 | + | ||
| 23 | +using namespace ge; | ||
| 24 | +namespace ops { | ||
| 25 | + | ||
| 26 | +static bool InferShapeForTopKV2DCommon(gert::InferShapeContext* context, int64_t k, const int64_t* dim) | ||
| 27 | +{ | ||
| 28 | + const gert::Shape* input_x_shape = context->GetInputShape(INPUT_X_INDEX); | ||
| 29 | + OP_CHECK_NULL_WITH_CONTEXT(context, input_x_shape); | ||
| 30 | + size_t dim_size = input_x_shape->GetDimNum(); | ||
| 31 | + if (dim_size <= 0) { | ||
| 32 | + OP_LOGE(context->GetNodeName(), "The dims_in size should more than 0!"); | ||
| 33 | + return GRAPH_FAILED; | ||
| 34 | + } | ||
| 35 | + int64_t sorted_axis = dim_size - 1; | ||
| 36 | + | ||
| 37 | + if (dim != nullptr) { | ||
| 38 | + sorted_axis = *dim; | ||
| 39 | + if (sorted_axis < 0) { | ||
| 40 | + sorted_axis += dim_size; | ||
| 41 | + } | ||
| 42 | + if (sorted_axis >= static_cast<int64_t>(dim_size)) { | ||
| 43 | + OP_LOGE(context->GetNodeName(), "Dim is out of shape size."); | ||
| 44 | + return GRAPH_FAILED; | ||
| 45 | + } | ||
| 46 | + } | ||
| 47 | + | ||
| 48 | + gert::Shape* output_values_shape = context->GetOutputShape(OUTPUT_VALUES_INDEX); | ||
| 49 | + OP_CHECK_NULL_WITH_CONTEXT(context, output_values_shape); | ||
| 50 | + gert::Shape* output_indices_shape = context->GetOutputShape(OUTPUT_INDICES_INDEX); | ||
| 51 | + OP_CHECK_NULL_WITH_CONTEXT(context, output_indices_shape); | ||
| 52 | + | ||
| 53 | + output_values_shape->SetDimNum(dim_size); | ||
| 54 | + output_indices_shape->SetDimNum(dim_size); | ||
| 55 | + for (size_t i = 0; i < dim_size; i++) { | ||
| 56 | + if (static_cast<int64_t>(i) == sorted_axis) { | ||
| 57 | + output_values_shape->SetDim(i, k); | ||
| 58 | + output_indices_shape->SetDim(i, k); | ||
| 59 | + continue; | ||
| 60 | + } | ||
| 61 | + output_values_shape->SetDim(i, input_x_shape->GetDim(i)); | ||
| 62 | + output_indices_shape->SetDim(i, input_x_shape->GetDim(i)); | ||
| 63 | + } | ||
| 64 | + return GRAPH_SUCCESS; | ||
| 65 | +} | ||
| 66 | + | ||
| 67 | +static graphStatus InferShapeForTopKV2D(gert::InferShapeContext* context) | ||
| 68 | +{ | ||
| 69 | + OP_LOGD(context->GetNodeName(), "Begin to do TopKV2D InferShape"); | ||
| 70 | + const gert::RuntimeAttrs* attrs = context->GetAttrs(); | ||
| 71 | + OP_CHECK_NULL_WITH_CONTEXT(context, attrs); | ||
| 72 | + const gert::Tensor* input_k_tensor = context->GetInputTensor(1U); | ||
| 73 | + OP_CHECK_NULL_WITH_CONTEXT(context, input_k_tensor); | ||
| 74 | + const int32_t* k = input_k_tensor->GetData<int32_t>(); | ||
| 75 | + OP_CHECK_NULL_WITH_CONTEXT(context, k); | ||
| 76 | + const int64_t* dim = attrs->GetInt(1U); | ||
| 77 | + return InferShapeForTopKV2DCommon(context, *k, dim); | ||
| 78 | +} | ||
| 79 | + | ||
| 80 | +IMPL_OP_INFERSHAPE(TopKV2D).InferShape(InferShapeForTopKV2D).InputsDataDependency({1}); | ||
| 81 | + | ||
| 82 | +} // namespace ops | ||
| @@ -0,0 +1,572 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +namespace { | ||
| 27 | +const char* const TOPKV2D = "TopKV2D"; | ||
| 28 | +constexpr int32_t kSmallN = 64; | ||
| 29 | +constexpr int32_t kSelectionMaxK = 8; | ||
| 30 | +constexpr int64_t kHugeRowcountThreshold = 1LL << 20; | ||
| 31 | +constexpr int32_t kMonotoneNonMonotone = 0; | ||
| 32 | +constexpr int32_t kMonotoneAscending = 1; | ||
| 33 | +constexpr int32_t kMonotoneDescending = 2; | ||
| 34 | +constexpr int32_t kMonotoneAllEqual = 3; | ||
| 35 | +constexpr int32_t kMonotoneProbeMaxN = 16; | ||
| 36 | +constexpr int32_t kMonotoneComparableMinN = 2; | ||
| 37 | +constexpr int64_t kParallelTargetShardBytes = 32LL * 1024LL; | ||
| 38 | +constexpr int64_t kSerialInputBytesThreshold = 8LL * 1024LL; | ||
| 39 | + | ||
| 40 | +inline int64_t ComputePerUnit(int64_t total_units, int64_t row_bytes, int64_t num_cores) | ||
| 41 | +{ | ||
| 42 | + if ((total_units <= 0) || (row_bytes <= 0) || (num_cores <= 0)) { | ||
| 43 | + return 1; | ||
| 44 | + } | ||
| 45 | + int64_t max_per_unit = kParallelTargetShardBytes / row_bytes; | ||
| 46 | + if (max_per_unit < 1) { | ||
| 47 | + max_per_unit = 1; | ||
| 48 | + } | ||
| 49 | + int64_t required_tasks = (total_units + max_per_unit - 1) / max_per_unit; | ||
| 50 | + int64_t aligned_tasks = num_cores * ((required_tasks + num_cores - 1) / num_cores); | ||
| 51 | + return (total_units + aligned_tasks - 1) / aligned_tasks; | ||
| 52 | +} | ||
| 53 | + | ||
| 54 | +template <typename T> | ||
| 55 | +struct ValueIndex { | ||
| 56 | + T value; | ||
| 57 | + int32_t index; | ||
| 58 | +}; | ||
| 59 | + | ||
| 60 | +template <typename T> | ||
| 61 | +inline bool LessDescTieIdx(const ValueIndex<T>& a, const ValueIndex<T>& b) | ||
| 62 | +{ | ||
| 63 | + if (aicpu::IsValueEqual<T>(a.value, b.value)) { | ||
| 64 | + return a.index < b.index; | ||
| 65 | + } | ||
| 66 | + return a.value > b.value; | ||
| 67 | +} | ||
| 68 | + | ||
| 69 | +template <typename T> | ||
| 70 | +inline bool LessAscTieIdx(const ValueIndex<T>& a, const ValueIndex<T>& b) | ||
| 71 | +{ | ||
| 72 | + if (aicpu::IsValueEqual<T>(a.value, b.value)) { | ||
| 73 | + return a.index < b.index; | ||
| 74 | + } | ||
| 75 | + return a.value < b.value; | ||
| 76 | +} | ||
| 77 | + | ||
| 78 | +template <typename T> | ||
| 79 | +inline int32_t DetectMonotone(const T* __restrict__ in, int32_t jump, int32_t n) | ||
| 80 | +{ | ||
| 81 | + if (n < kMonotoneComparableMinN) { | ||
| 82 | + return kMonotoneAllEqual; | ||
| 83 | + } | ||
| 84 | + if (n <= kMonotoneProbeMaxN) { | ||
| 85 | + const T a = in[0]; | ||
| 86 | + const T b = in[static_cast<int64_t>(n >> 1) * jump]; | ||
| 87 | + const T c = in[static_cast<int64_t>(n - 1) * jump]; | ||
| 88 | + if (aicpu::IsValueEqual<T>(a, b) && aicpu::IsValueEqual<T>(b, c)) { | ||
| 89 | + for (int32_t i = 1; i < n; ++i) { | ||
| 90 | + if (!aicpu::IsValueEqual<T>(in[static_cast<int64_t>(i) * jump], a)) { | ||
| 91 | + return kMonotoneNonMonotone; | ||
| 92 | + } | ||
| 93 | + } | ||
| 94 | + return kMonotoneAllEqual; | ||
| 95 | + } | ||
| 96 | + const bool p_asc = (a <= b) && (b <= c); | ||
| 97 | + const bool p_desc = (a >= b) && (b >= c); | ||
| 98 | + if (!p_asc && !p_desc) { | ||
| 99 | + return kMonotoneNonMonotone; | ||
| 100 | + } | ||
| 101 | + } | ||
| 102 | + bool asc = true; | ||
| 103 | + bool desc = true; | ||
| 104 | + for (int32_t i = 1; i < n; ++i) { | ||
| 105 | + const T prev = in[static_cast<int64_t>(i - 1) * jump]; | ||
| 106 | + const T cur = in[static_cast<int64_t>(i) * jump]; | ||
| 107 | + if (cur < prev) { | ||
| 108 | + asc = false; | ||
| 109 | + } else if (cur > prev) { | ||
| 110 | + desc = false; | ||
| 111 | + } | ||
| 112 | + if (!asc && !desc) { | ||
| 113 | + return kMonotoneNonMonotone; | ||
| 114 | + } | ||
| 115 | + } | ||
| 116 | + if (asc && desc) { | ||
| 117 | + return kMonotoneAllEqual; | ||
| 118 | + } | ||
| 119 | + return asc ? kMonotoneAscending : kMonotoneDescending; | ||
| 120 | +} | ||
| 121 | + | ||
| 122 | +template <typename T, bool largest_true> | ||
| 123 | +inline void EmitMonotoneTopK(const T* __restrict__ in, int32_t jump, T* __restrict__ out_v, int32_t* __restrict__ out_i, | ||
| 124 | + int32_t n, int32_t k, int32_t mono) | ||
| 125 | +{ | ||
| 126 | + if (mono == kMonotoneAllEqual) { | ||
| 127 | + const T v = in[0]; | ||
| 128 | + for (int32_t j = 0; j < k; ++j) { | ||
| 129 | + out_v[static_cast<int64_t>(j) * jump] = v; | ||
| 130 | + out_i[static_cast<int64_t>(j) * jump] = j; | ||
| 131 | + } | ||
| 132 | + return; | ||
| 133 | + } | ||
| 134 | + const bool head_is_largest = (mono == kMonotoneDescending); | ||
| 135 | + const bool take_head = (largest_true == head_is_largest); | ||
| 136 | + if (take_head) { | ||
| 137 | + for (int32_t j = 0; j < k; ++j) { | ||
| 138 | + out_v[static_cast<int64_t>(j) * jump] = in[static_cast<int64_t>(j) * jump]; | ||
| 139 | + out_i[static_cast<int64_t>(j) * jump] = j; | ||
| 140 | + } | ||
| 141 | + return; | ||
| 142 | + } | ||
| 143 | + int32_t j = 0; | ||
| 144 | + int32_t r = n - 1; | ||
| 145 | + while (j < k && r >= 0) { | ||
| 146 | + int32_t g_left = r; | ||
| 147 | + const T gv = in[static_cast<int64_t>(r) * jump]; | ||
| 148 | + while (g_left - 1 >= 0 && aicpu::IsValueEqual<T>(in[static_cast<int64_t>(g_left - 1) * jump], gv)) { | ||
| 149 | + --g_left; | ||
| 150 | + } | ||
| 151 | + const int32_t glen = r - g_left + 1; | ||
| 152 | + const int32_t take = (glen < (k - j)) ? glen : (k - j); | ||
| 153 | + for (int32_t t = 0; t < take; ++t) { | ||
| 154 | + out_v[static_cast<int64_t>(j) * jump] = gv; | ||
| 155 | + out_i[static_cast<int64_t>(j) * jump] = g_left + t; | ||
| 156 | + ++j; | ||
| 157 | + } | ||
| 158 | + r = g_left - 1; | ||
| 159 | + } | ||
| 160 | +} | ||
| 161 | + | ||
| 162 | +template <typename T, bool largest_true> | ||
| 163 | +inline void SelectionTopK(const T* __restrict__ in_base, int32_t jump, T* __restrict__ out_v, | ||
| 164 | + int32_t* __restrict__ out_i, int32_t n, int32_t k) | ||
| 165 | +{ | ||
| 166 | + const int32_t mono = DetectMonotone<T>(in_base, jump, n); | ||
| 167 | + if (mono != 0) { | ||
| 168 | + EmitMonotoneTopK<T, largest_true>(in_base, jump, out_v, out_i, n, k, mono); | ||
| 169 | + return; | ||
| 170 | + } | ||
| 171 | + T vs[kSmallN]; | ||
| 172 | + int32_t is_[kSmallN]; | ||
| 173 | + for (int32_t i = 0; i < n; ++i) { | ||
| 174 | + vs[i] = in_base[static_cast<int64_t>(i) * jump]; | ||
| 175 | + is_[i] = i; | ||
| 176 | + } | ||
| 177 | + for (int32_t j = 0; j < k; ++j) { | ||
| 178 | + int32_t best = j; | ||
| 179 | + T bv = vs[j]; | ||
| 180 | + int32_t bi = is_[j]; | ||
| 181 | + for (int32_t i = j + 1; i < n; ++i) { | ||
| 182 | + T v = vs[i]; | ||
| 183 | + int32_t ii = is_[i]; | ||
| 184 | + bool eq = aicpu::IsValueEqual<T>(v, bv); | ||
| 185 | + bool take = largest_true ? (v > bv || (eq && ii < bi)) : (v < bv || (eq && ii < bi)); | ||
| 186 | + if (take) { | ||
| 187 | + bv = v; | ||
| 188 | + bi = ii; | ||
| 189 | + best = i; | ||
| 190 | + } | ||
| 191 | + } | ||
| 192 | + if (best != j) { | ||
| 193 | + vs[best] = vs[j]; | ||
| 194 | + is_[best] = is_[j]; | ||
| 195 | + vs[j] = bv; | ||
| 196 | + is_[j] = bi; | ||
| 197 | + } | ||
| 198 | + out_v[static_cast<int64_t>(j) * jump] = bv; | ||
| 199 | + out_i[static_cast<int64_t>(j) * jump] = bi; | ||
| 200 | + } | ||
| 201 | +} | ||
| 202 | + | ||
| 203 | +template <typename T, bool largest_true> | ||
| 204 | +inline void FullSortTopK(const T* __restrict__ in_base, int32_t jump, T* __restrict__ out_v, | ||
| 205 | + int32_t* __restrict__ out_i, int32_t n) | ||
| 206 | +{ | ||
| 207 | + ValueIndex<T> aos[kSmallN]; | ||
| 208 | + for (int32_t i = 0; i < n; ++i) { | ||
| 209 | + aos[i].value = in_base[static_cast<int64_t>(i) * jump]; | ||
| 210 | + aos[i].index = i; | ||
| 211 | + } | ||
| 212 | + if (largest_true) { | ||
| 213 | + std::sort(aos, aos + n, LessDescTieIdx<T>); | ||
| 214 | + } else { | ||
| 215 | + std::sort(aos, aos + n, LessAscTieIdx<T>); | ||
| 216 | + } | ||
| 217 | + for (int32_t j = 0; j < n; ++j) { | ||
| 218 | + out_v[static_cast<int64_t>(j) * jump] = aos[j].value; | ||
| 219 | + out_i[static_cast<int64_t>(j) * jump] = aos[j].index; | ||
| 220 | + } | ||
| 221 | +} | ||
| 222 | + | ||
| 223 | +template <typename T> | ||
| 224 | +inline void SiftDownMin(T* val, int32_t* idx, int32_t pos, int32_t k, int32_t jump) | ||
| 225 | +{ | ||
| 226 | + while (true) { | ||
| 227 | + int32_t L = pos * 2 + 1; | ||
| 228 | + if (L >= k) { | ||
| 229 | + break; | ||
| 230 | + } | ||
| 231 | + int32_t R = L + 1; | ||
| 232 | + int64_t aL = static_cast<int64_t>(L) * jump; | ||
| 233 | + int32_t smallest = L; | ||
| 234 | + int64_t as = aL; | ||
| 235 | + if (R < k) { | ||
| 236 | + int64_t aR = static_cast<int64_t>(R) * jump; | ||
| 237 | + bool eqLR = aicpu::IsValueEqual<T>(val[aR], val[aL]); | ||
| 238 | + if (val[aR] < val[aL] || (eqLR && idx[aR] > idx[aL])) { | ||
| 239 | + smallest = R; | ||
| 240 | + as = aR; | ||
| 241 | + } | ||
| 242 | + } | ||
| 243 | + int64_t ap = static_cast<int64_t>(pos) * jump; | ||
| 244 | + bool eq = aicpu::IsValueEqual<T>(val[as], val[ap]); | ||
| 245 | + if (val[as] < val[ap] || (eq && idx[as] > idx[ap])) { | ||
| 246 | + T tv = val[ap]; | ||
| 247 | + val[ap] = val[as]; | ||
| 248 | + val[as] = tv; | ||
| 249 | + int32_t ti = idx[ap]; | ||
| 250 | + idx[ap] = idx[as]; | ||
| 251 | + idx[as] = ti; | ||
| 252 | + pos = smallest; | ||
| 253 | + } else { | ||
| 254 | + break; | ||
| 255 | + } | ||
| 256 | + } | ||
| 257 | +} | ||
| 258 | + | ||
| 259 | +template <typename T> | ||
| 260 | +inline void SiftDownMax(T* val, int32_t* idx, int32_t pos, int32_t k, int32_t jump) | ||
| 261 | +{ | ||
| 262 | + while (true) { | ||
| 263 | + int32_t L = pos * 2 + 1; | ||
| 264 | + if (L >= k) { | ||
| 265 | + break; | ||
| 266 | + } | ||
| 267 | + int32_t R = L + 1; | ||
| 268 | + int64_t aL = static_cast<int64_t>(L) * jump; | ||
| 269 | + int32_t largest = L; | ||
| 270 | + int64_t as = aL; | ||
| 271 | + if (R < k) { | ||
| 272 | + int64_t aR = static_cast<int64_t>(R) * jump; | ||
| 273 | + bool eqLR = aicpu::IsValueEqual<T>(val[aR], val[aL]); | ||
| 274 | + if (val[aR] > val[aL] || (eqLR && idx[aR] > idx[aL])) { | ||
| 275 | + largest = R; | ||
| 276 | + as = aR; | ||
| 277 | + } | ||
| 278 | + } | ||
| 279 | + int64_t ap = static_cast<int64_t>(pos) * jump; | ||
| 280 | + bool eq = aicpu::IsValueEqual<T>(val[as], val[ap]); | ||
| 281 | + if (val[as] > val[ap] || (eq && idx[as] > idx[ap])) { | ||
| 282 | + T tv = val[ap]; | ||
| 283 | + val[ap] = val[as]; | ||
| 284 | + val[as] = tv; | ||
| 285 | + int32_t ti = idx[ap]; | ||
| 286 | + idx[ap] = idx[as]; | ||
| 287 | + idx[as] = ti; | ||
| 288 | + pos = largest; | ||
| 289 | + } else { | ||
| 290 | + break; | ||
| 291 | + } | ||
| 292 | + } | ||
| 293 | +} | ||
| 294 | + | ||
| 295 | +template <typename T, bool largest_true> | ||
| 296 | +inline void HeapTopK(const T* __restrict__ in_base, int32_t jump, T* __restrict__ val, int32_t* __restrict__ idx, | ||
| 297 | + int32_t n, int32_t k) | ||
| 298 | +{ | ||
| 299 | + for (int32_t i = 0; i < k; ++i) { | ||
| 300 | + int64_t a = static_cast<int64_t>(i) * jump; | ||
| 301 | + val[a] = in_base[a]; | ||
| 302 | + idx[a] = i; | ||
| 303 | + } | ||
| 304 | + for (int32_t i = k / 2 - 1; i >= 0; --i) { | ||
| 305 | + if (largest_true) { | ||
| 306 | + SiftDownMin(val, idx, i, k, jump); | ||
| 307 | + } else { | ||
| 308 | + SiftDownMax(val, idx, i, k, jump); | ||
| 309 | + } | ||
| 310 | + } | ||
| 311 | + for (int32_t i = k; i < n; ++i) { | ||
| 312 | + int64_t a = static_cast<int64_t>(i) * jump; | ||
| 313 | + const T v = in_base[a]; | ||
| 314 | + const T root = val[0]; | ||
| 315 | + bool better = largest_true ? (v > root) : (v < root); | ||
| 316 | + if (better) { | ||
| 317 | + val[0] = v; | ||
| 318 | + idx[0] = i; | ||
| 319 | + if (largest_true) { | ||
| 320 | + SiftDownMin(val, idx, 0, k, jump); | ||
| 321 | + } else { | ||
| 322 | + SiftDownMax(val, idx, 0, k, jump); | ||
| 323 | + } | ||
| 324 | + } | ||
| 325 | + } | ||
| 326 | +} | ||
| 327 | + | ||
| 328 | +template <typename T, bool largest_true> | ||
| 329 | +inline void InPlaceHeapSort(T* __restrict__ val, int32_t* __restrict__ idx, int32_t k, int32_t jump) | ||
| 330 | +{ | ||
| 331 | + for (int32_t end = k - 1; end > 0; --end) { | ||
| 332 | + int64_t a0 = 0; | ||
| 333 | + int64_t ae = static_cast<int64_t>(end) * jump; | ||
| 334 | + T tv = val[a0]; | ||
| 335 | + val[a0] = val[ae]; | ||
| 336 | + val[ae] = tv; | ||
| 337 | + int32_t ti = idx[a0]; | ||
| 338 | + idx[a0] = idx[ae]; | ||
| 339 | + idx[ae] = ti; | ||
| 340 | + if (largest_true) { | ||
| 341 | + SiftDownMin(val, idx, 0, end, jump); | ||
| 342 | + } else { | ||
| 343 | + SiftDownMax(val, idx, 0, end, jump); | ||
| 344 | + } | ||
| 345 | + } | ||
| 346 | +} | ||
| 347 | + | ||
| 348 | +} // namespace | ||
| 349 | + | ||
| 350 | +namespace aicpu { | ||
| 351 | + | ||
| 352 | +uint32_t TopkV2DCpuKernel::Compute(CpuKernelContext& ctx) | ||
| 353 | +{ | ||
| 354 | + KernelStatus res = GetInputAndCheck(ctx); | ||
| 355 | + if (res != KERNEL_STATUS_OK) { | ||
| 356 | + return static_cast<uint32_t>(res); | ||
| 357 | + } | ||
| 358 | + KERNEL_LOG_INFO("[TopKV2D] Compute begin, dtype=%d, head=%d, n=%d, k=%d, tail=%d, " | ||
| 359 | + "largest=%d, sorted=%d.", | ||
| 360 | + static_cast<int>(data_type_), head_, n_, k_, tail_, static_cast<int>(largest_), | ||
| 361 | + static_cast<int>(sorted_)); | ||
| 362 | + res = DispatchByDtype(ctx); | ||
| 363 | + if (res != KERNEL_STATUS_OK) { | ||
| 364 | + return static_cast<uint32_t>(res); | ||
| 365 | + } | ||
| 366 | + return static_cast<uint32_t>(KERNEL_STATUS_OK); | ||
| 367 | +} | ||
| 368 | + | ||
| 369 | +KernelStatus TopkV2DCpuKernel::DispatchByDtype(const CpuKernelContext& ctx) | ||
| 370 | +{ | ||
| 371 | + switch (data_type_) { | ||
| 372 | + case DT_FLOAT16: | ||
| 373 | + return DoCompute<Eigen::half>(ctx); | ||
| 374 | + case DT_FLOAT: | ||
| 375 | + return DoCompute<float>(ctx); | ||
| 376 | + case DT_DOUBLE: | ||
| 377 | + return DoCompute<double>(ctx); | ||
| 378 | + case DT_UINT8: | ||
| 379 | + return DoCompute<uint8_t>(ctx); | ||
| 380 | + case DT_INT8: | ||
| 381 | + return DoCompute<int8_t>(ctx); | ||
| 382 | + case DT_UINT16: | ||
| 383 | + return DoCompute<uint16_t>(ctx); | ||
| 384 | + case DT_INT16: | ||
| 385 | + return DoCompute<int16_t>(ctx); | ||
| 386 | + case DT_UINT32: | ||
| 387 | + return DoCompute<uint32_t>(ctx); | ||
| 388 | + case DT_INT32: | ||
| 389 | + return DoCompute<int32_t>(ctx); | ||
| 390 | + case DT_UINT64: | ||
| 391 | + return DoCompute<uint64_t>(ctx); | ||
| 392 | + case DT_INT64: | ||
| 393 | + return DoCompute<int64_t>(ctx); | ||
| 394 | + default: | ||
| 395 | + KERNEL_LOG_ERROR("[TopKV2D] input tensor dtype=%d not supported.", static_cast<int>(data_type_)); | ||
| 396 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 397 | + } | ||
| 398 | +} | ||
| 399 | + | ||
| 400 | +template <typename T> | ||
| 401 | +KernelStatus TopkV2DCpuKernel::DoCompute(const CpuKernelContext& ctx) | ||
| 402 | +{ | ||
| 403 | + T* in = PtrToPtr<void, T>(input_tensor_->GetData()); | ||
| 404 | + T* val = PtrToPtr<void, T>(output_values_->GetData()); | ||
| 405 | + int32_t* indice = PtrToPtr<void, int32_t>(output_indices_->GetData()); | ||
| 406 | + KERNEL_CHECK_NULLPTR(in, KERNEL_STATUS_PARAM_INVALID, "[TopKV2D] input data pointer is null."); | ||
| 407 | + KERNEL_CHECK_NULLPTR(val, KERNEL_STATUS_PARAM_INVALID, "[TopKV2D] output values pointer is null."); | ||
| 408 | + KERNEL_CHECK_NULLPTR(indice, KERNEL_STATUS_PARAM_INVALID, "[TopKV2D] output indices pointer is null."); | ||
| 409 | + | ||
| 410 | + const int64_t total_units = static_cast<int64_t>(head_) * static_cast<int64_t>(tail_); | ||
| 411 | + if (total_units <= 0 || n_ <= 0 || k_ <= 0) { | ||
| 412 | + KERNEL_LOG_INFO("[TopKV2D] empty work (head*tail=%ld, n=%d, k=%d), skip.", total_units, n_, k_); | ||
| 413 | + return KERNEL_STATUS_OK; | ||
| 414 | + } | ||
| 415 | + | ||
| 416 | + const int64_t row_bytes_in = static_cast<int64_t>(n_) * static_cast<int64_t>(sizeof(T)); | ||
| 417 | + const int64_t num_cores = std::max(static_cast<int64_t>(1), | ||
| 418 | + static_cast<int64_t>(aicpu::CpuKernelUtils::GetCPUNum(ctx))); | ||
| 419 | + const int64_t per_unit = ComputePerUnit(total_units, row_bytes_in, num_cores); | ||
| 420 | + const int64_t total_input_bytes = total_units * static_cast<int64_t>(n_) * static_cast<int64_t>(sizeof(T)); | ||
| 421 | + const bool go_serial = (total_input_bytes < kSerialInputBytesThreshold) || (total_units == 1); | ||
| 422 | + | ||
| 423 | + auto shard = [this, in, val, indice](size_t start, size_t end) { | ||
| 424 | + TopKForNVectorImpl<T>(in, val, indice, static_cast<int64_t>(start), static_cast<int64_t>(end)); | ||
| 425 | + }; | ||
| 426 | + | ||
| 427 | + if (go_serial) { | ||
| 428 | + KERNEL_LOG_INFO("[TopKV2D] Serial path: total_units=%ld, total_input_bytes=%ld < %ld.", total_units, | ||
| 429 | + total_input_bytes, kSerialInputBytesThreshold); | ||
| 430 | + shard(0u, static_cast<size_t>(total_units)); | ||
| 431 | + return KERNEL_STATUS_OK; | ||
| 432 | + } | ||
| 433 | + | ||
| 434 | + KERNEL_LOG_INFO("[TopKV2D] Parallel path: total_units=%ld, per_unit=%ld, row_bytes=%ld.", total_units, per_unit, | ||
| 435 | + row_bytes_in); | ||
| 436 | + uint32_t ret = CpuKernelUtils::ParallelFor(ctx, total_units, per_unit, shard); | ||
| 437 | + if (ret != static_cast<uint32_t>(KERNEL_STATUS_OK)) { | ||
| 438 | + KERNEL_LOG_ERROR("[TopKV2D] CpuKernelUtils::ParallelFor failed, rc=%u, total_units=%ld, per_unit=%ld.", ret, | ||
| 439 | + total_units, per_unit); | ||
| 440 | + return KERNEL_STATUS_INNER_ERROR; | ||
| 441 | + } | ||
| 442 | + return KERNEL_STATUS_OK; | ||
| 443 | +} | ||
| 444 | + | ||
| 445 | +template <typename T> | ||
| 446 | +void TopkV2DCpuKernel::TopKForNVectorImpl(T* in, T* val, int32_t* indice, int64_t start, int64_t end) const | ||
| 447 | +{ | ||
| 448 | + const int64_t total_units = static_cast<int64_t>(head_) * static_cast<int64_t>(tail_); | ||
| 449 | + const bool k_almost_n = (k_ >= n_ - 1); | ||
| 450 | + const bool huge_rowcount_almost_full_sort = k_almost_n && | ||
| 451 | + (total_units * static_cast<int64_t>(n_) >= kHugeRowcountThreshold); | ||
| 452 | + const bool use_small_n_path = (n_ <= kSmallN) && (k_ < n_) && (k_ <= kSelectionMaxK) && | ||
| 453 | + !huge_rowcount_almost_full_sort; | ||
| 454 | + const bool use_full_sort_path = (n_ <= kSmallN) && (k_ == n_); | ||
| 455 | + | ||
| 456 | + for (int64_t u = start; u < end; ++u) { | ||
| 457 | + const int32_t head = static_cast<int32_t>(u / tail_); | ||
| 458 | + const int32_t tail = static_cast<int32_t>(u % tail_); | ||
| 459 | + const int64_t in_off = static_cast<int64_t>(head) * n_ * tail_ + tail; | ||
| 460 | + const int64_t out_off = static_cast<int64_t>(head) * k_ * tail_ + tail; | ||
| 461 | + const T* in_base = in + in_off; | ||
| 462 | + T* val_base = val + out_off; | ||
| 463 | + int32_t* idx_base = indice + out_off; | ||
| 464 | + | ||
| 465 | + if (use_small_n_path) { | ||
| 466 | + if (largest_) { | ||
| 467 | + SelectionTopK<T, true>(in_base, tail_, val_base, idx_base, n_, k_); | ||
| 468 | + } else { | ||
| 469 | + SelectionTopK<T, false>(in_base, tail_, val_base, idx_base, n_, k_); | ||
| 470 | + } | ||
| 471 | + } else if (use_full_sort_path) { | ||
| 472 | + if (largest_) { | ||
| 473 | + FullSortTopK<T, true>(in_base, tail_, val_base, idx_base, n_); | ||
| 474 | + } else { | ||
| 475 | + FullSortTopK<T, false>(in_base, tail_, val_base, idx_base, n_); | ||
| 476 | + } | ||
| 477 | + } else { | ||
| 478 | + if (largest_) { | ||
| 479 | + HeapTopK<T, true>(in_base, tail_, val_base, idx_base, n_, k_); | ||
| 480 | + if (sorted_) { | ||
| 481 | + InPlaceHeapSort<T, true>(val_base, idx_base, k_, tail_); | ||
| 482 | + } | ||
| 483 | + } else { | ||
| 484 | + HeapTopK<T, false>(in_base, tail_, val_base, idx_base, n_, k_); | ||
| 485 | + if (sorted_) { | ||
| 486 | + InPlaceHeapSort<T, false>(val_base, idx_base, k_, tail_); | ||
| 487 | + } | ||
| 488 | + } | ||
| 489 | + } | ||
| 490 | + } | ||
| 491 | +} | ||
| 492 | + | ||
| 493 | +template <typename T> | ||
| 494 | +void TopkV2DCpuKernel::TopKForNVector(size_t start, size_t end) | ||
| 495 | +{ | ||
| 496 | + T* in = PtrToPtr<void, T>(input_tensor_->GetData()); | ||
| 497 | + T* val = PtrToPtr<void, T>(output_values_->GetData()); | ||
| 498 | + int32_t* indice = PtrToPtr<void, int32_t>(output_indices_->GetData()); | ||
| 499 | + if (in == nullptr || val == nullptr || indice == nullptr) { | ||
| 500 | + return; | ||
| 501 | + } | ||
| 502 | + TopKForNVectorImpl<T>(in, val, indice, static_cast<int64_t>(start), static_cast<int64_t>(end)); | ||
| 503 | +} | ||
| 504 | + | ||
| 505 | +KernelStatus TopkV2DCpuKernel::ParseShapeAndDim(const CpuKernelContext& ctx) | ||
| 506 | +{ | ||
| 507 | + input_tensor_ = ctx.Input(0); | ||
| 508 | + KERNEL_CHECK_NULLPTR(input_tensor_, KERNEL_STATUS_PARAM_INVALID, "[TopKV2D] Get input[0] name[x] failed."); | ||
| 509 | + std::shared_ptr<TensorShape> input_shape = input_tensor_->GetTensorShape(); | ||
| 510 | + KERNEL_CHECK_NULLPTR(input_shape, KERNEL_STATUS_PARAM_INVALID, "[TopKV2D] Get shape of input[0] name[x] failed."); | ||
| 511 | + int32_t input_rank = input_shape->GetDims(); | ||
| 512 | + if (input_rank < 1) { | ||
| 513 | + KERNEL_LOG_ERROR("[TopKV2D] Input rank=%d must be >= 1.", input_rank); | ||
| 514 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 515 | + } | ||
| 516 | + input_rank_ = input_rank; | ||
| 517 | + AttrValue* dim = ctx.GetAttr("dim"); | ||
| 518 | + dim_ = static_cast<int32_t>(dim == nullptr ? -1 : (dim->GetInt())); | ||
| 519 | + dim_ = dim_ < 0 ? (input_rank + dim_) : dim_; | ||
| 520 | + KERNEL_CHECK_FALSE(((dim_ >= 0) && (dim_ < input_rank)), KERNEL_STATUS_PARAM_INVALID, | ||
| 521 | + "[TopKV2D] Invalid attr dim value=%d, must be in [%d, %d).", dim_, -input_rank, input_rank); | ||
| 522 | + head_ = 1; | ||
| 523 | + tail_ = 1; | ||
| 524 | + for (int32_t i = 0; i < input_rank; ++i) { | ||
| 525 | + if (i < dim_) { | ||
| 526 | + head_ *= static_cast<int32_t>(input_shape->GetDimSize(i)); | ||
| 527 | + } else if (i == dim_) { | ||
| 528 | + n_ = static_cast<int32_t>(input_shape->GetDimSize(i)); | ||
| 529 | + } else { | ||
| 530 | + tail_ *= static_cast<int32_t>(input_shape->GetDimSize(i)); | ||
| 531 | + KERNEL_CHECK_FALSE((tail_ != 0), KERNEL_STATUS_PARAM_INVALID, "[TopKV2D] input dim size cannot be 0."); | ||
| 532 | + } | ||
| 533 | + } | ||
| 534 | + data_type_ = static_cast<DataType>(input_tensor_->GetDataType()); | ||
| 535 | + return KERNEL_STATUS_OK; | ||
| 536 | +} | ||
| 537 | + | ||
| 538 | +KernelStatus TopkV2DCpuKernel::GetInputAndCheck(const CpuKernelContext& ctx) | ||
| 539 | +{ | ||
| 540 | + KernelStatus ret = ParseShapeAndDim(ctx); | ||
| 541 | + if (ret != KERNEL_STATUS_OK) { | ||
| 542 | + return ret; | ||
| 543 | + } | ||
| 544 | + | ||
| 545 | + Tensor* k_tensor = ctx.Input(1); | ||
| 546 | + KERNEL_CHECK_NULLPTR(k_tensor, KERNEL_STATUS_PARAM_INVALID, "[TopKV2D] Get input[1] name[k] failed."); | ||
| 547 | + KERNEL_CHECK_NULLPTR(k_tensor->GetData(), KERNEL_STATUS_PARAM_INVALID, | ||
| 548 | + "[TopKV2D] Get input[1] name[k] data failed."); | ||
| 549 | + k_ = *static_cast<int32_t*>(k_tensor->GetData()); | ||
| 550 | + if (k_ < 0) { | ||
| 551 | + KERNEL_LOG_ERROR("[TopKV2D] k=%d must be >= 0.", k_); | ||
| 552 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 553 | + } | ||
| 554 | + if (n_ < k_) { | ||
| 555 | + KERNEL_LOG_ERROR("[TopKV2D] Input must have at least k=%d elements along dim, got n=%d.", k_, n_); | ||
| 556 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 557 | + } | ||
| 558 | + | ||
| 559 | + AttrValue* sorted = ctx.GetAttr("sorted"); | ||
| 560 | + sorted_ = (sorted == nullptr) ? true : (sorted->GetBool()); | ||
| 561 | + AttrValue* largest = ctx.GetAttr("largest"); | ||
| 562 | + largest_ = (largest == nullptr) ? true : (largest->GetBool()); | ||
| 563 | + | ||
| 564 | + output_values_ = ctx.Output(0); | ||
| 565 | + KERNEL_CHECK_NULLPTR(output_values_, KERNEL_STATUS_PARAM_INVALID, "[TopKV2D] Get output[0] name[values] failed."); | ||
| 566 | + output_indices_ = ctx.Output(1); | ||
| 567 | + KERNEL_CHECK_NULLPTR(output_indices_, KERNEL_STATUS_PARAM_INVALID, "[TopKV2D] Get output[1] name[indices] failed."); | ||
| 568 | + return KERNEL_STATUS_OK; | ||
| 569 | +} | ||
| 570 | + | ||
| 571 | +REGISTER_CPU_KERNEL(TOPKV2D, TopkV2DCpuKernel); | ||
| 572 | +} // namespace aicpu | ||
| @@ -0,0 +1,50 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +namespace aicpu { | ||
| 18 | +class TopkV2DCpuKernel : public CpuKernel { | ||
| 19 | +public: | ||
| 20 | + ~TopkV2DCpuKernel() = default; | ||
| 21 | + | ||
| 22 | + uint32_t Compute(CpuKernelContext& ctx) override; | ||
| 23 | + | ||
| 24 | + KernelStatus GetInputAndCheck(const CpuKernelContext& ctx); | ||
| 25 | + template <typename T> | ||
| 26 | + KernelStatus DoCompute(const CpuKernelContext& ctx); | ||
| 27 | + template <typename T> | ||
| 28 | + void TopKForNVector(size_t start, size_t end); | ||
| 29 | + template <typename T> | ||
| 30 | + void TopKForNVectorImpl(T* in, T* val, int32_t* indice, int64_t start, int64_t end) const; | ||
| 31 | + KernelStatus DispatchByDtype(const CpuKernelContext& ctx); | ||
| 32 | + KernelStatus ParseShapeAndDim(const CpuKernelContext& ctx); | ||
| 33 | + | ||
| 34 | +private: | ||
| 35 | + int32_t k_ = 0; | ||
| 36 | + bool sorted_ = true; | ||
| 37 | + bool largest_ = true; | ||
| 38 | + int32_t dim_ = 0; | ||
| 39 | + int32_t input_rank_ = 0; | ||
| 40 | + DataType data_type_ = DT_DOUBLE; | ||
| 41 | + Tensor* input_tensor_ = nullptr; | ||
| 42 | + Tensor* output_values_ = nullptr; | ||
| 43 | + Tensor* output_indices_ = nullptr; | ||
| 44 | + int32_t head_ = 1; | ||
| 45 | + int32_t tail_ = 1; | ||
| 46 | + int32_t n_ = 1; | ||
| 47 | +}; | ||
| 48 | +} // namespace aicpu | ||
| 49 | + | ||
| 50 | + | ||
| @@ -0,0 +1,228 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +using namespace std; | ||
| 22 | +using namespace aicpu; | ||
| 23 | + | ||
| 24 | +namespace { | ||
| 25 | +template <typename T> | ||
| 26 | +struct ValueIndex { | ||
| 27 | + T value; | ||
| 28 | + int32_t index; | ||
| 29 | +}; | ||
| 30 | + | ||
| 31 | +template <typename T> | ||
| 32 | +bool CompareDescending(const ValueIndex<T>& one, const ValueIndex<T>& another) | ||
| 33 | +{ | ||
| 34 | + if (one.value == another.value) { | ||
| 35 | + return one.index < another.index; | ||
| 36 | + } | ||
| 37 | + return one.value > another.value; | ||
| 38 | +} | ||
| 39 | + | ||
| 40 | +template <typename T> | ||
| 41 | +bool CompareAscending(const ValueIndex<T>& one, const ValueIndex<T>& another) | ||
| 42 | +{ | ||
| 43 | + if (one.value == another.value) { | ||
| 44 | + return one.index < another.index; | ||
| 45 | + } | ||
| 46 | + return one.value < another.value; | ||
| 47 | +} | ||
| 48 | +} // namespace | ||
| 49 | + | ||
| 50 | +class TEST_TOPKV2D_UT : public testing::Test {}; | ||
| 51 | + | ||
| 52 | + | ||
| 53 | + auto node_def = CpuKernelUtils::CreateNodeDef(); \ | ||
| 54 | + NodeDefBuilder(node_def.get(), "TopKV2D", "TopKV2D") \ | ||
| 55 | + .Input({"x", data_types[0], shapes[0], datas[0]}) \ | ||
| 56 | + .Input({"k", data_types[1], shapes[1], datas[1]}) \ | ||
| 57 | + .Input({"assist_seq", data_types[4], shapes[4], datas[4]}) \ | ||
| 58 | + .Output({"values", data_types[2], shapes[2], datas[2]}) \ | ||
| 59 | + .Output({"indices", data_types[3], shapes[3], datas[3]}) \ | ||
| 60 | + .Attr("sorted", true) \ | ||
| 61 | + .Attr("largest", true) \ | ||
| 62 | + .Attr("dim", -1); | ||
| 63 | + | ||
| 64 | + | ||
| 65 | + auto node_def = CpuKernelUtils::CreateNodeDef(); \ | ||
| 66 | + NodeDefBuilder(node_def.get(), "TopKV2D", "TopKV2D") \ | ||
| 67 | + .Input({"x", data_types[0], shapes[0], datas[0]}) \ | ||
| 68 | + .Input({"k", data_types[1], shapes[1], datas[1]}) \ | ||
| 69 | + .Input({"assist_seq", data_types[4], shapes[4], datas[4]}) \ | ||
| 70 | + .Output({"values", data_types[2], shapes[2], datas[2]}) \ | ||
| 71 | + .Output({"indices", data_types[3], shapes[3], datas[3]}) \ | ||
| 72 | + .Attr("sorted", true) \ | ||
| 73 | + .Attr("largest", false) \ | ||
| 74 | + .Attr("dim", -1); | ||
| 75 | + | ||
| 76 | + | ||
| 77 | + auto node_def = CpuKernelUtils::CreateNodeDef(); \ | ||
| 78 | + NodeDefBuilder(node_def.get(), "TopKV2D", "TopKV2D") \ | ||
| 79 | + .Input({"x", data_types[0], shapes[0], datas[0]}) \ | ||
| 80 | + .Input({"k", data_types[1], shapes[1], datas[1]}) \ | ||
| 81 | + .Input({"assist_seq", data_types[4], shapes[4], datas[4]}) \ | ||
| 82 | + .Output({"values", data_types[2], shapes[2], datas[2]}) \ | ||
| 83 | + .Output({"indices", data_types[3], shapes[3], datas[3]}) \ | ||
| 84 | + .Attr("sorted", true) \ | ||
| 85 | + .Attr("largest", true) \ | ||
| 86 | + .Attr("dim", dim_val); | ||
| 87 | + | ||
| 88 | + | ||
| 89 | + TEST_F(TEST_TOPKV2D_UT, TestTopKV2D_##aicpu_type##_LARGEST) \ | ||
| 90 | + { \ | ||
| 91 | + vector<DataType> data_types = {aicpu_type, DT_INT32, aicpu_type, DT_INT32, DT_FLOAT16}; \ | ||
| 92 | + vector<vector<int64_t>> shapes = {{24}, {}, {7}, {7}, {24}}; \ | ||
| 93 | + base_type input[24]; \ | ||
| 94 | + SetRandomValue<base_type>(input, 24); \ | ||
| 95 | + vector<ValueIndex<base_type>> output_expect(24); \ | ||
| 96 | + for (int i = 0; i < 24; i++) { \ | ||
| 97 | + output_expect[i].index = i; \ | ||
| 98 | + output_expect[i].value = input[i]; \ | ||
| 99 | + } \ | ||
| 100 | + sort(output_expect.begin(), output_expect.end(), CompareDescending<base_type>); \ | ||
| 101 | + int32_t k = 7; \ | ||
| 102 | + base_type output_value[7] = {(base_type)0}; \ | ||
| 103 | + int32_t output_index[7] = {0}; \ | ||
| 104 | + Eigen::half assist_seq[24]; \ | ||
| 105 | + for (int i = 0; i < 24; i++) { \ | ||
| 106 | + assist_seq[i] = Eigen::half(0.0f); \ | ||
| 107 | + } \ | ||
| 108 | + vector<void*> datas = {(void*)input, (void*)&k, (void*)output_value, (void*)output_index, (void*)assist_seq}; \ | ||
| 109 | + CREATE_NODEDEF_V2D(shapes, data_types, datas); \ | ||
| 110 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); \ | ||
| 111 | + for (int i = 0; i < 7; i++) { \ | ||
| 112 | + EXPECT_EQ(output_value[i], output_expect[i].value); \ | ||
| 113 | + EXPECT_EQ(output_index[i], output_expect[i].index); \ | ||
| 114 | + } \ | ||
| 115 | + } \ | ||
| 116 | + TEST_F(TEST_TOPKV2D_UT, TestTopKV2D_##aicpu_type##_SMALLEST) \ | ||
| 117 | + { \ | ||
| 118 | + vector<DataType> data_types = {aicpu_type, DT_INT32, aicpu_type, DT_INT32, DT_FLOAT16}; \ | ||
| 119 | + vector<vector<int64_t>> shapes = {{24}, {}, {7}, {7}, {24}}; \ | ||
| 120 | + base_type input[24]; \ | ||
| 121 | + SetRandomValue<base_type>(input, 24); \ | ||
| 122 | + vector<ValueIndex<base_type>> output_expect(24); \ | ||
| 123 | + for (int i = 0; i < 24; i++) { \ | ||
| 124 | + output_expect[i].index = i; \ | ||
| 125 | + output_expect[i].value = input[i]; \ | ||
| 126 | + } \ | ||
| 127 | + sort(output_expect.begin(), output_expect.end(), CompareAscending<base_type>); \ | ||
| 128 | + int32_t k = 7; \ | ||
| 129 | + base_type output_value[7] = {(base_type)0}; \ | ||
| 130 | + int32_t output_index[7] = {0}; \ | ||
| 131 | + Eigen::half assist_seq[24]; \ | ||
| 132 | + for (int i = 0; i < 24; i++) { \ | ||
| 133 | + assist_seq[i] = Eigen::half(0.0f); \ | ||
| 134 | + } \ | ||
| 135 | + vector<void*> datas = {(void*)input, (void*)&k, (void*)output_value, (void*)output_index, (void*)assist_seq}; \ | ||
| 136 | + CREATE_NODEDEF_V2D_SMALLEST(shapes, data_types, datas); \ | ||
| 137 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); \ | ||
| 138 | + for (int i = 0; i < 7; i++) { \ | ||
| 139 | + EXPECT_EQ(output_value[i], output_expect[i].value); \ | ||
| 140 | + EXPECT_EQ(output_index[i], output_expect[i].index); \ | ||
| 141 | + } \ | ||
| 142 | + } \ | ||
| 143 | + TEST_F(TEST_TOPKV2D_UT, TestTopKV2D_##aicpu_type##_SECOND_LAST_DIM) \ | ||
| 144 | + { \ | ||
| 145 | + vector<DataType> data_types = {aicpu_type, DT_INT32, aicpu_type, DT_INT32, DT_FLOAT16}; \ | ||
| 146 | + vector<vector<int64_t>> shapes = {{2, 3, 4}, {}, {2, 2, 4}, {2, 2, 4}, {2, 3, 4}}; \ | ||
| 147 | + base_type input[24]; \ | ||
| 148 | + for (int i = 0; i < 24; i++) { \ | ||
| 149 | + input[i] = base_type(i + 1); \ | ||
| 150 | + } \ | ||
| 151 | + base_type output_value_expect[16] = {base_type(9), base_type(10), base_type(11), base_type(12), \ | ||
| 152 | + base_type(5), base_type(6), base_type(7), base_type(8), \ | ||
| 153 | + base_type(21), base_type(22), base_type(23), base_type(24), \ | ||
| 154 | + base_type(17), base_type(18), base_type(19), base_type(20)}; \ | ||
| 155 | + int32_t output_index_expect[16] = {2, 2, 2, 2, 1, 1, 1, 1, 2, 2, 2, 2, 1, 1, 1, 1}; \ | ||
| 156 | + int32_t k = 2; \ | ||
| 157 | + base_type output_value[16] = {(base_type)0}; \ | ||
| 158 | + int32_t output_index[16] = {0}; \ | ||
| 159 | + Eigen::half assist_seq[24]; \ | ||
| 160 | + for (int i = 0; i < 24; i++) { \ | ||
| 161 | + assist_seq[i] = Eigen::half(0.0f); \ | ||
| 162 | + } \ | ||
| 163 | + vector<void*> datas = {(void*)input, (void*)&k, (void*)output_value, (void*)output_index, (void*)assist_seq}; \ | ||
| 164 | + CREATE_NODEDEF_V2D_DIM(shapes, data_types, datas, -2); \ | ||
| 165 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); \ | ||
| 166 | + for (int i = 0; i < 16; i++) { \ | ||
| 167 | + EXPECT_EQ(output_value[i], output_value_expect[i]); \ | ||
| 168 | + EXPECT_EQ(output_index[i], output_index_expect[i]); \ | ||
| 169 | + } \ | ||
| 170 | + } | ||
| 171 | + | ||
| 172 | +TEST_F(TEST_TOPKV2D_UT, TestTopKV2D_KVALUE_EXCEPTION) | ||
| 173 | +{ | ||
| 174 | + vector<DataType> data_types = {DT_INT64, DT_INT32, DT_INT64, DT_INT32, DT_FLOAT16}; | ||
| 175 | + vector<vector<int64_t>> shapes = {{24}, {}, {7}, {7}, {24}}; | ||
| 176 | + int64_t input[24]; | ||
| 177 | + SetRandomValue<int64_t>(input, 24); | ||
| 178 | + int32_t k = -1; | ||
| 179 | + int64_t output_value[7] = {(int64_t)0}; | ||
| 180 | + int32_t output_index[7] = {0}; | ||
| 181 | + Eigen::half assist_seq[24]; | ||
| 182 | + for (int i = 0; i < 24; i++) { | ||
| 183 | + assist_seq[i] = Eigen::half(0.0f); | ||
| 184 | + } | ||
| 185 | + vector<void*> datas = {(void*)input, (void*)&k, (void*)output_value, (void*)output_index, (void*)assist_seq}; | ||
| 186 | + CREATE_NODEDEF_V2D(shapes, data_types, datas); | ||
| 187 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 188 | +} | ||
| 189 | + | ||
| 190 | +TEST_F(TEST_TOPKV2D_UT, TestTopKV2D_INPUT_DATADYPE_EXCEPTION) | ||
| 191 | +{ | ||
| 192 | + vector<DataType> data_types = {DT_BOOL, DT_INT32, DT_BOOL, DT_INT32, DT_FLOAT16}; | ||
| 193 | + vector<vector<int64_t>> shapes = {{24}, {}, {7}, {7}, {24}}; | ||
| 194 | + bool input[24]; | ||
| 195 | + SetRandomValue<bool>(input, 24); | ||
| 196 | + int32_t k = 7; | ||
| 197 | + bool output_value[7] = {(bool)0}; | ||
| 198 | + int32_t output_index[7] = {0}; | ||
| 199 | + Eigen::half assist_seq[24]; | ||
| 200 | + for (int i = 0; i < 24; i++) { | ||
| 201 | + assist_seq[i] = Eigen::half(0.0f); | ||
| 202 | + } | ||
| 203 | + vector<void*> datas = {(void*)input, (void*)&k, (void*)output_value, (void*)output_index, (void*)assist_seq}; | ||
| 204 | + CREATE_NODEDEF_V2D(shapes, data_types, datas); | ||
| 205 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 206 | +} | ||
| 207 | + | ||
| 208 | +ADD_CASE_V2D(Eigen::half, DT_FLOAT16) | ||
| 209 | + | ||
| 210 | +ADD_CASE_V2D(float, DT_FLOAT) | ||
| 211 | + | ||
| 212 | +ADD_CASE_V2D(double, DT_DOUBLE) | ||
| 213 | + | ||
| 214 | +ADD_CASE_V2D(int8_t, DT_INT8) | ||
| 215 | + | ||
| 216 | +ADD_CASE_V2D(int16_t, DT_INT16) | ||
| 217 | + | ||
| 218 | +ADD_CASE_V2D(int32_t, DT_INT32) | ||
| 219 | + | ||
| 220 | +ADD_CASE_V2D(int64_t, DT_INT64) | ||
| 221 | + | ||
| 222 | +ADD_CASE_V2D(uint8_t, DT_UINT8) | ||
| 223 | + | ||
| 224 | +ADD_CASE_V2D(uint16_t, DT_UINT16) | ||
| 225 | + | ||
| 226 | +ADD_CASE_V2D(uint32_t, DT_UINT32) | ||
| 227 | + | ||
| 228 | +ADD_CASE_V2D(uint64_t, DT_UINT64) | ||
| @@ -0,0 +1,11 @@ | |||
| 1 | +# ---------------------------------------------------------------------------- | ||
| 2 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | +# CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | +# Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | +# See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | +# ---------------------------------------------------------------------------- | ||
| 10 | + | ||
| 11 | +add_all_modules_sources(OPTYPE zeta ACLNNTYPE aclnn_exclude) | ||
| @@ -0,0 +1,72 @@ | |||
| 1 | +# Zeta | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 产品 | 是否支持 | | ||
| 6 | +| :------------------------------------------ | :------: | | ||
| 7 | +| Ascend 950PR/Ascend 950DT | √ | | ||
| 8 | +| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ | | ||
| 9 | +| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ | | ||
| 10 | +| Atlas 200I/500 A2 推理产品 | √ | | ||
| 11 | +| Atlas 推理系列产品 | √ | | ||
| 12 | +| Atlas 训练系列产品 | √ | | ||
| 13 | + | ||
| 14 | +## 功能说明 | ||
| 15 | + | ||
| 16 | +- 算子功能：计算Hurwitz zeta函数，即 ζ(x, q) = Σ(n=0 to ∞) 1/(q+n)^x。 | ||
| 17 | + | ||
| 18 | +- 计算公式： | ||
| 19 | + | ||
| 20 | +$$z = \zeta(x, q) = \sum_{n=0}^{\infty} \frac{1}{(q+n)^x}$$ | ||
| 21 | + | ||
| 22 | +## 参数说明 | ||
| 23 | + | ||
| 24 | +<table style="undefined;table-layout: fixed; width: 1005px"><colgroup> | ||
| 25 | +<col style="width: 140px"> | ||
| 26 | +<col style="width: 140px"> | ||
| 27 | +<col style="width: 180px"> | ||
| 28 | +<col style="width: 213px"> | ||
| 29 | +<col style="width: 100px"> | ||
| 30 | +</colgroup> | ||
| 31 | +<thead> | ||
| 32 | + <tr> | ||
| 33 | + <th>参数名</th> | ||
| 34 | + <th>输入/输出/属性</th> | ||
| 35 | + <th>描述</th> | ||
| 36 | + <th>数据类型</th> | ||
| 37 | + <th>数据格式</th> | ||
| 38 | + </tr></thead> | ||
| 39 | +<tbody> | ||
| 40 | + <tr> | ||
| 41 | + <td>x</td> | ||
| 42 | + <td>输入</td> | ||
| 43 | + <td>输入张量，Hurwitz zeta函数的指数参数。</td> | ||
| 44 | + <td>FLOAT、DOUBLE</td> | ||
| 45 | + <td>ND</td> | ||
| 46 | + </tr> | ||
| 47 | + <tr> | ||
| 48 | + <td>q</td> | ||
| 49 | + <td>输入</td> | ||
| 50 | + <td>输入张量，Hurwitz zeta函数的偏移参数，必须与x具有相同的数据类型。</td> | ||
| 51 | + <td>FLOAT、DOUBLE</td> | ||
| 52 | + <td>ND</td> | ||
| 53 | + </tr> | ||
| 54 | + <tr> | ||
| 55 | + <td>z</td> | ||
| 56 | + <td>输出</td> | ||
| 57 | + <td>输出张量，Hurwitz zeta函数的计算结果，与x具有相同的类型和形状。</td> | ||
| 58 | + <td>FLOAT、DOUBLE</td> | ||
| 59 | + <td>ND</td> | ||
| 60 | + </tr> | ||
| 61 | +</tbody></table> | ||
| 62 | + | ||
| 63 | +## 约束说明 | ||
| 64 | + | ||
| 65 | +- 输入x和q的数据类型必须一致。 | ||
| 66 | +- 输入x和q的数据大小必须一致（不支持broadcast）。 | ||
| 67 | + | ||
| 68 | +## 调用说明 | ||
| 69 | + | ||
| 70 | +| 调用方式 | 样例代码 | 说明 | | ||
| 71 | +| --------- | ------------------------------------------------------------ | ------------------------------------------------------------ | | ||
| 72 | +| 图模式接口 | [test_geir_zeta](examples/test_geir_zeta.cpp) | 通过[算子IR](op_graph/zeta_proto.h)接口方式调用Zeta算子。 | | ||
| @@ -0,0 +1,283 @@ | |||||||||||||||
| 1 | +/** | ||||||||||||||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||||||||||||||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||||||||||||||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||||||||||||||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||||||||||||||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||||||||||||||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||||||||||||||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||||||||||||||
| 9 | + */ | ||||||||||||||
| 10 | + | ||||||||||||||
| 11 | +/*! | ||||||||||||||
| 12 | + * \file test_geir_zeta.cpp | ||||||||||||||
| 13 | + * \brief GE IR test for Zeta operator | ||||||||||||||
| 14 | + */ | ||||||||||||||
| 15 | + | ||||||||||||||
| 16 | + | ||||||||||||||
| 17 | + | ||||||||||||||
| 18 | + | ||||||||||||||
| 19 | + | ||||||||||||||
| 20 | + | ||||||||||||||
| 21 | + | ||||||||||||||
| 22 | + | ||||||||||||||
| 23 | + | ||||||||||||||
| 24 | + | ||||||||||||||
| 25 | + | ||||||||||||||
| 26 | + | ||||||||||||||
| 27 | + | ||||||||||||||
| 28 | + | ||||||||||||||
| 29 | + | ||||||||||||||
| 30 | + | ||||||||||||||
| 31 | + | ||||||||||||||
| 32 | + | ||||||||||||||
| 33 | + | ||||||||||||||
| 34 | + | ||||||||||||||
| 35 | + | ||||||||||||||
| 36 | + | ||||||||||||||
| 37 | + | ||||||||||||||
| 38 | + | ||||||||||||||
| 39 | + | ||||||||||||||
| 40 | +using namespace ge; | ||||||||||||||
| 41 | +using std::map; | ||||||||||||||
| 42 | +using std::string; | ||||||||||||||
| 43 | +using std::vector; | ||||||||||||||
| 44 | + | ||||||||||||||
| 45 | +string GetTime() | ||||||||||||||
| 46 | +{ | ||||||||||||||
| 47 | + time_t timep; | ||||||||||||||
| 48 | + time(&timep); | ||||||||||||||
| 49 | + char tmp[64]; | ||||||||||||||
| 50 | + strftime(tmp, sizeof(tmp), "%Y-%m-%d %H:%M:%S,000", localtime(&timep)); | ||||||||||||||
| 51 | + return tmp; | ||||||||||||||
| 52 | +} | ||||||||||||||
| 53 | + | ||||||||||||||
| 54 | +uint32_t GetDataTypeSize(DataType dt) | ||||||||||||||
| 55 | +{ | ||||||||||||||
| 56 | + switch (dt) { | ||||||||||||||
| 57 | + case ge::DT_BOOL: | ||||||||||||||
| 58 | + return 1U; | ||||||||||||||
| 59 | + case ge::DT_INT8: | ||||||||||||||
| 60 | + case ge::DT_UINT8: | ||||||||||||||
| 61 | + return 1U; | ||||||||||||||
| 62 | + case ge::DT_FLOAT16: | ||||||||||||||
| 63 | + case ge::DT_INT16: | ||||||||||||||
| 64 | + case ge::DT_UINT16: | ||||||||||||||
| 65 | + return 2U; | ||||||||||||||
| 66 | + case ge::DT_FLOAT: | ||||||||||||||
| 67 | + case ge::DT_INT32: | ||||||||||||||
| 68 | + case ge::DT_UINT32: | ||||||||||||||
| 69 | + return 4U; | ||||||||||||||
| 70 | + case ge::DT_DOUBLE: | ||||||||||||||
| 71 | + case ge::DT_INT64: | ||||||||||||||
| 72 | + case ge::DT_UINT64: | ||||||||||||||
| 73 | + return 8U; | ||||||||||||||
| 74 | + default: | ||||||||||||||
| 75 | + return 0U; | ||||||||||||||
| 76 | + } | ||||||||||||||
| 77 | +} | ||||||||||||||
| 78 | + | ||||||||||||||
| 79 | +int32_t GenData(vector<int64_t> shapes, Tensor& tensor, TensorDesc& tensor_desc, DataType data_type, double value) | ||||||||||||||
| 80 | +{ | ||||||||||||||
| 81 | + tensor_desc.SetRealDimCnt(shapes.size()); | ||||||||||||||
| 82 | + size_t size = 1; | ||||||||||||||
| 83 | + for (uint32_t i = 0; i < shapes.size(); i++) { | ||||||||||||||
| 84 | + size *= static_cast<size_t>(shapes[i]); | ||||||||||||||
| 85 | + } | ||||||||||||||
| 86 | + uint32_t type_size = GetDataTypeSize(data_type); | ||||||||||||||
| 87 | + if (type_size == 0U) { | ||||||||||||||
| 88 | + printf("%s - ERROR - [XIR]: GenData: unsupported data type\n", GetTime().c_str()); | ||||||||||||||
| 89 | + return FAILED; | ||||||||||||||
| 90 | + } | ||||||||||||||
| 91 | + uint32_t data_len = static_cast<uint32_t>(size * type_size); | ||||||||||||||
| 92 | + uint8_t* buf = new (std::nothrow) uint8_t[data_len]; | ||||||||||||||
| 93 | + if (buf == nullptr) { | ||||||||||||||
| 94 | + printf("%s - ERROR - [XIR]: GenData: allocate memory failed\n", GetTime().c_str()); | ||||||||||||||
| 95 | + return FAILED; | ||||||||||||||
| 96 | + } | ||||||||||||||
| 97 | + | ||||||||||||||
| 98 | + if (data_type == ge::DT_FLOAT) { | ||||||||||||||
| 99 | + float* p = reinterpret_cast<float*>(buf); | ||||||||||||||
| 100 | + for (size_t i = 0; i < size; ++i) { | ||||||||||||||
| 101 | + p[i] = static_cast<float>(value); | ||||||||||||||
| 102 | + } | ||||||||||||||
| 103 | + } else if (data_type == ge::DT_DOUBLE) { | ||||||||||||||
| 104 | + double* p = reinterpret_cast<double*>(buf); | ||||||||||||||
| 105 | + for (size_t i = 0; i < size; ++i) { | ||||||||||||||
| 106 | + p[i] = static_cast<double>(value); | ||||||||||||||
| 107 | + } | ||||||||||||||
| 108 | + } | ||||||||||||||
| 109 | + | ||||||||||||||
| 110 | + tensor = Tensor(tensor_desc, buf, data_len); | ||||||||||||||
| 111 | + delete[] buf; | ||||||||||||||
| 112 | + return SUCCESS; | ||||||||||||||
| 113 | +} | ||||||||||||||
| 114 | + | ||||||||||||||
| 115 | +int32_t WriteDataToFile(string bin_file, uint64_t data_size, uint8_t* inputData) | ||||||||||||||
| 116 | +{ | ||||||||||||||
| 117 | + FILE* fp = fopen(bin_file.c_str(), "wb"); | ||||||||||||||
| 118 | + if (fp == nullptr) { | ||||||||||||||
| 119 | + return FAILED; | ||||||||||||||
| 120 | + } | ||||||||||||||
| 121 | + fwrite(inputData, sizeof(uint8_t), data_size, fp); | ||||||||||||||
| 122 | + fclose(fp); | ||||||||||||||
| 123 | + return SUCCESS; | ||||||||||||||
🟡 Medium Priority 文件 建议：与 改动建议
| |||||||||||||||
| 124 | +} | ||||||||||||||
| 125 | + | ||||||||||||||
| 126 | +void ProcessInputData(vector<Tensor>& input) | ||||||||||||||
| 127 | +{ | ||||||||||||||
| 128 | + for (size_t i = 0; i < input.size(); i++) { | ||||||||||||||
| 129 | + string input_file = "./tc_ge_irrun_test_0008_npu_input_" + std::to_string(i) + ".bin"; | ||||||||||||||
| 130 | + uint8_t* input_data_i = input[i].GetData(); | ||||||||||||||
| 131 | + int64_t input_shape = input[i].GetTensorDesc().GetShape().GetShapeSize(); | ||||||||||||||
| 132 | + uint32_t type_size = GetDataTypeSize(input[i].GetTensorDesc().GetDataType()); | ||||||||||||||
| 133 | + if (type_size == 0U) { | ||||||||||||||
| 134 | + printf("ERROR: input %zu has unsupported dtype\n", i); | ||||||||||||||
| 135 | + continue; | ||||||||||||||
| 136 | + } | ||||||||||||||
| 137 | + uint32_t data_size = static_cast<uint32_t>(input_shape * type_size); | ||||||||||||||
| 138 | + WriteDataToFile(input_file.c_str(), data_size, input_data_i); | ||||||||||||||
| 139 | + } | ||||||||||||||
| 140 | +} | ||||||||||||||
| 141 | + | ||||||||||||||
| 142 | +void ProcessOutputData(vector<Tensor>& output) | ||||||||||||||
| 143 | +{ | ||||||||||||||
| 144 | + for (size_t i = 0; i < output.size(); i++) { | ||||||||||||||
| 145 | + string output_file = "./tc_ge_irrun_test_0008_npu_output_" + std::to_string(i) + ".bin"; | ||||||||||||||
| 146 | + uint8_t* output_data_i = output[i].GetData(); | ||||||||||||||
| 147 | + int64_t output_shape = output[i].GetTensorDesc().GetShape().GetShapeSize(); | ||||||||||||||
| 148 | + uint32_t type_size = GetDataTypeSize(output[i].GetTensorDesc().GetDataType()); | ||||||||||||||
| 149 | + if (type_size == 0U) { | ||||||||||||||
| 150 | + printf("ERROR: output %zu has unsupported dtype\n", i); | ||||||||||||||
| 151 | + continue; | ||||||||||||||
| 152 | + } | ||||||||||||||
| 153 | + uint32_t data_size = static_cast<uint32_t>(output_shape * type_size); | ||||||||||||||
| 154 | + WriteDataToFile(output_file.c_str(), data_size, output_data_i); | ||||||||||||||
| 155 | + } | ||||||||||||||
| 156 | +} | ||||||||||||||
| 157 | + | ||||||||||||||
| 158 | +int CreateOppInGraph(DataType inDtype, vector<Tensor>& input, vector<Operator>& inputs, vector<Operator>& outputs, | ||||||||||||||
| 159 | + Graph& graph) | ||||||||||||||
| 160 | +{ | ||||||||||||||
| 161 | + vector<int64_t> shape = {2, 3}; | ||||||||||||||
| 162 | + | ||||||||||||||
| 163 | + // Input x | ||||||||||||||
| 164 | + auto x_data = op::Data("x_data").set_attr_index(0); | ||||||||||||||
| 165 | + TensorDesc x_desc = TensorDesc(ge::Shape(shape), FORMAT_ND, inDtype); | ||||||||||||||
| 166 | + x_desc.SetPlacement(ge::kPlacementHost); | ||||||||||||||
| 167 | + x_desc.SetFormat(FORMAT_ND); | ||||||||||||||
| 168 | + Tensor x_tensor; | ||||||||||||||
| 169 | + int32_t ret = GenData(shape, x_tensor, x_desc, inDtype, 2.0); | ||||||||||||||
| 170 | + if (ret != SUCCESS) { | ||||||||||||||
| 171 | + printf("%s - ERROR - [XIR]: Generate x data failed\n", GetTime().c_str()); | ||||||||||||||
| 172 | + return FAILED; | ||||||||||||||
| 173 | + } | ||||||||||||||
| 174 | + x_data.update_input_desc_x(x_desc); | ||||||||||||||
| 175 | + x_data.update_output_desc_y(x_desc); | ||||||||||||||
| 176 | + graph.AddOp(x_data); | ||||||||||||||
| 177 | + input.push_back(x_tensor); | ||||||||||||||
| 178 | + inputs.push_back(x_data); | ||||||||||||||
| 179 | + | ||||||||||||||
| 180 | + // Input q | ||||||||||||||
| 181 | + auto q_data = op::Data("q_data").set_attr_index(1); | ||||||||||||||
| 182 | + TensorDesc q_desc = TensorDesc(ge::Shape(shape), FORMAT_ND, inDtype); | ||||||||||||||
| 183 | + q_desc.SetPlacement(ge::kPlacementHost); | ||||||||||||||
| 184 | + q_desc.SetFormat(FORMAT_ND); | ||||||||||||||
| 185 | + Tensor q_tensor; | ||||||||||||||
| 186 | + ret = GenData(shape, q_tensor, q_desc, inDtype, 1.5); | ||||||||||||||
| 187 | + if (ret != SUCCESS) { | ||||||||||||||
| 188 | + printf("%s - ERROR - [XIR]: Generate q data failed\n", GetTime().c_str()); | ||||||||||||||
| 189 | + return FAILED; | ||||||||||||||
| 190 | + } | ||||||||||||||
| 191 | + q_data.update_input_desc_x(q_desc); | ||||||||||||||
| 192 | + q_data.update_output_desc_y(q_desc); | ||||||||||||||
| 193 | + graph.AddOp(q_data); | ||||||||||||||
| 194 | + input.push_back(q_tensor); | ||||||||||||||
| 195 | + inputs.push_back(q_data); | ||||||||||||||
| 196 | + | ||||||||||||||
| 197 | + // Zeta operator | ||||||||||||||
| 198 | + auto zeta_op = op::Zeta("zeta_op"); | ||||||||||||||
| 199 | + zeta_op.set_input_x(x_data); | ||||||||||||||
| 200 | + zeta_op.set_input_q(q_data); | ||||||||||||||
| 201 | + TensorDesc z_desc = TensorDesc(ge::Shape(shape), FORMAT_ND, inDtype); | ||||||||||||||
| 202 | + zeta_op.update_output_desc_z(z_desc); | ||||||||||||||
| 203 | + graph.AddOp(zeta_op); | ||||||||||||||
| 204 | + outputs.push_back(zeta_op); | ||||||||||||||
| 205 | + | ||||||||||||||
| 206 | + return SUCCESS; | ||||||||||||||
| 207 | +} | ||||||||||||||
| 208 | + | ||||||||||||||
| 209 | +int main() | ||||||||||||||
| 210 | +{ | ||||||||||||||
| 211 | + const char* graph_name = "tc_ge_irrun_test"; | ||||||||||||||
| 212 | + Graph graph(graph_name); | ||||||||||||||
| 213 | + vector<Tensor> input; | ||||||||||||||
| 214 | + | ||||||||||||||
| 215 | + printf("%s - INFO - [XIR]: Start to initialize ge using ge global options\n", GetTime().c_str()); | ||||||||||||||
| 216 | + std::map<AscendString, AscendString> global_options = {{"ge.exec.deviceId", "0"}, {"ge.graphRunMode", "1"}}; | ||||||||||||||
| 217 | + Status ret = ge::GEInitialize(global_options); | ||||||||||||||
| 218 | + if (ret != SUCCESS) { | ||||||||||||||
| 219 | + printf("%s - INFO - [XIR]: Initialize ge using ge global options failed\n", GetTime().c_str()); | ||||||||||||||
| 220 | + return FAILED; | ||||||||||||||
| 221 | + } | ||||||||||||||
| 222 | + printf("%s - INFO - [XIR]: Initialize ge using ge global options success\n", GetTime().c_str()); | ||||||||||||||
| 223 | + | ||||||||||||||
| 224 | + vector<Operator> inputs{}; | ||||||||||||||
| 225 | + vector<Operator> outputs{}; | ||||||||||||||
| 226 | + | ||||||||||||||
| 227 | + DataType inDtype = DT_FLOAT; | ||||||||||||||
| 228 | + | ||||||||||||||
| 229 | + ret = CreateOppInGraph(inDtype, input, inputs, outputs, graph); | ||||||||||||||
| 230 | + if (ret != SUCCESS) { | ||||||||||||||
| 231 | + printf("%s - ERROR - [XIR]: Create graph failed\n", GetTime().c_str()); | ||||||||||||||
| 232 | + return FAILED; | ||||||||||||||
| 233 | + } | ||||||||||||||
| 234 | + | ||||||||||||||
| 235 | + if (!inputs.empty() && !outputs.empty()) { | ||||||||||||||
| 236 | + graph.SetInputs(inputs).SetOutputs(outputs); | ||||||||||||||
| 237 | + } | ||||||||||||||
| 238 | + | ||||||||||||||
| 239 | + std::map<AscendString, AscendString> build_options = {}; | ||||||||||||||
| 240 | + printf("%s - INFO - [XIR]: Start to create ir session\n", GetTime().c_str()); | ||||||||||||||
| 241 | + ge::Session* session = new Session(build_options); | ||||||||||||||
| 242 | + if (session == nullptr) { | ||||||||||||||
| 243 | + printf("%s - ERROR - [XIR]: Create session failed\n", GetTime().c_str()); | ||||||||||||||
| 244 | + return FAILED; | ||||||||||||||
| 245 | + } | ||||||||||||||
| 246 | + printf("%s - INFO - [XIR]: Create ir session success\n", GetTime().c_str()); | ||||||||||||||
| 247 | + | ||||||||||||||
| 248 | + uint32_t graph_id = 0; | ||||||||||||||
| 249 | + std::map<AscendString, AscendString> graph_options = {}; | ||||||||||||||
| 250 | + | ||||||||||||||
| 251 | + printf("%s - INFO - [XIR]: Add graph\n", GetTime().c_str()); | ||||||||||||||
| 252 | + ret = session->AddGraph(graph_id, graph, graph_options); | ||||||||||||||
| 253 | + if (ret != SUCCESS) { | ||||||||||||||
| 254 | + printf("%s - ERROR - [XIR]: Add graph failed\n", GetTime().c_str()); | ||||||||||||||
| 255 | + delete session; | ||||||||||||||
| 256 | + GEFinalize(); | ||||||||||||||
| 257 | + return FAILED; | ||||||||||||||
| 258 | + } | ||||||||||||||
| 259 | + | ||||||||||||||
| 260 | + printf("%s - INFO - [XIR]: Start to run graph\n", GetTime().c_str()); | ||||||||||||||
| 261 | + vector<Tensor> output; | ||||||||||||||
| 262 | + ret = session->RunGraph(graph_id, input, output); | ||||||||||||||
| 263 | + if (ret != SUCCESS) { | ||||||||||||||
| 264 | + printf("%s - ERROR - [XIR]: Run graph failed\n", GetTime().c_str()); | ||||||||||||||
| 265 | + delete session; | ||||||||||||||
| 266 | + GEFinalize(); | ||||||||||||||
| 267 | + return FAILED; | ||||||||||||||
| 268 | + } | ||||||||||||||
| 269 | + printf("%s - INFO - [XIR]: Run graph success\n", GetTime().c_str()); | ||||||||||||||
| 270 | + | ||||||||||||||
| 271 | + ProcessInputData(input); | ||||||||||||||
| 272 | + ProcessOutputData(output); | ||||||||||||||
| 273 | + | ||||||||||||||
| 274 | + printf("%s - INFO - [XIR]: Start to finalize ir graph session\n", GetTime().c_str()); | ||||||||||||||
| 275 | + delete session; | ||||||||||||||
| 276 | + ret = ge::GEFinalize(); | ||||||||||||||
| 277 | + if (ret != SUCCESS) { | ||||||||||||||
| 278 | + printf("%s - INFO - [XIR]: Finalize ir graph session failed\n", GetTime().c_str()); | ||||||||||||||
| 279 | + return FAILED; | ||||||||||||||
| 280 | + } | ||||||||||||||
| 281 | + printf("%s - INFO - [XIR]: Finalize ir graph session success\n", GetTime().c_str()); | ||||||||||||||
| 282 | + return SUCCESS; | ||||||||||||||
| 283 | +} | ||||||||||||||
| @@ -0,0 +1,41 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +namespace ge { | ||
| 17 | +/** | ||
| 18 | + *@brief Compute the Hurwitz zeta function. | ||
| 19 | + | ||
| 20 | + *@par Inputs: | ||
| 21 | + *The input q must be the same type as x. Inputs include: | ||
| 22 | + *@li x:A Tensor. Must be one of the following types: float32, double. | ||
| 23 | + *@li q:A Tensor. Must have the same type as x. \n | ||
| 24 | + | ||
| 25 | + *@par Outputs: | ||
| 26 | + *z:A Tensor. Has the same type as x. \n | ||
| 27 | + | ||
| 28 | + *@attention Constraints: | ||
| 29 | + *The implementation for Zeta on Ascend uses ai cpu, with bad performance. | ||
| 30 | + | ||
| 31 | + *@par Third-party framework compatibility. | ||
| 32 | + *Compatible with tensorflow Zeta operator. | ||
| 33 | + */ | ||
| 34 | +REG_OP(Zeta) | ||
| 35 | + .INPUT(x, TensorType({DT_DOUBLE, DT_FLOAT})) | ||
| 36 | + .INPUT(q, TensorType({DT_DOUBLE, DT_FLOAT})) | ||
| 37 | + .OUTPUT(z, TensorType({DT_DOUBLE, DT_FLOAT})) | ||
| 38 | + .OP_END_FACTORY_REG(Zeta) | ||
| 39 | +} // namespace ge | ||
| 40 | + | ||
| 41 | + | ||
| @@ -0,0 +1,133 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | +namespace { | ||
| 23 | +const uint32_t kInputNum = 2; | ||
| 24 | +const uint32_t kOutputNum = 1; | ||
| 25 | +const char* const kZeta = "Zeta"; | ||
| 26 | +const int64_t kZetaParallelNum = 64 * 1024; | ||
| 27 | +} // namespace | ||
| 28 | + | ||
| 29 | +namespace aicpu { | ||
| 30 | +namespace detail { | ||
| 31 | +template <typename T> | ||
| 32 | +inline T ScalarZeta(T a, T b) | ||
| 33 | +{ | ||
| 34 | + return Eigen::numext::zeta(a, b); | ||
| 35 | +} | ||
| 36 | + | ||
| 37 | +inline uint32_t ParallelForZeta(const CpuKernelContext& ctx, int64_t total, int64_t per_unit_size, | ||
| 38 | + const std::function<void(int64_t, int64_t)>& work) | ||
| 39 | +{ | ||
| 40 | + if (total > kZetaParallelNum) { | ||
| 41 | + return CpuKernelUtils::ParallelFor(ctx, total, per_unit_size, work); | ||
| 42 | + } | ||
| 43 | + work(0, total); | ||
| 44 | + return KERNEL_STATUS_OK; | ||
| 45 | +} | ||
| 46 | + | ||
| 47 | +template <typename T> | ||
| 48 | +inline uint32_t ComputeZetaKernel(const CpuKernelContext& ctx) | ||
| 49 | +{ | ||
| 50 | + T* input0 = static_cast<T*>(ctx.Input(0)->GetData()); | ||
| 51 | + T* input1 = static_cast<T*>(ctx.Input(1)->GetData()); | ||
| 52 | + T* output = static_cast<T*>(ctx.Output(0)->GetData()); | ||
| 53 | + int64_t total = ctx.Input(0)->NumElements(); | ||
| 54 | + if (total == 0) { | ||
| 55 | + return KERNEL_STATUS_OK; | ||
| 56 | + } | ||
| 57 | + uint32_t cores = CpuKernelUtils::GetCPUNum(ctx); | ||
| 58 | + int64_t per_unit_size = total / std::min(std::max(1L, static_cast<long>(cores) - 2L), total); | ||
🟠 High Priority
建议：在计算 🟠 High Priority 文件 建议：在计算 | |||
| 59 | + return ParallelForZeta(ctx, total, per_unit_size, [&](int64_t begin, int64_t end) { | ||
| 60 | + (void)std::transform(input0 + begin, input0 + end, input1 + begin, output + begin, ScalarZeta<T>); | ||
| 61 | + }); | ||
| 62 | +} | ||
| 63 | + | ||
| 64 | +template <typename T> | ||
| 65 | +inline uint32_t ComputeZeta(const CpuKernelContext& ctx) | ||
| 66 | +{ | ||
| 67 | + uint32_t result = ComputeZetaKernel<T>(ctx); | ||
| 68 | + if (result != KERNEL_STATUS_OK) { | ||
| 69 | + KERNEL_LOG_ERROR("Zeta compute failed."); | ||
| 70 | + } | ||
| 71 | + return result; | ||
| 72 | +} | ||
| 73 | + | ||
| 74 | +inline uint32_t ExtraCheckZeta(const CpuKernelContext& ctx) | ||
| 75 | +{ | ||
| 76 | + if (ctx.Input(0)->GetData() == nullptr) { | ||
| 77 | + KERNEL_LOG_ERROR("Get input data failed."); | ||
| 78 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 79 | + } | ||
| 80 | + if (ctx.Output(0)->GetData() == nullptr) { | ||
| 81 | + KERNEL_LOG_ERROR("Get output data failed."); | ||
| 82 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 83 | + } | ||
| 84 | + if (ctx.Input(0)->GetDataType() != ctx.Input(1)->GetDataType()) { | ||
| 85 | + KERNEL_LOG_ERROR("The data type of the first input [%s] need be the same as the second input [%s].", | ||
| 86 | + DTypeStr(ctx.Input(0)->GetDataType()).c_str(), DTypeStr(ctx.Input(1)->GetDataType()).c_str()); | ||
| 87 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 88 | + } | ||
| 89 | + if (ctx.Input(0)->GetDataType() != ctx.Output(0)->GetDataType()) { | ||
| 90 | + KERNEL_LOG_ERROR("The data type of the input [%s] need be the same as the output [%s].", | ||
| 91 | + DTypeStr(ctx.Input(0)->GetDataType()).c_str(), DTypeStr(ctx.Output(0)->GetDataType()).c_str()); | ||
| 92 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 93 | + } | ||
| 94 | + if (ctx.Input(0)->GetDataSize() != ctx.Input(1)->GetDataSize()) { | ||
| 95 | + KERNEL_LOG_ERROR("The data size of the first input [%lu] need be the same as the second input [%lu].", | ||
| 96 | + ctx.Input(0)->GetDataSize(), ctx.Input(1)->GetDataSize()); | ||
| 97 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 98 | + } | ||
| 99 | + if (ctx.Input(0)->GetDataSize() != ctx.Output(0)->GetDataSize()) { | ||
| 100 | + KERNEL_LOG_ERROR("The data size of the input [%lu] need be the same as the output [%lu].", | ||
| 101 | + ctx.Input(0)->GetDataSize(), ctx.Output(0)->GetDataSize()); | ||
| 102 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 103 | + } | ||
| 104 | + return KERNEL_STATUS_OK; | ||
| 105 | +} | ||
| 106 | + | ||
| 107 | +inline uint32_t CheckZeta(CpuKernelContext& ctx) | ||
| 108 | +{ | ||
| 109 | + return NormalCheck(ctx, kInputNum, kOutputNum) ? KERNEL_STATUS_PARAM_INVALID : ExtraCheckZeta(ctx); | ||
| 110 | +} | ||
| 111 | + | ||
| 112 | +inline uint32_t ComputeZeta(const CpuKernelContext& ctx) | ||
| 113 | +{ | ||
| 114 | + DataType input_type = ctx.Input(0)->GetDataType(); | ||
| 115 | + switch (input_type) { | ||
| 116 | + case DT_FLOAT: | ||
| 117 | + return ComputeZeta<float>(ctx); | ||
| 118 | + case DT_DOUBLE: | ||
| 119 | + return ComputeZeta<double>(ctx); | ||
| 120 | + default: | ||
| 121 | + KERNEL_LOG_ERROR("Unsupported input data type [%s].", DTypeStr(input_type).c_str()); | ||
| 122 | + return KERNEL_STATUS_PARAM_INVALID; | ||
| 123 | + } | ||
| 124 | +} | ||
| 125 | +} // namespace detail | ||
| 126 | + | ||
| 127 | +uint32_t ZetaCpuKernel::Compute(CpuKernelContext& ctx) | ||
| 128 | +{ | ||
| 129 | + return detail::CheckZeta(ctx) ? static_cast<uint32_t>(KERNEL_STATUS_PARAM_INVALID) : detail::ComputeZeta(ctx); | ||
| 130 | +} | ||
| 131 | + | ||
| 132 | +REGISTER_CPU_KERNEL(kZeta, ZetaCpuKernel); | ||
| 133 | +} // namespace aicpu | ||
| @@ -0,0 +1,26 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +namespace aicpu { | ||
| 17 | +class ZetaCpuKernel : public CpuKernel { | ||
| 18 | +public: | ||
| 19 | + ZetaCpuKernel() = default; | ||
| 20 | + ~ZetaCpuKernel() override = default; | ||
| 21 | + | ||
| 22 | +protected: | ||
| 23 | + uint32_t Compute(CpuKernelContext& ctx) override; | ||
| 24 | +}; | ||
| 25 | +} // namespace aicpu | ||
| 26 | + | ||
| @@ -0,0 +1,29 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +namespace ops { | ||
| 15 | +class Zeta : public OpDef { | ||
| 16 | +public: | ||
| 17 | + explicit Zeta(const char* name) : OpDef(name) | ||
| 18 | + { | ||
| 19 | + this->Input("x").DataType({ge::DT_DOUBLE, ge::DT_FLOAT}); | ||
| 20 | + this->Input("q").DataType({ge::DT_DOUBLE, ge::DT_FLOAT}); | ||
| 21 | + this->Output("z").DataType({ge::DT_DOUBLE, ge::DT_FLOAT}); | ||
| 22 | + ApplyMathAicpuDefaultCfg(*this); | ||
| 23 | + this->AICPU().ExtendCfgInfo(OP_INFO_FORMAT_AGNOSTIC.c_str(), TRUE_FORMAT_AGNOSTIC.c_str()); | ||
| 24 | + this->AICPU().ExtendCfgInfo(OP_INFO_OPS_FLAG.c_str(), OPEN_OPS_FLAG.c_str()); | ||
| 25 | + } | ||
| 26 | +}; | ||
| 27 | + | ||
| 28 | +OP_ADD(Zeta); | ||
| 29 | +} // namespace ops | ||
| @@ -0,0 +1,142 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +using namespace std; | ||
| 26 | +using namespace aicpu; | ||
| 27 | + | ||
| 28 | +class TEST_ZETA_UT : public testing::Test {}; | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + auto node_def = CpuKernelUtils::CreateNodeDef(); \ | ||
| 32 | + NodeDefBuilder node(node_def.get(), "Zeta", "Zeta"); \ | ||
| 33 | + node.Input({"x", data_types[0], shapes[0], datas[0]}) \ | ||
| 34 | + .Input({"q", data_types[1], shapes[1], datas[1]}) \ | ||
| 35 | + .Output({"z", data_types[2], shapes[2], datas[2]}); | ||
| 36 | + | ||
| 37 | +TEST_F(TEST_ZETA_UT, TestZeta_DT_FLOAT) | ||
| 38 | +{ | ||
| 39 | + vector<DataType> data_types = {DT_FLOAT, DT_FLOAT, DT_FLOAT}; | ||
| 40 | + vector<vector<int64_t>> shapes = {{2, 3}, {2, 3}, {2, 3}}; | ||
| 41 | + float input_x[6] = {2.0f, 3.0f, 4.0f, 2.0f, 3.0f, 4.0f}; | ||
| 42 | + float input_q[6] = {1.0f, 1.0f, 1.0f, 2.0f, 2.0f, 2.0f}; | ||
| 43 | + float output_z[6] = {0.0f}; | ||
| 44 | + vector<void*> datas = {(void*)input_x, (void*)input_q, (void*)output_z}; | ||
| 45 | + CREATE_NODEDEF_ZETA(shapes, data_types, datas); | ||
| 46 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 47 | + // zeta(2,1) = pi^2/6 ≈ 1.6449, zeta(3,1) ≈ 1.2021, zeta(4,1) ≈ 1.0823 | ||
| 48 | + // zeta(2,2) = pi^2/6 - 1 ≈ 0.6449, zeta(3,2) ≈ 0.2021, zeta(4,2) ≈ 0.0823 | ||
| 49 | + EXPECT_NEAR(output_z[0], Eigen::numext::zeta(2.0f, 1.0f), 1e-4f); | ||
| 50 | + EXPECT_NEAR(output_z[1], Eigen::numext::zeta(3.0f, 1.0f), 1e-4f); | ||
| 51 | + EXPECT_NEAR(output_z[2], Eigen::numext::zeta(4.0f, 1.0f), 1e-4f); | ||
| 52 | + EXPECT_NEAR(output_z[3], Eigen::numext::zeta(2.0f, 2.0f), 1e-4f); | ||
| 53 | + EXPECT_NEAR(output_z[4], Eigen::numext::zeta(3.0f, 2.0f), 1e-4f); | ||
| 54 | + EXPECT_NEAR(output_z[5], Eigen::numext::zeta(4.0f, 2.0f), 1e-4f); | ||
| 55 | +} | ||
| 56 | + | ||
| 57 | +TEST_F(TEST_ZETA_UT, TestZeta_DT_DOUBLE) | ||
| 58 | +{ | ||
| 59 | + vector<DataType> data_types = {DT_DOUBLE, DT_DOUBLE, DT_DOUBLE}; | ||
| 60 | + vector<vector<int64_t>> shapes = {{2, 3}, {2, 3}, {2, 3}}; | ||
| 61 | + double input_x[6] = {2.0, 3.0, 4.0, 2.0, 3.0, 4.0}; | ||
| 62 | + double input_q[6] = {1.0, 1.0, 1.0, 2.0, 2.0, 2.0}; | ||
| 63 | + double output_z[6] = {0.0}; | ||
| 64 | + vector<void*> datas = {(void*)input_x, (void*)input_q, (void*)output_z}; | ||
| 65 | + CREATE_NODEDEF_ZETA(shapes, data_types, datas); | ||
| 66 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK); | ||
| 67 | + EXPECT_NEAR(output_z[0], Eigen::numext::zeta(2.0, 1.0), 1e-10); | ||
| 68 | + EXPECT_NEAR(output_z[1], Eigen::numext::zeta(3.0, 1.0), 1e-10); | ||
| 69 | + EXPECT_NEAR(output_z[2], Eigen::numext::zeta(4.0, 1.0), 1e-10); | ||
| 70 | + EXPECT_NEAR(output_z[3], Eigen::numext::zeta(2.0, 2.0), 1e-10); | ||
| 71 | + EXPECT_NEAR(output_z[4], Eigen::numext::zeta(3.0, 2.0), 1e-10); | ||
| 72 | + EXPECT_NEAR(output_z[5], Eigen::numext::zeta(4.0, 2.0), 1e-10); | ||
| 73 | +} | ||
| 74 | + | ||
| 75 | +TEST_F(TEST_ZETA_UT, TestZeta_INPUT_NULL_EXCEPTION) | ||
| 76 | +{ | ||
| 77 | + vector<DataType> data_types = {DT_FLOAT, DT_FLOAT, DT_FLOAT}; | ||
| 78 | + vector<vector<int64_t>> shapes = {{2, 3}, {2, 3}, {2, 3}}; | ||
| 79 | + float* null_ptr = nullptr; | ||
| 80 | + float input_q[6] = {1.0f}; | ||
| 81 | + float output_z[6] = {0.0f}; | ||
| 82 | + vector<void*> datas = {(void*)null_ptr, (void*)input_q, (void*)output_z}; | ||
| 83 | + CREATE_NODEDEF_ZETA(shapes, data_types, datas); | ||
| 84 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 85 | +} | ||
| 86 | + | ||
| 87 | +TEST_F(TEST_ZETA_UT, TestZeta_OUTPUT_NULL_EXCEPTION) | ||
| 88 | +{ | ||
| 89 | + vector<DataType> data_types = {DT_FLOAT, DT_FLOAT, DT_FLOAT}; | ||
| 90 | + vector<vector<int64_t>> shapes = {{0}, {0}, {0}}; | ||
| 91 | + float input_x[1] = {2.0f}; | ||
| 92 | + float input_q[1] = {1.0f}; | ||
| 93 | + float* null_ptr = nullptr; | ||
| 94 | + vector<void*> datas = {(void*)input_x, (void*)input_q, (void*)null_ptr}; | ||
| 95 | + CREATE_NODEDEF_ZETA(shapes, data_types, datas); | ||
| 96 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 97 | +} | ||
| 98 | + | ||
| 99 | +TEST_F(TEST_ZETA_UT, TestZeta_SHAPE_MISMATCH_EXCEPTION) | ||
| 100 | +{ | ||
| 101 | + vector<DataType> data_types = {DT_FLOAT, DT_FLOAT, DT_FLOAT}; | ||
| 102 | + vector<vector<int64_t>> shapes = {{2, 6}, {2, 8}, {2, 6}}; | ||
| 103 | + float input_x[12] = {2.0f}; | ||
| 104 | + float input_q[16] = {1.0f}; | ||
| 105 | + float output_z[12] = {0.0f}; | ||
| 106 | + vector<void*> datas = {(void*)input_x, (void*)input_q, (void*)output_z}; | ||
| 107 | + CREATE_NODEDEF_ZETA(shapes, data_types, datas); | ||
| 108 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 109 | +} | ||
| 110 | + | ||
| 111 | +TEST_F(TEST_ZETA_UT, TestZeta_DTYPE_MISMATCH_EXCEPTION) | ||
| 112 | +{ | ||
| 113 | + vector<DataType> data_types = {DT_FLOAT, DT_DOUBLE, DT_FLOAT}; | ||
| 114 | + vector<vector<int64_t>> shapes = {{2, 3}, {2, 3}, {2, 3}}; | ||
| 115 | + float input_x[6] = {2.0f}; | ||
| 116 | + double input_q[6] = {1.0}; | ||
| 117 | + float output_z[6] = {0.0f}; | ||
| 118 | + vector<void*> datas = {(void*)input_x, (void*)input_q, (void*)output_z}; | ||
| 119 | + CREATE_NODEDEF_ZETA(shapes, data_types, datas); | ||
| 120 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 121 | +} | ||
| 122 | + | ||
| 123 | +TEST_F(TEST_ZETA_UT, TestZeta_UNSUPPORTED_TYPE_EXCEPTION) | ||
| 124 | +{ | ||
| 125 | + vector<DataType> data_types = {DT_INT32, DT_INT32, DT_INT32}; | ||
| 126 | + vector<vector<int64_t>> shapes = {{2, 3}, {2, 3}, {2, 3}}; | ||
| 127 | + int32_t input_x[6] = {2}; | ||
| 128 | + int32_t input_q[6] = {1}; | ||
| 129 | + int32_t output_z[6] = {0}; | ||
| 130 | + vector<void*> datas = {(void*)input_x, (void*)input_q, (void*)output_z}; | ||
| 131 | + CREATE_NODEDEF_ZETA(shapes, data_types, datas); | ||
| 132 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 133 | +} | ||
| 134 | + | ||
| 135 | +TEST_F(TEST_ZETA_UT, TestZeta_NO_OUTPUT_EXCEPTION) | ||
| 136 | +{ | ||
| 137 | + auto node_def = CpuKernelUtils::CreateNodeDef(); | ||
| 138 | + float input_x[6] = {2.0f}; | ||
| 139 | + NodeDefBuilder node(node_def.get(), "Zeta", "Zeta"); | ||
| 140 | + node.Input({"x", DT_FLOAT, {2, 3}, input_x}); | ||
| 141 | + RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID); | ||
| 142 | +} | ||
🟡 Medium Priority
tile_with_axis_proto.h第45-46行的REG_OP中 Input/Output 的TensorType列表包含DT_BF16，但tile_with_axis_aicpu_def.cpp第19-22行的 DataType 列表中未包含DT_BF16，且tile_with_axis_aicpu.cpp第168-192行的Compute函数 switch-case 中也没有DT_BF16分支。当用户通过 IR 图传入 BF16 类型数据时，会在 kernel 执行阶段落入 default 分支返回PARAM_INVALID，而不是在图编译阶段就得到明确的类型不支持错误。建议：要么从 proto 中移除
DT_BF16以保持与 def/kernel 一致，要么在 def 和 kernel 中增加 BF16 支持。