Buffer 与内存

Buffer 与内存#

Parameter 使用 T.Tensor 类型标注，kernel body 中的临时 buffer 则由下面介绍的两种 declaration API 创建。使用 A[i, j] 访问元素， A[m0:m0+BM, 0:BK] 取得 TensorRegion，A.ptr_to([i, j]) 取得 元素 pointer，A.data 取得原始 data pointer。

声明 Buffer#

创建 tensor 的两个基本 API 是：

  • T.alloc_tensor(ty_args=[T.Tensor(shape, dtype, scope=..., layout=...)])： 分配新的 storage，返回 type 为 TensorType 的 tensor value。 T.alloc_shared 和 T.alloc_local 是相应 scope 的分配简写。

  • T.decl_tensor(data, ty_args=[T.Tensor(...)])：在已有的 typed pointer 上**声明 view**，不会分配 storage。TMEM 则使用下文的 T.cuda.decl_tmem 和 column address。

Buffer 的 data property 是 immutable、带类型的 pointer Expr，不一定 是 Var。T.decl_tensor 可以直接接收类型兼容的 pointer 表达式，详见 数据类型与表达式。

两种 API 使用同一个 buffer descriptor，主要参数如下：

参数

含义

dtype

元素类型，例如 "float32"、"float16"、"float4_e2m1fn"

shape

逻辑 shape，也就是各维 extents 组成的 tuple

layout

物理映射，详见 TileLayout； "default" 表示 dense row-major

pointer / TMEM column

用 T.ptr_byte_offset 对 pointer 偏移后声明 view； T.cuda.decl_tmem 接收 TMEM column address

align

Data pointer 的 byte alignment

scope 参数选择 memory space：

Scope

简写

Memory

"global"

默认值

Device global memory

"shared"

T.alloc_shared

Static shared memory（__shared__）

"shared.dyn"

pool

Dynamic shared memory，使用 pool 管理

"local"

T.alloc_local

Per-thread registers

"tmem"

TMEM pool

Blackwell tensor memory

 
A = T.decl_tensor(A_ptr, ty_args=[T.Tensor((M, K), "float16", scope="global", align=16)])   # parameter buffer
As = T.alloc_shared((BM, BK), "float16")                 # new shared tile
acc = T.alloc_local((4,), "float32")                     # register accumulator
view = T.decl_tensor(As.data, ty_args=[T.Tensor((BM, BK), "float16", scope="shared", layout="default")])  # a view over As

对于非 TMEM buffer，基于 pointer 的 buffer 只是 pointer 加上一组 metadata。 访问元素时，编译器根据 layout 计算地址：

addr(buffer[coord]) = buffer.data + elem_offset + layout.apply(*coord, shape=shape)["m"]

layout.apply 返回各个 physical axis 的映射，其中 "m" 分量是 element offset。因此，同一个逻辑访问会因为 buffer metadata 不同而生成不同的地址 计算。下面都在 4×8 区域上执行 B[i, j] = A[i, j] + 1，但用四种 方式声明 B：

from tvm_tirx.tirx.layout import TileLayout, S
from tvm_tirx.backend.cuda.axis import laneid

B = T.decl_tensor(p, ty_args=[T.Tensor((4, 8), "float32", scope="global")])                                       # row-major
B = T.decl_tensor(p, ty_args=[T.Tensor((4, 8), "float32", scope="global", layout=TileLayout(S[(4, 8):(1, 4)]))])  # column-major
B = T.decl_tensor(T.ptr_byte_offset(p, (64) * 4, ty=T.Ptr("float32", "global")), ty_args=[T.Tensor((4, 8), "float32", scope="global")])                       # shifted view
B = T.decl_tensor(p, ty_args=[T.Tensor((4, 8), "float32", scope="global", layout=TileLayout(S[(4, 8):(16, 1)]))]) # row stride 16

生成的 CUDA 会为 B[i, j] 使用不同的 index。A[i, j] 的 load 始终 保持 i*8 + j，变化的只有 B 的 metadata：

B_ptr[((i * 8) + j)]        = ...;   // row-major:        i*8 + j
B_ptr[((j * 4) + i)]        = ...;   // column-major:     j*4 + i
B_ptr[(((i * 8) + j) + 64)] = ...;   // elem_offset=64:   i*8 + j + 64
B_ptr[((i * 16) + j)]       = ...;   // row stride 16:    i*16 + j

Shared Memory#

Shared memory 分为两种：static shared memory 的大小在编译期固定， dynamic shared memory 的大小在 launch 时确定。TIRx 还提供 pool helper 管理 dynamic shared memory。

Static Shared Memory#

T.alloc_shared``（也就是 ``scope="shared"）创建最简单的 static shared buffer，其大小在编译期确定。下面先将数据写入 shared memory，通过 cta_sync 确保整个 block 都能看到写入，再读出结果：

