CANN pyasc 教程:使用 LocalTensor.set_user_tag 为 Tensor 打标并与 get_user_tag 协同控制算子执行流 📅 发布时间:2026/9/19 19:52:25 👁 浏览次数: CANN pyasc 教程使用 LocalTensor.set_user_tag 为 Tensor 打标并与 get_user_tag 协同控制算子执行流【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc导读LocalTensor.set_user_tag是 CANN pyasc 编程框架中用于为本地Local MemoryTensor 添加用户自定义 Tag 信息的核心接口。它允许开发者在算子 Kernel 内部为某个 Tensor 绑定一个整数标签并通过配套的get_user_tag读取该标签进而在 Kernel 执行期实现按标签分支的灵活控制流常用于多路输入 Tensor 混用、条件式计算、调试与调度策略差异化等场景。读完本文你将掌握set_user_tag/get_user_tag的完整用法、与 Ascend C 底层SetUserTag/GetUserTag的对应关系、在仓库源码与单元测试中的实现证据以及一套可直接落地的实战示例。一、接口定位LocalTensor 的用户自定义信息通道在 CANN pyasc 中LocalTensor 用于存放 AI Core 中 Local Memory内部存储的数据支持逻辑位置TPosition为VECIN、VECOUT、VECCALC、A1、A2、B1、B2、CO1、CO2参见 tensor.py 的类定义。LocalTensor 除了承载数据本身还提供了一批描述性/辅助性接口例如get_size、get_position、get_length、set_shape_info等。其中set_user_tag为 Tensor 添加用户自定义信息用户可以根据需要设置对应的 Tagget_user_tag获取指定 Tensor 块的 Tag 信息用户可以根据 Tag 信息对 Tensor 进行不同操作。两者的语义与 Ascend C 原生接口一一对应__aicore__ inline void SetUserTag(const TTagType tag) __aicore__ inline TTagType GetUserTag() const也就是说pyasc 的 Python 接口在编译期会被翻译为 Ascend C 的成员函数调用Tag 类型TTagType对应int32_t。二、接口原型与参数说明2.1 set_user_tagLocalTensor.set_user_tag(tag: int 0) → None参数说明tag设置的 Tag 信息类型TTagType对应为int32_t默认为0。调用示例来自官方文档tag 10 size input_local.get_size(tag)2.2 get_user_tagLocalTensor.get_user_tag() → int参数说明无。返回值说明指定 Tensor 块的 Tag 信息整数。调用示例来自官方文档tensor1 que1.deque(asc.half) tag1 tensor1.get_user_tag() tensor2 que2.deque(asc.half) tag2 tensor2.get_user_tag() tensor3 que3.alloc_tensor(asc.half) # 使用Tag控制条件语句执行 if tag1 10 and tag2 9: asc.add(tensor3, tensor1, tensor2, counttile_length)从示例可见get_user_tag常与deque/alloc_tensor配合使用从队列取出的 Tensor 自带某个 Tag或为新建 Tensor 设置 Tag 后通过条件判断决定是否执行计算。三、仓库源码级实现解析3.1 Python 侧实现set_user_tag与get_user_tag的 Python 实现位于 python/asc/language/core/tensor.py# get_user_tag require_jit set_tensor_docstring(tensor_nameLocalTensor, api_nameget_user_tag) def get_user_tag(self) - RuntimeInt: builder global_builder.get_ir_builder() handle builder.create_asc_LocalTensorGetUserTagOp(builder.get_i32_type(), self.to_ir()) return PlainValue(handle)# set_user_tag require_jit set_tensor_docstring(tensor_nameLocalTensor, api_nameset_user_tag) def set_user_tag(self, tag: RuntimeInt 0) - None: global_builder.get_ir_builder().create_asc_LocalTensorSetUserTagOp(self.to_ir(), _mat(tag).to_ir())关键点两个方法都带有require_jit装饰器说明它们必须在asc.jit修饰的 Kernel 函数内调用属于编译期内联展开的接口get_user_tag通过 IR Builder 创建asc_LocalTensorGetUserTagOp返回类型为i32并包装成RuntimeInt/PlainValueset_user_tag创建asc_LocalTensorSetUserTagOptag 参数通过_mat转为 IR 值后传入默认值0与文档一致两者的 docstring 均由set_tensor_docstring统一生成保证与 docs/python-api/language/generated/ 下的自动生成文档同步。3.2 IR 层TableGen定义在 MLIR TableGen 描述文件 include/ascir/Dialect/Asc/IR/Core/Tensor.td 中两个算子被定义为APIOpdef AscendC_LocalTensorGetUserTagOp : APIOplocal_tensor.get_user_tag, GetUserTag, [AscMemberFunc] { let summary Call AscendC::LocalTensor::GetUserTag method; let arguments (ins AscendC_LocalTensor:$tensor); let results (outs AnySignlessIntegerOrIndex:$result); ... } def AscendC_LocalTensorSetUserTagOp : APIOplocal_tensor.set_user_tag, SetUserTag, [AscMemberFunc] { let summary Call AscendC::LocalTensor::SetUserTag method; let arguments (ins AscendC_LocalTensor:$tensor, I32:$tag); }从源码结构可以推断SetUserTag算子接收AscendC_LocalTensor与I32类型的 tag 参数与int32_t对应GetUserTag算子接收一个 LocalTensor输出任意无符号整数或索引类型实际为i32两者都标注为AscMemberFunc即最终会被翻译成对AscendC::LocalTensor成员函数的调用分别对应GetUserTag()与SetUserTag(tag)。这一层设计说明Tag 机制在编译期被完整建模为 IR 算子而非简单的 Python 运行时属性因此它既能在设备端 Kernel 中生效也能参与后续的优化与代码生成流水线。四、在真实算子中的组合使用仓库的泛化测试 python/test/generalization/basic/test_vadd_sw.py 给出了一个完整的实战组合先用set_user_tag打标再用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() tmp_dtype asc.DataType(int32) tmp_local tmp_queue.alloc_tensor(tmp_dtype) new_tmp tmp_local.reinterpret_cast(x_gm.dtype) if z_tag 13 and src_pos asc.TPosition.VECIN: asc.add(z_local, x_local, y_local, counttile_length)这段代码展示的典型模式是为x_local设置自定义 Tag13调用get_user_tag()读回z_tag在条件语句中同时校验 Tag 与逻辑位置TPosition.VECIN满足条件才执行asc.add向量加运算。由此可以看出 Tag 的典型用途之一在同一个 Kernel 中区分携带不同元信息的 Tensor并用条件分支决定对它们施加的计算避免为每种情况编写独立 Kernel。五、单元测试验证仓库单元测试 python/test/unit/language/core/test_local_tensor.py 对两个接口分别做了最小化验证def test_get_user_tag(mock_launcher_run): asc.jit def kernel_get_user_tag() - None: x_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512) tag x_local.get_user_tag() kernel_get_user_tag[1]() assert mock_launcher_run.call_count 1 def test_set_user_tag(mock_launcher_run): asc.jit def kernel_set_user_tag() - None: x_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512) x_local.set_user_tag(10) kernel_set_user_tag[1]() assert mock_launcher_run.call_count 1要点测试通过asc.jit构造 Kernel用asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512)直接创建本地 Tensorset_user_tag(10)与get_user_tag()均能完成一次完整的编译启动流程mock_launcher_run被调用一次即通过证明接口在编译链路中可正常展开为 IR 算子并生成启动代码这也佐证了set_user_tag的默认值语义——显式传10与文档中默认0一致地写入 tag 参数。六、典型应用场景与注意事项6.1 典型应用场景多队列 Tensor 区分从不同TQue取出的 Tensor 打上不同 TagKernel 内统一按 Tag 分支处理收敛代码规模条件式计算结合get_user_tag返回值与if语句控制算子如asc.add、asc.data_copy等是否执行实现运行时策略选择调试与追踪为参与计算的 Tensor 标记来源或阶段号便于在复杂流水线中跟踪数据流转配合 TQue 生命周期接口deque/alloc_tensor取出的 Tensor 同样可调用set_user_tag/get_user_tag且不影响内存分配语义。6.2 注意事项必须在asc.jit修饰的 Kernel 函数内调用两个方法均带require_jit不可在 Host 侧 Python 代码中直接使用tag 类型为int32_tI32超出 32 位有符号整数范围的值不会被正确编码set_user_tag默认值0未显式设置时get_user_tag返回的即为该默认值Tag 信息是用户自定义元数据不会改变 Tensor 的地址、长度、shape 等实际数据属性也不影响内存布局。七、相关接口速查set_user_tag/get_user_tag是 LocalTensor 接口族的一员与之相邻的常用接口包括接口说明get_user_tag获取指定 Tensor 块的 Tag 信息set_user_tag为 Tensor 添加用户自定义 Tag 信息get_size获取当前 LocalTensor Size 大小单位为元素get_position获取 LocalTensor 所在的 TPosition 逻辑位置set_buffer_len设置 Buffer 长度便于编译器对内存及同步进行自动优化set_size设置当前 LocalTensor Size 大小总结LocalTensor.set_user_tag与LocalTensor.get_user_tag构成了 CANN pyasc 中 LocalTensor 的用户自定义元信息通道set_user_tag写入int32_t类型的 Tagget_user_tag在 Kernel 内读回该值用于条件分支。二者在编译期被建模为asc_LocalTensorSetUserTagOp/asc_LocalTensorGetUserTagOp两个 IR 算子最终翻译为 Ascend C 的SetUserTag/GetUserTag成员函数调用。借助仓库中的泛化测试与单元测试开发者可以快速验证这一模式并将其应用在多队列 Tensor 区分、条件式计算、调试追踪等实际算子开发场景中。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考