已开启
【代码侦探Challenge04-DivCustomTemplate】新增DivCustomTemplate工程化算子开发实现 #2154
【代码侦探Challenge04-DivCustomTemplate】新增DivCustomTemplate工程化算子开发实现 #2154
已开启
luoxiaoyan2024创建于 10 天前
7 个文件变更+485-0
@@ -0,0 +1,71 @@
1+#include "div_custom_template_tiling.h"
2+#include "register/op_def_registry.h"
3+ 
4+ 
5+namespace optiling {
6+static ge::graphStatus TilingFunc(gert::TilingContext* context)
7+{
8+ DivCustomTemplateTilingData tiling;
9+ const gert::StorageShape* x1_shape = context->GetInputShape(0);
10+ int64_t data_sz = 1;
11+ for (int i = 0; i < x1_shape->GetStorageShape().GetDimNum(); i++)
12+ data_sz *= x1_shape->GetStorageShape().GetDim(i);
13+ tiling.set_size(static_cast<uint32_t>(data_sz));
14+ context->SetBlockDim(8);
15+ tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
16+ context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
17+ 
18+ return ge::GRAPH_SUCCESS;
19+}
20+}
21+ 
22+ 
23+namespace ge {
24+static ge::graphStatus InferShape(gert::InferShapeContext* context)
25+{
26+ const gert::Shape* x1_shape = context->GetInputShape(0);
27+ gert::Shape* y_shape = context->GetOutputShape(0);
28+ *y_shape = *x1_shape;
29+ return GRAPH_SUCCESS;
30+}
31+static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
32+{
33+const auto inputDataType = context->GetInputDataType(0);
34+context->SetOutputDataType(0, inputDataType);
35+return ge::GRAPH_SUCCESS;
36+}
37+}
38+ 
39+ 
40+namespace ops {
41+class DivCustomTemplate : public OpDef {
42+public:
43+ explicit DivCustomTemplate(const char* name) : OpDef(name)
44+ {
45+ this->Input("x")
46+ .ParamType(REQUIRED)
47+ .DataType({ge::DT_FLOAT16})
48+ .Format({ge::FORMAT_ND})
49+ .UnknownShapeFormat({ge::FORMAT_ND});
50+ this->Input("y")
51+ .ParamType(REQUIRED)
52+ .DataType({ge::DT_FLOAT16})
53+ .Format({ge::FORMAT_ND})
54+ .UnknownShapeFormat({ge::FORMAT_ND});
55+ this->Output("z")
56+ .ParamType(REQUIRED)
57+ .DataType({ge::DT_FLOAT16})
58+ .Format({ge::FORMAT_ND})
59+ .UnknownShapeFormat({ge::FORMAT_ND});
60+ 
61+ this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
62+ 
63+ this->AICore()
64+ .SetTiling(optiling::TilingFunc);
65+ this->AICore().AddConfig("ascend910");
66+ 
67+ }
68+};
69+ 
70+OP_ADD(DivCustomTemplate);
71+}
@@ -0,0 +1,10 @@
1+ 
2+#include "register/tilingdata_base.h"
3+ 
4+namespace optiling {
5+BEGIN_TILING_DATA_DEF(DivCustomTemplateTilingData)
6+ TILING_DATA_FIELD_DEF(uint32_t, size);
7+END_TILING_DATA_DEF;
8+ 
9+REGISTER_TILING_DATA_CLASS(DivCustomTemplate, DivCustomTemplateTilingData)
10+}
@@ -0,0 +1,104 @@
1+#include "kernel_operator.h"
2+ 
3+using namespace AscendC;
4+ 
5+constexpr int32_t BUFFER_NUM = 2;
6+constexpr int32_t TILE_LENGTH = 2048;
7+ 
8+template <typename T>
9+class KernelDiv {
10+public:
11+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling)
12+ {
13+ GET_TILING_DATA(tilingData, tiling);
14+ totalLength = tilingData.size;
15+ 
16+ uint32_t blockNum = GetBlockNum();
17+ uint32_t blockId = GetBlockIdx();
18+ 
19+ // 每个block处理的元素数,最后一个block处理余数
20+ elementsPerBlock = totalLength / blockNum;
21+ uint32_t remainder = totalLength % blockNum;
22+ if (blockId == blockNum - 1 && remainder != 0) {
23+ elementsPerBlock += remainder;
24+ }
25+ offset = blockId * (totalLength / blockNum);
26+ 
27+ // 计算tile数量和最后一个tile的长度
28+ tileNum = (elementsPerBlock + TILE_LENGTH - 1) / TILE_LENGTH;
29+ if (elementsPerBlock == 0) {
30+ tileNum = 0;
31+ return;
32+ }
33+ lastTileLength = elementsPerBlock % TILE_LENGTH;
34+ if (lastTileLength == 0) {
35+ lastTileLength = TILE_LENGTH;
36+ }
37+ 
38+ xGm.SetGlobalBuffer((__gm__ T *)x + offset, elementsPerBlock);
39+ yGm.SetGlobalBuffer((__gm__ T *)y + offset, elementsPerBlock);
40+ zGm.SetGlobalBuffer((__gm__ T *)z + offset, elementsPerBlock);
41+ 
42+ pipe.InitBuffer(inQueueX, BUFFER_NUM, TILE_LENGTH * sizeof(T));
43+ pipe.InitBuffer(inQueueY, BUFFER_NUM, TILE_LENGTH * sizeof(T));
44+ pipe.InitBuffer(outQueueZ, BUFFER_NUM, TILE_LENGTH * sizeof(T));
45+ }
46+ 
47+ __aicore__ inline void Process()
48+ {
49+ for (uint32_t i = 0; i < tileNum; i++) {
50+ uint32_t currentTileLength = (i == tileNum - 1) ? lastTileLength : TILE_LENGTH;
51+ CopyIn(i, currentTileLength);
52+ Compute(currentTileLength);
53+ CopyOut(i, currentTileLength);
54+ }
55+ }
56+ 
57+private:
58+ __aicore__ inline void CopyIn(uint32_t tileIdx, uint32_t tileLength)
59+ {
60+ LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
61+ LocalTensor<T> yLocal = inQueueY.AllocTensor<T>();
62+ DataCopy(xLocal, xGm[tileIdx * TILE_LENGTH], tileLength);
63+ DataCopy(yLocal, yGm[tileIdx * TILE_LENGTH], tileLength);
64+ inQueueX.EnQue<T>(xLocal);
65+ inQueueY.EnQue<T>(yLocal);
66+ }
67+ 
68+ __aicore__ inline void Compute(uint32_t tileLength)
69+ {
70+ LocalTensor<T> xLocal = inQueueX.DeQue<T>();
71+ LocalTensor<T> yLocal = inQueueY.DeQue<T>();
72+ LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();
73+ Div(zLocal, xLocal, yLocal, tileLength);
74+ outQueueZ.EnQue<T>(zLocal);
75+ inQueueX.FreeTensor(xLocal);
76+ inQueueY.FreeTensor(yLocal);
77+ }
78+ 
79+ __aicore__ inline void CopyOut(uint32_t tileIdx, uint32_t tileLength)
80+ {
81+ LocalTensor<T> zLocal = outQueueZ.DeQue<T>();
82+ DataCopy(zGm[tileIdx * TILE_LENGTH], zLocal, tileLength);
83+ outQueueZ.FreeTensor(zLocal);
84+ }
85+ 
86+private:
87+ uint32_t totalLength;
88+ uint32_t elementsPerBlock;
89+ uint32_t offset;
90+ uint32_t tileNum;
91+ uint32_t lastTileLength;
92+ TPipe pipe;
93+ TQue<TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
94+ TQue<TPosition::VECOUT, BUFFER_NUM> outQueueZ;
95+ GlobalTensor<T> xGm, yGm, zGm;
96+};
97+ 
98+extern "C" __global__ __aicore__ void div_custom_template(GM_ADDR x, GM_ADDR y, GM_ADDR z,
99+ GM_ADDR workspace, GM_ADDR tiling)
100+{
101+ KernelDiv<half> op;
102+ op.Init(x, y, z, workspace, tiling);
103+ op.Process();
104+}
@@ -0,0 +1,26 @@
1+{
2+ "op": "DivCustomTemplate",
3+ "language": "cpp",
4+ "input_desc": [
5+ {
6+ "name": "x",
7+ "param_type": "required",
8+ "format": ["ND"],
9+ "type": ["float16"]
10+ },
11+ {
12+ "name": "y",
13+ "param_type": "required",
14+ "format": ["ND"],
15+ "type": ["float16"]
16+ }
17+ ],
18+ "output_desc": [
19+ {
20+ "name": "z",
21+ "param_type": "required",
22+ "format": ["ND"],
23+ "type": ["float16"]
24+ }
25+ ]
26+}
@@ -0,0 +1,87 @@
1+#!/bin/bash
2+set -e
3+ 
4+# 获取当前脚本所在目录
5+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
6+cd "$SCRIPT_DIR"
7+ 
8+echo "=========================================="
9+echo " 1. Loading CANN Environment"
10+echo "=========================================="
11+if [ -n "${ASCEND_TOOLKIT_HOME:-}" ] && [ -f "${ASCEND_TOOLKIT_HOME}/set_env.sh" ]; then
12+ source "${ASCEND_TOOLKIT_HOME}/set_env.sh"
13+elif [ -n "${ASCEND_HOME_PATH:-}" ] && [ -f "${ASCEND_HOME_PATH}/set_env.sh" ]; then
14+ source "${ASCEND_HOME_PATH}/set_env.sh"
15+else
16+ if [ -f "/usr/local/Ascend/ascend-toolkit/set_env.sh" ]; then
17+ source /usr/local/Ascend/ascend-toolkit/set_env.sh
18+ elif [ -f "/usr/local/Ascend/cann-8.5.1/set_env.sh" ]; then
19+ source /usr/local/Ascend/cann-8.5.1/set_env.sh
20+ else
21+ echo "Error: Cannot find set_env.sh."
22+ exit 1
23+ fi
24+fi
25+echo "CANN Environment loaded."
26+ 
27+echo "=========================================="
28+echo " 2. Generating Operator Project (msopgen)"
29+echo "=========================================="
30+if [ -d "custom_op" ] && [ -f "custom_op/build.sh" ]; then
31+ echo ">>> custom_op exists, skip msopgen."
32+else
33+ echo ">>> Generating operator project..."
34+ SOC_VERSION=$(python3 -c "import acl; print(acl.get_soc_name())" 2>/dev/null || echo "Ascend910B")
35+ if echo "$SOC_VERSION" | grep -qi "Ascend910"; then
36+ SOC_PARAM="ai_core-ascend910a"
37+ else
38+ SOC_PARAM="ai_core-ascend910b1"
39+ fi
40+ echo ">>> SoC: $SOC_VERSION, param: $SOC_PARAM"
41+ msopgen gen -i div_custom_template.json -c $SOC_PARAM -lan cpp -out ./custom_op
42+ echo ">>> Generated. Please implement kernel in custom_op/op_kernel, then re-run."
43+ exit 0
44+fi
45+ 
46+echo "=========================================="
47+echo " 3. Compiling Operator"
48+echo "=========================================="
49+cd custom_op
50+bash build.sh 2>&1
51+cd ..
52+ 
53+echo "=========================================="
54+echo " 4. Installing Operator"
55+echo "=========================================="
56+RUN_FILE=$(ls custom_op/build_out/custom_*.run 2>/dev/null | head -1)
57+if [ -z "$RUN_FILE" ]; then
58+ echo "Error: No .run file found."
59+ exit 1
60+fi
61+$RUN_FILE --install-path=${HOME}/
62+ 
63+VENDOR_PATH=${HOME}/vendors/customize
64+TILING_LIB=$VENDOR_PATH/op_impl/ai_core/tbe/op_tiling/lib/linux/aarch64
65+PROTO_LIB=$VENDOR_PATH/op_proto/lib/linux/aarch64
66+if [ -d "$TILING_LIB" ]; then
67+ export LD_LIBRARY_PATH=$TILING_LIB:$LD_LIBRARY_PATH
68+fi
69+if [ -d "$PROTO_LIB" ]; then
70+ export LD_LIBRARY_PATH=$PROTO_LIB:$LD_LIBRARY_PATH
71+fi
72+ 
73+echo "=========================================="
74+echo " 5. Building Test"
75+echo "=========================================="
76+g++ -I"${ASCEND_TOOLKIT_HOME}/include" \
77+ -I"${VENDOR_PATH}/op_api/include" \
78+ -L"${ASCEND_TOOLKIT_HOME}/lib64" \
79+ -L"${VENDOR_PATH}/op_api/lib" \
80+ test/main.cpp \
81+ -lcust_opapi -lnnopbase -lacl_rt \
82+ -o execute_div_op
83+ 
84+echo "=========================================="
85+echo " 6. Running Test"
86+echo "=========================================="
87+./execute_div_op
@@ -0,0 +1,29 @@
1+# Copyright (c) Huawei Technologies Co., Ltd. 2020. All rights reserved.
2+cmake_minimum_required(VERSION 3.5.1)
3+project(acl_execute_div)
4+ 
5+add_compile_options(-std=c++11)
6+set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "./")
7+ 
8+set(INC_PATH $ENV{DDK_PATH})
9+if (NOT DEFINED ENV{DDK_PATH})
10+ set(INC_PATH "/usr/local/Ascend/ascend-toolkit/latest")
11+endif()
12+ 
13+set(CUST_PKG_PATH "${INC_PATH}/opp/vendors/customize/op_api")
14+set(LIB_PATH $ENV{NPU_HOST_LIB})
15+if (NOT DEFINED ENV{NPU_HOST_LIB})
16+ set(LIB_PATH "/usr/local/Ascend/ascend-toolkit/latest/${CMAKE_SYSTEM_PROCESSOR}-${CMAKE_SYSTEM_NAME_LOWER}/devlib")
17+endif()
18+ 
19+include_directories(
20+ ${INC_PATH}/include
21+ ${CUST_PKG_PATH}/include
22+)
23+link_directories(
24+ ${LIB_PATH}
25+ ${CUST_PKG_PATH}/lib
26+)
27+add_executable(execute_div_op main.cpp)
28+target_link_libraries(execute_div_op ascendcl cust_opapi acl_op_compiler nnopbase stdc++)
29+install(TARGETS execute_div_op DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY})
@@ -0,0 +1,158 @@
1+/**
2+ * @file main.cpp
3+ *
4+ * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
5+ */
6+#include <algorithm>
7+#include <cstdint>
8+#include <iostream>
9+#include <vector>
10+ 
11+#include "acl/acl.h"
12+#include "aclnn_div_custom_template.h"
13+ 
14+#define SUCCESS 0
15+#define FAILED 1
16+ 
17+#define CHECK_RET(cond, return_expr) \
18+ do { \
19+ if (!(cond)) { \
20+ return_expr; \
21+ } \
22+ } while (0)
23+ 
24+#define LOG_PRINT(message, ...) \
25+ do { \
26+ printf(message, ##__VA_ARGS__); \
27+ } while (0)
28+ 
29+int64_t GetShapeSize(const std::vector<int64_t> &shape)
30+{
31+ int64_t shapeSize = 1;
32+ for (auto i : shape) {
33+ shapeSize *= i;
34+ }
35+ return shapeSize;
36+}
37+ 
38+int Init(int32_t deviceId, aclrtStream *stream)
39+{
40+ auto ret = aclInit(nullptr);
41+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return FAILED);
42+ ret = aclrtSetDevice(deviceId);
43+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return FAILED);
44+ ret = aclrtCreateStream(stream);
45+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return FAILED);
46+ return SUCCESS;
47+}
48+ 
49+template <typename T>
50+int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
51+ aclDataType dataType, aclTensor **tensor)
52+{
53+ auto size = GetShapeSize(shape) * sizeof(T);
54+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
55+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return FAILED);
56+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
57+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return FAILED);
58+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, nullptr, 0, aclFormat::ACL_FORMAT_ND, shape.data(),
59+ shape.size(), *deviceAddr);
60+ return SUCCESS;
61+}
62+ 
63+void DestroyResources(std::vector<void *> tensors, std::vector<void *> deviceAddrs, aclrtStream stream,
64+ int32_t deviceId, void *workspaceAddr = nullptr)
65+{
66+ for (uint32_t i = 0; i < tensors.size(); i++) {
67+ if (tensors[i] != nullptr) {
68+ aclDestroyTensor(reinterpret_cast<aclTensor *>(tensors[i]));
69+ }
70+ if (deviceAddrs[i] != nullptr) {
71+ aclrtFree(deviceAddrs[i]);
72+ }
73+ }
74+ if (workspaceAddr != nullptr) {
75+ aclrtFree(workspaceAddr);
76+ }
77+ aclrtDestroyStream(stream);
78+ aclrtResetDevice(deviceId);
79+ aclFinalize();
80+}
81+ 
82+int main(int argc, char **argv)
83+{
84+ int32_t deviceId = 0;
85+ aclrtStream stream;
86+ auto ret = Init(deviceId, &stream);
87+ CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return FAILED);
88+ 
89+ std::vector<int64_t> inputXShape = {8, 2048};
90+ std::vector<int64_t> inputYShape = {8, 2048};
91+ std::vector<int64_t> outputZShape = {8, 2048};
92+ void *inputXDeviceAddr = nullptr;
93+ void *inputYDeviceAddr = nullptr;
94+ void *outputZDeviceAddr = nullptr;
95+ aclTensor *inputX = nullptr;
96+ aclTensor *inputY = nullptr;
97+ aclTensor *outputZ = nullptr;
98+ std::vector<aclFloat16> inputXHostData(inputXShape[0] * inputXShape[1]);
99+ std::vector<aclFloat16> inputYHostData(inputYShape[0] * inputYShape[1]);
100+ std::vector<aclFloat16> outputZHostData(outputZShape[0] * outputZShape[1]);
101+ for (int i = 0; i < inputXShape[0] * inputXShape[1]; ++i) {
102+ inputXHostData[i] = aclFloatToFloat16(1.0);
103+ inputYHostData[i] = aclFloatToFloat16(2.0);
104+ outputZHostData[i] = aclFloatToFloat16(0.0);
105+ }
106+ std::vector<void *> tensors = {inputX, inputY, outputZ};
107+ std::vector<void *> deviceAddrs = {inputXDeviceAddr, inputYDeviceAddr, outputZDeviceAddr};
108+ ret = CreateAclTensor(inputXHostData, inputXShape, &inputXDeviceAddr, aclDataType::ACL_FLOAT16, &inputX);
109+ CHECK_RET(ret == ACL_SUCCESS, DestroyResources(tensors, deviceAddrs, stream, deviceId); return FAILED);
110+ ret = CreateAclTensor(inputYHostData, inputYShape, &inputYDeviceAddr, aclDataType::ACL_FLOAT16, &inputY);
111+ CHECK_RET(ret == ACL_SUCCESS, DestroyResources(tensors, deviceAddrs, stream, deviceId); return FAILED);
112+ ret = CreateAclTensor(outputZHostData, outputZShape, &outputZDeviceAddr, aclDataType::ACL_FLOAT16, &outputZ);
113+ CHECK_RET(ret == ACL_SUCCESS, DestroyResources(tensors, deviceAddrs, stream, deviceId); return FAILED);
114+ 
115+ uint64_t workspaceSize = 0;
116+ aclOpExecutor *executor;
117+ ret = aclnnDivCustomTemplateGetWorkspaceSize(inputX, inputY, outputZ, &workspaceSize, &executor);
118+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnDivCustomTemplateGetWorkspaceSize failed. ERROR: %d\n", ret);
119+ DestroyResources(tensors, deviceAddrs, stream, deviceId); return FAILED);
120+ 
121+ void *workspaceAddr = nullptr;
122+ if (workspaceSize > 0) {
123+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
124+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret);
125+ DestroyResources(tensors, deviceAddrs, stream, deviceId, workspaceAddr); return FAILED);
126+ }
127+ ret = aclnnDivCustomTemplate(workspaceAddr, workspaceSize, executor, stream);
128+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnDivCustomTemplate failed. ERROR: %d\n", ret);
129+ DestroyResources(tensors, deviceAddrs, stream, deviceId, workspaceAddr); return FAILED);
130+ 
131+ ret = aclrtSynchronizeStream(stream);
132+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret);
133+ DestroyResources(tensors, deviceAddrs, stream, deviceId, workspaceAddr); return FAILED);
134+ 
135+ auto size = GetShapeSize(outputZShape);
136+ std::vector<aclFloat16> resultData(size, 0);
137+ ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outputZDeviceAddr,
138+ size * sizeof(aclFloat16), ACL_MEMCPY_DEVICE_TO_HOST);
139+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret);
140+ DestroyResources(tensors, deviceAddrs, stream, deviceId, workspaceAddr); return FAILED);
141+ 
142+ DestroyResources(tensors, deviceAddrs, stream, deviceId, workspaceAddr);
143+ 
144+ std::vector<aclFloat16> goldenData(size, aclFloatToFloat16(0.5));
145+ 
146+ LOG_PRINT("result is:\n");
147+ for (int64_t i = 0; i < 10; i++) {
148+ LOG_PRINT("%.1f ", aclFloat16ToFloat(resultData[i]));
149+ }
150+ LOG_PRINT("\n");
151+ if (std::equal(resultData.begin(), resultData.end(), goldenData.begin())) {
152+ LOG_PRINT("test pass\n");
153+ } else {
154+ LOG_PRINT("test failed\n");
155+ return FAILED;
156+ }
157+ return SUCCESS;
158+}