CANN PyPTO 系统变量解析：get_subblock_num() 获取 Subblock 总数（Task Ration）的用法与底层原理-平芜编程栈

人工智能
编译器
模型编译
高性能计算
深度学习
CANN

【免费下载链接】pypto

PyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。

导读

在 CANN PyPTO 的混合算子（Cube + Vector 融合）开发中，pypto_pro.language.get_subblock_num()用于获取当前逻辑 Block 关联的从核（subblock）总数，即 AscendC 中的 task ration。它是在 AIV 核上区分"逻辑 Block 编号"与"物理 AI Core 编号"、实现 Cube/Vector 两侧统一数据切分的关键系统变量。读完本文，你将掌握该 API 的产品支持范围、返回值语义（AIC 与 AIV 的不同表现）、典型调用方式，以及其从 IR 注册到 CCE 代码生成的完整底层链路。

一、产品支持情况

get_subblock_num()的系统变量能力与硬件平台强相关，当前仓库文档明确的支持矩阵如下：

产品形态	支持情况
Ascend 950PR / Ascend 950DT	支持
Atlas A3 训练系列产品 / Atlas A3 推理系列产品	不支持
Atlas A2 训练系列产品 / Atlas A2 推理系列产品	不支持

也就是说，目前仅 Ascend 950 系列产品提供 subblock 机制；在 Atlas A2/A3 产品上使用该 API 将不受支持。撰写或移植 Kernel 时，请务必按目标平台核对此支持矩阵（详见 get_subblock_num.md 与同目录下的 index.md）。

二、功能说明与函数原型

get_subblock_num()获取当前 Block 的 subblock 总数，即一个 Block 关联的从核数量，在 AscendC 语义中等价于 task ration。在混合算子中，一个逻辑 Block 可以由一个 AIC（Cube 核）与多个 AIV（Vector 从核）组成，subblock 就是这些从核的编号维度。

函数原型：

pypto_pro.language.get_subblock_num() -> int

该 API 定义在 PyPTO 语言前端 python/pypto_pro/language/_api.py 中，其 docstring 直接说明了语义：

Get the sub-block count per AI Core (task ration). Returns 1 on AIC binaries, get_subblockdim() on AIV binaries. Matches AscendC GetTaskRation().

即：AIC 二进制中恒返回 1；AIV 二进制中返回get_subblockdim()，与 AscendC 的GetTaskRation()语义一致。

参数与约束说明

参数：无。该 API 不接受任何参数。
约束：无。可在 Kernel 内的整数计算、数据索引中直接使用。

IR 层的类型推导也印证了这一点：在 framework/src/interface/ir/op/block_ops/memory.cpp 中，DeduceBlockGetBlockIdxType会显式校验args.size() == 0（不接受任何参数），并返回ScalarType(DataType::INT64)；get_subblock_num与get_subblock_idx、get_block_idx、get_block_num一起注册为无参 IR 算子（见 memory.cpp）。

三、返回值说明：按核类型与编译模式区分

返回值类型为DT_INT64，具体数值与核类型及编译模式有关：

AIC 核（Cube 侧）：始终返回 1。AIC 本身即 Block，没有 AIC 从核，因此不存在 subblock 划分。
AIV 核（Vector 侧）：
- 融合算子（mix，AIC:AIV = 1:2）：返回 2，即每个 AI Core 内含 2 个 AIV 从核（subblock 编号为 0 和 1）。
- 纯 Vector 算子（aiv-only）：返回 1，此时 AIV 即为 Block，同样没有 subblock 划分。

这一行为在 CCE 后端代码生成中可直接看到。在 framework/src/interface/pypto_pro/backend/backend_cce_ops.cpp 中，get_subblock_num的代码生成逻辑为：

// Matches AscendC GetTaskRation(): AIC returns 1, AIV returns get_subblockdim(). auto& cg = dynamic_cast<codegen::CCECodegen&>(codegen_base); const auto target = cg.GetTarget(); if (target == ir::SectionKind::Vector) { return std::string("(int64_t)(get_subblockdim())"); } return std::string("(int64_t)(1)");

即代码生成器根据当前代码段的目标类型（SectionKind::Vector与否）决定下发get_subblockdim()还是常量 1——这正是"返回值与核类型及编译模式有关"的底层来源。

与 get_subblock_idx() 的配合

get_subblock_num()通常与get_subblock_idx()（获取当前逻辑 AI Core 内 AIC 或 AIV 的 subblock 索引，取值范围为[0, get_subblock_num())）配合使用。在 1:2 的混合 Kernel 中，同一逻辑 Block 对应的两个 AIV 分别返回 0 和 1。后端实现中：

