TritonGPU 操作详解

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills

本文档详细描述 TritonGPU 方言（ttg）中所有操作的定义、签名和语义。所有操作签名均从 TableGen 源码精确提取。

源码参考：TritonGPUOps.td

1. 布局转换操作

1.1 ttg.convert_layout

将张量从一种布局编码转换为另一种。这是 TritonGPU 中最核心的操作之一，实际的数据移动（如共享内存读写）在此操作中发生。

%result = ttg.convert_layout %src : tensor<128xf32, #blocked> -> tensor<128xf32, #dot_op>

2. 异步拷贝操作

2.1 ttg.async_copy_global_to_local

将数据从全局内存异步拷贝到共享内存（local memory）。类似于tt.load，但数据写入共享内存描述符而非分布式张量。

支持的加载字节数（按计算能力）：

%token = ttg.async_copy_global_to_local %src, %dst mask %mask other %other cacheModifier = #tt<cache ca> : tensor<128x!tt.ptr<f32>> -> !ttg.memdesc<128xf32, #shared>

2.2 ttg.async_wait

等待异步拷贝操作完成。

计算能力要求：>= 80

%token2 = ttg.async_wait %token1 {num = 0 : i32}

2.3 ttg.async_commit_group

将当前未提交的异步操作标记为一个组。

计算能力要求：>= 80

%token = ttg.async_commit_group tokens %token1, %token2

3. 共享内存操作

3.1 ttg.local_alloc

在共享内存中分配缓冲区，返回内存描述符。

构建器：

额外方法：

• isSharedMemoryAlloc()：检查是否为共享内存分配
• getAlignmentOrDefault()：获取对齐值或默认值

%buf = ttg.local_alloc %src : () -> !ttg.memdesc<128xf32, #shared> %buf2 = ttg.local_alloc : () -> !ttg.memdesc<128xf32, #shared>

3.2 ttg.local_dealloc

释放共享内存缓冲区。可选操作，未显式释放的缓冲区在所有使用后的第一个后支配点被隐式释放。

ttg.local_dealloc %buf : !ttg.memdesc<128xf32, #shared>

3.3 ttg.local_load

从共享内存描述符加载到分布式张量。

%result = ttg.local_load %buf token %async_token : !ttg.memdesc<128xf32, #shared> -> tensor<128xf32, #blocked>

3.4 ttg.local_store

将分布式张量存储到共享内存。

ttg.local_store %tensor, %buf : tensor<128xf32, #blocked> -> !ttg.memdesc<128xf32, #shared>

4. 内存描述符视图操作

4.1 ttg.memdesc_index

取内存描述符沿第 0 维第 i 个元素的子视图。不影响底层内存。

例如：输入2x4x16xf16，index=1 → 输出4x16xf16

%sub = ttg.memdesc_index %src[%index] : !ttg.memdesc<2x4x16xf16, #shared> -> !ttg.memdesc<4x16xf16, #shared>

4.2 ttg.memdesc_subslice

取内存描述符的子视图，指定各维偏移。不影响底层内存。

例如：输入32x16xf16，offsets=[2,1] → 输出8x16xf16，覆盖input[2:10, 1:17]

%sub = ttg.memdesc_subslice %src[0, 0] : !ttg.memdesc<32x16xf16, #shared> -> !ttg.memdesc<8x16xf16, #shared>

4.3 ttg.memdesc_trans

转置内存描述符的视图。不影响底层内存。

%transposed = ttg.memdesc_trans %src {order = [1, 0]} : !ttg.memdesc<8x16xf16, #shared> -> !ttg.memdesc<16x8xf16, #shared>

4.4 ttg.memdesc_reshape

创建不同形状的内存描述符视图。不影响底层内存。

%reshaped = ttg.memdesc_reshape %src : !ttg.memdesc<128xf16, #shared> -> !ttg.memdesc<8x16xf16, #shared>

4.5 ttg.memdesc_reinterpret

将内存描述符重新解释为不同类型和形状。要求原始描述符是连续的。

%reinterpreted = ttg.memdesc_reinterpret %src : !ttg.memdesc<128xf16, #shared> -> !ttg.memdesc<64xf32, #shared>

5. 流水线操作

5.1 ttg.predicate_stage

流水线阶段谓词，用于软件流水线。

%pred = ttg.predicate_stage %iv, %ub, %step maxStage 2 stage 0 : index -> i1

5.2 ttg.mask

流水线掩码操作，包含一个区域。

5.3 ttg.mask.return

ttg.mask区域的终止操作。

6. 类型转换操作

6.1 ttg.fp4_to_fp

将打包为 i8 的 FP4 (E2M1) 数据上转换为浮点类型。每个 i8 的低 4 位表示第一个 FP4 元素，高 4 位表示第二个。

%result = ttg.fp4_to_fp %src {axis = 1 : i32} : tensor<16x8xi8> -> tensor<16x16xf32>

7. 全局内存操作

7.1 ttg.global_scratch_alloc

在全局内存中分配当前程序私有的缓冲区。

%ptr = ttg.global_scratch_alloc {nbytes = 1024 : i32, alignment = 16 : i32} : !tt.ptr<i8>

8. Warp 特化操作

8.1 ttg.warp_specialize

在不同 warp 组上异步执行不同代码。默认区域可隐式捕获，分区区域与上方隔离。

额外方法：

• getPartitionRegions()：获取分区区域
• getCaptureSizeAlign()：获取捕获列表的大小和对齐
• getTotalPartitionWarps()：获取额外 warp 总数

%0 = ttg.warp_specialize(%a, %b) default { %out = some_operation(%a) ttg.warp_yield %out : i32 } partition0(%arg0: i32, %arg1: i32) num_warps(8) { some_async_dispatch(%arg0, %arg1) ttg.warp_return } : (i32, i32) -> i32

8.2 ttg.warp_specialize.partitions

包含ttg.warp_specialize的隔离分区区域的容器操作。

8.3 ttg.warp_yield

ttg.warp_specialize默认区域的终止操作。操作数作为ttg.warp_specialize的 SSA 结果传递。

ttg.warp_yield %a, %b : i32, tensor<32xbf16, #blocked>

8.4 ttg.warp_return

ttg.warp_specialize分区区域的隐式终止操作。无操作数，因为分区区域不能返回任何值。

9. 操作分类速查表

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills