前言用 Python 写算子比用 Ascend C 写要快很多pyasc 就是 Python 算子开发工具链。这篇文章从环境安装到第一个自定义算子手把手走一遍。1. pyasc 工具链安装和配置1.1 环境准备pyasc 是 CANNCompute Architecture for Neural Networks提供的 Python 算子开发工具让开发者可以用 Python 语法定义算子逻辑工具链自动将其转换为高效的 Ascend C 代码并编译为算子库。系统要求Ubuntu 18.04/20.04/22.04 或 CentOS 7.x/8.xPython 3.7CANN 社区版 6.0.RC1 或更高版本Ascend 910/310/610 系列处理器安装步骤# 1. 安装 CANN 工具包以 Ubuntu 20.04 为例wgethttps://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/6.0.RC1/Ascend-cann-toolkit_6.0.RC1_linux-x86_64.runchmodx Ascend-cann-toolkit_6.0.RC1_linux-x86_64.run ./Ascend-cann-toolkit_6.0.RC1_linux-x86_64.run--install# 2. 设置环境变量source/usr/local/Ascend/ascend-toolkit/set_env.sh# 3. 安装 pyascpip3installpyasc-ihttps://pypi.tuna.tsinghua.edu.cn/simple1.2 验证安装安装完成后验证 pyasc 是否正常工作importpyascprint(pyasc.__version__)# 验证工具链frompyascimportOperator,compile_opprint(pyasc 工具链就绪)预期输出0.2.0 pyasc 工具链就绪2. Python 算子定义语法和接口2.1 算子定义基础pyasc 提供了一套类 Python 语法用于定义算子。核心是通过Operator装饰器和特定的 API 来描述算子的计算逻辑。基本结构frompyascimportOperator,Tensor,DataTypeOperator.register(CustomAdd)classCustomAddOp:自定义加法算子def__init__(self):self.nameCustomAdddefforward(self,input1:Tensor,input2:Tensor)-Tensor: 前向计算逻辑 Args: input1: 输入张量1 input2: 输入张量2 Returns: 输出张量 # 定义输出张量outputTensor(shapeinput1.shape,dtypeinput1.dtype)# 算子计算逻辑Python 语法output[:]input1input2returnoutputdefinfer_shape(self,input1_shape,input2_shape):形状推理assertinput1_shapeinput2_shape,输入形状必须一致returninput1_shapedefinfer_dtype(self,input1_dtype,input2_dtype):数据类型推理assertinput1_dtypeinput2_dtype,输入数据类型必须一致returninput1_dtype2.2 关键 API 说明API说明示例Tensor张量定义Tensor(shape(32, 32), dtypeDataType.FLOAT32)DataType数据类型DataType.FLOAT32,DataType.INT32Operator.register算子注册装饰器Operator.register(MyOp)Operator.compile编译算子Operator.compile(op_instance)2.3 完整算子示例下面是一个完整的自定义 ReLU 算子定义frompyascimportOperator,Tensor,DataType,compile_opOperator.register(CustomReLU)classCustomReLUOp: 自定义 ReLU 激活函数算子 ReLU(x) max(0, x) def__init__(self,inplaceFalse):self.inplaceinplace self.nameCustomReLUdefforward(self,x:Tensor)-Tensor: ReLU 前向计算 Args: x: 输入张量 Returns: 激活后的张量 ifself.inplace:# 原地操作x[x0]0returnxelse:# 创建新张量outputTensor(shapex.shape,dtypex.dtype)output[:]x.maximum(0)returnoutputdefbackward(self,grad_output:Tensor,x:Tensor)-Tensor: ReLU 反向传播 Args: grad_output: 输出梯度 x: 前向输入 Returns: 输入梯度 grad_inputTensor(shapex.shape,dtypex.dtype)grad_input[:]grad_output*(x0)returngrad_inputdefinfer_shape(self,x_shape):形状推理输出形状等于输入形状returnx_shapedefinfer_dtype(self,x_dtype):数据类型推理输出类型等于输入类型returnx_dtypedefget_attrs(self):获取算子属性return{inplace:self.inplace}# 实例化算子relu_opCustomReLUOp(inplaceFalse)# 编译算子compiled_opcompile_op(relu_op)print(f算子{compiled_op.name}编译成功)3. 算子编译Python → Ascend C → 算子库3.1 编译流程概述pyasc 的编译流程分为三个阶段Python 解析解析 Python 算子定义生成中间表示IR代码生成将 IR 转换为 Ascend C 代码编译链接调用 CANN 编译器将 Ascend C 代码编译为算子库3.2 编译配置通过compile_op函数的参数控制编译行为frompyascimportcompile_op,CompileConfig# 创建编译配置configCompileConfig(targetascend910,# 目标芯片debugTrue,# 启用调试信息optimize_level2,# 优化等级0-3output_dir./build,# 输出目录gen_ascend_cTrue,# 生成 Ascend C 代码verboseTrue# 详细输出)# 编译算子compiled_opcompile_op(relu_op,configconfig)3.3 查看生成的 Ascend C 代码编译后可以查看生成的 Ascend C 代码# 获取生成的 Ascend C 代码路径ascend_c_code_pathcompiled_op.ascend_c_pathprint(fAscend C 代码路径:{ascend_c_code_path})# 查看生成的代码withopen(ascend_c_code_path,r)asf:ascend_c_codef.read()print(生成的 Ascend C 代码)print(ascend_c_code[:500])# 打印前500字符生成的 Ascend C 代码结构// 自动生成的 Ascend C 代码 - CustomReLU#includekernel_operator.husingnamespaceAscendC;classCustomReLU{public:__aicore__staticvoidforward(GM_ADDR x,GM_ADDR y,constCustomReLUAttrsattrs){// 初始化KernelRelufloatop;LocalTensorfloatxLocalop.AllocTensorfloat();LocalTensorfloatyLocalop.AllocTensorfloat();// 数据搬运op.CopyIn(x,xLocal);// 计算op.Forward(yLocal,xLocal);// 写回op.CopyOut(y,yLocal);}};3.4 编译为算子库编译完成后会生成libcustom_relu.so算子动态库custom_relu.json算子描述文件custom_relu.h算子头文件供 C 调用# 查看编译产物ls-lh./build/# 输出# -rw-r--r-- 1 user user 12K May 28 22:00 custom_relu.json# -rw-r--r-- 1 user user 45K May 28 22:00 libcustom_relu.so# -rw-r--r-- 1 user user 2.3K May 28 22:00 custom_relu.h4. 注册与调用在模型中使用自定义算子4.1 算子注册编译后的算子需要注册到 CANN 算子库才能被模型调用frompyascimportOpRegistry# 注册算子到 CANNregistryOpRegistry()# 注册自定义 ReLU 算子registry.register(op_nameCustomReLU,op_lib_path./build/libcustom_relu.so,op_desc_path./build/custom_relu.json,frameworkpytorch# 或 tensorflow, mindspore)print(算子注册成功)4.2 在 PyTorch 模型中使用注册后可以在 PyTorch 模型中直接调用importtorchimporttorch_npu# 华为 NPU 后端frompyasc.torchimportCustomReLU# 创建模型classMyModel(torch.nn.Module):def__init__(self):super().__init__()self.convtorch.nn.Conv2d(3,64,kernel_size3)self.reluCustomReLU(inplaceFalse)# 使用自定义 ReLUself.fctorch.nn.Linear(64*222*222,1000)defforward(self,x):xself.conv(x)xself.relu(x)# 调用自定义算子xx.view(x.size(0),-1)xself.fc(x)returnx# 实例化模型modelMyModel()modelmodel.to(npu:0)# 移动到 NPU# 测试input_tensortorch.randn(1,3,224,224).to(npu:0)outputmodel(input_tensor)print(f输出形状:{output.shape})4.3 性能对比自定义算子 vs 原生算子性能测试importtimedefbenchmark(op,input_tensor,num_iterations100):性能基准测试# 预热for_inrange(10):op(input_tensor)# 正式测试torch.npu.synchronize()starttime.time()for_inrange(num_iterations):outputop(input_tensor)torch.npu.synchronize()endtime.time()avg_time(end-start)/num_iterations*1000# msreturnavg_time# 创建测试输入test_inputtorch.randn(32,64,112,112).to(npu:0)# 测试自定义 ReLUcustom_reluCustomReLU().to(npu:0)custom_timebenchmark(custom_relu,test_input)print(f自定义 ReLU 平均耗时:{custom_time:.2f}ms)# 测试原生 ReLUnative_relutorch.nn.ReLU().to(npu:0)native_timebenchmark(native_relu,test_input)print(f原生 ReLU 平均耗时:{native_time:.2f}ms)print(f性能比:{custom_time/native_time:.2f}x)5. 调试Python 级别 vs C 级别调试5.1 Python 级别调试pyasc 提供了 Python 级别的调试工具可以在不生成 Ascend C 代码的情况下验证算子逻辑。使用debug模式frompyascimportcompile_op,CompileConfig# 启用 Python 调试模式configCompileConfig(debugTrue,python_onlyTrue,# 只运行 Python 逻辑不生成 C 代码validate_resultsTrue# 验证结果正确性)# 编译仅 Pythondebug_opcompile_op(relu_op,configconfig)# 测试数据test_inputtorch.randn(10,10).to(npu:0)test_outputdebug_op.forward(test_input)# 验证expectedtorch.relu(test_input)asserttorch.allclose(test_output,expected,atol1e-5),结果不正确print(Python 级别调试通过)5.2 C 级别调试当需要深入底层优化时可以进行 C 级别调试。生成调试版本的算子库# 生成带调试信息的算子库configCompileConfig(targetascend910,debugTrue,debug_level2,# 详细调试信息optimize_level0,# 关闭优化以方便调试output_dir./build_debug)debug_libcompile_op(relu_op,configconfig)使用 GDB 调试# 1. 安装 GDB for NPUsudoaptinstallgdb-aarch64-linux-gnu# 2. 启动调试gdb--argspython test_custom_op.py# 3. 设置断点(gdb)breakCustomReLU::forward(gdb)run# 4. 查看变量(gdb)print x.shape(gdb)print *x.data5.3 常见问题排查问题原因解决方案编译失败CANN 环境未正确配置检查source set_env.sh是否执行算子运行报错形状推理错误实现infer_shape和infer_dtype性能不如预期未启用优化设置optimize_level2或3NPU 内存溢出张量过大使用inplaceTrue减少内存5.4 调试技巧技巧 1逐步验证# 分步验证算子逻辑opCustomReLUOp()# 1. 验证形状推理shapeop.infer_shape((32,32))print(f形状推理结果:{shape})# 2. 验证数据类型推理dtypeop.infer_dtype(DataType.FLOAT32)print(f数据类型推理结果:{dtype})# 3. 小数据集测试small_inputTensor(shape(2,2),dtypeDataType.FLOAT32,data[-1,2,-3,4])outputop.forward(small_input)print(f小数据集测试结果:{output.data})# 应为 [0, 2, 0, 4]技巧 2对比测试defcompare_with_native(pyasc_op,native_op,test_input):对比自定义算子和原生算子结果pyasc_outputpyasc_op.forward(test_input)native_outputnative_op(test_input)# 数值对比max_diff(pyasc_output-native_output).abs().max()print(f最大差异:{max_diff})# 精度对比ifmax_diff1e-5:print(✓ 精度符合要求)else:print(✗ 精度不符合要求)returnmax_diff总结pyasc 工具链大幅降低了 Ascend 算子开发门槛让熟悉 Python 的开发者可以快速上手。关键要点安装简单pip 一键安装依赖 CANN 工具包语法友好类 Python 语法无需学习 Ascend C调试方便支持 Python 和 C 双级别调试性能可观自动优化接近手写 Ascend C 性能完整的代码示例和文档可以在仓库中找到https://atomgit.com/cann/pyasc无论是自定义激活函数、损失函数还是实现研究中的新算子pyasc 都是一个高效的选择。开始用 Python 写你的第一个 Ascend 算子吧
CANN pyasc 工具——Python 接口的算子开发
发布时间：2026/5/28 23:26:53
前言用 Python 写算子比用 Ascend C 写要快很多pyasc 就是 Python 算子开发工具链。这篇文章从环境安装到第一个自定义算子手把手走一遍。1. pyasc 工具链安装和配置1.1 环境准备pyasc 是 CANNCompute Architecture for Neural Networks提供的 Python 算子开发工具让开发者可以用 Python 语法定义算子逻辑工具链自动将其转换为高效的 Ascend C 代码并编译为算子库。系统要求Ubuntu 18.04/20.04/22.04 或 CentOS 7.x/8.xPython 3.7CANN 社区版 6.0.RC1 或更高版本Ascend 910/310/610 系列处理器安装步骤# 1. 安装 CANN 工具包以 Ubuntu 20.04 为例wgethttps://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/6.0.RC1/Ascend-cann-toolkit_6.0.RC1_linux-x86_64.runchmodx Ascend-cann-toolkit_6.0.RC1_linux-x86_64.run ./Ascend-cann-toolkit_6.0.RC1_linux-x86_64.run--install# 2. 设置环境变量source/usr/local/Ascend/ascend-toolkit/set_env.sh# 3. 安装 pyascpip3installpyasc-ihttps://pypi.tuna.tsinghua.edu.cn/simple1.2 验证安装安装完成后验证 pyasc 是否正常工作importpyascprint(pyasc.__version__)# 验证工具链frompyascimportOperator,compile_opprint(pyasc 工具链就绪)预期输出0.2.0 pyasc 工具链就绪2. Python 算子定义语法和接口2.1 算子定义基础pyasc 提供了一套类 Python 语法用于定义算子。核心是通过Operator装饰器和特定的 API 来描述算子的计算逻辑。基本结构frompyascimportOperator,Tensor,DataTypeOperator.register(CustomAdd)classCustomAddOp:自定义加法算子def__init__(self):self.nameCustomAdddefforward(self,input1:Tensor,input2:Tensor)-Tensor: 前向计算逻辑 Args: input1: 输入张量1 input2: 输入张量2 Returns: 输出张量 # 定义输出张量outputTensor(shapeinput1.shape,dtypeinput1.dtype)# 算子计算逻辑Python 语法output[:]input1input2returnoutputdefinfer_shape(self,input1_shape,input2_shape):形状推理assertinput1_shapeinput2_shape,输入形状必须一致returninput1_shapedefinfer_dtype(self,input1_dtype,input2_dtype):数据类型推理assertinput1_dtypeinput2_dtype,输入数据类型必须一致returninput1_dtype2.2 关键 API 说明API说明示例Tensor张量定义Tensor(shape(32, 32), dtypeDataType.FLOAT32)DataType数据类型DataType.FLOAT32,DataType.INT32Operator.register算子注册装饰器Operator.register(MyOp)Operator.compile编译算子Operator.compile(op_instance)2.3 完整算子示例下面是一个完整的自定义 ReLU 算子定义frompyascimportOperator,Tensor,DataType,compile_opOperator.register(CustomReLU)classCustomReLUOp: 自定义 ReLU 激活函数算子 ReLU(x) max(0, x) def__init__(self,inplaceFalse):self.inplaceinplace self.nameCustomReLUdefforward(self,x:Tensor)-Tensor: ReLU 前向计算 Args: x: 输入张量 Returns: 激活后的张量 ifself.inplace:# 原地操作x[x0]0returnxelse:# 创建新张量outputTensor(shapex.shape,dtypex.dtype)output[:]x.maximum(0)returnoutputdefbackward(self,grad_output:Tensor,x:Tensor)-Tensor: ReLU 反向传播 Args: grad_output: 输出梯度 x: 前向输入 Returns: 输入梯度 grad_inputTensor(shapex.shape,dtypex.dtype)grad_input[:]grad_output*(x0)returngrad_inputdefinfer_shape(self,x_shape):形状推理输出形状等于输入形状returnx_shapedefinfer_dtype(self,x_dtype):数据类型推理输出类型等于输入类型returnx_dtypedefget_attrs(self):获取算子属性return{inplace:self.inplace}# 实例化算子relu_opCustomReLUOp(inplaceFalse)# 编译算子compiled_opcompile_op(relu_op)print(f算子{compiled_op.name}编译成功)3. 算子编译Python → Ascend C → 算子库3.1 编译流程概述pyasc 的编译流程分为三个阶段Python 解析解析 Python 算子定义生成中间表示IR代码生成将 IR 转换为 Ascend C 代码编译链接调用 CANN 编译器将 Ascend C 代码编译为算子库3.2 编译配置通过compile_op函数的参数控制编译行为frompyascimportcompile_op,CompileConfig# 创建编译配置configCompileConfig(targetascend910,# 目标芯片debugTrue,# 启用调试信息optimize_level2,# 优化等级0-3output_dir./build,# 输出目录gen_ascend_cTrue,# 生成 Ascend C 代码verboseTrue# 详细输出)# 编译算子compiled_opcompile_op(relu_op,configconfig)3.3 查看生成的 Ascend C 代码编译后可以查看生成的 Ascend C 代码# 获取生成的 Ascend C 代码路径ascend_c_code_pathcompiled_op.ascend_c_pathprint(fAscend C 代码路径:{ascend_c_code_path})# 查看生成的代码withopen(ascend_c_code_path,r)asf:ascend_c_codef.read()print(生成的 Ascend C 代码)print(ascend_c_code[:500])# 打印前500字符生成的 Ascend C 代码结构// 自动生成的 Ascend C 代码 - CustomReLU#includekernel_operator.husingnamespaceAscendC;classCustomReLU{public:__aicore__staticvoidforward(GM_ADDR x,GM_ADDR y,constCustomReLUAttrsattrs){// 初始化KernelRelufloatop;LocalTensorfloatxLocalop.AllocTensorfloat();LocalTensorfloatyLocalop.AllocTensorfloat();// 数据搬运op.CopyIn(x,xLocal);// 计算op.Forward(yLocal,xLocal);// 写回op.CopyOut(y,yLocal);}};3.4 编译为算子库编译完成后会生成libcustom_relu.so算子动态库custom_relu.json算子描述文件custom_relu.h算子头文件供 C 调用# 查看编译产物ls-lh./build/# 输出# -rw-r--r-- 1 user user 12K May 28 22:00 custom_relu.json# -rw-r--r-- 1 user user 45K May 28 22:00 libcustom_relu.so# -rw-r--r-- 1 user user 2.3K May 28 22:00 custom_relu.h4. 注册与调用在模型中使用自定义算子4.1 算子注册编译后的算子需要注册到 CANN 算子库才能被模型调用frompyascimportOpRegistry# 注册算子到 CANNregistryOpRegistry()# 注册自定义 ReLU 算子registry.register(op_nameCustomReLU,op_lib_path./build/libcustom_relu.so,op_desc_path./build/custom_relu.json,frameworkpytorch# 或 tensorflow, mindspore)print(算子注册成功)4.2 在 PyTorch 模型中使用注册后可以在 PyTorch 模型中直接调用importtorchimporttorch_npu# 华为 NPU 后端frompyasc.torchimportCustomReLU# 创建模型classMyModel(torch.nn.Module):def__init__(self):super().__init__()self.convtorch.nn.Conv2d(3,64,kernel_size3)self.reluCustomReLU(inplaceFalse)# 使用自定义 ReLUself.fctorch.nn.Linear(64*222*222,1000)defforward(self,x):xself.conv(x)xself.relu(x)# 调用自定义算子xx.view(x.size(0),-1)xself.fc(x)returnx# 实例化模型modelMyModel()modelmodel.to(npu:0)# 移动到 NPU# 测试input_tensortorch.randn(1,3,224,224).to(npu:0)outputmodel(input_tensor)print(f输出形状:{output.shape})4.3 性能对比自定义算子 vs 原生算子性能测试importtimedefbenchmark(op,input_tensor,num_iterations100):性能基准测试# 预热for_inrange(10):op(input_tensor)# 正式测试torch.npu.synchronize()starttime.time()for_inrange(num_iterations):outputop(input_tensor)torch.npu.synchronize()endtime.time()avg_time(end-start)/num_iterations*1000# msreturnavg_time# 创建测试输入test_inputtorch.randn(32,64,112,112).to(npu:0)# 测试自定义 ReLUcustom_reluCustomReLU().to(npu:0)custom_timebenchmark(custom_relu,test_input)print(f自定义 ReLU 平均耗时:{custom_time:.2f}ms)# 测试原生 ReLUnative_relutorch.nn.ReLU().to(npu:0)native_timebenchmark(native_relu,test_input)print(f原生 ReLU 平均耗时:{native_time:.2f}ms)print(f性能比:{custom_time/native_time:.2f}x)5. 调试Python 级别 vs C 级别调试5.1 Python 级别调试pyasc 提供了 Python 级别的调试工具可以在不生成 Ascend C 代码的情况下验证算子逻辑。使用debug模式frompyascimportcompile_op,CompileConfig# 启用 Python 调试模式configCompileConfig(debugTrue,python_onlyTrue,# 只运行 Python 逻辑不生成 C 代码validate_resultsTrue# 验证结果正确性)# 编译仅 Pythondebug_opcompile_op(relu_op,configconfig)# 测试数据test_inputtorch.randn(10,10).to(npu:0)test_outputdebug_op.forward(test_input)# 验证expectedtorch.relu(test_input)asserttorch.allclose(test_output,expected,atol1e-5),结果不正确print(Python 级别调试通过)5.2 C 级别调试当需要深入底层优化时可以进行 C 级别调试。生成调试版本的算子库# 生成带调试信息的算子库configCompileConfig(targetascend910,debugTrue,debug_level2,# 详细调试信息optimize_level0,# 关闭优化以方便调试output_dir./build_debug)debug_libcompile_op(relu_op,configconfig)使用 GDB 调试# 1. 安装 GDB for NPUsudoaptinstallgdb-aarch64-linux-gnu# 2. 启动调试gdb--argspython test_custom_op.py# 3. 设置断点(gdb)breakCustomReLU::forward(gdb)run# 4. 查看变量(gdb)print x.shape(gdb)print *x.data5.3 常见问题排查问题原因解决方案编译失败CANN 环境未正确配置检查source set_env.sh是否执行算子运行报错形状推理错误实现infer_shape和infer_dtype性能不如预期未启用优化设置optimize_level2或3NPU 内存溢出张量过大使用inplaceTrue减少内存5.4 调试技巧技巧 1逐步验证# 分步验证算子逻辑opCustomReLUOp()# 1. 验证形状推理shapeop.infer_shape((32,32))print(f形状推理结果:{shape})# 2. 验证数据类型推理dtypeop.infer_dtype(DataType.FLOAT32)print(f数据类型推理结果:{dtype})# 3. 小数据集测试small_inputTensor(shape(2,2),dtypeDataType.FLOAT32,data[-1,2,-3,4])outputop.forward(small_input)print(f小数据集测试结果:{output.data})# 应为 [0, 2, 0, 4]技巧 2对比测试defcompare_with_native(pyasc_op,native_op,test_input):对比自定义算子和原生算子结果pyasc_outputpyasc_op.forward(test_input)native_outputnative_op(test_input)# 数值对比max_diff(pyasc_output-native_output).abs().max()print(f最大差异:{max_diff})# 精度对比ifmax_diff1e-5:print(✓ 精度符合要求)else:print(✗ 精度不符合要求)returnmax_diff总结pyasc 工具链大幅降低了 Ascend 算子开发门槛让熟悉 Python 的开发者可以快速上手。关键要点安装简单pip 一键安装依赖 CANN 工具包语法友好类 Python 语法无需学习 Ascend C调试方便支持 Python 和 C 双级别调试性能可观自动优化接近手写 Ascend C 性能完整的代码示例和文档可以在仓库中找到https://atomgit.com/cann/pyasc无论是自定义激活函数、损失函数还是实现研究中的新算子pyasc 都是一个高效的选择。开始用 Python 写你的第一个 Ascend 算子吧
