ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

CANN pyasc 算子开发:LocalTensor.get_position() 获取片上存储逻辑位置(TPosition)实战指南

2026/9/19 23:51:31 拓冰建站 浏览量
CANN pyasc 算子开发:LocalTensor.get_position() 获取片上存储逻辑位置(TPosition)实战指南 CANN pyasc 算子开发LocalTensor.get_position() 获取片上存储逻辑位置TPosition实战指南【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyascLocalTensor.get_position()是 CANN pyasc 中用于查询 LocalTensor 所在 TPosition片上 Local Memory 逻辑位置的核心内建方法在算子 Kernel 内可用于实现“依据张量存放位置分发不同计算/搬运逻辑”的动态分支。本指南基于 官方 API 文档 展开并结合 tensor.py、enums.py、Tensor.td 及单元/泛化测试完整讲解接口签名、返回值的枚举语义、底层实现链路与真实算子中的分支应用帮助读者在 pyasc Kernel 中正确、安全地使用该接口。一、接口速览与适用场景在昇腾 AI Core 上片上 Local Memory 被划分为多个逻辑位置TPosition例如用于矩阵运算的 A1/A2/B1/B2、用于向量计算的 VECIN/VECOUT/VECCALC、以及累加输出 CO1/CO2 等。Kernel 编程中同一段数据可能因指令类型不同而被搬运到不同位置get_position()正是为此设计的“查询型”接口函数签名LocalTensor.get_position() → int功能获取调用它的 LocalTensor 当前所在的 TPosition 逻辑位置。参数无。返回值一个整数值对应 TPosition 枚举中的某个逻辑位置。对应 Ascend C 原生原型__aicore__ inline int32_t GetPosition() const在 pyasc 中通过 JIT 编译等价映射为该 C 内联方法。典型应用场景包括位置感知的算子分发在 Kernel 内根据src_pos asc.TPosition.VECIN等条件决定走向量计算分支还是矩阵搬运分支通用内核generic kernel适配同一段内核代码兼容不同摆放位置的数据运行时动态判断调试与自检在开发阶段打印或断言张量位置是否符合预期。该接口是 LocalTensor 的只读内建方法不修改张量本身返回的整型值可参与if/elif/else分支判断。二、返回值语义TPosition 枚举完整对照get_position()返回的整数并非随意定义而是与 pyasc 暴露的asc.TPositionIntEnum一一对应。其完整定义位于 python/asc/language/core/enums.py现对照如下TPosition 成员整数值语义GM0Global Memory片外全局内存GlobalTensor 位置A11矩阵运算 A 侧输入位置A22矩阵运算 A 侧输入第二份/乒乓缓冲B13矩阵运算 B 侧输入位置B24矩阵运算 B 侧输入乒乓缓冲C15矩阵累加输出 C 位置C26矩阵累加输出 C乒乓缓冲CO17C 输出用于后续向量运算的扩展位置CO28C 输出扩展位置CO1 的乒乓/对偶VECIN9向量单元输入缓冲位置VECOUT10向量单元输出缓冲位置VECCALC11向量计算单元专用计算位置MAX12位置枚举上限哨兵值一般不作业务使用注意官方文档明确get_position()支持 TPosition 为VECIN、VECOUT、VECCALC、A1、A2、B1、B2、CO1、CO2从枚举定义看它实际可覆盖从GM到VECCALC的全部逻辑位置MAX仅作为边界哨兵。这也与 LocalTensor 的类注释 完全一致——LocalTensor 用于存放 AI Core 中 Local Memory 的数据支持逻辑位置为 VECIN、VECOUT、VECCALC、A1、A2、B1、B2、CO1、CO2。由于TPosition继承自 Python 的IntEnum返回值可以直接与枚举成员做比较如src_pos asc.TPosition.VECIN也可以当作普通整数参与运算。在 Kernel 中推荐始终通过asc.TPosition.XXX常量比较避免硬编码魔数、提高可读性。三、Python 侧实现原理从 pyasc API 到 Ascend C 调用链在 pyasc 中get_position()是一个“文档字符串由基类模板注入、实际行为由 JIT 前端生成 IR”的方法。其 Python 定义位于 python/asc/language/core/tensor.pyoverload def get_position(self) - int: ... require_jit set_tensor_docstring(tensor_nameLocalTensor, api_nameget_position) def get_position(self) - RuntimeInt: builder global_builder.get_ir_builder() handle builder.create_asc_LocalTensorGetPositionOp(builder.get_i32_type(), self.to_ir()) return PlainValue(handle)关键点解读require_jit该接口只能在 Kernel 函数被asc.jit装饰中调用编译期由前端解释执行而不是普通 Python 运行时行为IR 构建调用create_asc_LocalTensorGetPositionOp在中间表示IR中创建一条local_tensor.get_position算子指令输入是当前 LocalTensor输出类型为i3232 位整型与 Ascend C 的int32_t GetPosition()返回类型对齐返回值封装返回PlainValue(handle)即一个运行时整数值RuntimeInt可直接参与 Kernel 内的分支比较。其 IR 算子的 TableGen 定义位于 include/ascir/Dialect/Asc/IR/Core/Tensor.tddef AscendC_LocalTensorGetPositionOp : APIOplocal_tensor.get_position, GetPosition, [AscMemberFunc] { let summary Call AscendC::LocalTensor::GetPosition method; let arguments (ins AscendC_LocalTensor:$tensor); let results (outs I32:$result); let assemblyFormat [{ $tensor attr-dict : qualified(type($tensor)) , type($result) }]; }从该定义可以清晰看出完整调用链pyasc Python 方法 → pybind/IR Builder 生成local_tensor.get_positionOp → Lower 到 Ascend C 的LocalTensor::GetPosition()内联方法最终在 AI Core 上执行。这一“IR 即桥”的设计使 Python 侧保持了原生语法同时底层无缝复用 Ascend C 的全部能力。文档字符串docstring同样由 python/asc/language/core/utils.py 中的LocalTensorDocstring.get_position_docstring()模板生成其中的函数介绍、Ascend C 原型、参数/返回值说明与调用示例即官方 API 文档 的原始来源。四、标准调用示例与分支分发写法官方文档给出了get_position()的典型用法——把返回位置作为分发条件为不同位置执行不同处理逻辑src_pos input_local.get_position() if src_pos asc.TPosition.VECCALC: # 处理逻辑1例如按向量计算位置执行向量指令 elif src_pos asc.TPosition.A1: # 处理逻辑2例如按矩阵 A 侧输入位置执行矩阵搬运/计算 else: # 处理逻辑3其他位置兜底处理使用要点务必通过asc.TPosition.XXX枚举常量比较避免硬编码整数虽然TPosition是IntEnum直接用整数也能比较但可读性和可维护性差该接口无参数、无副作用可安全地在任意 Kernel 位置调用不影响流水返回值是编译期可知的常量或运行时值取决于张量来源若 LocalTensor 由asc.LocalTensor(dtype..., posasc.TPosition.VECIN, addr0, tile_size512)显式创建则其位置在创建时即已确定此时分支通常可被编译器常量折叠若张量来自队列如TQue.alloc_tensor或经过reinterpret_cast、set_addr_with_offset等操作位置语义需结合具体指令上下文判断。五、仓库实证单元测试与泛化测试中的真实用法5.1 单元测试验证接口可编译可运行pyasc 在 python/test/unit/language/core/test_local_tensor.py 中为get_position提供了专门的单元测试def test_get_position(mock_launcher_run): asc.jit def kernel_get_position() - None: x_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512) pos x_local.get_position() kernel_get_position[1]() assert mock_launcher_run.call_count 1该测试说明显式以posasc.TPosition.VECIN创建 LocalTensor 后直接调用get_position()可以正常完成 JIT 编译与启动mock_launcher_run被调用一次。这也是TPosition作为LocalTensor构造参数直接使用的官方佐证。5.2 泛化测试位置感知分支在真实 vadd 内核中的落地在 python/test/generalization/basic/test_vadd_sw.py 的软件流水线 vadd 内核中get_position()与get_length()、get_size()、get_user_tag()等接口组合使用展示了“查询位置 → 条件分支 → 执行计算”的完整实战模式src_pos x_local.get_position() x_len x_local.get_length() x_size x_local.get_size() x_tag 13 # set custom tag x_local.set_user_tag(x_tag) z_tag x_local.get_user_tag() if z_tag 13 and src_pos asc.TPosition.VECIN: asc.add(z_local, x_local, y_local, counttile_length)这里x_local来自in_queue_x.alloc_tensor(x_gm.dtype)由 TPipe 队列分配在加法指令前先校验其位置是否为VECIN只有位置匹配才执行asc.add向量加法——这正是位置感知分支在真实算子内核中的典型落点可作为开发者编写自定义内核时的参考范式。六、与其他 LocalTensor 内建接口的组合使用get_position()通常与 LocalTensor 的其余查询/设置接口配合构成完整的“张量元信息自省”能力全部定义于 python/asc/language/core/tensor.py接口功能返回类型get_position()获取 TPosition 逻辑位置inti32get_phy_addr(offset0)获取物理地址可带偏移intu64get_length()获取数据长度intu32get_size()获取元素个数inti32get_shape_info()获取形状信息ShapeInfoget_user_tag()/set_user_tag(tag)读写用户自定义标签intreinterpret_cast(dtype)按新数据类型重新解释LocalTensorset_addr_with_offset(base, offset)基于基地址设置偏移地址None典型组合先get_position()判断位置是否支持目标指令再用get_size()/get_length()确定计算边界最后通过get_user_tag()/set_user_tag()做跨流水段的自定义标记传递。测试文件 test_vadd_sw.py 即示范了该组合拳。七、注意事项与最佳实践调用上下文get_position()受require_jit约束只能在asc.jit装饰的 Kernel 函数内调用不能在纯 Python 主流程中使用比较方式优先使用asc.TPosition.XXX枚举成员比较避免魔法数字由于是IntEnum与int类型比较语义正确位置与指令的匹配不同指令对输入/输出张量的 TPosition 有硬性要求例如 Cube 矩阵指令要求 A1/A2/B1/B2/CO1/CO2向量指令要求 VECIN/VECOUT/VECCALC使用get_position()做前置校验可有效避免“位置不符导致指令行为异常”的隐性错误性能开销GetPosition()在 Ascend C 中是__aicore__ inline内联只读查询编译期往往可被常量折叠运行时开销可忽略可放心在热路径中使用文档一致性如需查阅最新语义官方 API 文档与 utils.py 中的 docstring 模板保持同步两者可交叉验证。八、总结LocalTensor.get_position()是 CANN pyasc 算子内核中查询张量片上逻辑位置的标准接口返回与asc.TPosition枚举一一对应的整型值覆盖VECIN/VECOUT/VECCALC/A1/A2/B1/B2/CO1/CO2等全部 Local Memory 逻辑位置。其实现链路清晰Python 方法在 JIT 前端生成local_tensor.get_positionIR 算子Tensor.td最终 lower 为 Ascend C 的LocalTensor::GetPosition()。通过结合单元测试与真实 vadd 泛化内核的分支应用开发者可以在自己的算子中实现“位置感知”的指令分发逻辑提升内核的通用性与健壮性。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考