跳到主要内容
TileLang 系统教程

tilelang-tutorial.huandu.me

第 24 章 NVIDIA 后端:代码生成与硬件特性

/code/24-nvidia-backend-codegen/

02-gemm-inst-selection/README.md1964 字节
# 02 · 复现 `T.gemm` 的指令选择

## 目的

把 `_vendor/tilelang-v0.1.15/src/cuda/op/gemm.cc:336-375` 的 `Gemm::SelectInst`
谓词链,以及 `_vendor/tilelang-v0.1.15/src/op/tcgen5_meta.h:22-177` 的
`GetTCGEN5MMAMeta` 判定,用标准库重写成可交互模型。用来回答:

* 同样一句 `T.gemm`,在 sm_70 / sm_75 / sm_80 / sm_90 / sm_100 上分别选什么?
* 「我的 GEMM 为什么没用 TCGEN05」——是哪一条不满足?
* fp8 的 WGMMA 为什么对 transpose 这么挑?

## 运行前提

Python 3.9+,只用标准库。**不需要 GPU,也不需要安装 TileLang。**

## 运行命令

```bash
python3 gemm_inst_selection.py
```

## 期望输出

脚本打印一张「配置 → 指令」表,最后做自检。参考片段:

```text
sm_90a  M=128 N=128 K=64   f16->f32  A=shared B=shared C=fragment  -> cuda.wgmma
sm_90a  M=32  N=128 K=64   f16->f32  A=shared B=shared C=fragment  -> cuda.mma
sm_100  M=128 N=128 K=64   f16->f32  A=shared B=shared C=tmem      -> cuda.tcgen05
sm_100  M=128 N=128 K=64   f16->f32  A=shared B=shared C=fragment  -> cuda.mma
sm_75   M=16  N=16  K=16   f16->f32  A=shared B=shared C=fragment  -> cuda.mma
sm_70   M=16  N=16  K=16   f16->f16  A=shared B=shared C=fragment  -> cuda.mma
sm_70   M=16  N=16  K=16   f16->f32  A=fragment B=shared C=fragment -> cuda.fma
self-check: ok
```

## 作者实际做到的验证级别

**CPU 逻辑验证**:在 macOS (Apple Silicon)、Python 3.11 上运行过。

## 简化了什么

1. 只实现 `GetTCGEN5MMAMeta` 里 **f16/bf16 → f32** 那一支。fp8/fp6/fp4、
   `kind::mxf8f6f4`、int8 等分支没有实现(它们的结构相同,只是 atom 取值不同)。
2. 不实现 `T.gemm` 的 warp 划分(`ComputeWarpPartition`),只回答「选哪条指令」。
3. `CheckWgmma` 只覆盖 f16、bf16、tf32、fp8、int8 的 dtype/转置/K 整除规则;
   B 必须在 shared 这一条也实现了。
4. 不实现 `use_2cta` 之后的 `LowerBlackwell2SM` 等 pass 行为。

## 与源码的对应关系

| 本脚本 | `gemm.cc` / `tcgen5_meta.h` |
| --- | --- |
| `select_inst()` | `Gemm::SelectInst`(`_vendor/tilelang-v0.1.15/src/cuda/op/gemm.cc:336-375`) |
| `check_wgmma()` | `CheckWgmma`(:40-78) |
| `allow_wgmma()` | `AllowWgmma`(:89-97) |
| `allow_tcgen5_mma()` | `AllowTcgen5Mma`(:80-87) |
| `tcgen5_meta_f16()` | `GetTCGEN5MMAMeta`(`_vendor/tilelang-v0.1.15/src/op/tcgen5_meta.h:22-177`) |
| `allow_volta_mma()` / `allow_turing_mma()` | :99-117 / :131-145 |

← 浏览本章全部代码