2026/9/20 13:06:27
FEATURED · 精选文章 ` 获取线程块三维网格坐标)
PyPTO SIMT 编程指南使用pl.simt.block_idx()获取线程块三维网格坐标【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptopl.simt.block_idx()是 PyPTOCANN 开源仓库SIMT 编程范式pypto_pro.language.simt命名空间中的核心上下文查询接口用于在 SIMT 内核中获取当前线程块在线程块网格block grid中的三维坐标从而让不同线程块各自负责不同的数据分片。读完本文你将掌握该接口的函数原型、返回值语义、使用约束并结合仓库源码理解其从 Python 前端解析到 CCE 后端生成的完整调用链能够在 Ascend 950 系列产品上编写多线程块并行协作的 SIMT 内核。一、产品支持情况根据 block_idx.md 的声明pl.simt.block_idx()在不同产品上的支持情况如下产品系列支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持仓库中的端到端测试也印证了这一点test_multicore.py 在运行前会检查设备名仅当torch.npu.get_device_name()返回Ascend950A5 平台时才继续执行否则直接pytest.skip测试用例本身也通过pytest.mark.soc(950)标记了目标芯片。因此使用该接口前请先确认目标设备为 Ascend 950 系列。二、功能说明线程块网格中的身份标识block_idx()的作用是获取当前 SIMT 线程块在线程块网格中的三维坐标。在线程块网格模型中内核被划分为多个线程块block每个线程块又包含多个线程thread。block_idx()返回的坐标用于区分网格中的不同线程块同一线程块内的所有线程读取到相同的block_idx()结果不同线程块通过不同的坐标各自处理属于自己的数据分片从而实现数据并行。这一多线程块分工的典型用法可以从仓库的端到端测试中直接看到test_multicore.py 将一个 SIMT 函数同时启动在 4 个 AIV 核GRID_BLOCKS 4上每个线程块内的线程使用pl.simt.block_idx().x拿到自己所属块的编号配合线程编号tid计算出全局唯一的写入位置core_id * THREADS tid最终验证输出与torch.arange(GRID_BLOCKS * THREADS)完全一致——这正是不同线程块通过不同坐标处理各自数据的标准数据并行模式。三、函数原型与参数说明pypto_pro.language.simt.block_idx() - Any参数无。调用时不接受任何位置参数或关键字参数这一点由前端解析器强制校验见下文源码实现一节。返回类型标注为Any实际返回一个三维线程块坐标对象IR 层以dim3_context命名元组表示。在 Python 层该接口的声明位于 _simt_api.py 的Simt类中docstring 为 Return the current block coordinates within the outer kernel grid.与thread_idx()、block_dim()、grid_dim()等同属pl.simt.*上下文查询族。这些声明仅供 IDE 跳转、类型检查与拼写校验使用并不在运行时真正执行——在内核中每一次pl.simt.xxx()调用都会在 Python 真正执行前被 AST 解析器拦截并翻译为 IR在内核之外调用则会抛出RuntimeError。四、返回值说明block_idx()返回一个三维线程块坐标对象通过.x、.y、.z三个字段读取各维坐标分量类型每个分量均为DT_UINT32类型的Scalar无符号 32 位整数标量可直接参与算术运算、与DT_INT32等类型进行转换后用于寻址。取值范围当前线程块网格仅使用 X 维。block.x的取值范围为[0, grid_dim().x)即从 0 到网格 X 维大小减 1block.y和block.z恒为 0。语义由于同一线程块内所有线程读取到相同的坐标block_idx()更适合作为线程块级的公共信息参与寻址计算而线程级差异应使用pl.simt.thread_idx()或pl.simt.linear_thread_idx()获取。IR 层的实现直接体现了三维 逐分量访问的设计simt_ops.py 中定义了_DIM3_FIELDS (x, y, z)_make_dim3_components会为每个轴生成一次独立的simt.block_idx调用通过axis关键字参数区分 0/1/2 三个维度最终拼装成一个名为dim3_context的命名元组返回给上层。五、约束说明block_idx()具有严格的调用上下文限制只能在由pypto_pro.language.vector_function(modesimt)定义的 SIMT 入口函数或辅助函数中调用不能在普通的pl.jit()内核主体、mode非simt的向量函数或宿主host代码中调用。该约束在前端解析阶段被强制执行simt_ops.py 中的_validate_simt_body_op会检查当前函数类型是否为FunctionType.SimtVFSIMT 入口函数或FunctionType.SimtCalleeSIMT 辅助函数否则抛出InvalidOperation并附带提示Move SIMT-context-dependent logic into pl.vector_function(modesimt).请把依赖 SIMT 上下文的逻辑移入modesimt的向量函数中。同时该校验还会拒绝任何位置参数与关键字参数。在代码生成阶段同样有兜底校验backend_cce_simt_ops.cpp 中的MakeSimtContextComponentCodegenCCE在生成代码前会检查codegen.IsInSimtContext()若simt.block_idx在非 SIMT 函数中到达 CCE 代码生成阶段会直接报INVALID_OPERATION错误防止生成非法设备代码。六、调用示例以下示例来自 block_idx.md完整展示了定义 SIMT 向量函数 → 在内核中启动的标准流程import pypto_pro.language as pl pl.vector_function(modesimt, max_threads256) def write_block_idx_x( output: pl.Tensor[[1, 256], pl.DT_UINT32], ): tid pl.simt.linear_thread_idx() block pl.simt.block_idx() output[0, tid] block.x pl.jit() def simt_block_idx_kernel( output: pl.Tensor[[1, 256], pl.DT_UINT32], ): with pl.section_vector(): write_block_idx_x256对该示例的逐段解读SIMT 入口函数声明pl.vector_function(modesimt, max_threads256)声明这是一个 SIMT 模式的向量函数max_threads256指定线程块内最多 256 个线程线程级坐标pl.simt.linear_thread_idx()返回当前线程在线程块内的展平索引x 主序用于区分同一线程块内的不同线程线程块级坐标pl.simt.block_idx()返回当前线程块在网格中的坐标block.x即当前线程块在 X 维的编号写入目标output[0, tid] block.x将线程块编号写入输出张量的第tid列。由于同一线程块内所有线程的block.x相同该列会被写成同一个值——这正是块内一致、块间不同语义的直观体现启动方式在pl.jit()内核中通过with pl.section_vector():将启动操作放入向量Vector流水段并以write_block_idx_x256语法启动该 SIMT 函数256为启动的线程数如果需要同时启动多个线程块可以像 test_multicore.py 那样使用simt_multicoreNone, GRID_BLOCKS形式的网格维度参数。UT 中的上下文探针函数进一步演示了block_idx()三维分量的组合使用方式见 test_cce_simt.py通过block_id.x block_id.y block_id.z同时访问三个分量并累加验证了每个分量在解析与代码生成链路中均可独立工作。七、源码实现纵深从 AST 解析到 CCE 代码生成pl.simt.block_idx()的完整调用链可以拆解为四个环节仓库源码均有明确落点API 声明层python/pypto_pro/language/_simt_api.pySimt.block_idx()静态方法声明供 IDE、类型检查器与文档生成使用前端解析层python/pypto_pro/ir/op/simt_ops.pyop_impl(simt.block_idx)注册的_parse_block_idx处理器被调用经由_parse_dim3_context完成三项工作——校验调用上下文与参数合法性、为 x/y/z 三个轴各生成一次带axis关键字参数的操作调用、将其组装为dim3_context命名元组MakeTuple作为表达式返回IR 中间表示每个轴对应一个simt.block_idx操作axis属性取 0/1/2后续的流水线属性、调度与代码生成均基于该 IR 进行CCE 后端代码生成framework/src/interface/pypto_pro/backend/backend_cce_simt_ops.cppMakeSimtBlockIdxCodegenCCE通过MakeSimtContextComponentCodegenCCE将 IR 操作映射为 CCE 设备代码中的blockIdx内建变量并根据axis属性输出blockIdx.x/blockIdx.y/blockIdx.z该文件同时以同样的机制实现了threadIdx、blockDim、gridDim见 同文件 L55-L73四个上下文查询共享同一套轴名生成逻辑x/y/z见GetSimtAxisName。操作注册处backend_cce_simt_ops.cpp L1021将simt.block_idx的流水线类型声明为PipeType::S标量流水。八、典型应用多线程块数据并行将block_idx()与线程索引、网格维度配合可以构建经典的数据并行内核。以 test_multicore.py 为例其分工模式为pl.vector_function(modesimt, max_threadsTHREADS) def write_multicore_result(dst): tid pl.simt.linear_thread_idx() core_id pl.simt.block_idx().x # 线程块级我属于哪个块 dst[0, tid] core_id * THREADS tid # 全局唯一索引 块号 * 线程数 线程号with pl.section_vector(): core_id pl.get_block_idx() # 内核级当前向量核编号 write_multicore_resultTHREADS # 每个核都启动一个线程块 # 同步与写出 ... pl.store(out, dst, [core_id, 0])这里值得注意的工程细节是两级索引的分工内核内使用pl.get_block_idx()Vector 编程范式下的向量核编号与pl.get_subblock_idx()配合可定位 AIV 子块见 backend_cce_ops.cpp决定把当前线程块的计算结果存储到输出张量的哪一行SIMT 函数内使用pl.simt.block_idx().x决定线程块内每个线程写到本地 Tile 的哪个位置。最终测试断言输出矩阵第core_id行的值恰好是[core_id*THREADS, (core_id1)*THREADS)的连续递增序列从而端到端验证了网格中不同线程块通过不同block_idx()坐标处理各自数据分片的正确性。九、与相关 SIMT 上下文接口的对照block_idx()属于pl.simt.*上下文查询族理解它与邻近接口的区别有助于正确选型接口返回内容典型用途pl.simt.block_idx()当前线程块在网格中的三维坐标block.x/y/z区分网格中不同线程块实现块级数据并行pl.simt.grid_dim()线程块网格的三维尺寸gridDim.x/y/z与block_idx()配合计算块总数、做网格同步范围判断pl.simt.thread_idx()当前线程在线程块内的三维坐标threadIdx.x/y/z区分块内不同线程决定单线程处理的数据元素pl.simt.block_dim()线程块的三维尺寸blockDim.x/y/z与thread_idx()配合做边界判断、计算线性索引pl.simt.linear_thread_idx()x 主序展平的块内线程索引threadIdx.x threadIdx.y * blockDim.x threadIdx.z * blockDim.x * blockDim.y见 backend_cce_simt_ops.cpp L81一维数组访问时最常用的线程索引UT 用例 test_cce_simt.py 的_context_probe函数同时使用了上述全部接口并将所有分量求和可以视为对上下文查询族完整性的一个冒烟验证。十、注意事项小结设备适配仅 Ascend 950PR/950DT 支持A2/A3 系列不支持运行前可参考 test_multicore.py 的设备名检查逻辑做前置判断调用位置仅限pl.vector_function(modesimt)的入口函数或辅助函数前端解析与 CCE 代码生成两道校验均会拦截非法调用无参调用不接受任何位置/关键字参数当前实现网格仅使用 X 维block.x ∈ [0, grid_dim().x)block.y与block.z恒为 0类型语义三个分量均为DT_UINT32的Scalar与DT_INT32索引混用时注意显式转换。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
PyPTO SIMT 编程指南：使用 `pl.simt.block_idx()` 获取线程块三维网格坐标
PyPTO SIMT 编程指南使用pl.simt.block_idx()获取线程块三维网格坐标【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptopl.simt.block_idx()是 PyPTOCANN 开源仓库SIMT 编程范式pypto_pro.language.simt命名空间中的核心上下文查询接口用于在 SIMT 内核中获取当前线程块在线程块网格block grid中的三维坐标从而让不同线程块各自负责不同的数据分片。读完本文你将掌握该接口的函数原型、返回值语义、使用约束并结合仓库源码理解其从 Python 前端解析到 CCE 后端生成的完整调用链能够在 Ascend 950 系列产品上编写多线程块并行协作的 SIMT 内核。一、产品支持情况根据 block_idx.md 的声明pl.simt.block_idx()在不同产品上的支持情况如下产品系列支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持仓库中的端到端测试也印证了这一点test_multicore.py 在运行前会检查设备名仅当torch.npu.get_device_name()返回Ascend950A5 平台时才继续执行否则直接pytest.skip测试用例本身也通过pytest.mark.soc(950)标记了目标芯片。因此使用该接口前请先确认目标设备为 Ascend 950 系列。二、功能说明线程块网格中的身份标识block_idx()的作用是获取当前 SIMT 线程块在线程块网格中的三维坐标。在线程块网格模型中内核被划分为多个线程块block每个线程块又包含多个线程thread。block_idx()返回的坐标用于区分网格中的不同线程块同一线程块内的所有线程读取到相同的block_idx()结果不同线程块通过不同的坐标各自处理属于自己的数据分片从而实现数据并行。这一多线程块分工的典型用法可以从仓库的端到端测试中直接看到test_multicore.py 将一个 SIMT 函数同时启动在 4 个 AIV 核GRID_BLOCKS 4上每个线程块内的线程使用pl.simt.block_idx().x拿到自己所属块的编号配合线程编号tid计算出全局唯一的写入位置core_id * THREADS tid最终验证输出与torch.arange(GRID_BLOCKS * THREADS)完全一致——这正是不同线程块通过不同坐标处理各自数据的标准数据并行模式。三、函数原型与参数说明pypto_pro.language.simt.block_idx() - Any参数无。调用时不接受任何位置参数或关键字参数这一点由前端解析器强制校验见下文源码实现一节。返回类型标注为Any实际返回一个三维线程块坐标对象IR 层以dim3_context命名元组表示。在 Python 层该接口的声明位于 _simt_api.py 的Simt类中docstring 为 Return the current block coordinates within the outer kernel grid.与thread_idx()、block_dim()、grid_dim()等同属pl.simt.*上下文查询族。这些声明仅供 IDE 跳转、类型检查与拼写校验使用并不在运行时真正执行——在内核中每一次pl.simt.xxx()调用都会在 Python 真正执行前被 AST 解析器拦截并翻译为 IR在内核之外调用则会抛出RuntimeError。四、返回值说明block_idx()返回一个三维线程块坐标对象通过.x、.y、.z三个字段读取各维坐标分量类型每个分量均为DT_UINT32类型的Scalar无符号 32 位整数标量可直接参与算术运算、与DT_INT32等类型进行转换后用于寻址。取值范围当前线程块网格仅使用 X 维。block.x的取值范围为[0, grid_dim().x)即从 0 到网格 X 维大小减 1block.y和block.z恒为 0。语义由于同一线程块内所有线程读取到相同的坐标block_idx()更适合作为线程块级的公共信息参与寻址计算而线程级差异应使用pl.simt.thread_idx()或pl.simt.linear_thread_idx()获取。IR 层的实现直接体现了三维 逐分量访问的设计simt_ops.py 中定义了_DIM3_FIELDS (x, y, z)_make_dim3_components会为每个轴生成一次独立的simt.block_idx调用通过axis关键字参数区分 0/1/2 三个维度最终拼装成一个名为dim3_context的命名元组返回给上层。五、约束说明block_idx()具有严格的调用上下文限制只能在由pypto_pro.language.vector_function(modesimt)定义的 SIMT 入口函数或辅助函数中调用不能在普通的pl.jit()内核主体、mode非simt的向量函数或宿主host代码中调用。该约束在前端解析阶段被强制执行simt_ops.py 中的_validate_simt_body_op会检查当前函数类型是否为FunctionType.SimtVFSIMT 入口函数或FunctionType.SimtCalleeSIMT 辅助函数否则抛出InvalidOperation并附带提示Move SIMT-context-dependent logic into pl.vector_function(modesimt).请把依赖 SIMT 上下文的逻辑移入modesimt的向量函数中。同时该校验还会拒绝任何位置参数与关键字参数。在代码生成阶段同样有兜底校验backend_cce_simt_ops.cpp 中的MakeSimtContextComponentCodegenCCE在生成代码前会检查codegen.IsInSimtContext()若simt.block_idx在非 SIMT 函数中到达 CCE 代码生成阶段会直接报INVALID_OPERATION错误防止生成非法设备代码。六、调用示例以下示例来自 block_idx.md完整展示了定义 SIMT 向量函数 → 在内核中启动的标准流程import pypto_pro.language as pl pl.vector_function(modesimt, max_threads256) def write_block_idx_x( output: pl.Tensor[[1, 256], pl.DT_UINT32], ): tid pl.simt.linear_thread_idx() block pl.simt.block_idx() output[0, tid] block.x pl.jit() def simt_block_idx_kernel( output: pl.Tensor[[1, 256], pl.DT_UINT32], ): with pl.section_vector(): write_block_idx_x256对该示例的逐段解读SIMT 入口函数声明pl.vector_function(modesimt, max_threads256)声明这是一个 SIMT 模式的向量函数max_threads256指定线程块内最多 256 个线程线程级坐标pl.simt.linear_thread_idx()返回当前线程在线程块内的展平索引x 主序用于区分同一线程块内的不同线程线程块级坐标pl.simt.block_idx()返回当前线程块在网格中的坐标block.x即当前线程块在 X 维的编号写入目标output[0, tid] block.x将线程块编号写入输出张量的第tid列。由于同一线程块内所有线程的block.x相同该列会被写成同一个值——这正是块内一致、块间不同语义的直观体现启动方式在pl.jit()内核中通过with pl.section_vector():将启动操作放入向量Vector流水段并以write_block_idx_x256语法启动该 SIMT 函数256为启动的线程数如果需要同时启动多个线程块可以像 test_multicore.py 那样使用simt_multicoreNone, GRID_BLOCKS形式的网格维度参数。UT 中的上下文探针函数进一步演示了block_idx()三维分量的组合使用方式见 test_cce_simt.py通过block_id.x block_id.y block_id.z同时访问三个分量并累加验证了每个分量在解析与代码生成链路中均可独立工作。七、源码实现纵深从 AST 解析到 CCE 代码生成pl.simt.block_idx()的完整调用链可以拆解为四个环节仓库源码均有明确落点API 声明层python/pypto_pro/language/_simt_api.pySimt.block_idx()静态方法声明供 IDE、类型检查器与文档生成使用前端解析层python/pypto_pro/ir/op/simt_ops.pyop_impl(simt.block_idx)注册的_parse_block_idx处理器被调用经由_parse_dim3_context完成三项工作——校验调用上下文与参数合法性、为 x/y/z 三个轴各生成一次带axis关键字参数的操作调用、将其组装为dim3_context命名元组MakeTuple作为表达式返回IR 中间表示每个轴对应一个simt.block_idx操作axis属性取 0/1/2后续的流水线属性、调度与代码生成均基于该 IR 进行CCE 后端代码生成framework/src/interface/pypto_pro/backend/backend_cce_simt_ops.cppMakeSimtBlockIdxCodegenCCE通过MakeSimtContextComponentCodegenCCE将 IR 操作映射为 CCE 设备代码中的blockIdx内建变量并根据axis属性输出blockIdx.x/blockIdx.y/blockIdx.z该文件同时以同样的机制实现了threadIdx、blockDim、gridDim见 同文件 L55-L73四个上下文查询共享同一套轴名生成逻辑x/y/z见GetSimtAxisName。操作注册处backend_cce_simt_ops.cpp L1021将simt.block_idx的流水线类型声明为PipeType::S标量流水。八、典型应用多线程块数据并行将block_idx()与线程索引、网格维度配合可以构建经典的数据并行内核。以 test_multicore.py 为例其分工模式为pl.vector_function(modesimt, max_threadsTHREADS) def write_multicore_result(dst): tid pl.simt.linear_thread_idx() core_id pl.simt.block_idx().x # 线程块级我属于哪个块 dst[0, tid] core_id * THREADS tid # 全局唯一索引 块号 * 线程数 线程号with pl.section_vector(): core_id pl.get_block_idx() # 内核级当前向量核编号 write_multicore_resultTHREADS # 每个核都启动一个线程块 # 同步与写出 ... pl.store(out, dst, [core_id, 0])这里值得注意的工程细节是两级索引的分工内核内使用pl.get_block_idx()Vector 编程范式下的向量核编号与pl.get_subblock_idx()配合可定位 AIV 子块见 backend_cce_ops.cpp决定把当前线程块的计算结果存储到输出张量的哪一行SIMT 函数内使用pl.simt.block_idx().x决定线程块内每个线程写到本地 Tile 的哪个位置。最终测试断言输出矩阵第core_id行的值恰好是[core_id*THREADS, (core_id1)*THREADS)的连续递增序列从而端到端验证了网格中不同线程块通过不同block_idx()坐标处理各自数据分片的正确性。九、与相关 SIMT 上下文接口的对照block_idx()属于pl.simt.*上下文查询族理解它与邻近接口的区别有助于正确选型接口返回内容典型用途pl.simt.block_idx()当前线程块在网格中的三维坐标block.x/y/z区分网格中不同线程块实现块级数据并行pl.simt.grid_dim()线程块网格的三维尺寸gridDim.x/y/z与block_idx()配合计算块总数、做网格同步范围判断pl.simt.thread_idx()当前线程在线程块内的三维坐标threadIdx.x/y/z区分块内不同线程决定单线程处理的数据元素pl.simt.block_dim()线程块的三维尺寸blockDim.x/y/z与thread_idx()配合做边界判断、计算线性索引pl.simt.linear_thread_idx()x 主序展平的块内线程索引threadIdx.x threadIdx.y * blockDim.x threadIdx.z * blockDim.x * blockDim.y见 backend_cce_simt_ops.cpp L81一维数组访问时最常用的线程索引UT 用例 test_cce_simt.py 的_context_probe函数同时使用了上述全部接口并将所有分量求和可以视为对上下文查询族完整性的一个冒烟验证。十、注意事项小结设备适配仅 Ascend 950PR/950DT 支持A2/A3 系列不支持运行前可参考 test_multicore.py 的设备名检查逻辑做前置判断调用位置仅限pl.vector_function(modesimt)的入口函数或辅助函数前端解析与 CCE 代码生成两道校验均会拦截非法调用无参调用不接受任何位置/关键字参数当前实现网格仅使用 X 维block.x ∈ [0, grid_dim().x)block.y与block.z恒为 0类型语义三个分量均为DT_UINT32的Scalar与DT_INT32索引混用时注意显式转换。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
