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

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

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto

SIMT（Single Instruction Multiple Threads，单指令多线程）是 PyPTO Pro 在 AIV（Vector Core）上提供的一种线程并行编程模型，以 Thread 为基本执行单元，适合表达不规则数据访问、逐线程分支控制和共享地址的原子更新。本文以 SIMT 编程范式模块为骨架，系统讲解 Grid / Thread Block / Thread 三级线程层次、Warp 执行与分支发散、线程索引查询接口、SIMT 入口函数与辅助函数的定义和启动方式，以及 SIMT 函数可操作的 Scalar / Tile / Tensor 内存层级与抽象硬件架构，并辅以仓库源码与实战示例，帮助读者掌握在 A5 架构上编写与启动 SIMT 计算的能力。

SIMT 与 SIMD 的定位：在并行执行模型中的角色

在进入 SIMT 细节之前，先明确它在 PyPTO Pro 并行体系中的位置。根据 programming_paradigm_overview.md，PyPTO 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 Block；pypto_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 维均为 1；X 维由外层 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_idx），IR 层的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(mode="simt", max_threads=N)`	定义每个 Thread 执行的完整计算，结果写入传入的 Tile 或 Tensor。	由外层 JIT Kernel 通过`simt_functhreads`调用。
SIMT 辅助函数	`@pypto_pro.language.vector_function(mode="simt")`	封装可复用的逐 Thread 计算，可以不返回值或返回一个 Scalar。	由 SIMT 入口函数或其他辅助函数调用。

Host 先启动外层@pypto_pro.language.jit(arch="a5")Kernel，再由外层 Kernel 在 Vector 执行域中启动 SIMT 入口函数；入口函数可以继续调用辅助函数。辅助函数在调用它的线程中执行，不创建新线程。

图中 Thread 0 至 Thread N-1 表示实际启动 N 个线程的情况。一般情况下，实际线程数由simt_functhreads中方括号内的 threads 决定，可以小于装饰器声明的 max_threads。当前不支持 Host 直接启动 SIMT 函数，也不支持在 SIMT 函数内嵌套调用 SIMT 入口函数。

内存层级与操作对象：Scalar、Tile 与 Tensor

SIMT 函数可以操作 Scalar、Tile 和 Tensor 三类对象，它们对应不同的内存层级和共享范围：

操作对象	内存层级	作用范围	主要用途
Scalar	通常映射到寄存器	Thread 私有	保存参数、索引、局部变量和标量计算的中间结果。
Tile	UB	Thread Block 内共享	保存块内数据，并在 Thread 之间交换中间结果。
Tensor	GM	Grid 范围内可访问	保存输入、输出和跨 Thread Block 访问的数据，支持基于运行时索引进行不规则访存。

抽象硬件架构：SIMT 在 AIV 上的执行环境

PyPTO Pro 的 SIMT 函数运行在 AIV（Vector 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(mode="simt") def affine(value: pl.DT_FP32, scale: pl.DT_FP32, bias: pl.DT_FP32) -> pl.DT_FP32: return value * scale + bias @pl.vector_function(mode="simt", max_threads=THREADS) 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(arch="a5") 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(): transformTHREADS

max_threads声明单个线程块的上限，threads指定本次调用的实际尺寸，圆括号内按位置传入实参。入口函数不能在 SIMT 函数中嵌套调用。

从 Host 启动外层 Kernel

在 A5 环境中，通过 Host 启动外层 JIT Kernel 函数：

import torch import torch_npu torch.npu.set_device(0) input_tensor = torch.arange(ELEMENTS, dtype=torch.float32, device="npu: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, rtol=0, atol=0)

本例在 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, 4；pypto_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(mode="simt", max_threads=256) 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 搬出：

load（MTE2） → MTE2/V同步 → SIMD或SIMT计算（V）→ V/MTE3同步 → store（MTE3）

普通 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_mutex=True不会为普通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(mode="simt", max_threads=THREADS) 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(arch="a5") 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设置为BLOCKS=48，表示实际使用 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，后者当前不支持。