已合并
[feat]add aicpu op: sqrt #2539
FengHaozhan创建于 4月28日
[feat]add aicpu op: sqrt #2539
已合并
FengHaozhan创建于 4月28日
6 个文件变更+643-0
@@ -68,3 +68,4 @@ $$
68| 调用方式 | 调用样例 | 说明 |68| 调用方式 | 调用样例 | 说明 |
69|--------------|------------------------------------------------------------------------|--------------------------------------------------------------|69|--------------|------------------------------------------------------------------------|--------------------------------------------------------------|
70| aclnn调用 | [test_aclnn_sqrt](./examples/test_aclnn_sqrt.cpp) | 通过[aclnnSqrt](docs/aclnnSqrt&aclnnInplaceSqrt.md)接口方式调用Sqrt算子。 |70| aclnn调用 | [test_aclnn_sqrt](./examples/test_aclnn_sqrt.cpp) | 通过[aclnnSqrt](docs/aclnnSqrt&aclnnInplaceSqrt.md)接口方式调用Sqrt算子。 |
71+| 图模式调用 | [test_geir_sqrt](./examples/test_geir_sqrt.cpp) | 通过[算子IR](./op_graph/sqrt_proto.h)构图方式调用Sqrt算子。 |
@@ -0,0 +1,332 @@
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+#include <iostream>
12+#include <fstream>
13+#include <string.h>
14+#include <stdint.h>
15+#include <vector>
16+#include <string>
17+#include <map>
18+#include "assert.h"
19+#include "graph.h"
20+#include "types.h"
21+#include "tensor.h"
22+#include "ge_error_codes.h"
23+#include "ge_api_types.h"
24+#include "ge_api.h"
25+#include "array_ops.h"
26+#include "ge_ir_build.h"
27+ 
28+#include "elewise_calculation_ops.h"
29+ 
30+#define FAILED -1
31+#define SUCCESS 0
32+ 
33+using namespace ge;
34+using std::map;
35+using std::string;
36+using std::vector;
37+ 
38+#define ADD_INPUT(inputIndex, inputName, inputDtype, inputShape) \
39+ vector<int64_t> placeholder##inputIndex##_shape = inputShape; \
40+ auto placeholder##inputIndex = \
41+ op::Data(std::string("placeholder") + std::to_string(inputIndex)).set_attr_index(0); \
42+ TensorDesc placeholder##inputIndex##_desc = \
43+ TensorDesc(ge::Shape(placeholder##inputIndex##_shape), FORMAT_ND, inputDtype); \
44+ placeholder##inputIndex##_desc.SetPlacement(ge::kPlacementHost); \
45+ Tensor tensor_placeholder##inputIndex; \
46+ ret = GenOnesData( \
47+ placeholder##inputIndex##_shape, tensor_placeholder##inputIndex, placeholder##inputIndex##_desc, inputDtype, \
48+ 4); \
49+ if (ret != SUCCESS) { \
50+ printf("%s - ERROR - [XIR]: Generate input data failed\n", GetTime().c_str()); \
51+ return FAILED; \
52+ } \
53+ placeholder##inputIndex.update_input_desc_x(placeholder##inputIndex##_desc); \
54+ input.push_back(tensor_placeholder##inputIndex); \
55+ graph.AddOp(placeholder##inputIndex); \
56+ sqrt1.set_input_##inputName(placeholder##inputIndex); \
57+ inputs.push_back(placeholder##inputIndex)
58+ 
59+#define ADD_OUTPUT(outputIndex, outputName, outputDtype, outputShape) \
60+ TensorDesc outputName##outputIndex##_desc = TensorDesc(ge::Shape(outputShape), FORMAT_ND, outputDtype); \
61+ sqrt1.update_output_desc_##outputName(outputName##outputIndex##_desc)
62+ 
63+#define LOG_PRINT(message, ...) \
64+ do { \
65+ printf(message, ##__VA_ARGS__); \
66+ } while (0)
67+ 
68+string GetTime()
69+{
70+ time_t timep;
71+ time(&timep);
72+ char tmp[64];
73+ struct tm tm_info;
74+ strftime(tmp, sizeof(tmp), "%Y-%m-%d %H:%M:%S,000", localtime_r(&timep, &tm_info));
75+ return tmp;
76+}
77+ 
78+uint32_t GetDataTypeSize(DataType dt)
79+{
80+ uint32_t dilation = 1;
81+ uint32_t oneByte = 1;
82+ uint32_t twoByte = 2;
83+ uint32_t fourByte = 4;
84+ uint32_t eightByte = 8;
85+ 
86+ if (dt == ge::DT_FLOAT) {
87+ dilation = fourByte;
88+ } else if (dt == ge::DT_FLOAT16) {
89+ dilation = twoByte;
90+ } else if (dt == ge::DT_BF16) {
91+ dilation = twoByte;
92+ } else if (dt == ge::DT_INT16) {
93+ dilation = twoByte;
94+ } else if (dt == ge::DT_UINT16) {
95+ dilation = twoByte;
96+ } else if (dt == ge::DT_INT32) {
97+ dilation = fourByte;
98+ } else if (dt == ge::DT_UINT32) {
99+ dilation = fourByte;
100+ } else if (dt == ge::DT_INT64) {
101+ dilation = eightByte;
102+ } else if (dt == ge::DT_DOUBLE) {
103+ dilation = eightByte;
104+ } else if (dt == ge::DT_COMPLEX64) {
105+ dilation = eightByte;
106+ } else if (dt == ge::DT_COMPLEX128) {
107+ dilation = eightByte * twoByte;
108+ } else if (dt == ge::DT_UINT64) {
109+ dilation = eightByte;
110+ } else if (dt == ge::DT_INT8) {
111+ dilation = oneByte;
112+ }
113+ return dilation;
114+}
115+ 
116+int32_t GenOnesData(
117+ vector<int64_t> shapes, Tensor& input_tensor, TensorDesc& input_tensor_desc, DataType data_type, int value)
118+{
119+ input_tensor_desc.SetRealDimCnt(shapes.size());
120+ size_t size = 1;
121+ for (uint32_t i = 0; i < shapes.size(); i++) {
122+ size *= shapes[i];
123+ }
124+ uint32_t data_len = size * GetDataTypeSize(data_type);
125+ uint8_t* pData = new (std::nothrow) uint8_t[data_len];
126+ if (pData == nullptr) {
127+ return FAILED;
128+ }
129+ for (uint32_t i = 0; i < data_len; ++i) {
130+ pData[i] = static_cast<uint8_t>(value);
131+ }
132+ input_tensor = Tensor(input_tensor_desc, reinterpret_cast<uint8_t*>(pData), data_len);
133+ return SUCCESS;
134+}
135+ 
136+int32_t WriteDataToFile(string bin_file, uint64_t data_size, uint8_t* inputData)
137+{
138+ FILE* fp = fopen(bin_file.c_str(), "wb");
139+ if (fp == nullptr) {
140+ printf("Failed to open file: %s\n", bin_file.c_str());
141+ return FAILED;
142+ }
143+ fwrite(inputData, sizeof(uint8_t), data_size, fp);
144+ fclose(fp);
145+ return SUCCESS;
146+}
147+ 
148+int CreateOppInGraph(
149+ DataType inDtype, std::vector<ge::Tensor>& input, std::vector<Operator>& inputs, std::vector<Operator>& outputs,
150+ Graph& graph)
151+{
152+ Status ret = SUCCESS;
153+ // 自定义代码:添加单算子定义到图中
154+ auto sqrt1 = op::Sqrt("sqrt");
155+ vector<vector<int64_t>> shapes = {{4, 4}, {4, 4}};
156+ 
157+ ADD_INPUT(1, x, inDtype, shapes[0]);
158+ ADD_OUTPUT(1, y, inDtype, shapes[1]);
159+ 
160+ outputs.push_back(sqrt1);
161+ // 添加完毕
162+ return SUCCESS;
163+}
164+ 
165+bool InitEnv()
166+{
167+ printf("%s - INFO - [XIR]: Start to initialize ge using ge global options\n", GetTime().c_str());
168+ std::map<AscendString, AscendString> global_options = {{"ge.exec.deviceId", "0"}, {"ge.graphRunMode", "1"}};
169+ Status ret = ge::GEInitialize(global_options);
170+ if (ret != SUCCESS) {
171+ printf("%s - INFO - [XIR]: Initialize ge using ge global options failed\n", GetTime().c_str());
172+ return false;
173+ }
174+ printf("%s - INFO - [XIR]: Initialize ge using ge global options success\n", GetTime().c_str());
175+ return true;
176+}
177+ 
178+bool CreateAndConfigGraph(Graph& graph, std::vector<ge::Tensor>& input)
179+{
180+ printf("%s - INFO - [XIR]: Start to CreateAndConfigGraph\n", GetTime().c_str());
181+ std::vector<Operator> inputs{};
182+ std::vector<Operator> outputs{};
183+ 
184+ // Use DT_COMPLEX64 to dispatch to AICPU kernel
185+ // (TBE Sqrt only supports real types on ascend910b)
186+ DataType inDtype = DT_COMPLEX64;
187+ std::cout << inDtype << std::endl;
188+ 
189+ Status ret = CreateOppInGraph(inDtype, input, inputs, outputs, graph);
190+ if (ret != SUCCESS) {
191+ printf("%s - ERROR - [XIR]: Create ir session using build options failed\n", GetTime().c_str());
192+ return false;
193+ }
194+ 
195+ if (!inputs.empty() && !outputs.empty()) {
196+ graph.SetInputs(inputs).SetOutputs(outputs);
197+ }
198+ return true;
199+}
200+ 
201+bool AddGraphToSession(ge::Session* session, Graph& graph, uint32_t graph_id)
202+{
203+ printf("%s - INFO - [XIR]: Create ir session using build options success\n", GetTime().c_str());
204+ printf("%s - INFO - [XIR]: Start to add compute graph to ir session\n", GetTime().c_str());
205+ 
206+ std::map<AscendString, AscendString> graph_options = {};
207+ Status ret = session->AddGraph(graph_id, graph, graph_options);
208+ if (ret != SUCCESS) {
209+ printf("%s - INFO - [XIR]: Add graph failed\n", GetTime().c_str());
210+ delete session;
211+ ge::GEFinalize();
212+ return false;
213+ }
214+ printf("%s - INFO - [XIR]: Session add ir compute graph to ir session success\n", GetTime().c_str());
215+ return true;
216+}
217+ 
218+bool DumpAndRunGraph(
219+ ge::Session* session, Graph& graph, std::vector<ge::Tensor>& input, std::vector<ge::Tensor>& output,
220+ uint32_t graph_id)
221+{
222+ printf("%s - INFO - [XIR]: dump graph to txt\n", GetTime().c_str());
223+ std::string file_path = "./dump";
224+ aclgrphDumpGraph(graph, file_path.c_str(), file_path.length());
225+ 
226+ printf("%s - INFO - [XIR]: Start to run ir compute graph\n", GetTime().c_str());
227+ 
228+ Status ret = session->RunGraph(graph_id, input, output);
229+ if (ret != SUCCESS) {
230+ printf("%s - INFO - [XIR]: Run graph failed\n", GetTime().c_str());
231+ delete session;
232+ ge::GEFinalize();
233+ return false;
234+ }
235+ printf("%s - INFO - [XIR]: Session run ir compute graph success\n", GetTime().c_str());
236+ return true;
237+}
238+ 
239+void ProcessInputData(std::vector<ge::Tensor>& input)
240+{
241+ int input_num = input.size();
242+ for (int i = 0; i < input_num; i++) {
243+ std::cout << "input " << i << " dtype : " << input[i].GetTensorDesc().GetDataType() << std::endl;
244+ string input_file = "./tc_ge_irrun_sqrt_input_" + std::to_string(i) + ".bin";
245+ uint8_t* input_data_i = input[i].GetData();
246+ int64_t input_shape = input[i].GetTensorDesc().GetShape().GetShapeSize();
247+ std::cout << "this is " << i << "th input, input shape size =" << input_shape << std::endl;
248+ uint32_t data_size = input_shape * GetDataTypeSize(input[i].GetTensorDesc().GetDataType());
249+ WriteDataToFile((const char*)input_file.c_str(), data_size, input_data_i);
250+ }
251+}
252+ 
253+void ProcessOutputData(std::vector<ge::Tensor>& output)
254+{
255+ int output_num = output.size();
256+ for (int i = 0; i < output_num; i++) {
257+ std::cout << "output " << i << " dtype : " << output[i].GetTensorDesc().GetDataType() << std::endl;
258+ string output_file = "./tc_ge_irrun_sqrt_output_" + std::to_string(i) + ".bin";
259+ uint8_t* output_data_i = output[i].GetData();
260+ int64_t output_shape = output[i].GetTensorDesc().GetShape().GetShapeSize();
261+ std::cout << "this is " << i << "th output, output shape size =" << output_shape << std::endl;
262+ uint32_t data_size = output_shape * GetDataTypeSize(output[i].GetTensorDesc().GetDataType());
263+ WriteDataToFile((const char*)output_file.c_str(), data_size, output_data_i);
264+ // complex64: each element = (float real, float imag)
265+ float* result = reinterpret_cast<float*>(output_data_i);
266+ for (int64_t j = 0; j < output_shape; j++) {
267+ LOG_PRINT("result[%ld] = (%f, %f)\n", j, result[2 * j], result[2 * j + 1]);
268+ }
269+ }
270+}
271+ 
272+int FinalizeRes()
273+{
274+ ge::AscendString error_msg = ge::GEGetErrorMsgV2();
275+ std::string error_str(error_msg.GetString());
276+ std::cout << "Error message: " << error_str << std::endl;
277+ ge::AscendString warning_msg = ge::GEGetWarningMsgV2();
278+ std::string warning_str(warning_msg.GetString());
279+ std::cout << "Warning message: " << warning_str << std::endl;
280+ printf("%s - INFO - [XIR]: Precision is ok\n", GetTime().c_str());
281+ printf("%s - INFO - [XIR]: Start to finalize ir graph session\n", GetTime().c_str());
282+ Status ret = ge::GEFinalize();
283+ if (ret != SUCCESS) {
284+ printf("%s - INFO - [XIR]: Finalize ir graph session failed\n", GetTime().c_str());
285+ return FAILED;
286+ }
287+ printf("%s - INFO - [XIR]: Finalize ir graph session success\n", GetTime().c_str());
288+ return SUCCESS;
289+}
290+ 
291+int main(int argc, char* argv[])
292+{
293+ // 初始化环境
294+ if (!InitEnv()) {
295+ return FAILED;
296+ }
297+ 
298+ // 创建计算图
299+ const char* graph_name = "tc_ge_irrun_sqrt_test";
300+ Graph graph(graph_name);
301+ std::vector<ge::Tensor> input;
302+ if (!CreateAndConfigGraph(graph, input)) {
303+ return FAILED;
304+ }
305+ 
306+ // 创建会话并添加图
307+ std::map<AscendString, AscendString> build_options = {};
308+ printf("%s - INFO - [XIR]: Start to create ir session using build options\n", GetTime().c_str());
309+ ge::Session* session = new (std::nothrow) Session(build_options);
310+ if (session == nullptr) {
311+ printf("%s - ERROR - [XIR]: Create ir session using build options failed\n", GetTime().c_str());
312+ return FAILED;
313+ }
314+ 
315+ uint32_t graph_id = 0;
316+ if (!AddGraphToSession(session, graph, graph_id)) {
317+ return FAILED;
318+ }
319+ 
320+ // 执行图计算
321+ std::vector<ge::Tensor> output;
322+ if (!DumpAndRunGraph(session, graph, input, output, graph_id)) {
323+ return FAILED;
324+ }
325+ 
326+ // 处理输入输出数据
327+ ProcessInputData(input);
328+ ProcessOutputData(output);
329+ 
330+ delete session;
331+ return FinalizeRes();
332+}
@@ -0,0 +1,23 @@
1+{
2+ "Sqrt":{
3+ "opInfo":{
4+ "computeCost":"100",
5+ "engine":"DNN_VM_AICPU",
6+ "flagAsync":"False",
7+ "flagPartial":"False",
8+ "functionName":"RunCpuKernel",
9+ "kernelSo":"libmath_aicpu_kernels.so",
10+ "opKernelLib":"CUSTAICPUKernel",
11+ "userDefined":"True",
12+ "workspaceSize":"100"
13+ },
14+ "input0": {
15+ "name": "x",
16+ "type": "DT_FLOAT16,DT_FLOAT,DT_DOUBLE,DT_COMPLEX64,DT_COMPLEX128"
17+ },
18+ "output0": {
19+ "name": "y",
20+ "type": "DT_FLOAT16,DT_FLOAT,DT_DOUBLE,DT_COMPLEX64,DT_COMPLEX128"
21+ }
22+ }
23+}
@@ -0,0 +1,108 @@
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+#include "sqrt_aicpu.h"
11+ 
12+#include <complex>
13+ 
14+#include "cpu_kernel_utils.h"
15+#include "utils/kernel_util.h"
16+ 
17+using namespace std;
18+ 
19+namespace {
20+const char *const kSqrt = "Sqrt";
21+const uint32_t kInputNum = 1;
22+const uint32_t kOutputNum = 1;
23+constexpr int64_t kParallelDataNums = 8 * 1024;
24+} // namespace
25+ 
26+namespace aicpu {
27+template <typename T>
28+uint32_t SqrtCpuKernel::DoComputeReal(const CpuKernelContext &ctx) {
29+ auto *input = reinterpret_cast<T *>(ctx.Input(0)->GetData());
30+ auto *output = reinterpret_cast<T *>(ctx.Output(0)->GetData());
31+ int64_t data_num = ctx.Input(0)->NumElements();
32+ 
33+ if (data_num <= kParallelDataNums) {
34+ Eigen::TensorMap<Eigen::Tensor<T, 1>, Eigen::Unaligned> tensor_x(input, data_num);
35+ Eigen::TensorMap<Eigen::Tensor<T, 1>, Eigen::Unaligned> tensor_y(output, data_num);
36+ tensor_y = tensor_x.sqrt();
37+ } else {
38+ uint32_t min_core_num = 1;
39+ int64_t max_core_num = std::max(min_core_num, aicpu::CpuKernelUtils::GetCPUNum(ctx) - kResvCpuNum);
40+ max_core_num = max_core_num > data_num ? data_num : max_core_num;
41+ auto shard_sqrt = [&input, &output](int64_t begin, int64_t end) {
42+ int64_t length = end - begin;
43+ Eigen::TensorMap<Eigen::Tensor<T, 1>, Eigen::Unaligned> tensor_x(input + begin, length);
44+ Eigen::TensorMap<Eigen::Tensor<T, 1>, Eigen::Unaligned> tensor_y(output + begin, length);
45+ tensor_y = tensor_x.sqrt();
46+ };
47+ KERNEL_HANDLE_ERROR(CpuKernelUtils::ParallelFor(ctx, data_num, data_num / max_core_num, shard_sqrt),
48+ "Sqrt Compute failed.");
49+ }
50+ return KERNEL_STATUS_OK;
51+}
52+ 
53+template <typename T>
54+uint32_t SqrtCpuKernel::DoComputeComplex(const CpuKernelContext &ctx) {
55+ auto *input = reinterpret_cast<T *>(ctx.Input(0)->GetData());
56+ auto *output = reinterpret_cast<T *>(ctx.Output(0)->GetData());
57+ int64_t data_num = ctx.Input(0)->NumElements();
58+ 
59+ auto shard_sqrt = [&input, &output](int64_t begin, int64_t end) {
60+ for (int64_t i = begin; i < end; ++i) {
61+ output[i] = std::sqrt(input[i]);
62+ }
63+ };
64+ 
65+ if (data_num <= kParallelDataNums) {
66+ shard_sqrt(0, data_num);
67+ } else {
68+ uint32_t min_core_num = 1;
69+ int64_t max_core_num = std::max(min_core_num, aicpu::CpuKernelUtils::GetCPUNum(ctx) - kResvCpuNum);
70+ max_core_num = max_core_num > data_num ? data_num : max_core_num;
71+ KERNEL_HANDLE_ERROR(CpuKernelUtils::ParallelFor(ctx, data_num, data_num / max_core_num, shard_sqrt),
72+ "Sqrt Compute failed.");
73+ }
74+ return KERNEL_STATUS_OK;
75+}
76+ 
77+uint32_t SqrtCpuKernel::Compute(CpuKernelContext &ctx) {
78+ KERNEL_HANDLE_ERROR(NormalCheck(ctx, kInputNum, kOutputNum), "Check Sqrt params failed.");
79+ 
80+ DataType input_type = ctx.Input(0)->GetDataType();
81+ DataType output_type = ctx.Output(0)->GetDataType();
82+ KERNEL_CHECK_FALSE((input_type == output_type), KERNEL_STATUS_PARAM_INVALID,
83+ "The data type of input [%s] must be the same as output [%s].",
84+ DTypeStr(input_type).c_str(), DTypeStr(output_type).c_str());
85+ KERNEL_CHECK_FALSE((ctx.Input(0)->GetDataSize() == ctx.Output(0)->GetDataSize()), KERNEL_STATUS_PARAM_INVALID,
86+ "The data size of input [%lu] must be the same as output [%lu].",
87+ ctx.Input(0)->GetDataSize(), ctx.Output(0)->GetDataSize());
88+ 
89+ KERNEL_LOG_DEBUG("%s op input[x] data type is [%s].", kSqrt, DTypeStr(input_type).c_str());
90+ switch (input_type) {
91+ case DT_FLOAT16:
92+ return DoComputeReal<Eigen::half>(ctx);
93+ case DT_FLOAT:
94+ return DoComputeReal<float>(ctx);
95+ case DT_DOUBLE:
96+ return DoComputeReal<double>(ctx);
97+ case DT_COMPLEX64:
98+ return DoComputeComplex<complex<float>>(ctx);
99+ case DT_COMPLEX128:
100+ return DoComputeComplex<complex<double>>(ctx);
101+ default:
102+ KERNEL_LOG_ERROR("Sqrt invalid input type [%s].", DTypeStr(input_type).c_str());
103+ return KERNEL_STATUS_PARAM_INVALID;
104+ }
105+}
106+ 
107+REGISTER_CPU_KERNEL(kSqrt, SqrtCpuKernel);
108+} // 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+#ifndef AICPU_KERNELS_NORMALIZED_SQRT_H
12+#define AICPU_KERNELS_NORMALIZED_SQRT_H
13+ 
14+#include <complex>
15+#include "cpu_kernel.h"
16+#include "unsupported/Eigen/CXX11/Tensor"
17+ 
18+namespace aicpu {
19+ 
20+class SqrtCpuKernel : public CpuKernel {
21+public:
22+ SqrtCpuKernel() = default;
23+ ~SqrtCpuKernel() override = default;
24+ 
25+ uint32_t Compute(CpuKernelContext &ctx) override;
26+ 
27+private:
28+ template <typename T>
29+ uint32_t DoComputeReal(const CpuKernelContext &ctx);
30+ 
31+ template <typename T>
32+ uint32_t DoComputeComplex(const CpuKernelContext &ctx);
33+};
34+ 
35+} // namespace aicpu
36+#endif
@@ -0,0 +1,143 @@
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+#include "gtest/gtest.h"
12+#ifndef private
13+#define private public
14+#define protected public
15+#endif
16+#include "utils/aicpu_test_utils.h"
17+#include "cpu_kernel_utils.h"
18+#include "node_def_builder.h"
19+#undef private
20+#undef protected
21+#include "Eigen/Core"
22+ 
23+using namespace std;
24+using namespace aicpu;
25+ 
26+class TEST_SQRT_UT : public testing::Test {};
27+ 
28+#define CREATE_NODEDEF(shapes, data_types, datas) \
29+ auto node_def = CpuKernelUtils::CreateNodeDef(); \
30+ NodeDefBuilder(node_def.get(), "Sqrt", "Sqrt") \
31+ .Input({"x", data_types[0], shapes[0], datas[0]}) \
32+ .Output({"y", data_types[1], shapes[1], datas[1]})
33+ 
34+// ---- float32 basic test ----
35+TEST_F(TEST_SQRT_UT, TestSqrt_FLOAT) {
36+ vector<DataType> data_types = {DT_FLOAT, DT_FLOAT};
37+ vector<vector<int64_t>> shapes = {{5, 5}, {5, 5}};
38+ // seed=23457, uniform(1,100), shape=[5,5]
39+ float input[25] = {81.52432f, 43.482735f, 67.64165f, 23.798012f, 36.49648f,
40+ 23.168648f, 11.171817f, 62.294262f, 83.37125f, 50.089794f,
41+ 69.03212f, 67.75925f, 46.901024f, 54.83537f, 98.68752f,
42+ 71.01781f, 92.65018f, 92.62283f, 99.74749f, 98.54601f,
43+ 47.414738f, 15.765511f, 78.63351f, 46.978813f, 94.61364f};
44+ float output[25] = {0.0f};
45+ vector<void *> datas = {(void *)input, (void *)output};
46+ CREATE_NODEDEF(shapes, data_types, datas);
47+ RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK);
48+ float output_exp[25] = {9.029082f, 6.594144f, 8.224454f, 4.8783207f, 6.0412316f,
49+ 4.813382f, 3.3424268f, 7.8926716f, 9.130786f, 7.0774145f,
50+ 8.308557f, 8.231601f, 6.848432f, 7.405091f, 9.934159f,
51+ 8.427206f, 9.625496f, 9.624076f, 9.987367f, 9.927034f,
52+ 6.885836f, 3.9705806f, 8.867554f, 6.8541093f, 9.726954f};
53+ EXPECT_EQ(CompareResult<float>(output, output_exp, 25), true);
54+}
55+ 
56+// ---- float64 basic test (use perfect squares for exact results) ----
57+TEST_F(TEST_SQRT_UT, TestSqrt_DOUBLE) {
58+ vector<DataType> data_types = {DT_DOUBLE, DT_DOUBLE};
59+ vector<vector<int64_t>> shapes = {{5, 5}, {5, 5}};
60+ // use perfect squares so expected output is exact
61+ double input[25] = {1.0, 4.0, 9.0, 16.0, 25.0,
62+ 36.0, 49.0, 64.0, 81.0, 100.0,
63+ 121.0, 144.0, 169.0, 196.0, 225.0,
64+ 256.0, 289.0, 324.0, 361.0, 400.0,
65+ 441.0, 484.0, 529.0, 576.0, 625.0};
66+ double output[25] = {0.0};
67+ vector<void *> datas = {(void *)input, (void *)output};
68+ CREATE_NODEDEF(shapes, data_types, datas);
69+ RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK);
70+ double output_exp[25] = {1.0, 2.0, 3.0, 4.0, 5.0,
71+ 6.0, 7.0, 8.0, 9.0, 10.0,
72+ 11.0, 12.0, 13.0, 14.0, 15.0,
73+ 16.0, 17.0, 18.0, 19.0, 20.0,
74+ 21.0, 22.0, 23.0, 24.0, 25.0};
75+ EXPECT_EQ(CompareResult<double>(output, output_exp, 25), true);
76+}
77+ 
78+// ---- float16 basic test ----
79+TEST_F(TEST_SQRT_UT, TestSqrt_FLOAT16) {
80+ vector<DataType> data_types = {DT_FLOAT16, DT_FLOAT16};
81+ vector<vector<int64_t>> shapes = {{5, 5}, {5, 5}};
82+ // seed=3457, randint(1,100), shape=[5,5]
83+ Eigen::half input[25] = {
84+ Eigen::half(15.0f), Eigen::half(46.0f), Eigen::half(8.0f), Eigen::half(95.0f), Eigen::half(72.0f),
85+ Eigen::half(54.0f), Eigen::half(20.0f), Eigen::half(59.0f), Eigen::half(98.0f), Eigen::half(23.0f),
86+ Eigen::half(98.0f), Eigen::half(63.0f), Eigen::half(76.0f), Eigen::half(37.0f), Eigen::half(72.0f),
87+ Eigen::half(13.0f), Eigen::half(59.0f), Eigen::half(70.0f), Eigen::half(77.0f), Eigen::half(88.0f),
88+ Eigen::half(85.0f), Eigen::half(96.0f), Eigen::half(28.0f), Eigen::half(2.0f), Eigen::half(7.0f)};
89+ Eigen::half output[25] = {Eigen::half(0.0f)};
90+ vector<void *> datas = {(void *)input, (void *)output};
91+ CREATE_NODEDEF(shapes, data_types, datas);
92+ RUN_KERNEL(node_def, HOST, KERNEL_STATUS_OK);
93+ Eigen::half output_exp[25] = {
94+ Eigen::half(3.873f), Eigen::half(6.78f), Eigen::half(2.828f), Eigen::half(9.75f), Eigen::half(8.484f),
95+ Eigen::half(7.348f), Eigen::half(4.473f), Eigen::half(7.68f), Eigen::half(9.9f), Eigen::half(4.797f),
96+ Eigen::half(9.9f), Eigen::half(7.938f), Eigen::half(8.72f), Eigen::half(6.082f), Eigen::half(8.484f),
97+ Eigen::half(3.605f), Eigen::half(7.68f), Eigen::half(8.37f), Eigen::half(8.77f), Eigen::half(9.38f),
98+ Eigen::half(9.22f), Eigen::half(9.8f), Eigen::half(5.293f), Eigen::half(1.414f), Eigen::half(2.646f)};
99+ EXPECT_EQ(CompareResult<Eigen::half>(output, output_exp, 25), true);
100+}
101+ 
102+// ---- exception: mismatched data type ----
103+TEST_F(TEST_SQRT_UT, TestSqrt_InputDtypeException) {
104+ vector<DataType> data_types = {DT_DOUBLE, DT_FLOAT};
105+ vector<vector<int64_t>> shapes = {{2, 3}, {2, 3}};
106+ double input[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
107+ float output[6] = {0.0f};
108+ vector<void *> datas = {(void *)input, (void *)output};
109+ CREATE_NODEDEF(shapes, data_types, datas);
110+ RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID);
111+}
112+ 
113+// ---- exception: mismatched shape ----
114+TEST_F(TEST_SQRT_UT, TestSqrt_InputShapeException) {
115+ vector<DataType> data_types = {DT_FLOAT, DT_FLOAT};
116+ vector<vector<int64_t>> shapes = {{2, 3}, {2, 4}};
117+ float input[6] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
118+ float output[8] = {0.0f};
119+ vector<void *> datas = {(void *)input, (void *)output};
120+ CREATE_NODEDEF(shapes, data_types, datas);
121+ RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID);
122+}
123+ 
124+// ---- exception: unsupported type (INT32) ----
125+TEST_F(TEST_SQRT_UT, TestSqrt_UnsupportedType) {
126+ vector<DataType> data_types = {DT_INT32, DT_INT32};
127+ vector<vector<int64_t>> shapes = {{2, 3}, {2, 3}};
128+ int32_t input[6] = {1, 2, 3, 4, 5, 6};
129+ int32_t output[6] = {0};
130+ vector<void *> datas = {(void *)input, (void *)output};
131+ CREATE_NODEDEF(shapes, data_types, datas);
132+ RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID);
133+}
134+ 
135+// ---- null input ----
136+TEST_F(TEST_SQRT_UT, TestSqrt_NullInput) {
137+ vector<DataType> data_types = {DT_FLOAT, DT_FLOAT};
138+ vector<vector<int64_t>> shapes = {{2, 3}, {2, 3}};
139+ float output[6] = {0.0f};
140+ vector<void *> datas = {(void *)nullptr, (void *)output};
141+ CREATE_NODEDEF(shapes, data_types, datas);
142+ RUN_KERNEL(node_def, HOST, KERNEL_STATUS_PARAM_INVALID);
143+}