@T.function
def smem_demo(A: T.Tensor((128,), "float32"), B: T.Tensor((128,), "float32")):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=[1], block=[128]))
    bx = T.cuda.block_idx("x")
    tx = T.cuda.thread_idx("x")
    sm = T.alloc_shared((128,), "float32")   # static shared memory
    sm[tx] = A[tx]
    T.cuda.cta_sync()
    B[tx] = sm[tx] * T.float32(2.0)

它会生成普通的 __shared__ array。下面省略了无关的 CUDA 代码：

extern "C" __global__ void __launch_bounds__(128)
smem_demo_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr) {
  int tx = ((int)threadIdx.x);
  __shared__ alignas(64) float sm_ptr[128];      // T.alloc_shared
  sm_ptr[tx] = A_ptr[tx];
  __syncthreads();                               // T.cuda.cta_sync()
  B_ptr[tx] = sm_ptr[tx] * 2.0f;
}

Dynamic Shared Memory#

Dynamic shared memory**（``scope="shared.dyn"``）的大小由 launch parameter ``sharedMemBytes`` 指定，而不是在编译期写死。一个 kernel 只能有**一个 dynamic-shared allocation，也就是 arena。因此需要先分配一次 arena，再用 T.decl_tensor 在 arena pointer 的不同 T.ptr_byte_offset 上声明 views：

arena = T.alloc_tensor(ty_args=[T.Tensor((128,), "float32", scope="shared.dyn")])   # the one arena
As = T.decl_tensor(arena.data, ty_args=[T.Tensor((64,), "float32", scope="shared.dyn")])                 # offset 0
Bs = T.decl_tensor(T.ptr_byte_offset(arena.data, (64) * 4, ty=T.Ptr("float32", "shared.dyn")), ty_args=[T.Tensor((64,), "float32", scope="shared.dyn")]) # offset 64
As[tx] = A[tx]
Bs[tx] = B[tx]
T.cuda.cta_sync()
C[tx] = As[tx] + Bs[tx]

两个 views 共用同一个 extern __shared__ arena。下面为便于阅读，将 arena 命名为 smem，并省略无关代码：

extern __shared__ __align__(64) float smem[];   // the one dynamic-shared arena
smem[tx]      = A_ptr[tx];                       // As — view at offset 0
smem[tx + 64] = B_ptr[tx];                       // Bs — view at offset 64
__syncthreads();
C_ptr[tx] = smem[tx] + smem[tx + 64];

分别调用两次 alloc_tensor``（``shared.dyn scope） 会报错，因为只允许一个 dynamic shared memory allocation。Static shared memory 的大小在编译期 确定，例如 __shared__ T x[N];；dynamic shared memory 则是 launch 时 指定大小的单个 arena，各个 buffer 是 arena 内不同 offset 上的 views。

备注

TVM 如何记录 dynamic-shared 大小。 Arena 的大小在编译期已知，本例中 是 128 个 float，也就是 512 bytes。Lowering 时，TVM 会在 device kernel 的 tirx.kernel_launch_params 中加入 "tirx.use_dyn_shared_memory" tag。Host launcher 计算总 byte 数，并将 它作为最后一个 launch argument：

# device kernel 属性：
"tirx.kernel_launch_params": ["blockIdx.x", "threadIdx.x", "tirx.use_dyn_shared_memory"]

# host 侧 launch 调用（..., gridDim.x, blockDim.x, dyn_shared_bytes）：
T.call_packed("dyn_kernel", A.data, B.data, C.data, 1, 64, 512)

运行时，这里的 512 会成为 cuLaunchKernelEx 调用中的 config.sharedMemBytes。用户不需要手工设置，它由 shared.dyn allocation 的大小推导得到。

SMEMPool#

SMEMPool 自动管理 arena。它使用 bump allocation 计算 offset，因此 不需要手工 decl 每个 view。除了 alloc 和 commit，它还支持 per-buffer align=、创建 MMA-compatible swizzle layout 的 alloc_tcgen05_mma_AB helper，以及将 cursor 回退以复用空间的 move_base_to：

from tvm_tirx.backend.cuda.lang import SMEMPool

pool = SMEMPool()                          # shared.dyn 上的 bump allocator
As = pool.alloc((BM, BK), "float16", align=128)   # 分配一个 tile
Bs = pool.alloc((BK, BN), "float16", align=128)
Cs = pool.alloc_tcgen05_mma_AB((BM, BN), "float16")  # 自动推导 MMA-compatible swizzle
pool.commit()                                 # 确定 pool 的最终大小
# pool.move_base_to(offset) 将 cursor 回退到可复用的位置

下方的 TMEM pool 建立在 SMEMPool 之上。

Registers#

