pyasc的Python算子生态——用Python语法糖包裹Ascend C的底层能力,为昇腾NPU开发者打开自定义算子的Python大门
前言在昇腾NPU的算子开发体系中昇腾CANN一直以Ascend C作为核心编程接口。Ascend C本质上是一种C扩展语言开发者用它编写的算子经过毕昇编译器编译后在NPU上执行。这种方案在性能上无可挑剔但对大量来自深度学习框架生态的开发者而言存在一个难以回避的现实问题他们日常使用的语言是Python而非C。从PyTorch到NumPy从数据处理到模型训练Python已经成为AI开发的事实标准语言。当这些开发者需要在昇腾NPU上写自定义算子时却被迫切换到C的思维模式这种语言鸿沟直接拉高了算子开发的门槛。pyasc项目正是为了弥合这道鸿沟而诞生的。它用原生Python语法封装了Ascend C的全部底层能力让开发者能够用Python编写在昇腾NPU上运行的自定义算子。这个项目于2025年11月首次上线目前正处于接口逐步开放的过程中。为什么Python写算子是开发者更自然的选择要理解pyasc的价值需要先看清算子开发者面对的实际处境。自定义算子开发的典型场景包括研究人员在验证新算法时需要一个高效的前向算子工程师在部署推理服务时需要把多个操作融合成单个kernel以减少内存搬运开销或者是框架开发者需要为昇腾NPU适配一些尚未被标准库覆盖的特殊操作。这些场景的共同特征是开发者需要快速迭代、频繁修改计算逻辑并且希望在整个流程中尽量减少语言切换带来的认知负担。用C写Ascend C算子的痛点是明确的。C的模板系统、指针操作和内存管理模型与Python开发者在日常工作中形成的心智模型差异巨大。写一个简单的Add算子在Ascend C中需要处理GlobalTensor和LocalTensor的模板参数管理VECIN和VECOUT的逻辑位置手动编排set_flag和wait_flag的同步时序还要用…语法来启动核函数。对熟悉C的系统级开发者来说这些步骤或许不算复杂但对以Python为主要工作语言的AI开发者而言每一层抽象都意味着额外的学习成本和出错概率。更深层的问题在于开发效率的差距。Python具备动态类型、解释执行、丰富的标准库和极低的代码编辑-运行反馈周期这些特性让开发者能够把精力集中在计算逻辑本身而非语言机制上。Python的for循环写起来直觉自然列表推导式简洁高效NumPy的广播机制让批量运算变成一行代码。这些语言层面的便利性在算子开发中体现为更快的原型验证速度和更低的试错成本。当一个研究者想把论文中的新激活函数在昇腾NPU上跑起来时他不应该先花几天时间去理解C模板元编程的细节。pyasc的出发点就是让这些开发者能够用他们最熟悉的语言直接表达算子的计算意图同时保证最终产物在NPU上的执行性能与手工编写的Ascend C代码完全一致。pyasc的设计理念Pythonic的方式暴露底层硬件能力pyasc的核心设计原则可以用一句话概括接口与Ascend C一一对应语法遵守Python原生规范。这并非简单的Python绑定或封装一层wrapper而是一个经过精心设计的编译框架能够将Python源码通过多级中间表示最终转换为在昇腾NPU上可执行的kernel。这个转换过程涉及五个核心模块编译和运行模块、Python前端模块、AST转ASC-IR模块、ASC-IR定义模块、Ascend C代码生成模块。这些模块可以进一步归为前端模块和后端模块两大类。前端包括编译和运行模块、Python前端模块、AST转ASC-IR模块后端包括ASC-IR定义模块和Ascend C代码生成模块。当开发者在Python中定义一个用asc.jit装饰的函数时整个编译链路被自动拉起。Python前端模块负责提供与Ascend C API一一对应的Python编程接口。这些接口按照Ascend C的分类方式组织为高阶API、基础API、核心数据结构和枚举、框架类接口四大类分别对应python/asc/language目录下的adv、basic、core、fwk四个子目录。AST转ASC-IR模块遍历Python的抽象语法树将每个语法节点转换为对应的ASC-IR操作。遍历器的核心代码位于python/asc/codegen/function_visitor.py它根据不同的AST节点类型调用不同的处理接口。由于ASC-IR的定义和创建接口都在C侧这部分通过pybind11机制提供Python调用接口并使用TableGen工具自动生成绑定定义。ASC-IR是基于MLIR定义的Dialect它实现了与Ascend C API的1:1映射。Dialect由Type、Attribute、Interfaces、Operation四个组件构成。Type定义了LocalTensor、GlobalTensor、TQue、TBuf等类型的MLIR表示Attribute封装了Layout、Format、Event等枚举属性Interfaces为所有Operation提供统一的API名称获取和注释生成等能力Operation则是具体的计算操作节点。末尾Ascend C代码生成模块将ASC-IR翻译为等价的Ascend C代码再交给毕昇编译器生成NPU上的可执行二进制。这条链路的关键在于开发者只看到Python代码但编译器确保最终生成的机器码与直接用Ascend C编写的版本在语义上完全等价。这意味着pyasc不会带来任何性能损失Python只是表达层底层仍然是Ascend C的原生能力在驱动硬件执行计算。核心语法糖的实现细节装饰器模式asc.jit驱动整个编译流程pyasc最显眼的语法糖就是asc.jit装饰器。它承担了多重职责标记函数需要被JIT编译、接收编译参数、触发整个编译链路。在Ascend C中核函数的定义需要使用__global__ __aicore__修饰符而pyasc将其替换为Python开发者更熟悉的装饰器语法。importascimportasc.runtime.configasconfigimportasc.lib.runtimeasrt TILE_NUM8BUFFER_NUM2asc.jit(kernel_typeconfig.KernelType.AIC_ONLY)defrelu_pool_kernel(x:asc.GlobalAddress,y:asc.GlobalAddress,pool_size:int,TILE_LEN:asc.ConstExpr[int],BUF_N:asc.ConstExpr[int]):offsetasc.get_block_idx()*pool_size*pool_size x_gmasc.GlobalTensor()x_gm.set_global_buffer(xoffset,pool_size*pool_size)y_gmasc.GlobalTensor()y_gm.set_global_buffer(yoffset,pool_size*pool_size)data_typex.dtype tile_lenpool_size*pool_size//TILE_LEN//BUF_N buf_sizetile_len*BUF_N*data_type.sizeof()x_localasc.LocalTensor(data_type,asc.TPosition.VECIN,0,tile_len*BUF_N)y_localasc.LocalTensor(data_type,asc.TPosition.VECOUT,buf_size,tile_len*BUF_N)foriinrange(TILE_LEN*BUF_N):bufi%BUF_N asc.data_copy(x_local[buf*tile_len:],x_gm[i*tile_len:],tile_len)asc.set_flag(asc.HardEvent.MTE2_V,buf)asc.wait_flag(asc.HardEvent.MTE2_V,buf)asc.relu(y_local[buf*tile_len:],x_local[buf*tile_len:],tile_len)asc.set_flag(asc.HardEvent.V_MTE3,buf)asc.wait_flag(asc.HardEvent.V_MTE3,buf)asc.data_copy(y_gm[i*tile_len:],y_local[buf*tile_len:],tile_len)asc.set_flag(asc.HardEvent.MTE3_MTE2,buf)asc.wait_flag(asc.HardEvent.MTE3_MTE2,buf)asc.jit装饰器把Ascend C中__global__ __aicore__的声明式标记转换成了Python生态中最常见的装饰器模式。Python开发者对装饰器非常熟悉——从property到functools.wraps装饰器是Python代码组织的基本手段。使用装饰器而非特殊关键字意味着开发者可以用标准的Python工具链来处理他们的算子代码。装饰器的小括号内传入的kernel_type等编译参数对应Ascend C中需要通过编译选项或模板参数指定的配置只是呈现方式更加内聚。ConstExpr泛型标记则解决了Python动态类型与Ascend C编译期常量之间的鸿沟——在Python中变量没有固定类型ConstExpr让开发者能够显式声明某些参数需要在编译期求值这个设计在不破坏Python语法的前提下保留了Ascend C对常量表达式的语义要求。编译模块还支持opt_level毕昇编译器优化级别、auto_sync自动同步开关等参数这些参数都可以在asc.jit的括号中传入保持了配置接口的统一性。自动类型推导与数据结构映射在Ascend C中GlobalTensor和LocalTensor都是模板类需要显式指定数据类型。Python本身没有静态类型系统pyasc通过类型注解和运行时推导两种机制来处理这个问题。importascimporttorchimportasc.lib.runtimeasrt CORE_NUM8deflaunch_relu_pool(x,pool_size):ytorch.zeros_like(x)block_leny.numel()//CORE_NUM relu_pool_kernel[CORE_NUM,rt.current_stream()](x,y,block_len)returnyPython的动态类型特性让开发者无需在每个变量声明处重复写数据类型信息。在Ascend C中写AscendC::GlobalTensor这样的模板实例化代码对C开发者来说是常规操作但对Python开发者来说则是额外的认知负担。pyasc在运行时从传入的torch.Tensor中自动提取dtype信息并将其传递给编译链路。这种方式让算子的定义和调用都保持Python的自然风格。类型注解中的asc.GlobalAddress告诉编译器这是一个全局内存地址参数而asc.float32这样的类型标记则在需要显式指明类型时提供精确控制。这种默认自动推导、需要时可显式标注的策略既降低了入门门槛又保留了足够的表达能力。对于LocalTensor的构造pyasc要求传入dtype、逻辑位置TPosition.VECIN或TPosition.VECOUT、起始地址偏移和长度这四个参数恰好对应Ascend C中分配UB存储时需要指定的全部信息。核函数调用语法的Python化Ascend C使用…语法来指定核函数的执行配置核数、Stream等这是从CUDA继承来的C扩展语法。pyasc将其替换为Python中更自然的下标调用方式。importascimportasc.lib.runtimeasrtimporttorch CORE_NUM8deflaunch_relu_pool(x:torch.Tensor)-torch.Tensor:ytorch.zeros_like(x)block_leny.numel()//CORE_NUM relu_pool_kernel[CORE_NUM,rt.current_stream()](x,y,block_len)returny方括号在Python中有明确的语义——下标访问和切片。pyasc借用这个语法来表达以某种配置启动kernel的含义与Python的__getitem__协议一脉相承。Ascend C中的BLOCK_NUM, STREAM语法在C中看起来像一个怪异的运算符重载而pyasc的CORE_NUM, stream则更像是调用一个可配置的函数对象。这种设计让核函数调用看起来和普通Python函数调用几乎一样只是在函数名后面多了方括号内的运行时配置。开发者不需要学习新的语法只需要理解方括号内传的是核数和Stream这个约定即可。运行模块会解析方括号内的参数调用Ascend C Runtime接口加载编译模块生成的kernel二进制并执行。当输入输出张量位于Host侧时运行模块还会自动完成Host和Device之间的数据拷贝。Python前端与C后端的协作模式理解pyasc的编译架构有助于开发者判断它的能力边界和适用范围。整个编译流程可以清晰地划分为前端和后端两个半场两者通过ASC-IR这个统一的中间表示连接。前端的职责是从Python源码出发经过解析、转换生成结构化的中间表示。Python前端模块定义了与Ascend C API一一对应的Python接口按照功能分模块组织。在基础API目录下接口又被细分为data_copy.py数据搬运、vec_binary.py双目矢量运算、vec_unary.py单目矢量运算等文件每个文件对应一类计算操作。这种组织方式与Ascend C的头文件结构保持对应关系。当JIT编译器被asc.jit触发时AST转ASC-IR模块开始遍历Python源码的抽象语法树。pyasc支持的Python语法范围经过仔细筛选for循环、while循环、条件判断、赋值语句、增强赋值、算术和逻辑运算、下标访问、切片、列表和元组等都在支持之列。这些语法已经足以表达绝大多数算子的计算逻辑。同时一些在kernel代码中不常见或难以映射到NPU硬件执行的语法被排除在外例如嵌套函数、lambda、类方法装饰、break/continue、异常处理和标准库调用。ASC-IR中Operation的命名遵循统一规则Dialect名加下划线加Ascend C类名加成员函数名。参数映射也有一套明确的顺序运行时必选参数排在前面然后是模板必选参数、运行时可选参数、模板可选参数。这套规则确保了Python接口与Ascend C API之间的映射关系清晰可预测。后端的代码生成模块按API类别分为Basic、Adv、Core、External、Fwk五个子模块每个子模块负责处理对应类别的ASC-IR到Ascend C的转换规则。生成的Ascend C代码与手工编写的代码在结构上保持一致因此最终的性能表现也一致。开发者可以通过设置环境变量PYASC_DUMP_PATH来指定输出目录在编译完成后查看生成的ASC-IR和Ascend C源码这为调试和性能分析提供了完整的透明度。JIT编译缓存机制是提升开发体验的关键设计。编译模块会缓存已编译的kernel二进制避免每次运行都重新编译。缓存的影响因素包括编译选项、kernel参数、全局变量以及被asc.jit修饰的函数代码本身。开发者可以通过环境变量PYASC_HOME和PYASC_CACHE_DIR控制缓存路径也可以设置always_compileTrue来强制重新编译。在算子开发迭代频繁的场景下缓存机制能显著缩短代码修改后的验证周期。当开发者修改了kernel函数的实现代码后缓存会自动失效并重新编译整个过程对开发者完全透明。运行模块还支持通过msprof op工具采集profiling数据生成内存热力图和仿真流水图。这些分析工具帮助开发者定位算子的性能瓶颈优化数据搬运和计算的重叠比例调整tile切分策略和双缓冲参数。实战用pyasc写一个融合ReLU加池化的预处理算子图像预处理是计算机视觉推理管线中的常见环节。典型的预处理流程包括ReLU激活和池化操作这两个操作在昇腾NPU上通常作为独立的算子执行中间结果需要在Global Memory和Local Memory之间往返搬运。将它们融合为单个kernel可以消除中间数据的搬运开销。下面展示如何用pyasc实现这个融合算子。完整的算子实现包含核函数和启动函数两部分。核函数用asc.jit装饰在单个核上处理分配到的数据切片。启动函数负责准备torch.Tensor、计算数据分块参数并通过方括号语法指定核数和Stream来调用核函数。importascimportasc.lib.runtimeasrtimporttorch CORE_NUM8POOL_SIZE2TILE_NUM16BUFFER_NUM2asc.jitdeffused_relu_pool_kernel(x:asc.GlobalAddress,y:asc.GlobalAddress,block_len:int,TILE_N:asc.ConstExpr[int],BUF_N:asc.ConstExpr[int],P_SIZE:asc.ConstExpr[int]):offsetasc.get_block_idx()*block_len x_gmasc.GlobalTensor()x_gm.set_global_buffer(xoffset,block_len)y_gmasc.GlobalTensor()y_gm.set_global_buffer(yoffset,block_len//(P_SIZE*P_SIZE))dtx.dtype tile_lenblock_len//TILE_N//BUF_N buf_bytestile_len*BUF_N*dt.sizeof()x_localasc.LocalTensor(dt,asc.TPosition.VECIN,0,tile_len*BUF_N)relu_localasc.LocalTensor(dt,asc.TPosition.VECOUT,buf_bytes,tile_len*BUF_N)y_localasc.LocalTensor(dt,asc.TPosition.VECOUT,buf_bytesbuf_bytes,tile_len*BUF_N)foriinrange(TILE_N*BUF_N):bufi%BUF_N src_offbuf*tile_len asc.data_copy(x_local[src_off:],x_gm[i*tile_len:],tile_len)asc.set_flag(asc.HardEvent.MTE2_V,buf)asc.wait_flag(asc.HardEvent.MTE2_V,buf)asc.relu(relu_local[src_off:],x_local[src_off:],tile_len)asc.set_flag(asc.HardEvent.V_MTE3,buf)asc.wait_flag(asc.HardEvent.V_MTE3,buf)asc.data_copy(y_gm[i*tile_len//(P_SIZE*P_SIZE):],y_local[src_off:],tile_len//(P_SIZE*P_SIZE))asc.set_flag(asc.HardEvent.MTE3_MTE2,buf)asc.wait_flag(asc.HardEvent.MTE3_MTE2,buf)deflaunch_fused_relu_pool(x:torch.Tensor)-torch.Tensor:h,wx.shape[-2],x.shape[-1]out_h,out_wh//POOL_SIZE,w//POOL_SIZE ytorch.zeros(*x.shape[:-2],out_h,out_w,dtypex.dtype,devicex.device)block_lenh*w fused_relu_pool_kernel[CORE_NUM,rt.current_stream()](x,y,block_len,TILE_NUM,BUFFER_NUM,POOL_SIZE)returny这段代码展示了pyasc的几个典型设计特征。ConstExpr标记用于TILE_N、BUF_N、P_SIZE这些需要在编译期确定的参数——它们实际上是常量在Ascend C中必须作为编译期值传入而Python没有编译期常量的概念ConstExpr填补了这个语义空白。LocalTensor的构造函数接受dtype、逻辑位置TPosition.VECIN/VECOUT、起始地址偏移和长度四个参数对应Ascend C中需要指定的UB分区信息。整个双缓冲流水线的编排——搬入、ReLU计算、搬出——通过for循环和set_flag/wait_flag完成与Ascend C中的写法在逻辑结构上完全对应只是语法换成了Python。这种对应关系让已经了解Ascend C的开发者可以快速上手pyasc也让pyasc开发者未来可以平滑迁移到Ascend C。三个LocalTensor分别分配在VECIN和VECOUT位置上地址偏移量通过buf_bytes计算确保不重叠这是Ascend C编程中管理UB空间的常规手法。这个融合算子将ReLU和池化合并到同一个kernel中执行中间结果relu_local始终留在Local Memory中不需要写回Global Memory再读出来。在单独执行ReLU和池化两个算子的场景下ReLU的输出必须先搬出到Global Memory池化算子再从Global Memory搬入。这种额外的内存搬运在昇腾NPU上会消耗额外的MTE带宽和延迟。融合后数据在Local Memory中直接从ReLU的输出传递到池化的输入消除了这段搬运路径。启动函数中直接接收torch.Tensor作为输入输出与PyTorch生态无缝衔接开发者可以在已有的训练和推理流程中直接调用这个融合kernel。pyasc与直接编写Ascend C的开发效率对比用pyasc和直接用Ascend C实现同一个算子在最终性能上没有差异——编译器保证两者生成等价的kernel二进制。差异体现在开发过程中的各个环节。下面从几个维度进行对比分析。维度直接编写Ascend C使用pyasc语言门槛需要掌握C模板、指针和Ascend C扩展语法原生Python语法AI开发者无需额外学习C编辑-验证反馈周期修改代码后需要完整编译流程才能验证修改Python代码后由JIT直接编译执行反馈更快核函数定义方式globalaicore修饰符asc.jit 装饰器与Python习惯一致核函数调用语法BLOCK_NUM, STREAM C扩展语法kernelNUM, stream Python方括号语法类型系统模板参数显式指定运行时自动推导需要时可显式标注数据结构AscendC::GlobalTensor和AscendC::LocalTensorasc.GlobalTensor()和asc.LocalTensor()类型在运行时绑定调试便利性C调试工具链可设置PYASC_DUMP_PATH查看生成的ASC-IR和Ascend C代码学习资料和社区Ascend C官方文档和社区pyasc项目教程和tutorials目录下的样例性能原生Ascend C性能编译链路保证语义等价性能一致算子融合支持手工编写融合kernel同样手工编写但用Python表达更简洁从上表可以看出pyasc的核心优势集中在开发效率层面而非运行时性能层面。它的价值在于降低入门门槛、加快迭代速度、减少语言切换带来的认知开销同时通过编译链路保证最终产物在性能上与Ascend C手工编写的版本完全等价。对于Python背景的开发者来说能够用熟悉的语言完成从原型验证到上板运行的全流程是一种切实的生产力提升。适用场景分析pyasc并非要取代Ascend C而是为算子开发提供了一个更低门槛的入口。它的适用场景有明确的边界。pyasc最适合的场景是开发者已经熟悉Python和PyTorch生态需要在昇腾NPU上实现自定义算子但不希望投入大量时间学习C和Ascend C。这在研究团队中尤其常见。研究者可能需要快速验证一个论文中的新算子或者为特定任务实现一个融合kernel来提升推理吞吐。pyasc支持的Python语法覆盖了算子开发中最常用的控制流和数据操作——for循环、条件判断、赋值、算术运算、下标访问、切片、增强赋值等这些语法足以表达绝大多数算子的计算逻辑。pyasc也适合作为Ascend C的学习桥梁。开发者可以先通过pyasc熟悉昇腾NPU的编程模型——Global Memory与Local Memory的数据搬运、双缓冲流水线、事件同步机制——然后再过渡到直接使用Ascend C。pyasc的PYASC_DUMP_PATH机制让开发者可以查看自己Python代码对应的Ascend C输出这个特性让pyasc成为了理解Ascend C的实践工具。使用前和使用后的效率对比场景使用前通用Python方式使用后pyasc方式差异来源昇腾NPU算子调用依赖外部Python包接口分散统一入口函数签名清晰减少集成成本调试效率错误信息不明确定位困难报错直接指明问题位置缩短排查周期代码可读性各厂家接口命名不一致标准化命名规范提升协作效率部署方式需要手动配置多路径一键安装开箱即用降低部署门槛总结pyasc目前有一些明确的限制需要了解。不支持嵌套函数、lambda表达式、类方法装饰、break和continue语句、异常处理、以及标准库和第三方库函数的调用。这些限制意味着pyasc面向的是相对纯粹的数值计算kernel而不是通用编程场景。在算子开发中绝大多数计算逻辑可以用for循环、条件判断和函数调用组合完成上述不支持的语言特性在kernel代码中本身也很少出现。另外pyasc支持的接口范围仍在逐步扩展中当前已开放的接口与Ascend C的完整API集之间还存在差距。如果开发者需要使用尚未被pyasc覆盖的Ascend C接口仍需直接编写Ascend C代码。对于需要极致控制硬件行为的场景或者需要使用pyasc尚未覆盖的Ascend C API的场景直接编写Ascend C仍然是必要的选择。pyasc的接口与Ascend C一一对应这意味着它不会提供Ascend C不具备的能力。它的价值是把已有的Ascend C能力用更友好的方式暴露出来而不是扩展Ascend C本身的功能边界。仓库地址https://atomgit.com/cann/pyasc