TME/TCE API Reference：MP31（mp_31）

本参考介绍 MTCC MP31（__MUSA_ARCH__ == 310）设备端 builtin 及其参数约束。 编译目标为 --offload-arch=mp_31。未列出的后续架构 builtin 和其他平台 wrapper 不属于本页。

1. WMMA builtin

__musa_wmma_m16n8k4_mma
__musa_wmma_m16n8k8_mma
__musa_wmma_m16n8k16_mma
__musa_wmma_m8n16k16_mma
__musa_wmma_m16n16k16_mma
__musa_wmma_m16n16k32_mma
__musa_wmma_m16n16k64_mma

控制字段包括 layout、mma_type 和 A/B major mode。直接调用时必须严格使用 对应 shape 的参数顺序和寄存器布局；不要以 CUDA WMMA 或其他架构签名替换。

2. SQMMA builtin

__musa_sqmma_m16n64_mma       __musa_sqmma_m32n32_mma
__musa_sqmma_m32n64_mma       __musa_sqmma_m32n128_mma
__musa_sqmma_m64n16_mma       __musa_sqmma_m64n32_mma
__musa_sqmma_m64n64_mma       __musa_sqmma_m64n128_mma
__musa_sqmma_m128n32_mma      __musa_sqmma_m128n64_mma
__musa_sqmma_m128n128_mma

descriptor、scale、A/B major mode 与 accumulator 组合必须匹配相应 shape。 不要从 WMMA 参数顺序推导 SQMMA raw 调用。

2.1 TceABDtype

Enum	Value	A	B	Accumulator
U8U8	0	u8	u8	u32
S8S8	1	s8	s8	s32
F16F16	2	f16	f16	f32
BF16BF16	3	bf16	bf16	f32
TF32TF32	4	tf32	tf32	f32
E4M3E4M3	5	e4m3	e4m3	f32
E5M2E5M2	6	e5m2	e5m2	f32
F16S4	7	f16	s4	f32
F16S8	8	f16	s8	f32
S4F16	9	s4	f16	f32
S8F16	10	s8	f16	f32
E4M3E5M2	11	e4m3	e5m2	f32
E5M2E4M3	12	e5m2	e4m3	f32
BF16S4	13	bf16	s4	f32
BF16S8	14	bf16	s8	f32
S4BF16	15	s4	bf16	f32
S8BF16	16	s8	bf16	f32

3. Async barrier builtin

void    __musa_async_bar_record(int32_t count);
void    __musa_async_init_arrival(int32_t barrier_id,
                                  int32_t arrival_count,
                                  int32_t initial_phase);
void    __musa_async_add_trans(int32_t barrier_id, int32_t transaction_bytes);
int32_t __musa_async_arrive(int32_t barrier_id);
void    __musa_async_wait(int32_t barrier_id, int32_t phase);
void    __musa_async_decrease_trans(int32_t barrier_id,
                                    int32_t transaction_bytes);

barrier_id 是 block-local，__musa_async_arrive 的合法输入为 1–63；0 不可用。add_trans 增加待完成字节数，arrive 返回 generation phase，wait 等待该 phase 完成，decrease_trans 释放此前添加的字节数。

4. TME tile load/store builtin

实际参数包括 descriptor、rank-specific 整数向量、坐标、shared 指针和 barrier。 使用对应 MP31 operation 可避免改变参数顺序。

__musa_tme_ld_tile_1d  __musa_tme_ld_tile_2d  __musa_tme_ld_tile_3d
__musa_tme_ld_tile_4d  __musa_tme_ld_tile_5d  __musa_tme_ld_blk
__musa_tme_ld_im2col_3d __musa_tme_ld_im2col_4d __musa_tme_ld_im2col_5d
__musa_tme_st_1d  __musa_tme_st_2d  __musa_tme_st_3d
__musa_tme_st_4d  __musa_tme_st_5d  __musa_tme_st_blk
__musa_tme_store_commit  __musa_tme_store_read_wait

load 的 transaction bytes 必须等于预期写入 shared memory 的字节数；store 的 commit/read-wait 是顺序控制标记，不是通用 barrier 替代品。

5. Prefetch-only 与 im2col

__musa_tme_tile_prefetch_only_1d ... __musa_tme_tile_prefetch_only_5d
__musa_tme_im2col_prefetch_only_3d
__musa_tme_im2col_prefetch_only_4d
__musa_tme_im2col_prefetch_only_5d

这些 builtin 只发出 TME prefetch hint，不写 shared memory，也不产生可等待的 完成事件；真正使用数据仍须发起对应 load 并管理 barrier。im2col 参数还包括 filter、output dimension 和 convolution parameter 向量。

6. 使用清单

• barrier ID 传入 __musa_async_arrive 时必须为 1–63；
• descriptor、tile 对齐、rank、stride 和 dtype 必须彼此匹配；
• 每个 async load 都要有 transaction accounting、arrival 和 phase wait；
• 不支持的 shape、架构或边界输入必须保留普通 load/store 回退。