Per-thread 临时数据使用 local scope。通过 T.alloc_local(shape, dtype) 分配后，这些数据只属于当前 thread。使用 静态索引的 local arrays 通常会被 scalarize 到 registers；使用动态索引，或 register pressure 较高时，也可能进入 local memory。

r = T.alloc_local((4,), "float32")   # 每个 thread 私有的 register array
for k in T.unroll(4):
    r[k] = A[tx, k]
# ... 使用 r[0..3] 计算 ...

alignas(64) float r_ptr[4];          // per-thread, register-resident
r_ptr[0] = A_ptr[tx * 4 + 0];
r_ptr[1] = A_ptr[tx * 4 + 1];
// ...

备注

alignas(64) 来自 buffer 的默认 alignment。data_alignment 默认为 runtime::kAllocAlignment，也就是 64 bytes；CUDA codegen 会把它加到 每个 allocation 上，包括 per-thread local array。对于保存在 registers 中的 array，这个 alignment 不会影响性能：只要 index 能在 编译期解析，nvcc/ptxas 就会通过 scalar replacement of aggregates （SROA）将 thread-local array 提升为 registers，因而不会使用可寻址的 local memory。只有使用动态 index 并 spill 到 local memory 的 array 才会实际受到 over-alignment 影响，这种情况并不常见。Register local 的 over-alignment 是一个已知问题，未来应为 local scope 使用 dtype 的 natural alignment。

Scalar#

Scalar 本质上是只有**一个元素**的 register array。可以直接分配大小为 1 的 local buffer，并通过 [0] 访问：

phase = T.alloc_local((1,), "int32")   # 单元素 register array
phase[0] = 0
while phase[0] < 4:
    acc = acc + A[tx, phase[0]]
    phase[0] += 1

每次都写 phase[0] 比较繁琐，因此 TIRx 提供 scalar 语法，表示同一个 单元素 register buffer，并允许直接通过名称读写：

phase: T.int32 = 0                 # mutable scalar，是上一种写法的语法糖
while phase < 4:
    acc = acc + A[tx, phase]
    phase += 1

s = T.local_scalar("int32")        # 显式形式；通过名称赋值，而不是 s[0]
acc: T.float32 = 0.0               # 带类型注解的赋值也会创建 scalar

两种写法在 parse 后会得到结构完全相同的 TIRx。Parser 会将 phase: T.int32 解析为单元素 local buffer，将 phase 和 phase += 1 解析为 phase[0] 和 phase[0] += 1。对两个 kernels 调用 tvm_tirx.ir.assert_structural_equal 会通过。Scalar 访问同一个单元素 local tensor，只是省去了 [0]。T.local_scalar、T.shared_scalar 和 T.alloc_scalar 可以显式选择 scope。

备注

为什么不使用 scalar-typed Var ？ Typed TIRx 中两种形式都 具有 Var identity。Buffer value 是带 TensorType 的 Var；store 修改的是 buffer contents，并不会重新绑定该 Var。下方 T.let 创建的则是只有一次 immutable binding 的 scalar-typed Var。因此 mutable scalar 必须由可反复 store contents 的单元素 buffer 支撑。

let#

T.let 是 immutable binding，对应一个 Bind node。它表示命名后的 值，不是 buffer，适合保存派生常量：

n: T.let = M * K               # immutable binding（Bind）
half: T.let[T.int32] = N // 2  # 显式指定类型

它会生成普通的 C scalar variable，而不是 array，也不需要 [0]。例如， 运行时变量 m 上的 half: T.let = m * 2 会生成：

int half = m * 2;     // `let` 生成类似 const 的 local variable

由于值不会改变，simplifier 可以自由执行 propagation 和 common subexpression elimination。因此在使用位置可能直接看到 m * 2，也可能 看到共享的 common-subexpression temporary，而不一定保留 half。

备注

为什么需要 immutable binding？ 因为值不会改变，arithmetic analyzer 在简化 Bind 时内部可以调用 arith::Analyzer::Bind，将关于这个 值的结论传播到所有使用位置，包括 constant bounds、表示 divisibility 和 alignment 的 modular set，以及 ranges。这些信息可用于简化 index、 消除 bounds check，以及决定 alignment 和 vectorization。Mutable scalar 是一次 memory load（buf[0]），analyzer 不能假设它保持不变，因此无法 传播这些性质。let 还是一个不需要 allocation 的 pure value，可以自由 inline、substitute 或 CSE；scalar 则是具有 load/store 语义的单元素 buffer。

Tensor Memory#

Blackwell tensor memory 不能像普通 scratch scope 那样直接分配。Kernel 必须通过 warp-uniform T.ptx.tcgen05.alloc 和 tcgen05.dealloc intrinsics 显式申请和释放，再用 T.cuda.decl_tmem(column, ty_args=[T.Tensor(..., scope="tmem", layout=...)]) 在其中声明 tensor view。

