PyPTOget_spr特殊寄存器读取 API 详解读取 AddrReg 字节数实现压缩数据长度感知【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读pypto_pro.language.get_spr是 CANN PyPTOParallel Tensor/Tile Operation 编程范式提供的一个系统级工具 API用于在 Vector 流水线中读取特殊目的寄存器SPR的值并返回标量。当前实现仅支持 AddrRegAR地址寄存器其典型应用场景是与vf.squeeze配合vf.squeeze将有效元素的总字节数写入 AR 寄存器随后在核函数体中通过get_spr读出该值从而感知压缩squeeze后有效数据的实际长度。读完本文你将掌握get_spr的函数原型、产品支持范围、源码级实现原理get_ar()指令下发链路以及它与vf.squeeze/vf.squeeze_store_unalign的完整配合套路并可直接复用仓库中的可运行示例。功能说明读取指定特殊寄存器的值get_spr用于读取指定特殊寄存器的值返回一个标量值。根据官方文档 get_spr API 文档当前仅支持 AddrReg 寄存器。AddrReg 是一个特殊的地址寄存器其核心使用模式是vf.squeeze执行数据压缩操作将有效元素的总字节数存储到 AddrReg 寄存器中随后通过pypto_pro.language.get_spr读取该字节数值供后续逻辑如非对齐存储、长度判断、动态步长计算使用。也就是说get_spr本质上是一个读出 Squeeze 结果长度的通道它把 Vector 指令内部产生、通常对程序员不可见的硬件状态暴露为可用的标量值。产品支持情况get_spr的产品支持范围非常明确与当前仓库的硬件后端密切相关Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持即该 API 目前仅在 Ascend 950 系列硬件上可用。在编写可移植内核时需要注意对 A2/A3 平台做能力判断或分支处理避免在不支持的硬件上调用导致编译或运行失败。函数原型与参数说明get_spr() - ar_value参数无参数。当前仅支持读取 AR 寄存器因此无需传入寄存器标识。返回值返回ar_value标量值即 AR 寄存器中保存的有效元素总字节数。约束无从 API 使用角度没有额外约束但底层指令存在作用域限制见下文注意事项。源码级实现从 Python API 到get_ar()指令get_spr虽是一个轻量 API但它在仓库中贯穿了 Python 前端声明、IR 注册、CCE 后端代码生成三层实现理解这条链路有助于掌握其真实行为。Python 前端声明在 python/pypto_pro/language/_api.py 中get_spr被声明为_api_decl def get_spr() - int: Read a special purpose register value (get_ar instruction). Currently only the AR register is supported. The AR register stores the total byte count of valid elements produced by Squeeze. get_ar() is an __aicore__ instruction and cannot be used inside pl.vector_function. Call this in the pl.jit kernel body. Returns: int64_t scalar value from the SPR 从这段 docstring 可以确认三个关键事实底层下发的指令是get_ar即 get address registerAR 寄存器中保存的是Squeeze 产生的有效元素的总字节数total byte countget_ar属于__aicore__指令不能在pl.vector_function内部调用必须写在pl.jit核函数体kernel body中——这与文档示例中get_spr出现在pl.section_vector()块之后的写法一致。该 API 会通过 python/pypto_pro/language/init.py 的导出清单暴露为pl.get_spr供用户在核函数中直接调用。IR 层注册在 python/pypto_pro/ir/op/system_ops.py 中get_spr被注册为 IR 算子get_spr: OpSpec(ir_nameget_spr, parse_argsFalse, parse_kwargsFalse),parse_argsFalse, parse_kwargsFalse表明该算子不接受任何参数这与无参数的 API 定义一致。而在 C 侧的 IR 校验实现 framework/src/interface/ir/op/block_ops/memory.cpp 中REGISTER_OP(get_spr)同样会执行args.size() 0的参数个数校验任何带参数的调用都会触发INVALID_ARGUMENT报错从编译早期就杜绝了误用。CCE 后端代码生成真正下发指令的位置在 framework/src/interface/pypto_pro/backend/backend_cce_ops.cpp// Helper function for get_spr (reads AR special purpose register via get_ar()) static std::string MakeGetSprCodegenCCE(const ir::CallPtr op, codegen::CodegenBase codegen_base) { (void)codegen_base; PRO_CODEGEN_CHECK(ExternalError::INVALID_ARGUMENT, op-args_.size() 0) get_spr requires no arguments; return get_ar(); } REGISTER_BACKEND_OP(BackendCCE, get_spr) .set_pipe(ir::PipeType::V) .f_codegen([](const ir::CallPtr op, codegen::CodegenBase codegen) { return MakeGetSprCodegenCCE(op, codegen); });这里有两个值得注意的实现细节代码生成直接输出get_ar()指令算子被挂载到V 流水线ir::PipeType::V说明get_spr的语义与 Vector 流水线执行单元绑定这也解释了它为什么必须与vf.squeeze等 Vector 指令配合使用、并应在 Vector section 相关的核函数体上下文中调用。与vf.squeeze的配合压缩长度感知的完整链路get_spr的价值只有在与vf.squeeze配合时才能体现。在 python/pypto_pro/language/_vf_api.py 中vf.squeeze被定义为staticmethod _api_decl def squeeze(src, preg, gather_mode: Optional[SqueezeMode] None): Squeeze mask to index register (vsqz instruction). Converts active mask bits into a packed index sequence in the destination register. For each active lane in mask, the corresponding element of src is compressed (squeezed) into a contiguous region at the start of dst. Args: src: Source register preg: Predicate mask register Kwargs: gather_mode: pl.SqueezeMode.STORE_REG or pl.SqueezeMode.NO_STORE_REG Returns: Destination register (RegTensor) with packed/squeezed elements. 它通过vsqz指令把谓词掩码predicate mask中激活的通道active lane对应的元素压缩到目标寄存器起始的连续区域并把有效数据的总字节数写入 AR 寄存器。gather_mode参数支持pl.SqueezeMode.STORE_REG写入 AR 寄存器与pl.SqueezeMode.NO_STORE_REG不写两种模式。AR 寄存器中的字节数还会被其他 API 复用。例如 python/pypto_pro/language/_vf_api.py 中的vf.squeeze_store_unalign文档明确说明Reads the valid byte count from the AR register (written byvf.squeezewithgather_modeSTORE_REG) as the implicit stride. Must be called aftervf.squeeze(STORE_REG)and paired withvf.squeeze_store_unalign_post.这说明 AR 寄存器是Squeeze → 后续处理这条数据链路上的共享硬件状态vf.squeeze写、get_spr读供宿主逻辑判断长度、vf.squeeze_store_unalign读作为隐式步长执行非对齐存储。三者共同构成了压缩数据从计算到存储的完整闭环。调用示例在核函数中读取 AR 寄存器以下是文档给出的完整可运行示例它在pl.jit核函数中执行加载 → 掩码压缩 → 非对齐存储的 Vector 流程并在pl.section_vector()块结束后调用pl.get_spr()读取压缩后的有效字节数import os import pypto_pro.language as pl import torch import torch_npu pl.vector_function def example_vf(src_tile, dst_tile): preg vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_FP32) reg vf.load_align(src_tile, 0) reg_sq vf.squeeze(reg, preg) vf.store_align(dst_tile, reg_sq, preg) pl.jit() def example_kernel( a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], out: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], ): tf pl.TileType(shape[1, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) in_a_grp pl.make_tile_group(typetf, addrs0x0, mutex_ids[0]) in_a in_a_grp.current() t_out_grp pl.make_tile_group(typetf, addrs0x100, mutex_ids[1]) t_out t_out_grp.current() with pl.section_vector(): pl.load(in_a, a, [0, 0]) example_vf(in_a, t_out) pl.store(out, t_out, [0, 0]) ar_value pl.get_spr() def test_example(): device_id int(os.environ.get(TILE_FWK_DEVICE_ID, 0)) device fnpu:{device_id} core_nums 1 torch.npu.set_device(device) a torch.randn([1, 64], devicedevice, dtypetorch.float32) out torch.empty([1, 64], devicedevice, dtypetorch.float32) example_kernelNone, core_nums torch.npu.synchronize() torch.testing.assert_close(out, a, rtol1e-5, atol1e-5) if __name__ __main__: test_example() print(PASSED)对示例的逐步拆解Vector 函数example_vf内先通过vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_FP32)创建全 1 谓词掩码vf.load_align对齐加载源 tile 到寄存器vf.squeeze(reg, preg)按掩码压缩数据默认不传gather_mode即使用缺省模式最后vf.store_align对齐存储到目标 tile核函数example_kernel中创建TileType(shape[1, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec)类型的 tile group通过addrs0x0与addrs0x100显式指定输入/输出在 Vec 内存UB中的地址偏移在pl.section_vector()上下文内完成 load → vector function → store随后在 Vector section 之外调用pl.get_spr()读取 AR 寄存器值存入ar_value宿主侧验证test_example通过TILE_FWK_DEVICE_ID环境变量默认 0选择 NPU 设备以example_kernel[None, core_nums]stream 为 None、1 个 AI Core启动核函数torch.npu.synchronize()等待执行完成后用torch.testing.assert_close(out, a, rtol1e-5, atol1e-5)校验输出与输入一致最终打印PASSED。注意事项与使用边界结合文档与源码使用get_spr时有以下边界需要牢记作用域限制get_ar是__aicore__指令禁止在pl.vector_function内部调用必须放在pl.jit核函数体如pl.section_vector()块之后中否则会因指令作用域不合法而失败调用时机必须确保在vf.squeeze或写入 AR 寄存器的相关 Vector 指令执行之后再读取读取的是最近一次 Squeeze 写入的有效元素总字节数int64 标量参数约束get_spr不接受任何参数带参数调用会在 IR 层与后端代码生成阶段被双重校验拦截INVALID_ARGUMENT硬件约束仅 Ascend 950PR / Ascend 950DT 支持Atlas A2/A3 系列不支持跨平台移植时需注意能力判断。延伸阅读API 文档入口特殊寄存器访问 API 索引Python 前端声明与语义注释python/pypto_pro/language/_api.pyvf.squeeze与vf.squeeze_store_unalign定义python/pypto_pro/language/_vf_api.py、python/pypto_pro/language/_vf_api.pyIR 算子注册与参数校验python/pypto_pro/ir/op/system_ops.py、framework/src/interface/ir/op/block_ops/memory.cppCCE 后端指令下发get_ar()framework/src/interface/pypto_pro/backend/backend_cce_ops.cpp相关测试用例python/tests/ut/pypto_pro/language/parser/test_vf_scope.py、python/tests/st/pypto_pro/frontend/vf_api/test_vf_basic_ops.py【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考