PyPTO 系统访问变量实战指南get_block_idx / get_block_num / get_subblock_idx / get_subblock_num 多核编程完全解析【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读PyPTOParallel Tensor/Tile Operation的 SIMDVector/Tile编程模型中多核数据切分、跨核同步与 Cube/Vector 混合调度都依赖一组系统访问变量System Variables——get_block_idx、get_block_num、get_subblock_idx、get_subblock_num。本指南以 system_variables/index.md 文档为骨架逐一拆解这 4 个接口的功能、返回语义、产品支持情况与典型使用场景并结合仓库中 Python API 声明 与 CCE 后端代码生成实现 说明其底层原理。读完本文你将掌握在纯 Vector Kernel 与混合 KernelAIC:AIV1:2中正确完成多核数据分片、条件执行与跨域协同的核心方法。一、系统访问变量总览四个接口的分工在 PyPTO 中Kernel 以逻辑 Block为执行单元在多个 AI Core逻辑 AI Core上并行运行。为了在 Kernel 内感知我在哪、一共多少人、我这一份怎么分PyPTO 提供了 4 个零参数的系统级接口接口函数原型语义典型用途get_block_idx()- int当前执行域中逻辑 AI Core 的全局索引多核数据分片、偏移计算get_block_num()- int实际启动的逻辑 Block 数量限核后循环步长、工作总量计算get_subblock_idx()- int当前逻辑 AI Core 内 AIC/AIV 的子核subblock索引区分同一 Block 内的不同 AIV条件执行get_subblock_num()- int当前 Block 的 subblock 总数task ration换算物理核号、还原 core_id从源码看这 4 个接口在 python/pypto_pro/language/_api.py 中以_api_decl声明对应 IR 层注册于 python/pypto_pro/ir/op/system_ops.py真正落地到设备侧 CCE 代码的是 backend_cce_ops.cpp 中REGISTER_BACKEND_OP注册的一系列后端算子。其中get_subblock_num在 Python 侧注释中明确标注Matches AscendC GetTaskRation()即与 AscendC 的GetTaskRation()语义对齐。二、产品支持情况仅限 Ascend 950 系列上述 4 个接口的支持范围在各自文档中完全一致Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持也就是说系统访问变量是 950 系列设备上多核 Vector/Tile 编程的基础能力在使用 A2/A3 平台时多核切分需要采用其他等价手段如 Host 侧分核传入不能直接依赖这组接口。三、get_block_idx()获取当前逻辑 AI Core 的全局索引3.1 功能与返回值get_block_idx()获取当前执行域中逻辑 AI Core 的全局索引用于多核控制和数据偏移计算。无参数、无约束返回设备运行时产生的整型标量值可用于 Kernel 内整数运算和索引。返回值取值范围与 Kernel 的执行域Section密切相关仅启动 CubeAIC或仅启动 VectorAIV时范围为[0, get_block_num())同时启动 AIC 与 AIV 时AIC 侧范围为[0, get_block_num())AIV 侧范围为[0, get_subblock_num() × get_block_num())。当前 AIC:AIV1:2 配置下即[0, 2 × get_block_num())。在 Vector 段中该接口返回的是全局 AIV 逻辑索引编号方式为物理核号 × get_subblock_num() 子核号可直接用于数据分片和偏移计算。混合 Kernel 在 Vector 段做跨步stride切分时工作单元总数为get_block_num() × get_subblock_num()。3.2 底层实现Vector 段如何映射为全局 AIV 索引CCE 后端的代码生成逻辑印证了上述文档语义见 backend_cce_ops.cppstatic std::string MakeBlockGetBlockIdxCodegenCCE(const ir::CallPtr op, codegen::CodegenBase codegen_base) { CHECK(op-args_.size() 0) get_block_idx requires no arguments; auto cg dynamic_castcodegen::CCECodegen(codegen_base); const auto target cg.GetTarget(); if (target ir::SectionKind::Vector) { return (int32_t)(get_block_idx() * get_subblockdim() get_subblockid()); } return (int32_t)(get_block_idx()); }可以看到当编译目标为 Vector 段时生成的设备代码是get_block_idx() * get_subblockdim() get_subblockid()——即物理核号 × 子核数 子核号而 Cube 段则直接返回get_block_idx()。这从实现层面保证了文档所述的两种取值范围。3.3 调用示例一纯 Vector Kernel 多核数据分片以Kernel[None, NUM_CORES]形式启动 2 个逻辑 Block每个 AIV 用get_block_idx()获取全局逻辑索引并处理 64 行逐元素加法用 printf 打印索引值import os import pypto_pro.language as pl import torch NUM_CORES 2 pl.jit(auto_mutexTrue) def multicore_add_kernel( x: pl.Tensor[[128, 128], pl.DT_FP16], y: pl.Tensor[[128, 128], pl.DT_FP16], z: pl.Tensor[[128, 128], pl.DT_FP16], ): tt pl.TileType(shape[64, 128], dtypepl.DT_FP16, target_memorypl.MemorySpace.Vec) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_b pl.make_tile_group(typett, addrs0x4000, mutex_ids[1]) tile_c pl.make_tile_group(typett, addrs0x8000, mutex_ids[2]) with pl.section_vector(): vidx pl.get_block_idx() # 当前AIV的全局逻辑索引 num_blocks pl.get_block_num() # 实际启动的Block数 pl.printf(block_idx %d, block_num %d\n, vidx, num_blocks) for tile_idx in pl.range(vidx, 2, num_blocks): offset tile_idx * 64 cur_a tile_a.current() cur_b tile_b.current() cur_c tile_c.current() pl.load(cur_a, x, [offset, 0]) pl.load(cur_b, y, [offset, 0]) pl.add(cur_c, cur_a, cur_b) pl.store(z, cur_c, [offset, 0]) if __name__ __main__: device fnpu:{int(os.environ.get(TILE_FWK_DEVICE_ID, 0))} torch.npu.set_device(device) torch.manual_seed(42) x torch.rand([128, 128], devicedevice, dtypetorch.float16) y torch.rand([128, 128], devicedevice, dtypetorch.float16) z torch.zeros([128, 128], devicedevice, dtypetorch.float16) multicore_add_kernelNone, NUM_CORES torch.npu.synchronize() torch.testing.assert_close(z, x y, rtol1e-2, atol1e-2) print(fmax diff {(z - (x y)).abs().max().item()})回显 Vec后为核号多核间输出顺序不固定 Vec 0 block_idx 0, block_num 2 Vec 1 block_idx 1, block_num 23.4 调用示例二混合 Kernel 中 AIC 与 AIV 的返回值差异混合 KernelAIC:AIV1:2中Cube 段返回物理核号[0, get_block_num())Vector 段返回全局 AIV 逻辑编号[0, get_block_num() × get_subblock_num())import os import pypto_pro.language as pl import torch pl.jit() def block_idx_mix_kernel(out: pl.Tensor[[1], pl.DT_INT32]): with pl.section_cube(): aic_idx pl.get_block_idx() pl.printf([cube] block_idx %d\n, aic_idx) pl.setval(out, 0, 1) with pl.section_vector(): aiv_idx pl.get_block_idx() pl.printf([vector] block_idx %d\n, aiv_idx) if __name__ __main__: device fnpu:{int(os.environ.get(TILE_FWK_DEVICE_ID, 0))} torch.npu.set_device(device) out torch.zeros(1, devicedevice, dtypetorch.int32) block_idx_mix_kernelNone, 2 torch.npu.synchronize()回显启动 2 个逻辑 BlockAIC 侧返回 0/1AIV 侧返回 0~3 Cube 0 [cube] block_idx 0 Vec 0 [vector] block_idx 0 Vec 1 [vector] block_idx 1 Cube 1 [cube] block_idx 1 Vec 2 [vector] block_idx 2 Vec 3 [vector] block_idx 3四、get_block_num()读取实际启动的 Block 数并作为循环步长4.1 功能与返回值get_block_num()获取本次实际启动的逻辑 Block 数量用于多核控制和数据偏移计算。返回值是设备运行时产生的整型标量等于限核后实际启动的逻辑 Block 数。关键语义JIT 每次启动会通过 C 启动器查询实际 Stream 的有效资源限制按 Kernel 执行域和配对比例限制 Host 请求的block_dim。因此返回值可能小于 Host 请求值——数据切分应使用本接口返回值作为循环步长避免遗漏任务。这一点在 Python 侧 API 注释中同样强调This can be smaller than the hosts requestedkernel[stream, block_dim]. Use this runtime count as the work-distribution stride so limiting cores does not leave tiles unprocessed.见 language/_api.py。分执行域的取值规则仅启动 CubeAIC或仅启动 VectorAIV时该值等于执行域逻辑核数在 AIC:AIV 为 1:2 的混合 Kernel 中该值表示逻辑 Block 数AIC 逻辑核数为get_block_num()AIV 逻辑核数为get_block_num() * get_subblock_num()。4.2 调用示例一按实际 Block 数循环切分用Kernel[None, NUM_CORES]请求最多 2 个逻辑 Block每个 AIV 以get_block_num()返回值为步长跨步处理 64 行 Tile完整 Kernel 定义与 3.3 节相同此处省略重复部分with pl.section_vector(): vidx pl.get_block_idx() # 当前AIV的全局逻辑索引 num_blocks pl.get_block_num() # 实际启动的Block数 pl.printf(block_idx %d, block_num %d\n, vidx, num_blocks) for tile_idx in pl.range(vidx, 2, num_blocks): offset tile_idx * 64 ...回显 Vec 0 block_idx 0, block_num 2 Vec 1 block_idx 1, block_num 24.3 调用示例二限核启动下的返回值启动 Block 数是请求的上界实际启动数可能更小。复用 3.3 节的 Kernel 与__main__中的 x/y仅把启动 Block 数改为 1z torch.zeros([128, 128], devicedevice, dtypetorch.float16) multicore_add_kernelNone, 1 torch.npu.synchronize() torch.testing.assert_close(z, x y, rtol1e-2, atol1e-2) # 仍覆盖全部128行回显 Vec 0 block_idx 0, block_num 1该示例同时验证了 4.1 节的核心结论即使只启动 1 个 Block由于循环步长来自get_block_num()此处为 1单核也会完整遍历所有 Tile128 行全部被覆盖assert_close依然通过。五、get_subblock_idx()区分同一 Block 内的不同 AIV5.1 功能与返回值get_subblock_idx()获取当前逻辑 AI Core 内 AIC 或 AIV 的 subblock 索引。返回整型标量取值范围为[0, get_subblock_num())在 AIC 与 AIV 比例为 1:2 的混合 Kernel 中同一逻辑 Block 对应的两个 AIV 分别返回 0 和 1。5.2 典型使用场景文档给出两种典型模式insert Cube 模式每个子核计算部分结果用 insert 拼入 L1 Buffer 中的 NZTileCube 侧读取合并后的完整数据条件执行根据子核号决定是否执行某段代码例如只让 sub-core 0 执行指定操作。[!CAUTION] 注意 纯 Vector Kernel 中的两个子核共享 MTE 搬运管道不能由每个子核分别使用pypto_pro.language.store向 GM 的不同区域写入数据。按子核切分数据搬运时应使用 insert Cube 模式。一句话总结两个索引接口的分工get_block_idx()用于 Vector 段的全局 AIV 数据分片get_subblock_idx()用于区分同一逻辑 Block 内的不同 AIV。5.3 调用示例一纯 Vector Kernel 中读取子核号纯 Vector Kernel 中每个 Block 即一个 AIVget_subblock_num()返回 1get_subblock_idx()恒为 0import os import pypto_pro.language as pl import torch pl.jit(auto_mutexTrue) def subblock_add_kernel( x: pl.Tensor[[64, 64], pl.DT_FP32], y: pl.Tensor[[64, 64], pl.DT_FP32], out: pl.Tensor[[64, 64], pl.DT_FP32], ): tt pl.TileType(shape[64, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tile_x pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_y pl.make_tile_group(typett, addrs0x4000, mutex_ids[1]) tile_sum pl.make_tile_group(typett, addrs0x8000, mutex_ids[2]) with pl.section_vector(): sub_idx pl.get_subblock_idx() pl.printf(subblock_idx %d, subblock_num %d\n, sub_idx, pl.get_subblock_num()) cur_x tile_x.current() cur_y tile_y.current() cur_sum tile_sum.current() pl.load(cur_x, x, [0, 0]) pl.load(cur_y, y, [0, 0]) pl.add(cur_sum, cur_x, cur_y) pl.store(out, cur_sum, [0, 0]) if __name__ __main__: device fnpu:{int(os.environ.get(TILE_FWK_DEVICE_ID, 0))} torch.npu.set_device(device) torch.manual_seed(42) x torch.randn([64, 64], devicedevice, dtypetorch.float32) y torch.randn([64, 64], devicedevice, dtypetorch.float32) out torch.zeros([64, 64], devicedevice, dtypetorch.float32) subblock_add_kernel(x, y, out) torch.npu.synchronize() torch.testing.assert_close(out, x y, rtol1e-5, atol1e-5) print(fmax diff {(out - (x y)).abs().max().item()})回显 Vec 0 subblock_idx 0, subblock_num 15.4 调用示例二混合 Kernel 中的子核号与条件执行混合 Kernel 的 Vector 段中get_block_idx()返回全局 AIV 逻辑索引用于数据分片get_subblock_idx()用于区分同一逻辑 Block 内的两个 AIV。根据子核号可做条件执行例如只让 sub-core 0 写结果import os import pypto_pro.language as pl import torch pl.jit() def subblock_cond_kernel(out: pl.Tensor[[1], pl.DT_INT32]): with pl.section_cube(): pass with pl.section_vector(): block_idx pl.get_block_idx() # 全局AIV逻辑索引0~3 sub_idx pl.get_subblock_idx() # 块内AIV编号0或1 pl.printf(block_idx %d, subblock_idx %d\n, block_idx, sub_idx) # 条件执行只让每个逻辑Block的sub-core 0写结果 if sub_idx 0: pl.printf(sub-core 0 of core %d writes output\n, block_idx // pl.get_subblock_num()) pl.setval(out, 0, 1) if __name__ __main__: device fnpu:{int(os.environ.get(TILE_FWK_DEVICE_ID, 0))} torch.npu.set_device(device) out torch.zeros(1, devicedevice, dtypetorch.int32) subblock_cond_kernelNone, 2 torch.npu.synchronize()回显启动 2 个逻辑 Block Cube条目为 Cube 段本例无打印 Cube 0 Vec 0 block_idx 0, subblock_idx 0 sub-core 0 of core 0 writes output Vec 1 block_idx 1, subblock_idx 1 Cube 1 Vec 2 block_idx 2, subblock_idx 0 sub-core 0 of core 1 writes output Vec 3 block_idx 3, subblock_idx 1注意此处block_idx // pl.get_subblock_num()正是 6.2 节所述还原物理核号的用法。六、get_subblock_num()获取 task ration 并统一 Cube/Vector 的 core_id6.1 功能与返回值get_subblock_num()获取当前 block 的 subblock 总数即一个 block 关联的从核数量又称 task ration返回整型Expr。返回值与核类型及编译模式有关AIC 核始终返回 1AIC 为 block无 AIC 从核AIV 核融合算子mixAIC:AIV 1:2返回 2每个 AI Core 含 2 个 AIV 从核纯 Vector 算子aiv-only返回 1AIV 为 block无 subblock 划分。Python 侧 API 注释与文档一致Returns 1 on AIC binaries, get_subblockdim() on AIV binaries. Matches AscendC GetTaskRation().见 language/_api.py。6.2 调用示例用除法还原物理核号统一两侧数据切分在融合算子中get_block_idx()在 AIV 核上返回的是逻辑编号block_idx * subblock_num subblock_idx通过除以get_subblock_num()可还原物理 AI Core 编号使 Cube 与 Vector 两侧用统一的core_id切分数据import pypto_pro.language as pl NUM_CORES 2 pl.jit(auto_mutexTrue) def matmul_example( a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP16], b: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP16], out: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP16], ): num_cores pl.get_block_num() # AIC/AIV两侧得到相同的物理核号详见下方NOTE core_id pl.get_block_idx() // pl.get_subblock_num() with pl.section_cube(): for i in pl.range(core_id, a.shape[0] // 128, num_cores): ... # Cube侧按行块i执行load/matmul/store with pl.section_vector(): for i in pl.range(core_id, a.shape[0] // 128, num_cores): ... # Vector侧用同一core_id切分与Cube侧对齐 matmul_exampleNone, NUM_CORES[!NOTE] 说明 该除法在 AIC 核上为block_idx // 1在 AIV 核上为(block_idx * 2 subblock_idx) // 2两者均得到相同的 AI Core 编号因此 Cube 与 Vector 可共享同一core_id做数据切分。七、综合使用建议与注意事项7.1 四接口组合速查场景推荐组合纯 Vector Kernel 多核切分get_block_idx()分片索引get_block_num()步长混合 Kernel Vector 段切分get_block_idx()全局 AIV 逻辑索引get_subblock_num()工作单元换算区分块内 AIV / 条件执行get_subblock_idx()统一 Cube 与 Vector 的物理核号get_block_idx() // get_subblock_num()7.2 关键注意事项步长必须使用get_block_num()实际启动数可能小于 Host 请求的block_dim若以请求值硬编码为步长限核时会产生数据空洞见 4.3 节Vector 段的索引语义与 Cube 段不同Vector 段返回的是全局 AIV 逻辑索引含子核维度CCE 后端通过get_block_idx() * get_subblockdim() get_subblockid()实现见 backend_cce_ops.cpp纯 Vector Kernel 中禁止子核分别 store 到 GM两个子核共享 MTE 搬运管道按子核切分数据搬运时应走 insert Cube 模式平台限制本组接口仅受 Ascend 950PR/950DT 支持A2/A3 系列不支持移植时需做平台判断多核输出顺序不固定使用 printf 调试时各核打印顺序无法保证应以内容而非顺序判断结果。7.3 进一步阅读系统访问变量索引文档insert子核结果拼入 NZTileprintfKernel 内调试打印Python 侧 API 声明python/pypto_pro/language/_api.pyIR 算子注册python/pypto_pro/ir/op/system_ops.pyCCE 代码生成实现framework/src/interface/pypto_pro/backend/backend_cce_ops.cpp【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考