Star902
1820
代码
代码
Issues228
Pull Requests306
Actions
项目讨论
Wiki
分析
Star902
1820
已合并
更新极简样例资料和工程 #620
lileizheng创建于 1月4日
更新极简样例资料和工程 #620
已合并
lileizheng创建于 1月4日
从lileizheng/ops-nnremote_update_fast_kernel合入到cann/ops-nnmaster
共 21 个文件变更+919-388
Mexamples/CMakeLists.txt+5-0
@@ -10,6 +10,11 @@
10file(GLOB SUBDIRECTORIES LIST_DIRECTORIES true RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)10file(GLOB SUBDIRECTORIES LIST_DIRECTORIES true RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
11# 遍历子目录11# 遍历子目录
12foreach(SUBDIR ${SUBDIRECTORIES})12foreach(SUBDIR ${SUBDIRECTORIES})
13+ # 检查是否跳过目录
14+ if("${SUBDIR}" STREQUAL "fast_kernel_launch_example")
15+ continue()
16+ endif()
17+
13 # 检查子目录中是否存在 CMakeLists.txt18 # 检查子目录中是否存在 CMakeLists.txt
14 if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${SUBDIR}/CMakeLists.txt)19 if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${SUBDIR}/CMakeLists.txt)
15 add_subdirectory(${SUBDIR})20 add_subdirectory(${SUBDIR})
Mexamples/fast_kernel_launch_example/CMakeLists.txt+43-122
@@ -9,107 +9,42 @@
9# ----------------------------------------------------------------------------9# ----------------------------------------------------------------------------
10 10 
11cmake_minimum_required(VERSION 3.16)11cmake_minimum_required(VERSION 3.16)
12-project(AscendOps)12+set(PKG_NAME AscendOps)
13-option(BUILD_TORCH_OPS "Build torch ops project (PyTorch extension style)" OFF)13+project(${PKG_NAME} VERSION 1.0.0)
14-message(STATUS "Build torch ops: ${BUILD_TORCH_OPS}")
15- 
16-if(NOT BUILD_TORCH_OPS)
17- message(STATUS "BUILD_TORCH_OPS is OFF, exiting CMake configuration")
18- return()
19-endif()
20- 
21-#================================
22-# Torch Ops 项目特定配置
23-#================================
24 14 
25set(CMAKE_EXPORT_COMPILE_COMMANDS ON)15set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
26set(CMAKE_CXX_STANDARD 17)16set(CMAKE_CXX_STANDARD 17)
27set(CMAKE_CXX_STANDARD_REQUIRED ON)17set(CMAKE_CXX_STANDARD_REQUIRED ON)
28set(CMAKE_POSITION_INDEPENDENT_CODE ON)18set(CMAKE_POSITION_INDEPENDENT_CODE ON)
29 19 
30-# 设置默认构建类型20+include(cmake/ascend.cmake)
31-if(NOT CMAKE_BUILD_TYPE)21+include(cmake/python.cmake)
32- set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type (Release/Debug)" FORCE)22+include(cmake/torch.cmake)
33-endif()23+include(cmake/torch_npu.cmake)
24+include(cmake/func.cmake)
34 25 
35-#================================26+# set NPU architecture, default to ascend910b
36-# Ascend 环境配置27+set(NPU_ARCH "ascend910b" CACHE STRING "NPU architecture")
37-#================================28+# set EXTENSION_MODULE_NAME to `ascend_ops`
38-message(STATUS "$ENV{ASCEND_HOME_PATH}")29+set(EXTENSION_MODULE_NAME "ascend_ops" CACHE STRING "Extension module name")
39-if(DEFINED ENV{ASCEND_HOME_PATH})
40- set(ASCEND_HOME "$ENV{ASCEND_HOME_PATH}" CACHE PATH "Ascend installation path" FORCE)
41-else()
42- set(ASCEND_HOME "/usr/local/Ascend/cann" CACHE PATH "Ascend installation path")
43-endif()
44 30 
45-set(BISHENG "${ASCEND_HOME}/tools/ccec_compiler/bin/bisheng" CACHE FILEPATH "Path to Bisheng compiler")31+# include directories
46-message(STATUS "ASCEND_HOME = ${ASCEND_HOME}")32+set(INCLUDE_DIRECTORIES
47-message(STATUS "BISHENG = ${BISHENG}")
48-# 设置编译器为 bisheng
49-set(CMAKE_C_COMPILER ${BISHENG})
50-set(CMAKE_CXX_COMPILER ${BISHENG})
51-set(CMAKE_LINKER ${BISHENG})
52- 
53-# 设置Ascend Log模块
54-set(ASCEND_DIR $ENV{ASCEND_HOME_PATH})
55-set(CMAKE_PREFIX_PATH ${ASCEND_DIR}/)
56-include(../../cmake/modules/Finddlog.cmake)
57- 
58-# Python 配置
59-find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
60-message(STATUS "Python3_EXECUTABLE = ${Python3_EXECUTABLE}")
61-message(STATUS "Python3_INCLUDE_DIRS = ${Python3_INCLUDE_DIRS}")
62-message(STATUS "Python3_LIBRARIES = ${Python3_LIBRARIES}")
63- 
64-# Torch 配置
65-execute_process(
66- COMMAND ${Python3_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)"
67- OUTPUT_VARIABLE TORCH_CMAKE_PATH
68- OUTPUT_STRIP_TRAILING_WHITESPACE
69-)
70-set(Torch_DIR "${TORCH_CMAKE_PATH}/Torch")
71-find_package(Torch REQUIRED)
72-message(STATUS "Torch_DIR = ${TORCH_DIR}")
73-message(STATUS "Torch_LIBRARIES = ${TORCH_LIBRARIES}")
74-message(STATUS "Torch_INCLUDES = ${TORCH_INCLUDE_DIRS}")
75- 
76-# Torch-NPU 配置
77-execute_process(
78- COMMAND ${Python3_EXECUTABLE} -c "import torch_npu, os; print(os.path.dirname(torch_npu.__file__))"
79- OUTPUT_VARIABLE TORCH_NPU_PATH
80- OUTPUT_STRIP_TRAILING_WHITESPACE
81-)
82-set(TORCH_NPU_INCLUDE_PATH "${TORCH_NPU_PATH}/include")
83-set(TORCH_NPU_LIB_PATH "${TORCH_NPU_PATH}/lib")
84- 
85-message(STATUS "TORCH_NPU_PATH = ${TORCH_NPU_PATH}")
86-message(STATUS "TORCH_NPU_INCLUDE_DIR = ${TORCH_NPU_INCLUDE_PATH}")
87-message(STATUS "TORCH_NPU_LIB_DIR = ${TORCH_NPU_LIB_PATH}")
88- 
89-# Torch Ops 公共配置
90-set(COMMON_INCLUDE_DIRS
91 ${Python3_INCLUDE_DIRS}33 ${Python3_INCLUDE_DIRS}
92 ${TORCH_INCLUDE_DIRS}34 ${TORCH_INCLUDE_DIRS}
93 ${TORCH_NPU_INCLUDE_PATH}35 ${TORCH_NPU_INCLUDE_PATH}
36+ ${ASCEND_INCLUDE_DIRS}
94 ${CMAKE_CURRENT_SOURCE_DIR}/../..37 ${CMAKE_CURRENT_SOURCE_DIR}/../..
95- ${CMAKE_CURRENT_SOURCE_DIR}/../../common/inc
96- ${ASCEND_HOME}/include
97- ${ASCEND_HOME}/${CMAKE_SYSTEM_PROCESSOR}-linux/ascendc/include/basic_api/impl
98- ${ASCEND_HOME}/${CMAKE_SYSTEM_PROCESSOR}-linux/ascendc/include/basic_api
99- ${ASCEND_HOME}/include/ascendc
100- ${ASCEND_HOME}/include/ascendc/basic_api/interface
101- ${ASCEND_HOME}/include/ascendc/basic_api
102- ${ASCEND_HOME}/pkg_inc/op_common
103- ${ASCEND_HOME}/pkg_inc/base
104- ${ASCEND_HOME}/pkg_inc
105)38)
106 39 
107-set(COMMON_LINK_DIRS40+# link directories
41+set(LINK_DIRECTORIES
108 ${TORCH_NPU_LIB_PATH}42 ${TORCH_NPU_LIB_PATH}
109- ${ASCEND_HOME}/lib6443+ ${ASCEND_DIR}/lib64
110)44)
111 45 
112-set(COMMON_LINK_LIBS46+# link libraries
47+set(LINK_LIBRARIES
113 ${TORCH_LIBRARIES}48 ${TORCH_LIBRARIES}
114 torch_npu49 torch_npu
115 ascendcl50 ascendcl
@@ -117,58 +52,44 @@ set(COMMON_LINK_LIBS
117 register52 register
118 tiling_api53 tiling_api
119 runtime54 runtime
120- ${Python3_LIBRARIES}
121)55)
122 56 
123-set(COMMON_COMPILE_OPTIONS57+# compile options
58+set(COMPILE_OPTIONS
124 ${TORCH_CXX_FLAGS}59 ${TORCH_CXX_FLAGS}
125- -O260+ -O3
126 -fdiagnostics-color=always61 -fdiagnostics-color=always
127- -DPy_LIMITED_API=0x03090000
128 -w62 -w
129- -DTORCH_MODE63+ -DEXTENSION_MODULE_NAME=${EXTENSION_MODULE_NAME}
130)64)
131 65 
132-message(STATUS "Building Torch Ops project...")66+# define a LIST variable to hold object files
67+set(OBJECTS_LIST "" CACHE INTERNAL "List of objects")
68+message(STATUS "OBJECTS_LIST before adding sources: ${OBJECTS_LIST}")
69+add_subdirectory(csrc)
70+message(STATUS "OBJECTS_LIST after adding sources: ${OBJECTS_LIST}")
133 71 
134-file(GLOB OPERATOR_DIRS "${CMAKE_CURRENT_SOURCE_DIR}/ascend_ops/csrc/*")72+# create shared library
135-set(OPERATOR_TARGETS "")73+set("EXTENSION_CPP" ${CMAKE_CURRENT_SOURCE_DIR}/csrc/extension.cpp)
136- 
137-foreach(OP_DIR ${OPERATOR_DIRS})
138- if(IS_DIRECTORY ${OP_DIR})
139- # 获取目录名
140- get_filename_component(OP_NAME ${OP_DIR} NAME)
141- 
142- # 检查目录中是否有 CMakeLists.txt
143- if(EXISTS "${OP_DIR}/CMakeLists.txt")
144- message(STATUS "Found operator: ${OP_NAME}")
145- add_subdirectory(ascend_ops/csrc/${OP_NAME})
146- list(APPEND OPERATOR_TARGETS $<TARGET_OBJECTS:${OP_NAME}_objects>)
147- endif()
148- endif()
149-endforeach()
150- 
151-message(STATUS "Discovered torch operator targets: ${OPERATOR_TARGETS}")
152- 
153-# 主库源文件
154-set(MAIN_LIBRARY_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/ascend_ops/csrc/npu_ops_def.cpp")
155- 
156-# 创建主库
157add_library(_C SHARED74add_library(_C SHARED
158- ${MAIN_LIBRARY_SOURCE}75+ ${EXTENSION_CPP}
159- ${OPERATOR_TARGETS}76+ ${OBJECTS_LIST}
160)77)
161- 
162set_target_properties(_C PROPERTIES78set_target_properties(_C PROPERTIES
163 POSITION_INDEPENDENT_CODE ON79 POSITION_INDEPENDENT_CODE ON
164- LINK_FLAGS "--cce-fatobj-link"
165 PREFIX ""80 PREFIX ""
166 SUFFIX ".abi3.so"81 SUFFIX ".abi3.so"
167 OUTPUT_NAME "_C"82 OUTPUT_NAME "_C"
168)83)
84+target_compile_definitions(_C PRIVATE Py_LIMITED_API=0x03080000)
85+target_compile_options(_C PRIVATE ${COMPILE_OPTIONS})
86+target_include_directories(_C PRIVATE ${INCLUDE_DIRECTORIES})
87+target_link_directories(_C PRIVATE ${LINK_DIRECTORIES})
88+target_link_libraries(_C PRIVATE ${LINK_LIBRARIES})
169 89 
170-target_compile_options(_C PRIVATE ${COMMON_COMPILE_OPTIONS})90+add_custom_command(TARGET _C POST_BUILD
171-target_include_directories(_C PRIVATE ${COMMON_INCLUDE_DIRS})91+ COMMAND ${CMAKE_COMMAND} -E copy
172-target_link_directories(_C PRIVATE ${COMMON_LINK_DIRS})92+ $<TARGET_FILE:_C>
173-target_link_libraries(_C PRIVATE ${COMMON_LINK_LIBS})93+ ${CMAKE_CURRENT_SOURCE_DIR}/${EXTENSION_MODULE_NAME}/$<TARGET_FILE_NAME:_C>
174-return()94+ COMMENT "Copying compiled extension $<TARGET_FILE_NAME:_C> to ${CMAKE_CURRENT_SOURCE_DIR}/${EXTENSION_MODULE_NAME}/"
95+)
Mexamples/fast_kernel_launch_example/README.md+152-164
@@ -1,200 +1,188 @@
1# AscendOps1# AscendOps
2 2 
3-**AscendOps** - 一个轻量级，高性能的算子开发工程模板3+## 环境部署 | Prerequisites
4 4 
5-## 项目简介 | Introduction5+- 请先参考[环境部署](../../docs/zh/context/quick_install.md)完成基础环境搭建
6-AscendOps 是一个轻量级，高性能的算子开发工程模板，它集成了PyTorch、PyBind11和昇腾CANN工具链，提供了从算子内核编写，编译到Python封装的完整工具链。6+- GCC 9.4.0+
7+- Python 3.8+
8+- PyTorch>=2.6.0
9+- 对应版本的[TorchNPU](https://gitcode.com/Ascend/pytorch/releases)
7 10 
8-## 核心特性 | Features11+## 安装步骤 | Installation Steps
9-🚀 开箱即用 (Out-of-the-Box): 预置完整的昇腾NPU算子开发环境配置，克隆后即可开始开发。
10 12 
11-🧩 极简设计 (Minimalist Design): 代码结构清晰直观，专注于核心算子开发流程。13+1. 安装依赖 | Install Dependencies:
12- 14+ ```sh
13-⚡ 高性能 (High Performance): 基于AscendC编程模型，充分发挥昇腾NPU硬件能力。15+ python3 -m pip install -r requirements.txt
14- 
15-📦 一键部署 (One-Click Deployment): 集成setuptools构建系统，支持一键编译和安装。
16- 
17-🔌 PyTorch集成 (PyTorch Integration): 无缝集成PyTorch张量操作，支持自动微分和GPU/NPU统一接口。
18- 
19-## 核心交付件 | Core Deliverables
20-1. `csrc/xxx/xxx_torch.cpp` 算子Kernel实现
21-2. `csrc/xxx/CMakeLists.txt` 算子cmake配置
22-3. `csrc/npu_ops_def.cpp` 注册算子接口
23- 
24-## 环境要求 | Prerequisites
25-* Python: 3.8+
26-* CANN Ascend Toolkit
27-* CANN Ascend Legacy
28-* PyTorch: 2.1.0+
29-* PyTorchAdapter 7.1.0+
30- 
31-## 环境准备 | Preparation
32- 
33-1. **安装社区版CANN toolkit包和社区版CANN legacy包**
34- 
35- 根据实际环境，安装社区版CANN toolkit包和社区版CANN legacy包，社区包的安装部署参考[算子调用](../../docs/zh/invocation/quick_op_invocation.md)环境准备章节。
36- 
37-2. **配置环境变量**
38- 
39- 根据实际场景，选择合适的命令。
40- 
41- ```bash
42- # 默认路径安装，以root用户为例（非root用户，将/usr/local替换为${HOME}）
43- source /usr/local/Ascend/set_env.sh
44- # 指定路径安装
45- # source ${install-path}/set_env.sh
46 ```16 ```
47 17 
48-3. **安装torch与torch_npu包**18+2. 构建Wheel包 | Build the Wheel:
49- 
50- 根据实际环境，下载对应 torch 包，常见的 wheel 文件名示例如下：
51- 
52- - x86_64 Linux 版本：`torch-${torch_version}+cpu-${python_version}-linux_x86_64.whl`
53- - ARM Linux（旧版本通常无 `+cpu` 后缀）：`torch-${torch_version}-${python_version}-linux_aarch64.whl`
54- 
55- 说明：上面示例中的版本号及 `+cpu` 后缀仅作为示例。不同操作系统和架构（尤其是 ARM Linux）对应的 wheel 名称可能没有 `+cpu` 后缀，请以 PyTorch 官方下载页面或 `pip` 实际可用的包名为准。
56- 
57- 下载链接为：[官网地址](http://download.pytorch.org/whl/torch)
58- 
59- 安装命令如下（将 `<torch_whl>` 替换为实际下载的文件名）：
60- 
61- ```sh
62- pip3 install <torch_whl>
63- ```
64- 
65- 根据实际环境，安装对应 torch-npu 包，例如：`torch_npu-${torch_version}-${python_version}-linux_${arch}.whl`，下载链接：[官网地址](https://www.hiascend.com/document/detail/zh/Pytorch/710/configandinstg/instg/insg_0004.html)
66- 
67- 可以直接使用 pip 命令下载安装（将 `<torch_npu_whl>` 替换为实际下载的文件名），命令如下：
68- 
69- ```sh
70- pip3 install <torch_npu_whl>
71- ```
72- 
73- - ${torch_version}：表示 torch 包版本号。
74- - ${python_version}：表示 python 版本号。
75- - ${arch}：表示 CPU 架构，如 aarch64、x86_64。
76- 
77-## 安装步骤 | Installation
78- 
79-1. 进入目录，安装依赖
80- ```sh
81- cd ./examples/fast_kernel_launch_example
82- pip3 install -r requirements.txt
83- ```
84- 
85-2. 从源码构建.whl包
86 ```sh19 ```sh
20+ # -n: non-isolated build (uses existing environment)
87 python3 -m build --wheel -n21 python3 -m build --wheel -n
88 ```22 ```
89 23 
90-3. 安装构建好的.whl包24+3. 安装 | Install Package:
91 ```sh25 ```sh
92- pip3 install dist/xxx.whl26+ python3 -m pip install dist/*.whl --force-reinstall --no-deps
93 ```27 ```
94 28 
95- 重新安装请使用以下命令覆盖已安装过的版本：29+4. （可选）再次构建前建议先执行以下命令清理编译缓存
96- ```sh30+ ```sh
97- pip3 install dist/xxx.whl --force-reinstall --no-deps31+ python3 setup.py clean
98- ```32+ ```
99 33 
100-4. (可选)再次构建前建议先执行以下命令清理编译缓存34+## 快速开始 | Quick Start
101- ```sh35+安装完成后，您可以像使用普通PyTorch操作一样使用NPU算子
102- python3 setup.py clean
103- ```
104 36 
105-## 开发模式构建 | Developing Mode37+```python
38+import torch
39+import torch_npu
40+import ascend_ops
106 41 
107-此命令实现即时生效的开发环境配置，执行后即可使源码修改生效，省略了构建完整whl包和安装的过程，适用于需要多次修改验证算子的场景：42+# Initialize data on NPU
108- ```sh43+x = torch.randn(10, 32, dtype=torch.float32).npu()
109- pip3 install --no-build-isolation -e .44+y = torch.randn(10, 32, dtype=torch.float32).npu()
110- ```
111 45 
112-## 使用示例 | Usage Example46+# Call the custom NPU operator
47+npu_result = torch.ops.ascend_ops.add(x, y)
113 48 
114-安装完成后，您可以像使用普通PyTorch操作一样使用NPU算子，以conv3d算子为例，您可以在`ascend_ops/csrc/conv3d_custom/test`目录下找到并执行这个脚本:49+# Verify against CPU ATen implementation
50+cpu_x = x.cpu()
51+cpu_y = y.cpu()
52+cpu_result = cpu_x + cpu_y
115 53 
116-```bash54+assert torch.allclose(cpu_result, npu_result.cpu(), rtol=1e-6)
117-cd ./ascend_ops/csrc/conv3d_custom/test/55+print("Verification successful!")
118-python3 test_conv3d_custom.py
119```56```
120 57 
121-最终看到如下输出，即为执行成功：58+## 开发指南：新增一个算子 | Developer Guide: Adding a New Operator
122-```bash
123-compare CPU Result vs NPU Result: True
124-```
125 59 
60+为了实现一个新算子(如`add`)，你只需要提供一个C++实现即可。
126 61 
127-## 开发新算子 | Developing New Operators62+1. 首先你需要在csrc目录下使用算子名`add`建立一个文件夹，在此文件夹内使用你当前想要开发的soc名建立一个子文件夹`ascend910b`.
128-1. 编写算子调用文件
129 63 
130- 在 `ascend_ops/csrc/` 目录下添加新的算子目录 `mykernel`，在 `mykernel` 目录下添加新的算子调用文件 `mykernel_torch.cpp`64+2. 在soc目录下新建一个`CMakeLists.txt`
131- ```c++65+ ```
132- __global__ __aicore__ void mykernel(GM_ADDR input, GM_ADDR output, int64_t num_element) {66+ add_sources("--npu-arch=dav-2201")
133- // 您的算子kernel实现67+ ```
134- }68+ 这里`dav-2201`为ascend910b芯片对应的编译参数
135 69 
136- void mykernel_api(aclrtStream stream, const at::Tensor& x, const at::Tensor& y) {70+3. 在soc目录下新建一个`add.cpp`(建议使用算子名为文件名)。这个文件包含了开发一个AICORE算子所需要的全部模块。
137- // 您的算子入口实现，在该方法中使用<<<>>>的方式调用算子kernel71+ - 算子Schema注册
138- mykernel<<<blockDim, nullptr, stream>>>(x, y, num_element);72+ - 算子Meta Function实现 & 注册
139- }73+ - 算子Kernel实现 (AscendC)
74+ - 算子NPU调用实现 & 注册
140 75 
141- torch::Tensor mykernel_npu(torch::Tensor x, torch::Tensor y) {76+ ```cpp
142- // 您的算子wrapper接口，用于向pytorch注册自定义接口77+ #include <ATen/Operators.h>
143- AT_DISPATCH_FLOATING_TYPES_AND2(78+ #include <torch/all.h>
144- at::kHalf, at::kBFloat16, x.scalar_type(), "mykernel_npu", [&] { mykernel_api(stream, x, y); });79+ #include <torch/library.h>
145- }80+ #include "torch_npu/csrc/core/npu/NPUStream.h"
81+ #include "torch_npu/csrc/framework/OpCommand.h"
82+ #include "kernel_operator.h"
83+ #include "platform/platform_ascendc.h"
84+ #include <type_traits>
146 85 
147- // PyTorch提供的宏，用于在特定后端注册算子86+ namespace ascend_ops { // 当前项目为一个命名空间
148- TORCH_LIBRARY_IMPL(ascend_ops, PrivateUse1, m)87+ namespace Add { // 建议每个算子自己有一个独立的namespace，防止全局变量污染
88+ 
89+ /**
90+ * 将算子schema注册给PyTorch框架
91+ * 框架知道有这样一个算子
92+ */
93+ // Register the operator's schema
94+ TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m)
149 {95 {
150- m.impl("mykernel", mykernel_npu);96+ m.def("add(Tensor x, Tensor y) -> Tensor");
151 }97 }
152- ```
153 98 
154-2. 在`mykernel`目录下创建`CMakeLists.txt`99+ /**
155- 100+ * 实现算子的Meta函数，即InferShape+InferDtype
156- 将如下样例中的mykernel，替换为自己的算子名称101+ * 根据输入推导出这个算子的输出是什么样子，需要多少空间，不需要实际计算这个算子
157- ```cmake102+ */
158- message(STATUS "BUILD_TORCH_OPS ON in mykernel")103+ // Meta function implementation of Add
159- # MYKERNEL operation sources104+ torch::Tensor add_meta(const torch::Tensor &x, const torch::Tensor &y)
160- file(GLOB MYKERNEL_NPU_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")105+ {
161- 106+ TORCH_CHECK(x.sizes() == y.sizes(), "The shapes of x and y must be the same.");
162- set(MYKERNEL_SOURCES ${MYKERNEL_NPU_SOURCES})107+ auto z = torch::empty_like(x);
163- # Mark .cpp files with special properties108+ return z;
164- set_source_files_properties(
165- ${MYKERNEL_NPU_SOURCES} PROPERTIES
166- LANGUAGE CXX
167- COMPILE_FLAGS "--cce-soc-version=Ascend910B1 --cce-soc-core-type=CubeCore --cce-auto-sync -xcce"
168- )
169- 
170- # Create object library
171- add_library(mykernel_objects OBJECT ${MYKERNEL_SOURCES})
172- 
173- target_compile_options(mykernel_objects PRIVATE ${COMMON_COMPILE_OPTIONS})
174- target_include_directories(mykernel_objects PRIVATE ${COMMON_INCLUDE_DIRS})
175- return()
176- ```
177- 
178-3. 在 `ascend_ops/csrc/npu_ops_def.cpp`中添加TORCH_LIBRARY_IMPL定义
179- 
180- ```c++
181- TORCH_LIBRARY(ascend_ops, m) {
182- m.def("mykernel(Tensor x) -> Tensor");
183 }109 }
184- ```
185 110 
186-4. (可选)在 `ascend_ops/ops.py`中封装自定义接口111+ /**
187- ```python112+ * 将算子的Meta函数注册给框架
188- def mykernel(x: Tensor) -> Tensor:113+ * 框架可以调用这个Meta函数，在真正执行这个算子计算前知道需要多大空间
189- return torch.ops.ascend_ops.mykernel(x)114+ * 后续可以支持torch.compile/AutoGrad/AclGraph等图加速
190- ```115+ */
116+ // Register the Meta implementation
117+ TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, Meta, m)
118+ {
119+ m.impl("add", add_meta);
120+ }
191 121 
192-5. 使用开发模式进行编译122+ /**
193- ```bash123+ * NPU算子Kernel实现，使用AscendC API，面向当前的soc编写
194- pip install --no-build-isolation -e .124+ */
195- ```125+ template <typename T>
126+ __global__ __aicore__ void add_kernel(GM_ADDR x, GM_ADDR y, GM_ADDR z, int64_t totalLength, int64_t blockLength, uint32_t tileSize)
127+ {
128+ // kernel implementation
129+ }
130+ 
131+ /**
132+ * 实现算子调用接口
133+ * 在这个接口中, 需要完成NPU Kernel的调用
134+ * 1. 计算出输出的Tensor的个数/Shape/Dtype(可以调用Meta函数实现，也可以直接实现)
135+ * 2. 计算Tiling：根据Shape得到如何分块计算
136+ * 3. 调用NPU Kernel
137+ *
138+ */
139+ torch::Tensor add_npu(const torch::Tensor &x, const torch::Tensor &y)
140+ {
141+ auto z = add_meta(x, y);
142+ auto stream = c10_npu::getCurrentNPUStream().stream(false);
143+ int64_t totalLength, blockDim, blockLength, tileSize;
144+ totalLength = x.numel();
145+ std::tie(blockDim, blockLength, tileSize) = calc_tiling_params(totalLength);
146+ auto x_ptr = (GM_ADDR)x.data_ptr();
147+ auto y_ptr = (GM_ADDR)y.data_ptr();
148+ auto z_ptr = (GM_ADDR)z.data_ptr();
149+ auto acl_call = [=]() -> int {
150+ AT_DISPATCH_SWITCH(
151+ x.scalar_type(), "add_npu",
152+ // 根据不同的数据类型，调用不同的NPU Kernel
153+ AT_DISPATCH_CASE(torch::kFloat32, [&] {
154+ using scalar_t = float;
155+ add_kernel<scalar_t><<<blockDim, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize);
156+ })
157+ AT_DISPATCH_CASE(torch::kFloat16, [&] {
158+ using scalar_t = half;
159+ add_kernel<scalar_t><<<blockDim, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize);
160+ })
161+ AT_DISPATCH_CASE(torch::kInt32, [&] {
162+ using scalar_t = int32_t;
163+ add_kernel<scalar_t><<<blockDim, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize);
164+ })
165+ );
166+ return 0;
167+ };
168+ // 需要使用RunOpApi/RunOpApiV2接口调用，保证时序与TorchNPU调用aclnn接口一致。
169+ at_npu::native::OpCommand::RunOpApi("Add", acl_call);
170+ return z;
171+ }
172+ 
173+ /**
174+ * 将算子的调用函数注册给框架，Device为PrivateUse1
175+ * 框架知道当输入均在NPU Device上时，Dispatch到这个算子实现
176+ */
177+ // Register the NPU implementation
178+ TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, PrivateUse1, m)
179+ {
180+ m.impl("add", add_npu);
181+ }
182+ 
183+ } // namespace Add
184+ } // namespace ascend_ops
196 185 
197-6. 编写测试脚本并测试新算子
198- ```python
199- torch.ops.ascend_ops.mykernel(x)
200 ```186 ```
187+4. 使用[安装步骤](#安装步骤--installation-steps)章节构建Wheel包，安装并测试
188+5. 测试算子API请参考[test_add.py](tests/add/test_add.py)的实现
Mexamples/fast_kernel_launch_example/ascend_ops/__init__.py+12-4
@@ -1,5 +1,5 @@
1#!/usr/bin/env python31#!/usr/bin/env python3
2-# -*- coding: UTF-8 -*-2+# -*- coding: utf-8 -*-
CANN-robot
CANN-robotCANN-robot1月4日

代码可读性: 编码声明从'UTF-8'改为'utf-8'，虽然两者在功能上等价，但在Python中推荐使用小写'utf-8'，因为这是Python官方文档和PEP 263中使用的标准形式。这种不一致性可能会影响代码的统一性和可读性。

问题类型: 代码可读性 文件路径: examples/fast_kernel_launch_example/ascend_ops/__init__.py 行号: 2 问题代码:

# -*- coding: utf-8 -*-

修改建议:

建议保持编码声明的一致性，如果项目中其他文件使用小写'utf-8'，则此修改是正确的；如果项目中其他文件使用大写'UTF-8'，则应该统一为大写形式。建议检查项目中其他文件的编码声明，确保一致性。

此评论由代码审查工具自动生成

likedislike
3# ----------------------------------------------------------------------------3# ----------------------------------------------------------------------------
4# Copyright (c) 2025 Huawei Technologies Co., Ltd.4# Copyright (c) 2025 Huawei Technologies Co., Ltd.
5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
@@ -10,7 +10,15 @@
10# See LICENSE in the root of the software repository for the full text of the License.10# See LICENSE in the root of the software repository for the full text of the License.
11# ----------------------------------------------------------------------------11# ----------------------------------------------------------------------------
12 12 
13-from pathlib import Path13+"""Ascend Ops"""
CANN-robot
CANN-robotCANN-robot1月4日

注释与文档: 新增了模块文档字符串'"""Ascend Ops"""'，这是一个好的实践，但文档字符串内容过于简单，没有提供足够的模块信息。对于Ascend Ops这样的重要模块，应该提供更详细的文档说明。

问题类型: 注释与文档 文件路径: examples/fast_kernel_launch_example/ascend_ops/__init__.py 行号: 13 问题代码:

"""Ascend Ops"""

修改建议:

建议扩展模块文档字符串，至少包含以下内容：1) 模块的主要功能和用途；2) 主要提供的类和函数；3) 使用示例或快速入门指南；4) 相关的依赖和安装要求。例如："""Ascend Ops - Huawei Ascend AI处理器的高性能算子库。提供针对Ascend芯片优化的PyTorch算子扩展。"""

此评论由代码审查工具自动生成

likedislike
14+__version__ = "1.0.0"
CANN-robot
CANN-robotCANN-robot1月4日

变量命名: 新增了__version__ = "1.0.0"，这是一个好的实践，但版本号硬编码为字符串字面量。如果项目使用集中化的版本管理（如通过setup.py或pyproject.toml），这种硬编码可能导致版本不一致。

问题类型: 变量命名 文件路径: examples/fast_kernel_launch_example/ascend_ops/__init__.py 行号: 14 问题代码:

__version__ = "1.0.0"

修改建议:

建议考虑从单一源获取版本号，例如：1) 从setup.py或pyproject.toml导入；2) 使用importlib.metadata读取已安装包的版本；3) 或者至少确保此处的版本号与项目其他地方的版本声明保持一致。

