尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

在 pyasc 中使用 asc.language.adv.power 实现按元素幂运算

在 pyasc 中使用 asc.language.adv.power 实现按元素幂运算 在 pyasc 中使用 asc.language.adv.power 实现按元素幂运算【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc本文围绕 CANN pyasc 项目为 Python 用户提供与 Ascend C 一一对应的算子编程接口中的asc.language.adv.power高阶 API系统讲解其函数签名、参数语义、对应的 Ascend C 函数原型、底层 IR 构建链路与约束条件并结合仓库源码与单元测试给出可直接运行的双输入幂运算 kernel 写法。读完本文你将能够在昇腾 AI 处理器上正确、高效地编写按元素幂Power运算算子。功能概述asc.language.adv.power用于按元素做幂运算即对两个源操作数src0与src1逐元素计算dst[i] src0[i] ** src1[i]其中src0是底数、src1是指数结果写回目的操作数dst。它是 pyasc 高阶 APIasc.language.adv命名空间中少有的二元数学运算接口区别于sin、cos、exp等一元接口与按位异或xor在形式上相似但语义完全不同。该接口在仓库中的官方文档位于 docs/python-api/language/generated/asc.language.adv.power.md实现位于 python/asc/language/adv/math.py与 Ascend C 的Power算子原语一一对应。函数签名与参数说明asc.language.adv.power(dst: LocalTensor, src0: LocalTensor, src1: LocalTensor, count: int | None None, temp_buffer: LocalTensor | None None, is_reuse_source: bool False) - None参数类型说明dstLocalTensor目的操作数支持的 TPosition 为 VECIN / VECCALC / VECOUT。src0LocalTensor源操作数底数支持的 TPosition 为 VECIN / VECCALC / VECOUT数据类型需与dst保持一致。src1LocalTensor源操作数指数支持的 TPosition 为 VECIN / VECCALC / VECOUT数据类型需与dst保持一致。countint | None参与计算的元素个数默认None表示对整块缓冲区执行计算。temp_bufferLocalTensor | None临时内存空间类型为 LocalTensor支持的 TPosition 为 VECIN / VECCALC / VECOUT通常以uint8类型的 VECCALC 张量提供。is_reuse_sourcebool是否允许修改源操作数默认值为false。其中LocalTensor的完整定义可参考 asc.language.core.LocalTensor相关文档索引见 asc.language.adv 总览。从源码看power同时提供了静态类型版本count: Optional[int]、is_reuse_source: bool与 JIT 运行时版本count: Optional[RuntimeInt]、is_reuse_source: RuntimeBool两个重载前者用于 IDE 类型提示与静态检查后者在asc.jit编译场景下实际生效源码见 math.py#L392-L446overload def power(dst: LocalTensor, src0: LocalTensor, src1: LocalTensor, count: Optional[int] None, temp_buffer: Optional[LocalTensor] None, is_reuse_source: bool False) - None: ... require_jit def power(dst: LocalTensor, src0: LocalTensor, src1: LocalTensor, count: Optional[RuntimeInt] None, temp_buffer: Optional[LocalTensor] None, is_reuse_source: RuntimeBool False) - None: ... math_op_impl((dst, src0, src1), count, temp_buffer, is_reuse_source, create_asc_PowerOp)对应的 Ascend C 函数原型asc.language.adv.power与 Ascend C 的Power模板函数一一对应。pyasc 会根据是否传入count即calCount和temp_buffer即sharedTmpBuffer自动选择底层函数重载对应的四个 C 原型如下template typename T, bool isReuseSource false __aicore__ inline void Power(const LocalTensorT dstTensor, const LocalTensorT src0Tensor, const LocalTensorT src1Tensor, const LocalTensoruint8_t sharedTmpBuffer, uint32_t calCount) template typename T, bool isReuseSource false __aicore__ inline void Power(const LocalTensorT dstTensor, const LocalTensorT src0Tensor, const LocalTensorT src1Tensor, const LocalTensoruint8_t sharedTmpBuffer) template typename T, bool isReuseSource false __aicore__ inline void Power(const LocalTensorT dstTensor, const LocalTensorT src0Tensor, const LocalTensorT src1Tensor, uint32_t calCount) template typename T, bool isReuseSource false __aicore__ inline void Power(const LocalTensorT dstTensor, const LocalTensorT src0Tensor, const LocalTensorT src1Tensor)注意几个关键对应关系dst→dstTensorsrc0→src0Tensorsrc1→src1Tensorcount→calCountuint32_t参与计算的元素个数temp_buffer→sharedTmpBufferLocalTensoruint8_t即共享临时缓冲区这正是测试中临时张量使用uint8类型的原因is_reuse_source→ 模板参数isReuseSource默认false。底层实现原理从 Python 调用到 IR 构建power的实现非常精简真正的逻辑都收敛在math_op_impl中。该函数是所有数学类高阶 APIsin、cos、exp、power、xor等共用的底层实现源码见 math.py#L19-L29def math_op_impl(tensors: Tuple[LocalTensor], count: Optional[RuntimeInt], temp_buffer: Optional[LocalTensor], is_reuse_source: RuntimeBool, build_method: str) - None: if count is not None: check_type(count, count, RuntimeInt) count _mat(count, KnownTypes.int32).to_ir() if temp_buffer is not None: check_type(temp_buffer, temp_buffer, LocalTensor) temp_buffer temp_buffer.to_ir() is_reuse_source _mat(is_reuse_source, KnownTypes.bit).to_ir() getattr(global_builder.get_ir_builder(), build_method)(*(t.to_ir() for t in tensors), sharedTmpBuffertemp_buffer, calCountcount, isReuseSourceis_reuse_source)其核心处理流程可以概括为四条类型校验count必须是RuntimeInt对应int32类型temp_buffer必须是LocalTensor常量物化通过materialize_ir_value_mat将count物化为int32、将is_reuse_source物化为bit类型的 IR 常量张量转换把所有LocalTensor参数通过to_ir()转换为 IR 值算子构建动态调用 IR builder 上的create_asc_PowerOp传入sharedTmpBuffer、calCount、isReuseSource三个关键字参数完成算子 IR 节点的创建。也就是说一次asc.adv.power(dst, src0, src1, count512, temp_buffertmp)调用最终会在 IR 中构建出一个create_asc_PowerOp节点之后再经过代码生成流水线翻译为对应的 Ascend C 代码。在发射端PowerOp已在 lib/Target/AscendC/Translation.cpp#L128 中注册进翻译映射表ascendc::PowerOp确保 IR 能正确下译为 Ascend C 的Power原语并最终落盘为昇腾可执行的算子代码。约束说明使用asc.language.adv.power时需遵守以下约束地址不允许重叠不支持源操作数与目的操作数地址重叠即dst不能与src0、src1指向同一段缓冲区地址对齐操作数地址对齐要求请参见《Ascend C 算子开发接口》中的通用说明和约束-通用地址对齐约束文档对应链接见 asc.language.adv.power 官方文档数据类型一致src0、src1的数据类型需要与dst保持一致TPosition 限制所有张量均须位于 VECIN / VECCALC / VECOUT 位置向量计算相关位置。调用示例官方文档给出的最小调用形式为asc.adv.power(dst, src0, src1)即省略count、temp_buffer、is_reuse_source三个可选参数此时按整块缓冲区进行全量幂运算。仓库单元测试中的完整用法在仓库的单元测试 python/test/unit/language/adv/test_ops.py#L352-L364 中test_power_kernel给出了一个完整、可运行参考的 kernel 写法展示了temp_buffer传与不传两种形态def test_power_kernel(mock_launcher_run): asc.jit def power_kernel(): x_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512) y_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512) z_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECOUT, addr0, tile_size512) tmp asc.LocalTensor(dtypeasc.uint8, posasc.TPosition.VECCALC, addr0, tile_size512) asc.adv.power(z_local, x_local, y_local, count512, temp_buffertmp) asc.adv.power(z_local, x_local, y_local, count512) power_kernel[1]() assert mock_launcher_run.call_count 1从该测试可以提炼出编写powerkernel 的实操要点张量声明使用asc.LocalTensor(dtype..., posasc.TPosition.X, addr..., tile_size...)声明片上张量dtype统一为asc.float16源/目的位置分别用VECIN与VECOUT临时缓冲区temp_buffer声明为dtypeasc.uint8、posasc.TPosition.VECCALC的 LocalTensor与 Ascend C 原型中的LocalTensoruint8_t sharedTmpBuffer对应两种调用形态带temp_buffer的调用内部自动匹配带sharedTmpBuffer的 C 重载与不带temp_buffer的调用匹配不带临时缓冲区的重载可以同时出现在同一 kernel 中count语义count512表示本次计算参与幂运算的元素个数为 512对应 C 侧的calCount512JIT 编译kernel 函数需用asc.jit装饰并通过power_kernel[1]()以指定核数1 核启动随后由 pyasc 运行时完成编译与mock_launcher_run启动。常见问题与使用建议power与xor的区别二者都是双输入逐元素接口但power是数学幂运算dst[i] src0[i] ** src1[i]xor是按位异或运算dst[i] src0[i] ^ src1[i]且xor支持整型如int16输入务必按业务语义选择is_reuse_source的取舍默认false表示不允许修改源操作数更安全若确知源数据后续不再使用、且希望释放临时空间或减少缓冲占用可显式置True但需自行承担源数据被覆盖的风险地址重叠约束由于不支持源与目的地址重叠实践中应保证dst使用独立缓冲区或将同一数据先拷贝到新的VECCALC缓冲后再作为源输入性能考量幂运算涉及指数计算需要临时中间缓冲建议复用temp_buffer并在循环外统一申请避免反复分配count精确指定参与元素数可避免对尾部无效数据的多余计算。小结asc.language.adv.power是 pyasc 高阶 API 中用于按元素幂运算的标准接口它通过math_op_impl统一完成类型校验、常量物化与create_asc_PowerOp算子节点构建并在翻译阶段映射到 Ascend C 的Power原语。结合本文给出的参数语义、四个 C 原型对应关系、约束条件以及源自仓库单元测试的完整 kernel 写法你可以直接在昇腾 AI 处理器上编写正确的幂运算算子。更多数学类高阶 APIsin、cos、exp、sqrt、xor等均可参考同一实现模式见 python/asc/language/adv/math.py 与 asc.language.adv 文档总览。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表