CANN PyPTO Pro SIMT 编程范式：线程架构、SIMT 函数与抽象硬件详解

CANN PyPTO Pro SIMT 编程范式线程架构、SIMT 函数与抽象硬件详解【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptoSIMTSingle Instruction Multiple Threads单指令多线程是 PyPTO Pro 在 AIVVector Core上提供的一种线程并行编程模型以 Thread 为基本执行单元适合表达不规则数据访问、逐线程分支控制和共享地址的原子更新。本文以 SIMT 编程范式模块为骨架系统讲解 Grid / Thread Block / Thread 三级线程层次、Warp 执行与分支发散、线程索引查询接口、SIMT 入口函数与辅助函数的定义和启动方式以及 SIMT 函数可操作的 Scalar / Tile / Tensor 内存层级与抽象硬件架构并辅以仓库源码与实战示例帮助读者掌握在 A5 架构上编写与启动 SIMT 计算的能力。SIMT 与 SIMD 的定位在并行执行模型中的角色在进入 SIMT 细节之前先明确它在 PyPTO Pro 并行体系中的位置。根据 programming_paradigm_overview.mdPyPTO Pro 采用 Host 与 Device 协同的异构编程方式Host 代码运行在 CPU 上负责设备资源管理、数据搬运、任务下发与结果同步Device 代码运行在 NPU 上提供 SIMD 和 SIMT 两种并行编程方式。SIMD单指令多数据是数据并行模型一条指令在同一个时钟周期内对多个数据元素执行完全相同的操作适合连续规整的数据访问与批量相同操作例如逐元素计算、归约、矩阵乘和融合计算SIMT单指令多线程是线程并行模型同一份程序由多个 Thread 并行执行每个 Thread 根据自身索引处理不同数据允许独立寻址并进入不同分支或循环适合离散索引、不规则数据访问、原子更新和动态数据依赖较强的局部计算。两者的层次关系是SPMD 负责组织多个逻辑 AI Core 之间的任务并行SIMD 与 SIMT 负责描述单个逻辑执行域内的计算方式。PyPTO Pro 在 AIV 上提供 SIMD 与 SIMT混合编程能力外层 Kernel 组织规则的 Tile 计算、数据搬运和 SIMT 启动SIMT 函数描述逐线程逻辑具体开发步骤见 SIMT计算。线程架构Grid、Thread Block 与 Thread 三级层次SIMT 采用 Grid、Thread Block 和 Thread 三级线程层次从顶层到底层逐级划分并行任务。Grid 和 Thread Block 的规模以及 Thread Block 和 Thread 的坐标均为包含 X、Y、Z 三个分量的 dim3 三维结构。Grid线程块网格Grid 是 SIMT 线程层次结构的最顶层由多个 Thread Block 组成。grid_dim (grid_x, grid_y, grid_z)表示 Grid 在各维度上的 Thread Block 数量各 Thread Block 通过自身坐标标识。在 PyPTO Pro 中Grid 具有以下特点每个执行到 SIMT 启动点的 Vector 核启动一个 Thread Blockpypto_pro.language.simt.grid_dim()描述外层 Vector 执行域规模由 Host 启动 Kernel 时实际生效的 Vector 核数决定不同 Thread Block 彼此独立不能依赖固定的执行顺序当前 Grid 仅使用 X 维即grid_dim (grid_x, 1, 1)对应的 Thread Block 坐标中 Y、Z 均为 0。Thread Block线程块Thread Block 是 Grid 的组成单元由若干 Thread 组成。block_dim (block_x, block_y, block_z)表示一个 Thread Block 在各维度上的 Thread 数量三个分量的乘积为 Block 内 Thread 总数当前不超过2048。Thread Block 具有以下特点同一 Thread Block 内的 Thread 执行相同的 SIMT 入口函数并具有相同的 Thread Block 尺寸块内 Thread 可以访问传入的共享 Tile定义 SIMT 入口函数时可以声明单个 Thread Block 允许启动的最大 Thread 数量max_threads实际的 Thread 数由simt_functhreads中方括号内的 threads 决定。Thread线程Thread 是 SIMT 结构中的最小编程单元。每个 Thread 具有独立的局部变量和执行状态并通过自身在 Thread Block 内的三维坐标处理不同数据。当每个外层 Vector 核都调用同一simt_functhreads一次时启动的 Thread 总数为$$ \text{total_threads} grid_x \times grid_y \times grid_z \times block_x \times block_y \times block_z $$Warp 执行与分支Warp 是硬件在线程块内组织执行的单位。当前 A5 的 Warp Size 为 32线程按块内线性编号划分到 Warp线程总数不是 32 的整数倍时最后一个 Warp 只有部分线程有效。同一 Warp 中的线程可以根据数据进入不同分支但分支发散会降低执行效率因此线程数取 32 的整数倍是可考虑的性能选择。从源码看python/pypto_pro/language/_simt_api.py 也暴露了pypto_pro.language.simt.warp_size()接口用于返回目标 SIMT Warp 尺寸与文档中 Warp Size 为 32 的说明相互印证。线程索引层级查询接口每个 Thread 都有对应的三维坐标。开发者通过线程层级查询接口获取 Grid、Thread Block 和 Thread 的信息从而确定当前 Thread 负责处理的数据。下表汇总了 PyPTO Pro 的 SIMT 查询接口及其约束PyPTO Pro 接口说明返回形式约束pypto_pro.language.simt.grid_dim()Grid 在各维度上的 Thread Block 数量。dim3 形式的三维对象各分量为 DT_UINT32 Scalar。当前仅使用 X 维Y、Z 维均为 1X 维由外层 Kernel 实际使用的 Vector 核数决定。pypto_pro.language.simt.block_dim()Thread Block 在各维度上的 Thread 数量。dim3 形式的三维对象各分量为 DT_UINT32 Scalar。三个分量的乘积不能超过入口函数的 max_threads且不能超过 2048。pypto_pro.language.simt.block_idx()当前 Thread Block 在 Grid 中的三维坐标。dim3 形式的三维对象各分量为 DT_UINT32 Scalar。X 坐标范围由 Grid 的 X 维大小决定当前 Y、Z 坐标均为 0。pypto_pro.language.simt.thread_idx()当前 Thread 在 Thread Block 内的三维坐标。dim3 形式的三维对象各分量为 DT_UINT32 Scalar。各维坐标范围由 Thread Block 对应维度的大小决定。pypto_pro.language.simt.linear_thread_idx()当前 Thread 在块内按 X 维优先展开的编号。DT_UINT32 Scalar。范围为 [0, 块内线程总数)不含 Block 偏移。这些接口在 python/pypto_pro/language/_simt_api.py 中有对应的 API 声明thread_idx、block_dim、block_idx、grid_dim、linear_thread_idxIR 层的linear_thread_idx操作定义见 python/pypto_pro/ir/op/simt_ops.py。块内线性索引和线程索引的具体使用方法参见 SIMT计算。SIMT 函数入口函数与辅助函数PyPTO Pro 支持两类 SIMT 函数入口函数描述 Thread Block 中每个 Thread 执行的计算逻辑辅助函数用于复用逐 Thread 逻辑在调用它的 Thread 中执行。函数类型定义方式作用调用方式SIMT 入口函数pypto_pro.language.vector_function(modesimt, max_threadsN)定义每个 Thread 执行的完整计算结果写入传入的 Tile 或 Tensor。由外层 JIT Kernel 通过simt_functhreads调用。SIMT 辅助函数pypto_pro.language.vector_function(modesimt)封装可复用的逐 Thread 计算可以不返回值或返回一个 Scalar。由 SIMT 入口函数或其他辅助函数调用。Host 先启动外层pypto_pro.language.jit(archa5)Kernel再由外层 Kernel 在 Vector 执行域中启动 SIMT 入口函数入口函数可以继续调用辅助函数。辅助函数在调用它的线程中执行不创建新线程。图中 Thread 0 至 Thread N-1 表示实际启动 N 个线程的情况。一般情况下实际线程数由simt_functhreads中方括号内的 threads 决定可以小于装饰器声明的 max_threads。当前不支持 Host 直接启动 SIMT 函数也不支持在 SIMT 函数内嵌套调用 SIMT 入口函数。内存层级与操作对象Scalar、Tile 与 TensorSIMT 函数可以操作 Scalar、Tile 和 Tensor 三类对象它们对应不同的内存层级和共享范围操作对象内存层级作用范围主要用途Scalar通常映射到寄存器Thread 私有保存参数、索引、局部变量和标量计算的中间结果。TileUBThread Block 内共享保存块内数据并在 Thread 之间交换中间结果。TensorGMGrid 范围内可访问保存输入、输出和跨 Thread Block 访问的数据支持基于运行时索引进行不规则访存。抽象硬件架构SIMT 在 AIV 上的执行环境PyPTO Pro 的 SIMT 函数运行在 AIVVector Core上。AI 处理器内部有多个 Vector Core每个 Vector Core 包含计算单元、Shared Memory位于 UB和寄存器。核外的 GM 是全局内存空间被所有 Vector Core 共享L2 Cache 位于 GM 与各 Vector Core 之间也由多个 Vector Core 共享。SIMT 计算涉及的主要硬件资源如下计算单元执行 SIMT 函数中的标量计算、地址计算和控制逻辑寄存器和栈空间每个 Thread 独立使用用于保存函数参数、线程索引、局部变量和中间结果Shared Memory使用 UB 中的部分空间供同一 Thread Block 内的 Thread 交换数据和复用中间结果Data Cache使用 UB 中的部分空间缓存 SIMT 线程访问的 GM 数据L2 Cache缓存 GM 数据并降低访问延迟由硬件管理GM保存输入、输出以及不同 Thread Block 访问的数据。SIMT 线程访问 GM 时数据经过 L2 Cache 和 Vector Core 内的 Data Cache 后进入 Thread 私有寄存器。Thread Block 内共享的数据可以保存在 Shared Memory 中由块内 Thread 直接访问。实战定义、启动与数据映射定义入口函数和辅助函数以下示例计算 1000 个 FP32 元素的output input_tensor * scale bias。辅助函数复用逐线程计算入口函数通过 Block 和 Thread 索引定位元素并检查尾部边界import pypto_pro.language as pl ELEMENTS 1000 THREADS 256 pl.vector_function(modesimt) def affine(value: pl.DT_FP32, scale: pl.DT_FP32, bias: pl.DT_FP32) - pl.DT_FP32: return value * scale bias pl.vector_function(modesimt, max_threadsTHREADS) def transform( output: pl.Tensor[[1, ELEMENTS], pl.DT_FP32], input_tensor: pl.Tensor[[1, ELEMENTS], pl.DT_FP32], count: pl.DT_UINT32, scale: pl.DT_FP32, bias: pl.DT_FP32, ): index pl.simt.block_idx().x * pl.simt.block_dim().x pl.simt.thread_idx().x if index count: output[0, index] affine(input_tensor[0, index], scale, bias)参数类型可由实参推导Tensor 和 Scalar 参数可以写出注解以便阅读和类型校验Tile 参数不支持使用注解。在外层 Vector 执行域调用线程块外层 Kernel 的 vector section 中通过simt_functhreads调用 SIMT 入口函数import pypto_pro.language as pl pl.jit(archa5) def transform_kernel( input_tensor: pl.Tensor[[1, ELEMENTS], pl.DT_FP32], output: pl.Tensor[[1, ELEMENTS], pl.DT_FP32], count: pl.DT_UINT32, scale: pl.DT_FP32, bias: pl.DT_FP32, ): with pl.section_vector(): transformTHREADSmax_threads声明单个线程块的上限threads指定本次调用的实际尺寸圆括号内按位置传入实参。入口函数不能在 SIMT 函数中嵌套调用。从 Host 启动外层 Kernel在 A5 环境中通过 Host 启动外层 JIT Kernel 函数import torch import torch_npu torch.npu.set_device(0) input_tensor torch.arange(ELEMENTS, dtypetorch.float32, devicenpu:0).reshape(1, ELEMENTS) output torch.empty_like(input_tensor) scale 2.0 bias 1.0 blocks (ELEMENTS THREADS - 1) // THREADS transform_kernelNone, blocks torch.npu.synchronize() torch.testing.assert_close(output, input_tensor * scale bias, rtol0, atol0)本例在 Host 侧将block_dim设置为 4表示实际使用 4 个 Vector 核。每个 Vector 核调用一次transformTHREADS各启动一个包含 256 个线程的 Thread Block最后一个 Thread Block 只有 232 个线程访问数据其余 24 个线程被边界判断跳过。注意Host 启动参数block_dim用于配置核数SIMT 函数内的pypto_pro.language.simt.block_dim()表示 Thread Block 在各维度上的线程数两者含义不同。配置线程与映射数据索引方括号中的 threads 可以写成一至三个整数表达式例如simt_func256、simt_func16, 16、simt_func8, 8, 4pypto_pro.language.simt.block_dim()返回相应的三维尺寸未给出的维度补 1。各维必须为编译期正整数乘积不得超过 max_threads且不得超过 2048。例如simt_func8, 8, 4表示每个线程块有 256 个线程。一维线程块可使用全局索引。二维或三维线程块可以使用 X 维优先的pypto_pro.language.simt.linear_thread_idx()展开再结合线程块编号和每块线程总数计算全局索引import pypto_pro.language as pl pl.vector_function(modesimt, max_threads256) def copy_3d( source: pl.Tensor[[1, 1024], pl.DT_FP32], output: pl.Tensor[[1, 1024], pl.DT_FP32], ): dims pl.simt.block_dim() threads_per_block dims.x * dims.y * dims.z index pl.simt.block_idx().x * threads_per_block pl.simt.linear_thread_idx() if index 1024: output[0, index] source[0, index]若外层 Kernel 通过copy_3d8, 8, 4调用该 SIMT 入口函数pypto_pro.language.simt.linear_thread_idx()给出块内 0 至 255 的编号pypto_pro.language.simt.block_idx().x乘以 256 后提供块偏移。例如当实际有 4 个线程块时它们分别处理索引 0~255、256~511、512~767 和 768~1023。标量计算与数据依赖处理标量计算与类型转换SIMT 函数可以使用公共 Scalar 表达式并提供标量计算与类型转换 API支持以下功能对应接口在 python/pypto_pro/language/_simt_api.py 中均有声明类型转换pypto_pro.language.simt.cast用于数值转换pypto_pro.language.simt.bitcast用于重新解释二进制位模式基础数学运算simt.abs、simt.min、simt.max、simt.sqrt、simt.rsqrt和simt.fma分别提供绝对值、最值、平方根、平方根倒数和融合乘加指数、对数和三角函数simt.exp、simt.exp2、simt.log、simt.log2、simt.log1p、simt.sin、simt.cos和simt.tanh取整运算simt.rint、simt.round、simt.floor、simt.ceil和simt.trunc支持按不同规则取整浮点数值判断simt.isnan和simt.isinf分别判断数值是否为 NaN 或无穷。外层流水同步SIMT 入口函数调用在 V 流水异步执行。混合计算中MTE2 搬入的数据需要就绪后才能被 SIMT 访问SIMT 更新的 UB 数据需要计算完成后才能被 MTE3 搬出loadMTE2 → MTE2/V同步 → SIMD或SIMT计算V→ V/MTE3同步 → storeMTE3普通 Tile 通过成对的pypto_pro.language.system.sync_src和pypto_pro.language.system.sync_dst表达依赖。使用带mutex_ids的pypto_pro.language.make_tile_group时默认启用的 Auto Mutex 可管理pypto_pro.language.load、SIMT 入口函数调用和pypto_pro.language.store的缓冲区依赖仅开启auto_mutexTrue不会为普通pypto_pro.language.make_tile自动补全同步。原子更新多个线程操作同一元素时使用原子操作系列接口simt.atomic_add、simt.atomic_sub、simt.atomic_exch、simt.atomic_max、simt.atomic_min、simt.atomic_inc、simt.atomic_dec、simt.atomic_cas、simt.atomic_and、simt.atomic_or、simt.atomic_xor声明见 python/pypto_pro/language/_simt_api.py。原子操作不提供跨 Block 屏障也不保证其他地址的数据已经就绪。完整类型和返回值规则见原子操作 API。示例用标量索引实现 Gather以下示例从输入 Tensor 中按indices指定的行号读取数据表达语义为output[row, col] input_tensor[indices[0, row], col]每个线程处理一个或多个输出行行内使用循环复制这是索引表达示例实际性能还需结合数据布局和线程映射测量import pypto_pro.language as pl INPUT_ROWS 100000 WIDTH 128 OUTPUT_ROWS 12288 THREADS 256 BLOCKS (OUTPUT_ROWS THREADS - 1) // THREADS pl.vector_function(modesimt, max_threadsTHREADS) def gather_rows( output: pl.Tensor[[OUTPUT_ROWS, WIDTH], pl.DT_FP32], input_tensor: pl.Tensor[[INPUT_ROWS, WIDTH], pl.DT_FP32], indices: pl.Tensor[[1, OUTPUT_ROWS], pl.DT_INT32], row_count: pl.DT_UINT32, ): first_row pl.simt.block_idx().x * pl.simt.block_dim().x pl.simt.thread_idx().x row_stride pl.simt.grid_dim().x * pl.simt.block_dim().x for row in pl.range(first_row, row_count, row_stride): input_row indices[0, row] for col in pl.range(0, WIDTH, 1): output[row, col] input_tensor[input_row, col] pl.jit(archa5) def gather_kernel( input_tensor: pl.Tensor[[INPUT_ROWS, WIDTH], pl.DT_FP32], indices: pl.Tensor[[1, OUTPUT_ROWS], pl.DT_INT32], output: pl.Tensor[[OUTPUT_ROWS, WIDTH], pl.DT_FP32], row_count: pl.DT_UINT32, ): with pl.section_vector(): gather_rowsTHREADS调用方须保证0 row_count OUTPUT_ROWS且被读取的 indices 元素均处于0, INPUT_ROWS)。形状固定索引值可以在运行时变化。准备好满足注解的 NPU Tensor 后从 Host 启动gather_kernel[None, BLOCKS本例在 Host 侧将block_dim设置为BLOCKS48表示实际使用 48 个 Vector 核每个 Vector 核启动一个包含 256 个线程的 Thread Block。线程按pypto_pro.language.simt.grid_dim().x * pypto_pro.language.simt.block_dim().x跨步处理后续行覆盖全部 12288 行。各线程写入不同输出行行间没有数据依赖不需要线程块屏障。当前能力边界从 SIMT计算 的说明可以归纳出当前 SIMT 能力边界SIMT 入口必须由外层 A5 Vector 执行域调用不支持 Host 直接启动 SIMT 函数或在 SIMT 函数中嵌套调用 SIMT 入口函数SIMT 中不支持 Tile 创建、SIMD Tile 计算、Reg 计算或 System 流水操作不支持动态 GM Shape、Tile Subview、L1 Buffer Tile、DN/NZ 布局和通用指针参数未提供 Warp shuffle/vote/reduce、线程私有数组和显式 Cached GM 访问接口Tensor/Tile 须以完整对象传入不支持元素、Slice 或 Tile Subview 作为函数参数元素位宽不得小于 8 bit运行期索引不等于动态 Tensor Shape后者当前不支持。进一步阅读SIMT计算完整的接口约束、标量计算与数据依赖处理实战编程范式概述SIMD 与 SIMT 的定位与 AI Core 硬件基础SIMT API线程查询、标量计算、同步与原子操作的完整接口清单Add 算子快速入门SIMT从零开始的 SIMT 算子示例Kernel 核函数了解 Host 启动参数block_dim的含义与设置。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考