CANN pyasc 算子开发指南用原生 Python 从零实现并验证一个 Ascend C Add 算子【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc本篇指南以 pyasc 开源项目中的 Addz x y算子为例完整讲解基于 Python 原生语法开发昇腾 Ascend C 算子的全流程从环境准备、算子分析数学表达式、输入输出规格、接口选型到asc.jit核函数与 Launch 函数的实现再到仿真器/上板两种模式的编译运行与结果验证。读完本文你将掌握 pyasc 多核并行、双缓冲流水与set_flag/wait_flag手动同步的完整写法并能照此套路开发自己的自定义算子。总体开发流程在正式开发前需要先完成环境准备工作。pyasc 算子开发的基本流程如下其中每一步对应本文的一个章节环境准备安装 pyasc 与 CANN 软件→算子分析确定数学表达式、输入输出规格与所需接口→核函数开发编写asc.jit核函数与 Launch 函数→核函数运行验证编写完整验证程序并编译运行。环境准备开发 pyasc 算子前需要同时就绪两套环境pyasc 安装pyasc 支持通过 pip 快速安装和基于源码编译安装两种方式。快速安装直接执行pip install pyasc即可获得最新稳定版二进制 wheel 支持 CPython 3.9~3.12基于源码安装则需要先下载并安装 LLVM 19.1.7 预编译包、设置LLVM_INSTALL_PREFIX环境变量再执行python3 -m pip install .普通模式或python3 -m pip install -e .开发者模式。详细步骤请参考 pyasc 快速入门-环境准备与安装。CANN 软件安装开发算子前需安装 CANN 软件CANN toolkit 包与 CANN ops 包具体请参考 pyasc 快速入门-下载安装 CANN 包。安装 CANN 软件后需要配置运行环境变量默认路径安装时执行source /usr/local/Ascend/cann/set_env.sh若采用仿真器模式如 Ascend910B1 simulator运行还需追加设置export LD_LIBRARY_PATH$ASCEND_HOME_PATH/tools/simulator/Ascend910B1/lib:$LD_LIBRARY_PATH若仿真器模式下要运行接入 torch 的算子由于torch_npu默认只支持 NPU 上板且导入时会自动加载libruntime.so需提前预加载仿真动态库export LD_PRELOADlibruntime_camodel.so # 仿真器模式 unset LD_PRELOAD # NPU 上板模式完整的环境变量配置说明请参考 pyasc 快速入门-运行环境变量配置。算子分析开发算子的第一步是完成算子分析主要分析算子的数学表达式、输入输出的数量、Shape 范围以及计算逻辑的实现明确需要调用的 pyasc 接口。下文以 Add 算子为例介绍具体的分析过程。1. 明确算子的数学表达式及计算逻辑Add 算子的数学表达式为z x y计算逻辑是从外部存储 Global Memory 搬运数据至内部存储 Local Memory然后使用 pyasc 计算接口完成两个输入参数相加得到最终结果再搬运回 Global Memory。2. 明确输入和输出Add 算子有两个输入x与y输出为z本样例中算子输入支持的数据类型为float算子输出的数据类型与输入数据类型相同算子输入支持的 shape 为(8, 2048)输出 shape 与输入 shape 相同算子输入支持的 format 为ND。3. 确定核函数名称和参数本样例中核函数命名为vadd_kernel根据对算子输入输出的分析确定核函数有 3 个参数x、y、z其中x、y为输入参数z为输出参数。4. 确定算子实现所需接口实现涉及外部存储和内部存储间的数据搬运使用asc.data_copy接口来实现数据搬移本样例只涉及矢量计算的加法操作使用asc.add接口实现x y计算中使用到的 Tensor 数据结构使用asc.GlobalTensor、asc.LocalTensor进行管理并行流水任务之间使用asc.set_flag/asc.wait_flag接口完成同步。通过以上分析得到 pyasc Add 算子的设计规格如下项目内容算子类型OpTypeAdd算子输入xshape(8, 2048)数据类型floatformatND算子输入yshape(8, 2048)数据类型floatformatND算子输出zshape(8, 2048)数据类型floatformatND核函数名vadd_kernel使用的主要接口asc.data_copy数据搬运接口使用的主要接口asc.add矢量基础算术接口使用的主要接口asc.GlobalTensor/LocalTensor内存管理接口使用的主要接口asc.set_flag/wait_flag同步接口算子实现文件名称add.py在仓库中该样例的完整实现位于 examples/01_add/add.py对应的规格说明与执行说明可参考 examples/01_add/README.md。核函数开发完成环境准备和初步的算子分析后即可开始 pyasc 核函数的开发。多核并行与数据切分策略本样例使用多核并行计算即把数据进行分片分配到多个核上进行处理。pyasc 核函数是单个核上的处理函数所以每个核只处理部分数据。分配方案是假设共启用 8 个核数据整体长度为8 * 2048个元素平均分配到 8 个核上运行每个核上处理的数据大小为 2048 个元素。对于单核上的处理数据还可以进行数据切块实现对数据的流水并行处理。本样例使用以下参数控制数据切分USE_CORE_NUM 8启用 8 个核TILE_NUM 8每个核上数据分块个数BUFFER_NUM 2双缓冲注释中特别说明BUFFER_NUM只能取 1 或 2见 examples/01_add/add.py。由此可推算出每一层的数据量block_length total_length // USE_CORE_NUM每个核的数据量tile_length block_length // TILE_NUM // BUFFER_NUM每个 tile 的数据量。多核切分通过asc.get_block_idx()获取当前核的索引再用get_block_idx() * block_length计算该核在 Global Memory 中的起始偏移。核函数定义与实现使用asc.jit装饰器定义核函数并在核函数中实现算子逻辑import asc import asc.lib.runtime as rt USE_CORE_NUM 8 BUFFER_NUM 2 TILE_NUM 8 asc.jit def vadd_kernel(x: asc.GlobalAddress, y: asc.GlobalAddress, z: asc.GlobalAddress, block_length: int): # 获取当前核的索引计算数据偏移 offset asc.get_block_idx() * block_length # 创建 GlobalTensor 管理全局内存地址 x_gm asc.GlobalTensor() y_gm asc.GlobalTensor() z_gm asc.GlobalTensor() # 设置 Global Memory 起始地址和长度 x_gm.set_global_buffer(x offset, block_length) y_gm.set_global_buffer(y offset, block_length) z_gm.set_global_buffer(z offset, block_length) # 计算每个 tile 的长度考虑双缓冲 tile_length block_length // TILE_NUM // BUFFER_NUM # 获取数据类型信息 data_type x.dtype buffer_size tile_length * BUFFER_NUM * data_type.sizeof() # 创建 LocalTensor基于指定的逻辑位置/地址/长度 # x_local 和 y_local 放在 VECIN 位置 x_local asc.LocalTensor(data_type, asc.TPosition.VECIN, 0, tile_length * BUFFER_NUM) y_local asc.LocalTensor(data_type, asc.TPosition.VECIN, buffer_size, tile_length * BUFFER_NUM) # z_local 放在 VECOUT 位置 z_local asc.LocalTensor(data_type, asc.TPosition.VECOUT, buffer_size buffer_size, tile_length * BUFFER_NUM) # 流水循环处理双缓冲需要循环次数翻倍 for i in range(TILE_NUM * BUFFER_NUM): buf_id i % BUFFER_NUM # Step 1: 搬入 - 从 Global Memory 拷贝数据到 Local Memory asc.data_copy(x_local[buf_id * tile_length:], x_gm[i * tile_length:], tile_length) asc.data_copy(y_local[buf_id * tile_length:], y_gm[i * tile_length:], tile_length) # 同步等待 MTE2_V 事件确保数据搬入完成 asc.set_flag(asc.HardEvent.MTE2_V, buf_id) asc.wait_flag(asc.HardEvent.MTE2_V, buf_id) # Step 2: 计算 - 执行矢量加法 asc.add(z_local[buf_id * tile_length:], x_local[buf_id * tile_length:], y_local[buf_id * tile_length:], tile_length) # 同步等待 V_MTE3 事件确保计算完成 asc.set_flag(asc.HardEvent.V_MTE3, buf_id) asc.wait_flag(asc.HardEvent.V_MTE3, buf_id) # Step 3: 搬出 - 从 Local Memory 拷贝数据到 Global Memory asc.data_copy(z_gm[i * tile_length:], z_local[buf_id * tile_length:], tile_length) # 同步等待 MTE3_MTE2 事件确保数据搬出完成 asc.set_flag(asc.HardEvent.MTE3_MTE2, buf_id) asc.wait_flag(asc.HardEvent.MTE3_MTE2, buf_id)内部函数的调用关系示意图vadd_kernel ├── offset get_block_idx() * block_length ├── GlobalTensor 设置 │ ├── x_gm.set_global_buffer() │ ├── y_gm.set_global_buffer() │ └── z_gm.set_global_buffer() ├── LocalTensor 创建 └── for i in range(TILE_NUM * BUFFER_NUM): ├── CopyIn: data_copy (x_local, y_local - x_gm, y_gm) ├── Compute: add (z_local - x_local y_local) └── CopyOut: data_copy (z_gm - z_local)核函数代码逐段解析GlobalTensor 与全局缓冲区绑定。asc.GlobalTensor用于存放 Global Memory外部存储的全局数据其set_global_buffer(buffer, buffer_size)接口将张量与一段全局地址绑定。在 python/asc/language/core/tensor.py 中可以看到GlobalTensor的__getitem__切片操作只支持start偏移不支持step/stop返回一个从原起始地址偏移后的新GlobalTensor——这正是样例中x_gm[i * tile_length:]切片语义的来源。LocalTensor 的逻辑位置与地址分配。asc.LocalTensor用于存放 AI Core 中 Local Memory内部存储的数据支持逻辑位置TPosition为VECIN、VECOUT、VECCALC、A1、A2、B1、B2、CO1、CO2等。样例中x_local、y_local创建在VECIN矢量计算输入区地址分别为0与buffer_size各占tile_length * BUFFER_NUM个元素的空间z_local创建在VECOUT矢量计算输出区地址为buffer_size buffer_size与前两者在逻辑地址上错开避免数据覆盖。LocalTensor的构造函数签名LocalTensor(dtype, pos, addr, tile_size)与 Ascend C 的模板参数映射规则一致——在 pyasc 中DataType作为 Tensor 的成员变量以第一个参数传入详见 architecture_introduction.md 中“Python 前端模块”的参数映射说明。双缓冲流水循环。循环次数为TILE_NUM * BUFFER_NUM16 次每次通过buf_id i % BUFFER_NUM在 0/1 两个缓冲区之间交替使用。双缓冲机制使得搬入与计算、计算与搬出可以在不同 buffer 间流水叠加当第i个 tile 在 buffer 0 上计算时第i1个 tile 的数据可以同时搬入 buffer 1从而隐藏搬运时延。手动同步事件。同一核内不同流水线MTE2 数据搬入、V 矢量计算、MTE3 数据搬出之间使用set_flag/wait_flag成对同步本样例每轮迭代显式插入 3 对同步事件事件对含义作用MTE2_V数据搬入MTE2→ 矢量计算V确保data_copy搬入完成后才执行addV_MTE3矢量计算V→ 数据搬出MTE3确保add计算完成后才搬出结果MTE3_MTE2数据搬出MTE3→ 下一轮搬入MTE2确保本块搬出结束后才能复用缓冲区搬入新数据asc.HardEvent是 pyasc 提供的硬件同步事件枚举在 python/asc/language/core/enums.py 中定义了一整套事件类型如MTE2_V 4、V_MTE3 7、MTE3_MTE2 19等样例只使用其中与本流水最相关的三对。Launch 函数实现核函数只定义了单核上的处理逻辑Host 侧还需要一个 Launch 函数负责分配输出、计算分块并启动多核执行def vadd_launch(x: torch.Tensor, y: torch.Tensor) - torch.Tensor: z torch.zeros_like(x) total_length z.numel() block_length total_length // USE_CORE_NUM vadd_kernelUSE_CORE_NUM, rt.current_stream() return zvadd_kernel[USE_CORE_NUM, rt.current_stream()]使用内核调用符指定核数和流。其中方括号内第一个参数USE_CORE_NUM为运行核数必选项不能大于硬件实际可用核数第二个参数rt.current_stream()为 Kernel 执行流可选项这与 pyasc 编译运行模块“小括号传编译参数、中括号传运行时配置”的约定一致详见 architecture_introduction.md 中“编译和运行模块”一节(x, y, z, block_length)传递参数。pyasc 运行模块解析这些参数后对位于 Host 侧的输入输出张量会自动完成 Host 与 Device 之间的数据拷贝。核函数运行验证完成核函数开发后即可编写完整的核函数调用程序执行计算过程。完整的算子验证程序import logging import argparse import torch try: import torch_npu except ModuleNotFoundError: pass import asc import asc.runtime.config as config import asc.lib.runtime as rt USE_CORE_NUM 8 BUFFER_NUM 2 TILE_NUM 8 logging.basicConfig(levellogging.INFO) asc.jit def vadd_kernel(x: asc.GlobalAddress, y: asc.GlobalAddress, z: asc.GlobalAddress, block_length: int): offset asc.get_block_idx() * block_length x_gm asc.GlobalTensor() y_gm asc.GlobalTensor() z_gm asc.GlobalTensor() x_gm.set_global_buffer(x offset, block_length) y_gm.set_global_buffer(y offset, block_length) z_gm.set_global_buffer(z offset, block_length) tile_length block_length // TILE_NUM // BUFFER_NUM data_type x.dtype buffer_size tile_length * BUFFER_NUM * data_type.sizeof() x_local asc.LocalTensor(data_type, asc.TPosition.VECIN, 0, tile_length * BUFFER_NUM) y_local asc.LocalTensor(data_type, asc.TPosition.VECIN, buffer_size, tile_length * BUFFER_NUM) z_local asc.LocalTensor(data_type, asc.TPosition.VECOUT, buffer_size buffer_size, tile_length * BUFFER_NUM) for i in range(TILE_NUM * BUFFER_NUM): buf_id i % BUFFER_NUM asc.data_copy(x_local[buf_id * tile_length:], x_gm[i * tile_length:], tile_length) asc.data_copy(y_local[buf_id * tile_length:], y_gm[i * tile_length:], tile_length) asc.set_flag(asc.HardEvent.MTE2_V, buf_id) asc.wait_flag(asc.HardEvent.MTE2_V, buf_id) asc.add(z_local[buf_id * tile_length:], x_local[buf_id * tile_length:], y_local[buf_id * tile_length:], tile_length) asc.set_flag(asc.HardEvent.V_MTE3, buf_id) asc.wait_flag(asc.HardEvent.V_MTE3, buf_id) asc.data_copy(z_gm[i * tile_length:], z_local[buf_id * tile_length:], tile_length) asc.set_flag(asc.HardEvent.MTE3_MTE2, buf_id) asc.wait_flag(asc.HardEvent.MTE3_MTE2, buf_id) def vadd_launch(x: torch.Tensor, y: torch.Tensor) - torch.Tensor: z torch.zeros_like(x) total_length z.numel() block_length total_length // USE_CORE_NUM vadd_kernelUSE_CORE_NUM, rt.current_stream() return z # Backend 对应执行脚本时传入的 [RUN_MODE] # Platform 对应执行脚本时传入的 [SOC_VERSION] def vadd_custom(backend: config.Backend, platform: config.Platform): config.set_platform(backend, platform) device npu if config.Backend(backend) config.Backend.NPU else cpu size 8 * 2048 x torch.rand(size, dtypetorch.float32, devicedevice) y torch.rand(size, dtypetorch.float32, devicedevice) z vadd_launch(x, y) assert torch.allclose(z, x y) if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument(-r, typestr, defaultModel, helpbackend to run) parser.add_argument(-v, typestr, defaultNone, helpplatform to run) args parser.parse_args() backend args.r platform args.v if backend not in config.Backend.__members__: raise ValueError(Unsupported Backend! Supported: [Model, NPU]) backend config.Backend(backend) if platform is not None: platform_values [platform.value for platform in config.Platform] if platform not in platform_values: raise ValueError(fUnsupported Platform! Supported: {platform_values}) platform config.Platform(platform) logging.info([INFO] start process sample add.) vadd_custom(backend, platform) logging.info([INFO] Sample add run success.)验证程序的运行模式解析验证程序通过命令行参数-r/-v控制运行模式内部调用config.set_platform(backend, platform)完成运行环境配置。其底层实现在 python/asc/runtime/config.pyBackend运行后端枚举Model与NPU。Model为仿真器模式不需要 NPU 硬件NPU为上板模式需要真实的昇腾 AI 处理器。仿真器模式下若未指定平台默认使用Ascend910B1NPU 模式下若显式指定了soc_version会与rt.current_platform()获取的实际硬件平台比对不一致时抛出ValueErrorPlatform昇腾平台在 config.py 中定义了完整的平台枚举包括Ascend910B1、Ascend910B2、Ascend910B3、Ascend910B4、Ascend910_9362等 910 系列以及Ascend950PR_950z、Ascend950PR_9579等 950 系列型号正确性校验vadd_custom中随机生成x、y两个float32张量调用vadd_launch得到输出z再通过torch.allclose(z, x y)断言与 CPU 上的直接加法结果一致从而完成端到端功能验证。此外验证程序的assert torch.allclose(z, x y)与 python/test/kernels/test_vadd.py 等测试用例的校验思路一致说明该样例的验证方式与仓库泛化测试保持一致。编译和运行运行时使用以下命令python3 add.py -r [RUN_MODE] -v [SOC_VERSION]其中RUN_MODE编译执行方式可选择Model仿真或NPU上板。缺省时默认是仿真器模式SOC_VERSION昇腾 AI 处理器型号。如果无法确定具体的SOC_VERSION则在安装昇腾 AI 处理器的服务器上执行npu-smi info命令进行查询在查询到的 Name 前增加Ascend信息例如 Name 对应取值为xxxyy实际配置的SOC_VERSION值为Ascendxxxyy。仿真器模式下缺省默认是Ascend910B1环境NPU 上板模式下缺省自动检测。示例# 仿真器模式运行 python3 add.py -r Model -v Ascend910B1 # NPU 上板模式运行 python3 add.py -r NPU -v Ascend910B1用例执行完成打屏信息出现Sample add run success.说明样例执行成功[INFO] start process sample add. [INFO] Sample add run success.pyasc 与 Ascend C 算子开发接口/语法特性的对比特性Ascend C (.asc)pyasc (Python)编程语言C 扩展语法原生 Python核函数定义__global__ __aicore__asc.jit装饰器GlobalTensorAscendC::GlobalTensorTasc.GlobalTensor()LocalTensorAscendC::LocalTensorTasc.LocalTensor()数据搬运AscendC::DataCopy()asc.data_copy()矢量计算AscendC::Add()asc.add()同步事件AscendC::SetFlag()/WaitFlag()asc.set_flag()/wait_flag()核函数调用add_custom...()vadd_kernel[num_blocks, stream]()从对比可以看出pyasc 的核心设计目标是与 Ascend C API 一一对应Python 前端接口按高阶 API、基础 API、核心数据结构和枚举、同步与内存管理框架接口四类组织位于 python/asc/language 目录下而每个 Python 接口经过 AST 转 ASC-IR、Ascend C 代码生成后最终翻译为等价的 Ascend C 语法结构这一全链路架构在 architecture_introduction.md 中有详细说明。因此熟悉 Ascend C 的开发者几乎可以零成本迁移到 pyasc而 Python 用户则无需学习 C 扩展语法即可上手昇腾算子开发。后续引导如果想了解更多 pyasc 算子示例可以参考 examples 目录下的样例。除了本文讲解的 01_add仓库还提供了 02_add_framework基于框架类接口 TQue/TPipe 的 Add、03_matmul_mix、04_matmul_cube_only、05_matmul_leakyrelu、06_gelu、07_swiglu、08_rmsnorm、09_linear、10_fused_infer_attention 等从矢量到矩阵、从基础到融合算子的进阶样例如果想深入了解 pyasc 的 API 接口请参考 API 文档其中按 language高阶adv、基础basic、核心core、框架fwk与 libhost等分类收录了全部接口说明如果想了解 pyasc 的构建和调试方法请参考 快速入门。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN pyasc 算子开发指南：用原生 Python 从零实现并验证一个 Ascend C Add 算子
CANN pyasc 算子开发指南用原生 Python 从零实现并验证一个 Ascend C Add 算子【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc本篇指南以 pyasc 开源项目中的 Addz x y算子为例完整讲解基于 Python 原生语法开发昇腾 Ascend C 算子的全流程从环境准备、算子分析数学表达式、输入输出规格、接口选型到asc.jit核函数与 Launch 函数的实现再到仿真器/上板两种模式的编译运行与结果验证。读完本文你将掌握 pyasc 多核并行、双缓冲流水与set_flag/wait_flag手动同步的完整写法并能照此套路开发自己的自定义算子。总体开发流程在正式开发前需要先完成环境准备工作。pyasc 算子开发的基本流程如下其中每一步对应本文的一个章节环境准备安装 pyasc 与 CANN 软件→算子分析确定数学表达式、输入输出规格与所需接口→核函数开发编写asc.jit核函数与 Launch 函数→核函数运行验证编写完整验证程序并编译运行。环境准备开发 pyasc 算子前需要同时就绪两套环境pyasc 安装pyasc 支持通过 pip 快速安装和基于源码编译安装两种方式。快速安装直接执行pip install pyasc即可获得最新稳定版二进制 wheel 支持 CPython 3.9~3.12基于源码安装则需要先下载并安装 LLVM 19.1.7 预编译包、设置LLVM_INSTALL_PREFIX环境变量再执行python3 -m pip install .普通模式或python3 -m pip install -e .开发者模式。详细步骤请参考 pyasc 快速入门-环境准备与安装。CANN 软件安装开发算子前需安装 CANN 软件CANN toolkit 包与 CANN ops 包具体请参考 pyasc 快速入门-下载安装 CANN 包。安装 CANN 软件后需要配置运行环境变量默认路径安装时执行source /usr/local/Ascend/cann/set_env.sh若采用仿真器模式如 Ascend910B1 simulator运行还需追加设置export LD_LIBRARY_PATH$ASCEND_HOME_PATH/tools/simulator/Ascend910B1/lib:$LD_LIBRARY_PATH若仿真器模式下要运行接入 torch 的算子由于torch_npu默认只支持 NPU 上板且导入时会自动加载libruntime.so需提前预加载仿真动态库export LD_PRELOADlibruntime_camodel.so # 仿真器模式 unset LD_PRELOAD # NPU 上板模式完整的环境变量配置说明请参考 pyasc 快速入门-运行环境变量配置。算子分析开发算子的第一步是完成算子分析主要分析算子的数学表达式、输入输出的数量、Shape 范围以及计算逻辑的实现明确需要调用的 pyasc 接口。下文以 Add 算子为例介绍具体的分析过程。1. 明确算子的数学表达式及计算逻辑Add 算子的数学表达式为z x y计算逻辑是从外部存储 Global Memory 搬运数据至内部存储 Local Memory然后使用 pyasc 计算接口完成两个输入参数相加得到最终结果再搬运回 Global Memory。2. 明确输入和输出Add 算子有两个输入x与y输出为z本样例中算子输入支持的数据类型为float算子输出的数据类型与输入数据类型相同算子输入支持的 shape 为(8, 2048)输出 shape 与输入 shape 相同算子输入支持的 format 为ND。3. 确定核函数名称和参数本样例中核函数命名为vadd_kernel根据对算子输入输出的分析确定核函数有 3 个参数x、y、z其中x、y为输入参数z为输出参数。4. 确定算子实现所需接口实现涉及外部存储和内部存储间的数据搬运使用asc.data_copy接口来实现数据搬移本样例只涉及矢量计算的加法操作使用asc.add接口实现x y计算中使用到的 Tensor 数据结构使用asc.GlobalTensor、asc.LocalTensor进行管理并行流水任务之间使用asc.set_flag/asc.wait_flag接口完成同步。通过以上分析得到 pyasc Add 算子的设计规格如下项目内容算子类型OpTypeAdd算子输入xshape(8, 2048)数据类型floatformatND算子输入yshape(8, 2048)数据类型floatformatND算子输出zshape(8, 2048)数据类型floatformatND核函数名vadd_kernel使用的主要接口asc.data_copy数据搬运接口使用的主要接口asc.add矢量基础算术接口使用的主要接口asc.GlobalTensor/LocalTensor内存管理接口使用的主要接口asc.set_flag/wait_flag同步接口算子实现文件名称add.py在仓库中该样例的完整实现位于 examples/01_add/add.py对应的规格说明与执行说明可参考 examples/01_add/README.md。核函数开发完成环境准备和初步的算子分析后即可开始 pyasc 核函数的开发。多核并行与数据切分策略本样例使用多核并行计算即把数据进行分片分配到多个核上进行处理。pyasc 核函数是单个核上的处理函数所以每个核只处理部分数据。分配方案是假设共启用 8 个核数据整体长度为8 * 2048个元素平均分配到 8 个核上运行每个核上处理的数据大小为 2048 个元素。对于单核上的处理数据还可以进行数据切块实现对数据的流水并行处理。本样例使用以下参数控制数据切分USE_CORE_NUM 8启用 8 个核TILE_NUM 8每个核上数据分块个数BUFFER_NUM 2双缓冲注释中特别说明BUFFER_NUM只能取 1 或 2见 examples/01_add/add.py。由此可推算出每一层的数据量block_length total_length // USE_CORE_NUM每个核的数据量tile_length block_length // TILE_NUM // BUFFER_NUM每个 tile 的数据量。多核切分通过asc.get_block_idx()获取当前核的索引再用get_block_idx() * block_length计算该核在 Global Memory 中的起始偏移。核函数定义与实现使用asc.jit装饰器定义核函数并在核函数中实现算子逻辑import asc import asc.lib.runtime as rt USE_CORE_NUM 8 BUFFER_NUM 2 TILE_NUM 8 asc.jit def vadd_kernel(x: asc.GlobalAddress, y: asc.GlobalAddress, z: asc.GlobalAddress, block_length: int): # 获取当前核的索引计算数据偏移 offset asc.get_block_idx() * block_length # 创建 GlobalTensor 管理全局内存地址 x_gm asc.GlobalTensor() y_gm asc.GlobalTensor() z_gm asc.GlobalTensor() # 设置 Global Memory 起始地址和长度 x_gm.set_global_buffer(x offset, block_length) y_gm.set_global_buffer(y offset, block_length) z_gm.set_global_buffer(z offset, block_length) # 计算每个 tile 的长度考虑双缓冲 tile_length block_length // TILE_NUM // BUFFER_NUM # 获取数据类型信息 data_type x.dtype buffer_size tile_length * BUFFER_NUM * data_type.sizeof() # 创建 LocalTensor基于指定的逻辑位置/地址/长度 # x_local 和 y_local 放在 VECIN 位置 x_local asc.LocalTensor(data_type, asc.TPosition.VECIN, 0, tile_length * BUFFER_NUM) y_local asc.LocalTensor(data_type, asc.TPosition.VECIN, buffer_size, tile_length * BUFFER_NUM) # z_local 放在 VECOUT 位置 z_local asc.LocalTensor(data_type, asc.TPosition.VECOUT, buffer_size buffer_size, tile_length * BUFFER_NUM) # 流水循环处理双缓冲需要循环次数翻倍 for i in range(TILE_NUM * BUFFER_NUM): buf_id i % BUFFER_NUM # Step 1: 搬入 - 从 Global Memory 拷贝数据到 Local Memory asc.data_copy(x_local[buf_id * tile_length:], x_gm[i * tile_length:], tile_length) asc.data_copy(y_local[buf_id * tile_length:], y_gm[i * tile_length:], tile_length) # 同步等待 MTE2_V 事件确保数据搬入完成 asc.set_flag(asc.HardEvent.MTE2_V, buf_id) asc.wait_flag(asc.HardEvent.MTE2_V, buf_id) # Step 2: 计算 - 执行矢量加法 asc.add(z_local[buf_id * tile_length:], x_local[buf_id * tile_length:], y_local[buf_id * tile_length:], tile_length) # 同步等待 V_MTE3 事件确保计算完成 asc.set_flag(asc.HardEvent.V_MTE3, buf_id) asc.wait_flag(asc.HardEvent.V_MTE3, buf_id) # Step 3: 搬出 - 从 Local Memory 拷贝数据到 Global Memory asc.data_copy(z_gm[i * tile_length:], z_local[buf_id * tile_length:], tile_length) # 同步等待 MTE3_MTE2 事件确保数据搬出完成 asc.set_flag(asc.HardEvent.MTE3_MTE2, buf_id) asc.wait_flag(asc.HardEvent.MTE3_MTE2, buf_id)内部函数的调用关系示意图vadd_kernel ├── offset get_block_idx() * block_length ├── GlobalTensor 设置 │ ├── x_gm.set_global_buffer() │ ├── y_gm.set_global_buffer() │ └── z_gm.set_global_buffer() ├── LocalTensor 创建 └── for i in range(TILE_NUM * BUFFER_NUM): ├── CopyIn: data_copy (x_local, y_local - x_gm, y_gm) ├── Compute: add (z_local - x_local y_local) └── CopyOut: data_copy (z_gm - z_local)核函数代码逐段解析GlobalTensor 与全局缓冲区绑定。asc.GlobalTensor用于存放 Global Memory外部存储的全局数据其set_global_buffer(buffer, buffer_size)接口将张量与一段全局地址绑定。在 python/asc/language/core/tensor.py 中可以看到GlobalTensor的__getitem__切片操作只支持start偏移不支持step/stop返回一个从原起始地址偏移后的新GlobalTensor——这正是样例中x_gm[i * tile_length:]切片语义的来源。LocalTensor 的逻辑位置与地址分配。asc.LocalTensor用于存放 AI Core 中 Local Memory内部存储的数据支持逻辑位置TPosition为VECIN、VECOUT、VECCALC、A1、A2、B1、B2、CO1、CO2等。样例中x_local、y_local创建在VECIN矢量计算输入区地址分别为0与buffer_size各占tile_length * BUFFER_NUM个元素的空间z_local创建在VECOUT矢量计算输出区地址为buffer_size buffer_size与前两者在逻辑地址上错开避免数据覆盖。LocalTensor的构造函数签名LocalTensor(dtype, pos, addr, tile_size)与 Ascend C 的模板参数映射规则一致——在 pyasc 中DataType作为 Tensor 的成员变量以第一个参数传入详见 architecture_introduction.md 中“Python 前端模块”的参数映射说明。双缓冲流水循环。循环次数为TILE_NUM * BUFFER_NUM16 次每次通过buf_id i % BUFFER_NUM在 0/1 两个缓冲区之间交替使用。双缓冲机制使得搬入与计算、计算与搬出可以在不同 buffer 间流水叠加当第i个 tile 在 buffer 0 上计算时第i1个 tile 的数据可以同时搬入 buffer 1从而隐藏搬运时延。手动同步事件。同一核内不同流水线MTE2 数据搬入、V 矢量计算、MTE3 数据搬出之间使用set_flag/wait_flag成对同步本样例每轮迭代显式插入 3 对同步事件事件对含义作用MTE2_V数据搬入MTE2→ 矢量计算V确保data_copy搬入完成后才执行addV_MTE3矢量计算V→ 数据搬出MTE3确保add计算完成后才搬出结果MTE3_MTE2数据搬出MTE3→ 下一轮搬入MTE2确保本块搬出结束后才能复用缓冲区搬入新数据asc.HardEvent是 pyasc 提供的硬件同步事件枚举在 python/asc/language/core/enums.py 中定义了一整套事件类型如MTE2_V 4、V_MTE3 7、MTE3_MTE2 19等样例只使用其中与本流水最相关的三对。Launch 函数实现核函数只定义了单核上的处理逻辑Host 侧还需要一个 Launch 函数负责分配输出、计算分块并启动多核执行def vadd_launch(x: torch.Tensor, y: torch.Tensor) - torch.Tensor: z torch.zeros_like(x) total_length z.numel() block_length total_length // USE_CORE_NUM vadd_kernelUSE_CORE_NUM, rt.current_stream() return zvadd_kernel[USE_CORE_NUM, rt.current_stream()]使用内核调用符指定核数和流。其中方括号内第一个参数USE_CORE_NUM为运行核数必选项不能大于硬件实际可用核数第二个参数rt.current_stream()为 Kernel 执行流可选项这与 pyasc 编译运行模块“小括号传编译参数、中括号传运行时配置”的约定一致详见 architecture_introduction.md 中“编译和运行模块”一节(x, y, z, block_length)传递参数。pyasc 运行模块解析这些参数后对位于 Host 侧的输入输出张量会自动完成 Host 与 Device 之间的数据拷贝。核函数运行验证完成核函数开发后即可编写完整的核函数调用程序执行计算过程。完整的算子验证程序import logging import argparse import torch try: import torch_npu except ModuleNotFoundError: pass import asc import asc.runtime.config as config import asc.lib.runtime as rt USE_CORE_NUM 8 BUFFER_NUM 2 TILE_NUM 8 logging.basicConfig(levellogging.INFO) asc.jit def vadd_kernel(x: asc.GlobalAddress, y: asc.GlobalAddress, z: asc.GlobalAddress, block_length: int): offset asc.get_block_idx() * block_length x_gm asc.GlobalTensor() y_gm asc.GlobalTensor() z_gm asc.GlobalTensor() x_gm.set_global_buffer(x offset, block_length) y_gm.set_global_buffer(y offset, block_length) z_gm.set_global_buffer(z offset, block_length) tile_length block_length // TILE_NUM // BUFFER_NUM data_type x.dtype buffer_size tile_length * BUFFER_NUM * data_type.sizeof() x_local asc.LocalTensor(data_type, asc.TPosition.VECIN, 0, tile_length * BUFFER_NUM) y_local asc.LocalTensor(data_type, asc.TPosition.VECIN, buffer_size, tile_length * BUFFER_NUM) z_local asc.LocalTensor(data_type, asc.TPosition.VECOUT, buffer_size buffer_size, tile_length * BUFFER_NUM) for i in range(TILE_NUM * BUFFER_NUM): buf_id i % BUFFER_NUM asc.data_copy(x_local[buf_id * tile_length:], x_gm[i * tile_length:], tile_length) asc.data_copy(y_local[buf_id * tile_length:], y_gm[i * tile_length:], tile_length) asc.set_flag(asc.HardEvent.MTE2_V, buf_id) asc.wait_flag(asc.HardEvent.MTE2_V, buf_id) asc.add(z_local[buf_id * tile_length:], x_local[buf_id * tile_length:], y_local[buf_id * tile_length:], tile_length) asc.set_flag(asc.HardEvent.V_MTE3, buf_id) asc.wait_flag(asc.HardEvent.V_MTE3, buf_id) asc.data_copy(z_gm[i * tile_length:], z_local[buf_id * tile_length:], tile_length) asc.set_flag(asc.HardEvent.MTE3_MTE2, buf_id) asc.wait_flag(asc.HardEvent.MTE3_MTE2, buf_id) def vadd_launch(x: torch.Tensor, y: torch.Tensor) - torch.Tensor: z torch.zeros_like(x) total_length z.numel() block_length total_length // USE_CORE_NUM vadd_kernelUSE_CORE_NUM, rt.current_stream() return z # Backend 对应执行脚本时传入的 [RUN_MODE] # Platform 对应执行脚本时传入的 [SOC_VERSION] def vadd_custom(backend: config.Backend, platform: config.Platform): config.set_platform(backend, platform) device npu if config.Backend(backend) config.Backend.NPU else cpu size 8 * 2048 x torch.rand(size, dtypetorch.float32, devicedevice) y torch.rand(size, dtypetorch.float32, devicedevice) z vadd_launch(x, y) assert torch.allclose(z, x y) if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument(-r, typestr, defaultModel, helpbackend to run) parser.add_argument(-v, typestr, defaultNone, helpplatform to run) args parser.parse_args() backend args.r platform args.v if backend not in config.Backend.__members__: raise ValueError(Unsupported Backend! Supported: [Model, NPU]) backend config.Backend(backend) if platform is not None: platform_values [platform.value for platform in config.Platform] if platform not in platform_values: raise ValueError(fUnsupported Platform! Supported: {platform_values}) platform config.Platform(platform) logging.info([INFO] start process sample add.) vadd_custom(backend, platform) logging.info([INFO] Sample add run success.)验证程序的运行模式解析验证程序通过命令行参数-r/-v控制运行模式内部调用config.set_platform(backend, platform)完成运行环境配置。其底层实现在 python/asc/runtime/config.pyBackend运行后端枚举Model与NPU。Model为仿真器模式不需要 NPU 硬件NPU为上板模式需要真实的昇腾 AI 处理器。仿真器模式下若未指定平台默认使用Ascend910B1NPU 模式下若显式指定了soc_version会与rt.current_platform()获取的实际硬件平台比对不一致时抛出ValueErrorPlatform昇腾平台在 config.py 中定义了完整的平台枚举包括Ascend910B1、Ascend910B2、Ascend910B3、Ascend910B4、Ascend910_9362等 910 系列以及Ascend950PR_950z、Ascend950PR_9579等 950 系列型号正确性校验vadd_custom中随机生成x、y两个float32张量调用vadd_launch得到输出z再通过torch.allclose(z, x y)断言与 CPU 上的直接加法结果一致从而完成端到端功能验证。此外验证程序的assert torch.allclose(z, x y)与 python/test/kernels/test_vadd.py 等测试用例的校验思路一致说明该样例的验证方式与仓库泛化测试保持一致。编译和运行运行时使用以下命令python3 add.py -r [RUN_MODE] -v [SOC_VERSION]其中RUN_MODE编译执行方式可选择Model仿真或NPU上板。缺省时默认是仿真器模式SOC_VERSION昇腾 AI 处理器型号。如果无法确定具体的SOC_VERSION则在安装昇腾 AI 处理器的服务器上执行npu-smi info命令进行查询在查询到的 Name 前增加Ascend信息例如 Name 对应取值为xxxyy实际配置的SOC_VERSION值为Ascendxxxyy。仿真器模式下缺省默认是Ascend910B1环境NPU 上板模式下缺省自动检测。示例# 仿真器模式运行 python3 add.py -r Model -v Ascend910B1 # NPU 上板模式运行 python3 add.py -r NPU -v Ascend910B1用例执行完成打屏信息出现Sample add run success.说明样例执行成功[INFO] start process sample add. [INFO] Sample add run success.pyasc 与 Ascend C 算子开发接口/语法特性的对比特性Ascend C (.asc)pyasc (Python)编程语言C 扩展语法原生 Python核函数定义__global__ __aicore__asc.jit装饰器GlobalTensorAscendC::GlobalTensorTasc.GlobalTensor()LocalTensorAscendC::LocalTensorTasc.LocalTensor()数据搬运AscendC::DataCopy()asc.data_copy()矢量计算AscendC::Add()asc.add()同步事件AscendC::SetFlag()/WaitFlag()asc.set_flag()/wait_flag()核函数调用add_custom...()vadd_kernel[num_blocks, stream]()从对比可以看出pyasc 的核心设计目标是与 Ascend C API 一一对应Python 前端接口按高阶 API、基础 API、核心数据结构和枚举、同步与内存管理框架接口四类组织位于 python/asc/language 目录下而每个 Python 接口经过 AST 转 ASC-IR、Ascend C 代码生成后最终翻译为等价的 Ascend C 语法结构这一全链路架构在 architecture_introduction.md 中有详细说明。因此熟悉 Ascend C 的开发者几乎可以零成本迁移到 pyasc而 Python 用户则无需学习 C 扩展语法即可上手昇腾算子开发。后续引导如果想了解更多 pyasc 算子示例可以参考 examples 目录下的样例。除了本文讲解的 01_add仓库还提供了 02_add_framework基于框架类接口 TQue/TPipe 的 Add、03_matmul_mix、04_matmul_cube_only、05_matmul_leakyrelu、06_gelu、07_swiglu、08_rmsnorm、09_linear、10_fused_infer_attention 等从矢量到矩阵、从基础到融合算子的进阶样例如果想深入了解 pyasc 的 API 接口请参考 API 文档其中按 language高阶adv、基础basic、核心core、框架fwk与 libhost等分类收录了全部接口说明如果想了解 pyasc 的构建和调试方法请参考 快速入门。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考