REGISTER_BACKEND_OP(BackendCCE, "get_subblock_idx") ... return std::string("(int64_t)(get_subblockid())");

参见 backend_cce_ops.cpp。更完整的条件执行示例可参考姊妹文档 get_subblock_idx.md。

四、调用示例：混合算子中还原物理 AI Core 编号

在融合算子中，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_mutex=True) 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做数据切分。

这一示例的底层依据在get_block_idx的后端实现中清晰可见。在 backend_cce_ops.cpp 中：

// Helper function for get_block_idx (returns value expression). // Matches AscendC GetBlockIdx(): AIV returns global AIV index, AIC returns AIC index. static std::string MakeBlockGetBlockIdxCodegenCCE(const ir::CallPtr& op, codegen::CodegenBase& codegen_base) { ... if (target == ir::SectionKind::Vector) { return "(int64_t)(get_block_idx() * get_subblockdim() + get_subblockid())"; } return "(int64_t)(get_block_idx())"; }

AIV 侧返回的是get_block_idx() * get_subblockdim() + get_subblockid()（全局 AIV 逻辑索引），AIC 侧返回的是get_block_idx()（AIC 索引）。因此：

AIV 核：logical_idx // subblock_num = (block_idx * 2 + subblock_idx) // 2 = block_idx（物理 AI Core 编号）；
AIC 核：block_idx // 1 = block_idx。

两式结果一致，core_id即可作为 Cube/Vector 两侧统一的数据切分依据。get_block_num()则提供"经过流上 core 限制后实际生效的 worker block 数"，应以其作为切分步长，避免限制核数后留下未处理的 tile（其语义见 python/pypto_pro/language/_api.py 的 docstring）。

五、从 Python API 到 CCE 代码生成的完整链路

get_subblock_num()的调用并不只是 Python 层的一个普通函数，而是被注册为系统级 IR 算子，经过前端解析、IR 类型推导、后端代码生成三步完成下发：

Python 前端声明：在 python/pypto_pro/language/_api.py 中以@_api_decl声明，并在语言入口 python/pypto_pro/language/init.py 中导出为pl.get_subblock_num；
IR 算子注册与类型推导：在 python/pypto_pro/ir/op/system_ops.py 中注册为OpSpec(ir_name="get_subblock_num", parse_args=False, parse_kwargs=False)；IR 层由DeduceBlockGetBlockIdxType校验无参并推导返回INT64标量类型（memory.cpp）；
CCE 后端代码生成：在 backend_cce_ops.cpp 中按代码段类型（Vector 段下发get_subblockdim()，其他段下发常量 1）生成 CCE 代码，与 AscendC 的GetTaskRation()语义对齐。

此外，subblock 维度还参与了 Kernel 级 block 寻址。在 backend_cce_ops.cpp 中，sub-block 寻址按[block_num, block_num + block_num * subblockdim)（最多[N, 3N)）的区间扩展逻辑 block 索引，注释明确说明"Sub-block addressing follows the established get_block_idx() backend-op"，即get_subblock_num()返回的 subblock 总数直接决定了逻辑 Block 展开为物理 worker 的数量关系。

六、易混淆点与使用建议

不要用 AIC 侧语义推断 AIV 侧：AIC 恒返回 1、AIV 在 mix 模式下返回 2，两者不可混用；切分数据时应以get_block_idx() // get_subblock_num()得到物理核号，而不是直接使用get_block_idx()。
与get_block_num()配合使用：get_block_num()是经过 core 限制后的实际 block 数，用它作为循环步长可保证限核后所有 tile 仍被处理；get_subblock_num()用于在逻辑与物理编号之间换算。
平台兼容性：该能力仅 Ascend 950PR/Ascend 950DT 支持，Atlas A2/A3 系列不支持，跨平台移植时需要提供等价替代方案。
返回值类型：始终为DT_INT64标量，可直接参与 Kernel 内整数运算与数据索引，无需额外转换。

七、相关文档与源码索引

API 参考：get_subblock_num.md、get_subblock_idx.md、get_block_idx.md、get_block_num.md、index.md
Python 前端声明：python/pypto_pro/language/_api.py
IR 算子注册：python/pypto_pro/ir/op/system_ops.py
IR 类型推导：framework/src/interface/ir/op/block_ops/memory.cpp
CCE 后端代码生成：framework/src/interface/pypto_pro/backend/backend_cce_ops.cpp