Column address 是必需参数；T.alloc_tensor 不会预留 TMEM。Tensor memory 只能通过 tcgen05 的 mma、ld、st 和 cp 访问。

手工管理时，一个 warp 将 allocation 地址写入 shared slot，再在不同 column offset 上 decl tensor views，结束时由一个 warp 释放：

addr = T.alloc_shared((1,), "uint32")
if warp_id == alloc_warp:
    T.ptx[f"tcgen05.alloc.cta_group::{cta_group}.sync.aligned.shared::cta.b32"](
        addr.ptr_to([0]), 512)
T.cuda.cta_sync()
acc = T.cuda.decl_tmem(addr[0], ty_args=[T.Tensor(
    (CTA_M, 512), "float32", scope="tmem", layout=tmem_layout)])
# ... issue and wait for all MMA / TMEM transfers ...
T.cuda.cta_sync()
if warp_id == alloc_warp:
    T.ptx[f"tcgen05.relinquish_alloc_permit.cta_group::{cta_group}.sync.aligned"]()
    T.ptx[f"tcgen05.dealloc.cta_group::{cta_group}.sync.aligned.b32"](addr[0], 512)

此时 column offsets 和 tmem_layout （datapath D/F layout）都需要手工 管理。下面的 pool 会自动生成同样的步骤。

Views#

使用 view 对同一块 TMEM allocation 分区，不再分配新的 storage：

acc0 = acc.sub[:, :128]
acc1 = acc.sub[:, 128:256]

完整示例见 Part III 的 GEMM kernels。

Buffer API#

Tensor value 包含 pointer 和 TensorType 访问契约，详见上方 声明 Buffer。因此， 多数 methods 都是在编译期改变 shape、layout 或 index arithmetic，或者返回 pointer，本身不会生成运行时操作。常用 methods 如下：

Method

作用

B.data

带类型的 base-pointer 表达式，输出为 B_ptr

B.ptr_to([i, j])

指向某个元素的 typed pointer（address_of），输出为 &B_ptr[…]

vector-typed T.decl_tensor view

Vectorized load/store，输出为 *(float4*)(B_ptr + …)

B.view(*shape, layout=…)

用新的 shape/layout 重新解释同一 storage，不复制数据

T.decl_tensor(B.data, ty_args=[T.Tensor(...)])

local buffer 中属于当前 thread 的 private register slice

B.permute(*dims)

交换 axes 后的 view，也就是 transposed layout

Pointer：``ptr_to`` / ``data``。 ptr_to 将某个元素的地址传给 intrinsic 或 inline function；data 则是 base pointer：

B[tx] = T.cuda.func_call("ld", A.ptr_to([tx]), SRC, ty="float32")

B_ptr[tx] = ld(&A_ptr[tx]);          // ptr_to([tx]) -> &A_ptr[tx];  A.data -> A_ptr

Vectorized access：typed vector view。 一次 wide transfer 搬运多个 元素，另见 数据类型与表达式：

A4 = T.decl_tensor(A.data, ty_args=[T.Tensor((N // 4,), "float32x4")])
B4 = T.decl_tensor(B.data, ty_args=[T.Tensor((N // 4,), "float32x4")])
B4[tx] = A4[tx]

*(float4*)(B_ptr + tx * 4) = *(float4*)(A_ptr + tx * 4);

Reshape / reinterpret：``view`` / ``permute``。 两者都只修改 metadata， data pointer 保持不变，变化的是 index arithmetic。A.view(64, 4) 将包含 256 个元素的 buffer 看作 64×4；A2.permute(1, 0) 交换这两个 axes：

A2 = A.view(64, 4);     y = A2[tx, 0] + A2[tx, 3]   # A2[tx, j] -> A_ptr[tx*4 + j]
At = A2.permute(1, 0);  z = At[i, j]                # At[i, j]  -> A_ptr[j*4 + i]

A2_ptr[tx * 4]  /* +3 */                 // view: row-major 64x4 index
At_ptr[(j * 4) + i]                       // permute: swapped strides

Register：显式 local view。 Thread-axis layout 把逻辑元素分布到各个 thread。为当前 thread 的 8 个 physical registers 声明 dense view 后， 可以继续 reshape 这个 view：

R  = T.alloc_tensor(ty_args=[T.Tensor((32, 8), "float32", scope="local", layout=TileLayout(S[(32, 8) : (1 @ laneid, 1)]))])
R_flat = T.decl_tensor(R.data, ty_args=[T.Tensor((8,), "float32", scope="local", layout="default")])
R_2d = R_flat.view(2, 4)

alignas(64) float R_flat_ptr[8];         // the lane's private registers