此评论由代码审查工具自动生成

likedislike
14import torch15import torch
15-from . import _C16+ 
16-from . import ops17+try:
CANN-robot
CANN-robotCANN-robot1月4日

错误处理与异常管理: 新增了try-except块来捕获_C模块的导入错误，这是一个良好的错误处理实践。但是错误消息'Cannot import _C. Please make sure the ascend_ops is properly installed.'过于通用，没有提供具体的解决建议。

问题类型: 错误处理与异常管理 文件路径: examples/fast_kernel_launch_example/ascend_ops/__init__.py 行号: 17 问题代码:

try:
    from . import _C
except ImportError as e:
    raise ImportError(
        "Cannot import _C. Please make sure the `ascend_ops` is properly installed. "
    ) from e

修改建议:

建议改进错误消息，提供更具体的解决建议，例如：1) 检查是否已正确编译和安装扩展；2) 提供编译和安装的文档链接；3) 建议检查Python版本和PyTorch版本的兼容性。例如："Cannot import _C extension module. This usually means the ascend_ops package was not properly built. Please ensure you have built the extension with 'python setup.py build_ext --inplace' or installed the package via pip."

此评论由代码审查工具自动生成

likedislike
18+ from . import _C
19+except ImportError as e:
20+ raise ImportError(
21+ "Cannot import _C. Please make sure the `ascend_ops` is properly installed. "
22+ ) from e
23+ 
24+from . import ops
Mexamples/fast_kernel_launch_example/ascend_ops/ops.py+3-0
@@ -19,6 +19,9 @@ __all__ = ["conv3d_custom", ]
19 19 
20def conv3d_custom(input: Tensor, weight: Tensor, strides: list, pads: list, dilations: list,20def conv3d_custom(input: Tensor, weight: Tensor, strides: list, pads: list, dilations: list,
21 bias: Tensor = None, enable_hf32: bool = False) -> Tensor:21 bias: Tensor = None, enable_hf32: bool = False) -> Tensor:
22+ print(torch.ops.ascend_ops.conv3d_custom)
23+ assert hasattr(torch.ops.ascend_ops, "conv3d_custom"), "The 'conv3d_custom' operator is not registered in the 'torch.ops.ascend_ops' namespace."
24+
22 origin_input_shape = list(input.shape)25 origin_input_shape = list(input.shape)
23 origin_weight_shape = list(weight.shape)26 origin_weight_shape = list(weight.shape)
24 origin_cout = origin_weight_shape[0]27 origin_cout = origin_weight_shape[0]
Aexamples/fast_kernel_launch_example/build_and_test.sh+28-0
@@ -0,0 +1,28 @@
1+#!/bin/bash
2+# ----------------------------------------------------------------------------
3+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# 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,
8+# INCLUDING 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+set -e
13+cd "$(dirname "$0")"
14+ 
15+# Install dependencies
16+echo "Installing dependencies..."
17+pip install -r requirements.txt
18+ 
19+# Build the project
20+echo "Building the project..."
21+python3 setup.py clean
22+python3 -m build --wheel --no-isolation
23+python3 -m pip install dist/*.whl --force-reinstall --no-deps
24+ 
25+# Run tests
26+echo "Running tests..."
27+pytest tests/ -v
28+echo "execute samples success"
Aexamples/fast_kernel_launch_example/cmake/ascend.cmake+63-0
@@ -0,0 +1,63 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2025 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+# find ascend toolkit
12+if(UNIX)
13+ set(SYSTEM_PREFIX ${CMAKE_SYSTEM_PROCESSOR}-linux)
14+endif()
15+ 
16+if(DEFINED ENV{ASCEND_HOME_PATH})
17+ set(ASCEND_DIR $ENV{ASCEND_HOME_PATH})
18+else()
19+ message(WARNING "Environment variable ASCEND_HOME_PATH is not set. Using default path.")
20+ if("$ENV{USER}" STREQUAL "root")
21+ message(STATUS "Running as root user, checking default root paths for Ascend toolkit.")
22+ if(EXISTS /usr/local/Ascend/ascend-toolkit/latest)
23+ set(ASCEND_DIR /usr/local/Ascend/ascend-toolkit/latest)
24+ elseif(EXISTS /usr/local/Ascend/latest)
25+ set(ASCEND_DIR /usr/local/Ascend/latest)
26+ else()
27+ message(FATAL_ERROR "Ascend toolkit not found in default root paths. Please set ASCEND_HOME_PATH.")
28+ endif()
29+ else()
30+ message(STATUS "Running as non-root user, checking default user paths for Ascend toolkit.")
31+ if(EXISTS $ENV{HOME}/Ascend/ascend-toolkit/latest)
32+ set(ASCEND_DIR $ENV{HOME}/Ascend/ascend-toolkit/latest)
33+ elseif(EXISTS $ENV{HOME}/Ascend/latest)
34+ set(ASCEND_DIR $ENV{HOME}/Ascend/latest)
35+ else()
36+ message(FATAL_ERROR "Ascend toolkit not found in default user paths. Please set ASCEND_HOME_PATH.")
37+ endif()
38+ endif()
39+endif()
40+ 
41+message(STATUS "Using Ascend toolkit path: ${ASCEND_DIR}")
42+set(CMAKE_PREFIX_PATH ${ASCEND_DIR}/)
43+set(BISHENG "${ASCEND_DIR}/${SYSTEM_PREFIX}/ccec_compiler/bin/bisheng" CACHE FILEPATH "Path to Bisheng compiler")
44+message(STATUS "Bisheng compiler path: ${BISHENG}")
45+ 
46+# set the default compiler and linker to bisheng
47+set(CMAKE_C_COMPILER ${BISHENG})
48+set(CMAKE_CXX_COMPILER ${BISHENG})
49+set(CMAKE_LINKER ${BISHENG})
50+ 
51+# set ASCEND_INCLUDE_DIRS
52+set(ASCEND_INCLUDE_DIRS
53+ ${ASCEND_DIR}/include
54+ ${ASCEND_DIR}/compiler/tikcpp/include
55+ ${ASCEND_DIR}/compiler/ascendc/include/basic_api/impl
56+ ${ASCEND_DIR}/compiler/ascendc/include/basic_api/interface
57+ ${ASCEND_DIR}/compiler/ascendc/include/highlevel_api/impl
58+ ${ASCEND_DIR}/compiler/ascendc/include/highlevel_api/tiling
59+ ${ASCEND_DIR}/compiler/ascendc/impl/aicore/basic_api
60+ ${ASCEND_DIR}/pkg_inc/op_common
61+ ${ASCEND_DIR}/pkg_inc/base
62+ ${ASCEND_DIR}/pkg_inc
63+)
Aexamples/fast_kernel_launch_example/cmake/func.cmake+95-0
@@ -0,0 +1,95 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2025 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+# define functions
12+ 
13+# usage: recursive_add_subdirectory()
14+macro(recursive_add_subdirectory)
15+ file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
16+ foreach(SUB_DIR ${CURRENT_DIRS})
17+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/${NPU_ARCH}/CMakeLists.txt")
18+ add_subdirectory(${SUB_DIR}/${NPU_ARCH})
19+ endif()
20+ endforeach()
21+endmacro()
22+ 
23+# usage: add_sources("--npu-arch=dav-3101")
24+# usage: add_sources("--npu-arch=dav-3101" "file1.cpp;file2.cpp;file3.cpp")
25+macro(add_sources ARGS)
26+ # 解析参数
27+ set(COMPILE_ARGS "${ARGS}") # 第一个参数为编译参数
28+ set(CUSTOM_SOURCES "") # 第二个参数为自定义源文件列表（可选）
29+
30+ # 检查是否有第二个参数
31+ if(${ARGC} GREATER 1)
32+ set(CUSTOM_SOURCES "${ARGV1}")
33+ endif()
34+ 
35+ # clear CMAKE_CXX_FLAGS to avoid affecting bisheng compile
36+ unset(CMAKE_CXX_FLAGS)
37+ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
38+ set(CMAKE_CXX_STANDARD 17)
39+ set(CMAKE_CXX_STANDARD_REQUIRED ON)
40+ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
41+ set(CMAKE_C_COMPILER ${BISHENG})
42+ set(CMAKE_CXX_COMPILER ${BISHENG})
43+ set(CMAKE_LINKER ${BISHENG})
44+ 
45+ message(STATUS "CMAKE_CURRENT_SOURCE_DIR = ${CMAKE_CURRENT_SOURCE_DIR}")
46+ 
47+ # get parent dir name as OP_NAME
48+ get_filename_component(PARENT_DIR ${CMAKE_CURRENT_SOURCE_DIR} DIRECTORY)
49+ get_filename_component(OP_NAME ${PARENT_DIR} NAME)
50+ message(STATUS "OP_NAME: ${OP_NAME}")
51+ 
52+ # get compile flags for current op
53+ set(COMPILE_FLAGS "${COMPILE_ARGS} -xasc ")
54+ message(STATUS "COMPILE FLAGS: ${COMPILE_FLAGS}")
55+ 
56+ file(GLOB_RECURSE SOURCE_FILES RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp)
57+ 
58+ # 根据是否传入自定义源文件列表决定如何获取源文件
59+ if(CUSTOM_SOURCES)
60+ foreach(SRC ${CUSTOM_SOURCES})
61+ # 确保文件存在
62+ if(EXISTS ${SRC})
63+ # 获取相对于当前源目录的相对路径
64+ file(RELATIVE_PATH REL_SRC ${CMAKE_CURRENT_SOURCE_DIR} ${SRC})
65+ list(APPEND SOURCE_FILES ${REL_SRC})
66+ else()
67+ message(WARNING "Source file not found: ${SRC}")
68+ endif()
69+ endforeach()
70+ endif()
71+
72+ message(STATUS "SOURCE FILES: ${SOURCE_FILES}")
73+ if(SOURCE_FILES STREQUAL "")
74+ message(FATAL_ERROR "No source files found")
75+ endif()
76+ 
77+ # set_source_files_properties
78+ set_source_files_properties(
79+ ${SOURCE_FILES} PROPERTIES
80+ LANGUAGE CXX
81+ COMPILE_FLAGS "${COMPILE_FLAGS}"
82+ )
83+ 
84+ # set target name
85+ set(TARGET_NAME ${OP_NAME}_obj)
86+ add_library(${TARGET_NAME} OBJECT ${SOURCE_FILES})
87+ target_compile_options(${TARGET_NAME} PRIVATE ${COMPILE_OPTIONS})
88+ target_include_directories(${TARGET_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} ${INCLUDE_DIRECTORIES})
89+ 
90+ # add target obj to OBJECTS_LIST
91+ set(NEW_OBJECT_EXPRESSION $<TARGET_OBJECTS:${TARGET_NAME}>)
92+ set(TEMP_LIST ${OBJECTS_LIST})
93+ list(APPEND TEMP_LIST ${NEW_OBJECT_EXPRESSION})
94+ set(OBJECTS_LIST ${TEMP_LIST} CACHE INTERNAL "List of objects" FORCE)
95+endmacro()
Aexamples/fast_kernel_launch_example/cmake/python.cmake+15-0
@@ -0,0 +1,15 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2025 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+# find python
12+find_package(Python3 REQUIRED COMPONENTS Interpreter Development REQUIRED)
13+message(STATUS "Found Python3: ${Python3_EXECUTABLE} (found version ${Python3_VERSION})")
14+message(STATUS "Python3 include dir: ${Python3_INCLUDE_DIRS}")
15+message(STATUS "Python3 libraries: ${Python3_LIBRARIES}")
Aexamples/fast_kernel_launch_example/cmake/torch.cmake+16-0
@@ -0,0 +1,16 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2025 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+# find pytorch
12+message(STATUS "Using Torch path: ${Torch_DIR}")
13+find_package(Torch REQUIRED)
14+message(STATUS "Found Torch version: ${Torch_VERSION}")
15+message(STATUS "Torch include dirs: ${TORCH_INCLUDE_DIRS}")
16+message(STATUS "Torch libraries: ${TORCH_LIBRARIES}")
Aexamples/fast_kernel_launch_example/cmake/torch_npu.cmake+16-0
@@ -0,0 +1,16 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2025 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+# find torch_npu
12+set(TORCH_NPU_INCLUDE_PATH "${TORCH_NPU_PATH}/include")
13+set(TORCH_NPU_LIB_PATH "${TORCH_NPU_PATH}/lib")
14+message(STATUS "Using Torch NPU path: ${TORCH_NPU_PATH}")
15+message(STATUS "Torch NPU include path: ${TORCH_NPU_INCLUDE_PATH}")
16+message(STATUS "Torch NPU lib path: ${TORCH_NPU_LIB_PATH}")
Aexamples/fast_kernel_launch_example/csrc/CMakeLists.txt+11-0
@@ -0,0 +1,11 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2025 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+recursive_add_subdirectory()
Aexamples/fast_kernel_launch_example/csrc/add/ascend910b/CMakeLists.txt+11-0
@@ -0,0 +1,11 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2025 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_sources("--npu-arch=dav-2201")
Aexamples/fast_kernel_launch_example/csrc/add/ascend910b/add.cpp+186-0
@@ -0,0 +1,186 @@
1+/**
2+ * Copyright (c) 2025 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 add.cpp
13+ * \brief
14+ */
15+ 
16+#include <ATen/Operators.h>
17+#include <torch/all.h>
18+#include <torch/library.h>
19+#include "torch_npu/csrc/core/npu/NPUStream.h"
20+#include "torch_npu/csrc/framework/OpCommand.h"
21+#include "kernel_operator.h"
22+#include "platform/platform_ascendc.h"
23+#include <type_traits>
24+ 
25+namespace ascend_ops {
26+namespace Add {
27+ 
28+// Register the operator's schema
29+TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m)
30+{
31+ m.def("add(Tensor x, Tensor y) -> Tensor");
32+}
33+ 
34+// Meta function implementation of Add
35+torch::Tensor add_meta(const torch::Tensor &x, const torch::Tensor &y)
36+{
37+ TORCH_CHECK(x.sizes() == y.sizes(), "The shapes of x and y must be the same.");
CANN-robot
CANN-robotCANN-robot1月4日

错误处理: add_meta函数中只检查了x和y的sizes是否相等，但没有检查它们的dtype是否兼容。虽然kernel通过AT_DISPATCH_SWITCH处理不同类型，但如果x和y的dtype不同，结果可能不正确。

问题类型: 错误处理 文件路径: examples/fast_kernel_launch_example/csrc/add/ascend910b/add.cpp 行号: 37 问题代码:

TORCH_CHECK(x.sizes() == y.sizes(), "The shapes of x and y must be the same.");

修改建议:

添加dtype检查：`TORCH_CHECK(x.scalar_type() == y.scalar_type(), "The dtypes of x and y must be the same.");` 或者考虑类型提升逻辑。

此评论由代码审查工具自动生成

likedislike
38+ auto z = torch::empty_like(x);
39+ return z;
40+}
41+ 
42+// Register the Meta implementation
43+TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, Meta, m)
44+{
45+ m.impl("add", add_meta);
46+}
47+ 
48+std::tuple<int64_t, int64_t, int64_t> calc_tiling_params(int64_t totalLength)
49+{
50+ constexpr static int64_t MIN_ELEMS_PER_CORE = 1024;
51+ constexpr static int64_t PIPELINE_DEPTH = 2;
52+ constexpr static int64_t BUFFER_NUM = 3;
53+ auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
54+ uint64_t ubSize;
55+ ascendcPlatform->GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
56+ int64_t coreNum = ascendcPlatform->GetCoreNumAiv();
57+ TORCH_CHECK(coreNum > 0, "coreNum must be positive.");
58+ int64_t blockDim = std::min(coreNum, (totalLength + MIN_ELEMS_PER_CORE - 1) / MIN_ELEMS_PER_CORE);
59+ int64_t blockLength = (totalLength + blockDim - 1) / blockDim;
60+ int64_t tileSize = ubSize / PIPELINE_DEPTH / BUFFER_NUM;
61+ return std::make_tuple(blockDim, blockLength, tileSize);
62+}
63+ 
64+template <typename T>
65+__global__ __aicore__ void add_kernel(GM_ADDR x, GM_ADDR y, GM_ADDR z, int64_t totalLength, int64_t blockLength, uint32_t tileSize)
66+{
67+ constexpr static int64_t PIPELINE_DEPTH = 2;
68+ AscendC::TPipe pipe;
69+ AscendC::GlobalTensor<T> xGm, yGm, zGm;
70+ AscendC::TQue<AscendC::QuePosition::VECIN, PIPELINE_DEPTH> inQueueX;
71+ AscendC::TQue<AscendC::QuePosition::VECIN, PIPELINE_DEPTH> inQueueY;
72+ AscendC::TQue<AscendC::QuePosition::VECOUT, PIPELINE_DEPTH> outQueueZ;
73+ pipe.InitBuffer(inQueueX, PIPELINE_DEPTH, tileSize);
74+ pipe.InitBuffer(inQueueY, PIPELINE_DEPTH, tileSize);
75+ pipe.InitBuffer(outQueueZ, PIPELINE_DEPTH, tileSize);
76+ xGm.SetGlobalBuffer((__gm__ T *)x + blockLength * AscendC::GetBlockIdx());
77+ yGm.SetGlobalBuffer((__gm__ T *)y + blockLength * AscendC::GetBlockIdx());
78+ zGm.SetGlobalBuffer((__gm__ T *)z + blockLength * AscendC::GetBlockIdx());
79+ 
80+ int64_t currentBlockLength = totalLength - AscendC::GetBlockIdx() * blockLength;
81+ if (currentBlockLength > blockLength) {
82+ currentBlockLength = blockLength;
83+ }
84+ int64_t elementNumPerTile = tileSize / sizeof(T);
85+ int64_t tileNum = currentBlockLength / elementNumPerTile;
86+ int64_t tailTileElementNum = currentBlockLength - tileNum * elementNumPerTile;
87+ 
88+ for (int64_t i = 0; i < tileNum; ++i) {
89+ int64_t offset = i * elementNumPerTile;
90+ // CopyIn
91+ AscendC::DataCopyExtParams copyParams;
92+ copyParams.blockCount = 1;
93+ copyParams.blockLen = elementNumPerTile * sizeof(T);
94+ copyParams.srcStride = 0;
95+ copyParams.dstStride = 0;
96+ AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
97+ AscendC::LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
98+ AscendC::LocalTensor<T> yLocal = inQueueY.AllocTensor<T>();
99+ AscendC::DataCopyPad(xLocal, xGm[offset], copyParams, padParams);
100+ AscendC::DataCopyPad(yLocal, yGm[offset], copyParams, padParams);
101+ inQueueX.EnQue(xLocal);
102+ inQueueY.EnQue(yLocal);
103+ // Compute
104+ xLocal = inQueueX.DeQue<T>();
105+ yLocal = inQueueY.DeQue<T>();
106+ AscendC::LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();
107+ AscendC::Add(zLocal, xLocal, yLocal, elementNumPerTile);
108+ outQueueZ.EnQue(zLocal);
109+ inQueueX.FreeTensor(xLocal);
110+ inQueueY.FreeTensor(yLocal);
111+ // CopyOut
112+ zLocal = outQueueZ.DeQue<T>();
113+ AscendC::DataCopyPad(zGm[offset], zLocal, copyParams);
114+ outQueueZ.FreeTensor(zLocal);
115+ }
116+ 
117+ if (tailTileElementNum > 0) {
118+ int64_t offset = tileNum * elementNumPerTile;
119+ // CopyIn
120+ AscendC::DataCopyExtParams copyParams;
121+ copyParams.blockCount = 1;
122+ copyParams.blockLen = tailTileElementNum * sizeof(T);
123+ copyParams.srcStride = 0;
124+ copyParams.dstStride = 0;
125+ AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
126+ AscendC::LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
127+ AscendC::LocalTensor<T> yLocal = inQueueY.AllocTensor<T>();
128+ AscendC::DataCopyPad(xLocal, xGm[offset], copyParams, padParams);
129+ AscendC::DataCopyPad(yLocal, yGm[offset], copyParams, padParams);
130+ inQueueX.EnQue(xLocal);
131+ inQueueY.EnQue(yLocal);
132+ // Compute
133+ xLocal = inQueueX.DeQue<T>();
134+ yLocal = inQueueY.DeQue<T>();
135+ AscendC::LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();
136+ AscendC::Add(zLocal, xLocal, yLocal, tailTileElementNum);
137+ outQueueZ.EnQue(zLocal);
138+ inQueueX.FreeTensor(xLocal);
139+ inQueueY.FreeTensor(yLocal);
140+ // CopyOut
141+ zLocal = outQueueZ.DeQue<T>();
142+ AscendC::DataCopyPad(zGm[offset], zLocal, copyParams);
143+ outQueueZ.FreeTensor(zLocal);
144+ }
145+}
146+ 
147+torch::Tensor add_npu(const torch::Tensor &x, const torch::Tensor &y)
148+{
149+ auto z = add_meta(x, y);
150+ auto stream = c10_npu::getCurrentNPUStream().stream(false);
151+ int64_t totalLength, blockDim, blockLength, tileSize;
152+ totalLength = x.numel();
153+ std::tie(blockDim, blockLength, tileSize) = calc_tiling_params(totalLength);
154+ auto x_ptr = (GM_ADDR)x.data_ptr();
155+ auto y_ptr = (GM_ADDR)y.data_ptr();
156+ auto z_ptr = (GM_ADDR)z.data_ptr();
157+ auto acl_call = [=]() -> int {
158+ AT_DISPATCH_SWITCH(
159+ x.scalar_type(), "add_npu",
160+ AT_DISPATCH_CASE(torch::kFloat32, [&] {
161+ using scalar_t = float;
162+ add_kernel<scalar_t><<<blockDim, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize);
163+ })
164+ AT_DISPATCH_CASE(torch::kFloat16, [&] {
165+ using scalar_t = half;
166+ add_kernel<scalar_t><<<blockDim, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize);
167+ })
168+ AT_DISPATCH_CASE(torch::kInt32, [&] {
169+ using scalar_t = int32_t;
170+ add_kernel<scalar_t><<<blockDim, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize);
171+ })
172+ );
173+ return 0;
174+ };
175+ at_npu::native::OpCommand::RunOpApi("Add", acl_call);
176+ return z;
177+}
178+ 
179+// Register the NPU implementation
180+TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, PrivateUse1, m)
181+{
182+ m.impl("add", add_npu);
183+}
184+ 
185+} // namespace Add
186+} // namespace ascend_ops
Rexamples/fast_kernel_launch_example/ascend_ops/csrc/conv3d_custom/CMakeLists.txt→examples/fast_kernel_launch_example/csrc/conv3d_custom/ascend910b/CMakeLists.txt+1-16
@@ -8,9 +8,6 @@
8# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
9# ----------------------------------------------------------------------------9# ----------------------------------------------------------------------------
10 10 
11-message(STATUS "BUILD_TORCH_OPS ON in conv3d_custom")
12-# CONV3D_CUSTOM operation sources
13-file(GLOB CONV3D_CUSTOM_NPU_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
14set(CONV3D_TILING_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../../conv/conv3d_v2/op_host/op_tiling")11set(CONV3D_TILING_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../../conv/conv3d_v2/op_host/op_tiling")
15set(CONV3D_TILING_SOURCES 12set(CONV3D_TILING_SOURCES
16 "${CONV3D_TILING_DIR}/conv3d_api_tiling_algorithm.cpp"13 "${CONV3D_TILING_DIR}/conv3d_api_tiling_algorithm.cpp"
@@ -22,17 +19,5 @@ set(CONV3D_TILING_SOURCES
22 "${CONV3D_TILING_DIR}/conv3d_tiling_engine.cpp"19 "${CONV3D_TILING_DIR}/conv3d_tiling_engine.cpp"
23 "${CONV3D_TILING_DIR}/conv3d_tiling_utils.cpp"20 "${CONV3D_TILING_DIR}/conv3d_tiling_utils.cpp"
24)21)
25-set(CONV3D_CUSTOM_SOURCES ${CONV3D_CUSTOM_NPU_SOURCES} ${CONV3D_TILING_SOURCES})
26-# Mark .cpp files with special properties
27-set_source_files_properties(
28- ${CONV3D_CUSTOM_NPU_SOURCES} PROPERTIES
29- LANGUAGE CXX
30- COMPILE_FLAGS "--cce-soc-version=Ascend910B1 --cce-soc-core-type=CubeCore --cce-auto-sync -xcce"
31-)
32 22 
33-# Create object library23+add_sources("--npu-arch=dav-2201" "${CONV3D_TILING_SOURCES}")
34-add_library(conv3d_custom_objects OBJECT ${CONV3D_CUSTOM_SOURCES})
35- 
36-target_compile_options(conv3d_custom_objects PRIVATE ${COMMON_COMPILE_OPTIONS})
37-target_include_directories(conv3d_custom_objects PRIVATE ${COMMON_INCLUDE_DIRS})
38-return()
Rexamples/fast_kernel_launch_example/ascend_ops/csrc/conv3d_custom/conv3d_custom_torch.cpp→examples/fast_kernel_launch_example/csrc/conv3d_custom/ascend910b/conv3d_custom_torch.cpp+8-1
@@ -411,8 +411,15 @@ torch::Tensor Conv3dCustomNpu(
411 return output;411 return output;
412}412}
413 413 
414+TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m)
415+{
416+ m.def(
417+ "conv3d_custom(Tensor input, Tensor weight, int[3] stride, int[3] padding, int[3] dilation, int[5] "
418+ "oriInputShape, int[5] oriWeightShape, Tensor? bias, bool enable_hf32 = False) -> Tensor");
419+}
420+ 
414// Register Ascend implementations for conv3dv2421// Register Ascend implementations for conv3dv2
415-TORCH_LIBRARY_IMPL(ascend_ops, PrivateUse1, m)422+TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, PrivateUse1, m)
416{423{
417 m.impl("conv3d_custom", Conv3dCustomNpu);424 m.impl("conv3d_custom", Conv3dCustomNpu);
418}425}
Aexamples/fast_kernel_launch_example/csrc/extension.cpp+36-0
@@ -0,0 +1,36 @@
1+/**
2+ * Copyright (c) 2025 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 extension.cpp
13+ * \brief
14+ */
15+ 
16+#include <Python.h>
17+ 
18+extern "C"
19+{
20+ /* Creates a dummy empty _C module that can be imported from Python.
21+ The import from Python will load the .so consisting of this file
22+ in this extension, so that the TORCH_LIBRARY static initializers
23+ below are run. */
24+ PyObject *PyInit__C(void)
25+ {
26+ static struct PyModuleDef module_def = {
27+ PyModuleDef_HEAD_INIT,
28+ "_C", /* name of module */
29+ NULL, /* module documentation, may be NULL */
30+ -1, /* size of per-interpreter state of the module,
31+ or -1 if the module keeps state in global variables. */
32+ NULL, /* methods */
33+ };
34+ return PyModule_Create(&module_def);
35+ }
36+}
Mexamples/fast_kernel_launch_example/requirements.txt+5-7
@@ -1,8 +1,6 @@
1-pyyaml1+--extra-index-url https://download.pytorch.org/whl/cpu
2+ 
2build3build
3-decorator4+pyyaml
4-scipy5+numpy<2
5-attrs6+pytest
6-psutil
7-expecttest
8-wheel
Mexamples/fast_kernel_launch_example/setup.py+108-74
@@ -1,5 +1,5 @@
1#!/usr/bin/env python31#!/usr/bin/env python3
2-# -*- coding: UTF-8 -*-2+# -*- coding: utf-8 -*-
3# ----------------------------------------------------------------------------3# ----------------------------------------------------------------------------
4# Copyright (c) 2025 Huawei Technologies Co., Ltd.4# Copyright (c) 2025 Huawei Technologies Co., Ltd.
5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
@@ -11,18 +11,23 @@
11# ----------------------------------------------------------------------------11# ----------------------------------------------------------------------------
12 12 
13import os13import os
14-import sys
15-import glob
16import shutil14import shutil
17import subprocess15import subprocess
18-import sysconfig16+import logging
19-from pathlib import Path17+from setuptools import setup, find_packages, Distribution, Command
20-from setuptools import setup, find_packages, Extension18+from wheel.bdist_wheel import bdist_wheel
21-from setuptools.command.build_ext import build_ext19+ 
22-from setuptools import Command20+logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
21+PACKAGE_NAME = "ascend_ops"
22+VERSION = "1.0.0"
23+DESCRIPTION = "Example of PyTorch C++ and Ascend extensions"
23 24 
24 25 
25class CleanCommand(Command):26class CleanCommand(Command):
27+ """
28+ usage: python setup.py clean
29+ """
30+ description = "Clean build artifacts from the source tree"
26 user_options = []31 user_options = []
27 32
28 def initialize_options(self):33 def initialize_options(self):
@@ -32,89 +37,118 @@ class CleanCommand(Command):
32 pass37 pass
33 38
34 def run(self):39 def run(self):
35- # 删除构建目录40+ folders_to_remove = ['build', 'dist', f'{PACKAGE_NAME}.egg-info']
36- if os.path.exists('build'):41+ for folder in folders_to_remove:
37- shutil.rmtree('build')42+ if os.path.exists(folder):
38- print("Removed build/")43+ shutil.rmtree(folder)
39- 44+ logging.info(f"Removed folder: {folder}")
40- # 删除dist目录45+ for root, _, files in os.walk('.'):
41- if os.path.exists('dist'):
42- shutil.rmtree('dist')
43- print("Removed dist/")
44-
45- # 删除egg-info目录
46- egg_info_dir = f"{self.distribution.get_name().replace('-', '_')}.egg-info"
47- if os.path.exists(egg_info_dir):
48- shutil.rmtree(egg_info_dir)
49- print(f"Removed {egg_info_dir}/")
50-
51- # 删除.pyc文件和__pycache__目录
52- for root, dirs, files in os.walk('.'):
53 for file in files:46 for file in files:
54- if file.endswith('.pyc'):47+ if file.endswith(('.pyc', '.pyo')):
55- os.remove(os.path.join(root, file))48+ file_path = os.path.join(root, file)
56- print(f"Removed {os.path.join(root, file)}")49+ os.remove(file_path)
57- 50+ logging.info(f"Removed file: {file_path}")
58- for dir in dirs:51+ logging.info("Cleaned build artifacts.")
59- if dir == '__pycache__':
60- shutil.rmtree(os.path.join(root, dir))
61- print(f"Removed {os.path.join(root, dir)}/")
62 52 
63 53 
64-class CMakeExtension(Extension):54+class BinaryDistribution(Distribution):
65- def __init__(self, name, sourcedir=""):55+ """
66- super().__init__(name, sources=[])56+ Make this wheel not a pure python package
67- self.sourcedir = os.path.abspath(sourcedir)57+ """
58+ def is_pure(self):
59+ return False
60+ 
61+ def has_ext_modules(self):
62+ return True
68 63 
69 64 
70-class CMakeBuild(build_ext):65+class ABI3Wheel(bdist_wheel):
66+ """
67+ Force to use abi3 tag for wheel, this wheel supports multiple python versions >= 3.8
68+ """
69+ def get_tag(self):
70+ python, abi, plat = super().get_tag()
71+ python = "cp38"
72+ abi = "abi3"
73+ return python, abi, plat
74+ 
71 def run(self):75 def run(self):
72- try:76+ self.run_command('cmake_build')
73- subprocess.check_output(["cmake", "--version"])77+ super().run()
74- except OSError:
75- raise RuntimeError("CMake must be installed to build the extensions")
76 78 
77- for ext in self.extensions:
78- self.build_cmake(ext)
79 79 
80- def build_cmake(self, ext):80+class CMakeBuildCommand(Command):
81- extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))81+ """
82- python_include = sysconfig.get_path('include')82+ Custom command to build CMake extensions
83- python_libs = sysconfig.get_config_var('LIBDIR')83+ """
84+ description = "Build CMake extensions"
85+ user_options = []
84 86 
85- cmake_args = [87+ def initialize_options(self):
86- f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}",88+ pass
87- f"-DPYTHON_EXECUTABLE={sys.executable}",
88- f"-DBUILD_TORCH_OPS=ON",
89- f"-DPYTHON_EXTENSION_USE_ABI3=ON",
90- f"-DPYTHON_INCLUDE_DIR={python_include}",
91- f"-DPYTHON_LIBRARIES={python_libs}",
92- f"-DPy_LIMITED_API_VERSION=0x03080000",
93 89 
94- ]90+ def finalize_options(self):
91+ pass
95 92 
96- build_type = "Debug" if self.debug else "Release"93+ def run(self):
97- build_args = ["--config", build_type]94+ """
95+ This file `setup.py` and the CMakeLists.txt are in the same directory.
96+ Use multi-core to speed up compilation.
97+ """
98+ cpu_count = os.cpu_count() or 2
99+ num_jobs = str(cpu_count)
100+ # Get Torch and Torch NPU paths
101+ import torch
102+ TORCH_CMAKE_PATH = torch.utils.cmake_prefix_path
103+ Torch_DIR = os.path.join(TORCH_CMAKE_PATH, "Torch")
104+ logging.info(f"Using Torch path: {Torch_DIR}")
105+ import torch_npu
106+ TORCH_NPU_PATH = os.path.dirname(torch_npu.__file__)
107+ logging.info(f"Using Torch NPU path: {TORCH_NPU_PATH}")
98 108 
99- cpu_count = os.cpu_count() or 1109+ # Get NPU_ARCH from environment variable or set default
100- parallel_jobs = max(16, cpu_count // 2)110+ NPU_ARCH = os.environ.get('NPU_ARCH', 'ascend910b')
101- build_args += ["--", f"-j{parallel_jobs}"]111+ logging.info(f"Using NPU_ARCH: {NPU_ARCH}")
102 112 
103- build_temp = Path(self.build_temp) / ext.name113+ # Build the CMake project
104- build_temp.mkdir(parents=True, exist_ok=True)114+ build_temp = os.path.join(os.getcwd(), 'build')
115+ cmake_config_command = ['cmake', '-S', os.getcwd(), '-B', build_temp,
116+ '-DCMAKE_BUILD_TYPE=Release',
117+ f'-DTorch_DIR={Torch_DIR}',
118+ f'-DTORCH_NPU_PATH={TORCH_NPU_PATH}',
119+ f'-DNPU_ARCH={NPU_ARCH}'
120+ ]
121+ subprocess.check_call(cmake_config_command, cwd=os.getcwd())
122+ subprocess.check_call(['cmake', '--build', build_temp, '--config', 'Release', '--parallel', num_jobs], cwd=os.getcwd())
123+ logging.info("CMake extensions built successfully.")
105 124 
106- subprocess.check_call(["cmake", ext.sourcedir] + cmake_args, cwd=build_temp)125+ 
107- subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=build_temp)126+cmdclass = {
127+ 'clean': CleanCommand,
128+ 'bdist_wheel': ABI3Wheel,
129+ 'cmake_build': CMakeBuildCommand,
130+}
108 131 
109 132 
110setup(133setup(
111- name='ascend_ops',134+ name=PACKAGE_NAME,
112- version='0.0.1',135+ version=VERSION,
136+ description=DESCRIPTION,
113 packages=find_packages(),137 packages=find_packages(),
114- ext_modules=[CMakeExtension("ascend_ops._C", sourcedir=".")],138+ package_data={PACKAGE_NAME: ['*.so']},
115- cmdclass={'build_ext': CMakeBuild, 'clean': CleanCommand},139+ distclass=BinaryDistribution,
140+ cmdclass=cmdclass,
116 zip_safe=False,141 zip_safe=False,
117- install_requires=["torch"],142+ install_requires=[
118- options={"bdist_wheel": {"py_limited_api": "cp38"}},143+ "torch",
119- description="Example of PyTorch C++ and Ascend extensions (with CMake)",144+ "torch_npu"
145+ ],
146+ classifiers=[
147+ "Programming Language :: Python :: 3.8",
148+ "Programming Language :: Python :: 3.9",
149+ "Programming Language :: Python :: 3.10",
150+ "Programming Language :: Python :: 3.11",
151+ "Operating System :: POSIX :: Linux",
152+ ],
153+ python_requires='>=3.8',
120)154)
Aexamples/fast_kernel_launch_example/tests/add/test_add.py+105-0
@@ -0,0 +1,105 @@
1+#!/usr/bin/env python3
2+# -*- coding: utf-8 -*-
3+# ----------------------------------------------------------------------------
4+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# ----------------------------------------------------------------------------
12+ 
13+import torch
14+import torch_npu
15+import ascend_ops
16+import pytest
17+ 
18+ 
19+def test_add_interface_exist():
20+ """
21+ Test that the 'ascend_ops.add' operator is present in torch.ops.
22+ This existence test asserts that the custom operator registered under the
23+ 'ascend_ops' namespace is discoverable from Python via torch.ops.ascend_ops.add.
24+ It does not exercise operator functionality — only that the Python binding
25+ and registration are available.
26+ Rationale:
27+ The presence of this test guards against a common failure mode where an
28+ operator is implemented and registered in C++/ATen but is not exposed to
29+ the Python torch.ops namespace due to mismatches between the PyTorch
30+ operator schema and the C++ registration signature (argument names, types,
31+ or overloads). Such schema/signature inconsistencies can cause the
32+ operator to be hidden or not exported to Python, breaking consumers that
33+ expect to call torch.ops.ascend_ops.add. This test will fail loudly if the
34+ binding is missing, prompting investigation into schema/registration issues.
35+ """
36+ # This test specifically protects against discrepancies between the
37+ # PyTorch operator schema and the C++ signature/registration that can
38+ # prevent the operator from being visible in torch.ops.ascend_ops.
39+ print(torch.ops.ascend_ops.add)
40+ assert hasattr(torch.ops.ascend_ops, "add"), "The 'add' operator is not registered in the 'torch.ops.ascend_ops' namespace."
41+ 
42+ 
43+SHAPES = [
44+ (1,),
45+ (3,),
46+ (10,),
47+ (100,),
48+ (1024,),
49+ (10000,),
50+ (10, 10),
51+ (32, 32),
52+ (100, 100),
53+ (10, 100),
54+ (100, 10),
55+ (256, 512),
56+ (5, 10, 15),
57+ (16, 32, 64),
58+ (32, 64, 128),
59+ (1, 3, 32, 32),
60+ (4, 3, 64, 64),
61+ (8, 3, 128, 128),
62+ (1000, 1000),
63+]
64+ 
65+DTYPES = [
66+ torch.float32,
67+ torch.float16,
68+ torch.int32,
69+]
70+ 
71+ 
72+@pytest.mark.skipif(not torch.npu.is_available(), reason="NPU device not found")
73+@pytest.mark.parametrize("shape", SHAPES)
74+@pytest.mark.parametrize("dtype", DTYPES)
75+def test_add_operator(shape, dtype):
76+ """
77+ Test the functionality of the add operator, using concise but comprehensive combinations of shapes and data types.
78+ 
79+ Parameters:
80+ shape: Tensor shape
81+ dtype: Data type
82+ """
83+ if dtype in [torch.int32]:
84+ a = torch.randint(-100, 100, shape, dtype=dtype)
85+ b = torch.randint(-100, 100, shape, dtype=dtype)
86+ else:
87+ a = torch.randn(*shape, dtype=dtype)
88+ b = torch.randn(*shape, dtype=dtype)
89+ 
90+ expected = a + b
91+ a_npu = a.npu()
92+ b_npu = b.npu()
93+ result_npu = torch.ops.ascend_ops.add(a_npu, b_npu)
94+ result = result_npu.cpu()
95+ 
96+ if dtype in [torch.int32]:
97+ assert torch.equal(result, expected), \
98+ f"Add failed for shape {shape}, dtype {dtype}. " \
99+ f"Expected {expected}, but got {result}"
100+ else:
101+ assert torch.allclose(result, expected, rtol=1e-4, atol=1e-4), \
102+ f"Add failed for shape {shape}, dtype {dtype}. " \
103+ f"Max diff: {torch.max(torch.abs(result - expected)):.6f}"
104+ 
105+ print(f"✓ Test passed: shape={shape}, dtype={dtype}")
Rexamples/fast_kernel_launch_example/ascend_ops/csrc/conv3d_custom/test/test_conv3d_custom.py→examples/fast_kernel_launch_example/tests/conv3d_custom/test_conv3d_custom.py+0-0
文件重命名但无更改。