调试 TileLang 算子时布局老对不上？plot_layout 与推断 Pass 的双路径排错法

调试 TileLang 算子时布局老对不上plot_layout 与推断 Pass 的双路径排错法【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang写 Tile 级算子时最容易让结果静默出错的往往不是算法本身而是布局数据在 shared memory、寄存器片元Fragment与线程之间怎么摆决定了T.copy、T.gemm这些高层操作最终被编译成什么样的访存指令。布局写错数值可能依然正确比如对某些巧合形状但一旦形状变化、换一个 MMA/MFMA 微架构性能或正确性就崩掉。TileLang 提供了一条可视化排错双路径设计阶段用plot_layout直接绘制手头的T.Layout/T.Fragment编译阶段用布局推断 Pass 输出编译器真正选定的映射。本文结合仓库源码与示例给出这套双路径排错法的完整用法与边界。设计时绘图与编译期推断的职责边界在 docs/tools/layout_visualization.md 中官方把两条路径的分工讲得很清楚tilelang.tools.plot_layout绘制你已经拥有的T.Layout或T.Fragment——用于设计阶段验证你脑子里的映射。布局推断可视化Layout inference visualization打印并绘制编译器在layout_map注解中为每个 Fragment 推断出的布局——用于检查编译后的 kernel 里 buffer 真正被选定的映射。即设计时自证与编译后复核。前者回答我想让数据这样摆对吗后者回答编译器最后到底怎么摆的。二者结合恰好覆盖排错的完整闭环先在设计期用plot_layout排除显而易见的映射错误再打开推断 Pass 核对编译产物确认推断结果没有偏离你的预期。绘制需要 Matplotlib安装可视化扩展即可pip install tilelang[vis]纯文本输出txt不依赖 Matplotlib。4×4 转置映射与按线程着色的正确姿势最简单的起点是 examples/plot_layout/layout_transform.py 中的 4×4 转置。T.Layout的构造语义是输入索引 → 输出索引的映射函数lambda i, j: (j, i)即把(i, j)送到(j, i)import tilelang.language as T from tilelang.tools import plot_layout transpose T.Layout([4, 4], lambda i, j: (j, i)) plot_layout( transpose, save_directory./tmp, nametranspose_4x4, formatspng, )默认的viewinput视图里每个格子代表一个输入位置格内标注的是展平后的输出坐标切到viewoutput网格则变为输出空间格内标注回源坐标。绘制时默认按 Spectral 色带为输出坐标着色格子色块 坐标标注的组合能一眼看出转置的对角结构。对T.Fragment则走另一套绘制逻辑见 tilelang/tools/plot_layout.py 中的_plot_fragment_layoutplot_layout(fragment, namemma_load, formatspdf)Fragment 图以线程为单位着色前min(warp_size, num_threads)个线程即第一个 warp 的 32 条 lane使用 HSV 色带逐条区分其余线程沿用默认的 RdPu 色带每个格子内同时标注T{线程号}Thread ID与L{本地寄存器号}Local ID图例也明确标注T: Thread ID / L: Local ID。如果线程数小于 warp size32工具会发出 UserWarning提示至少 32 线程观感最佳——低于一个 warp 的片元图很难看出 lane 间的跨步规律。仓库在 examples/plot_layout/images/base_layout.png 存放了一张代表性片元图FP16 MMA 加载基础布局格子内的T/L双标签与按线程的 HSV 着色一目了然从 micro-tile 到 block 的片元组装以及 shared memory swizzle真实算子里没有裸的 4×4片元是从硬件微布局开始一层层 repeat / replicate 出来的。以 NVIDIA MMA 为例examples/plot_layout/fragment_mma_load_a.py 先用shared_16x16_to_mma_32x8_layout_sr_a这类IndexMap反推 thread 与 local index 的映射构造出与 dtype 微尺寸get_mma_micro_size匹配的基础T.Fragment再逐级放大# base layout 16x16 base_layout make_mma_load_base_layout(dtypeT.float16, matrixA, transposedFalse) plot_layout(base_layout, namebase_layout) # warp layout 32x16按线程方向 repeat 2 行再 replicate 2 列 warp_layout base_layout.repeat([block_rows, 1], repeat_on_threadTrue).replicate(block_cols) plot_layout(warp_layout, namewarp_layout) # block layout 128x32跨线程展开 4x2 份 block_layout warp_layout.repeat([warp_rows, chunk], repeat_on_threadFalse, lower_dim_firstFalse) plot_layout(block_layout, nameblock_layout)关键在repeat的两个参数repeat_on_thread决定放大是吃线程维度还是纯迭代空间维度lower_dim_first决定低维优先还是高维优先二者选错会直接产出线程映射错位的错误布局。绘制每一级布局就能肉眼验证 lane 在 M、K 方向上的跨步是否符合ldmatrix的取数模式。ROCm 侧同理examples/plot_layout/fragment_mfma_load_a.py 基于shared_16x16_to_local_64x4_layout_A等 MFMA 布局构造基础片元再按warp_rows × warp_cols与block_rows × block_cols逐级放大。对照两张图的 T 标注可以直观看出 CUDA MMA32 线程跨 16×16与 AMD MFMA64 线程跨 16×16的片元分布差异——这正是跨后端移植算子时最常踩的坑。shared memory 一侧的映射由 swizzle 布局刻画。tilelang/layout/swizzle.py 提供三档标准 swizzlemake_quarter_bank_swizzled_layout32B1-bit XOR、make_half_bank_swizzled_layout64B2-bit XOR、make_full_bank_swizzled_layout128B3-bit XOR。examples/plot_layout/layout_swizzle.py 对 fp16 分别绘制 8×16、8×32、8×64 以及多 tile 的 32×64 布局可以直接观察 XOR 位数增加时行内块如何逐步拆换。高维限制与诊断视图的适用边界可视化工具好用但有明确的边界提前知道能省去大量图怎么画歪了的排查时间高维 Layout 一律压成 2D 网格除最后一维作为列坐标外前面所有维度合并为行坐标见_flatten_to_2d。3D 的[2, 4, 8] → [8, 8]重排布局因此能正常绘制若想强制某种网格排布viewinput下可用grid_shape(rows, cols)覆盖但rows × cols必须严格等于输入元素总数。Fragment 绘制要求更高必须是二维输入形状且 thread 映射与 local-index 映射均为单值map_forward_thread/map_forward_index的结果长度断言为 1。带 replicate 的 Fragment 在格内用T分隔多个线程号但 local id 必须一致。output 视图是诊断视图而非校验工具输出空间形状是从映射坐标反推的取各维最大坐标 1非双射或稀疏映射可能覆盖已有格子或留下空格子。此时图会缺格或重影只能辅助观察不能当作布局合法性证明。大布局要谨慎绘制逻辑在 Python 里逐一枚举所有元素超大 shape 会生成巨型图并显著拖慢绘图、占用内存。高维 Fragment 以及编译器推断出的非二维 Fragment只打印文本、跳过绘图并给出 warning见 tilelang/analysis/layout_visual.py。pass_configs 启用自动推断的完整配置设计期验证完还需要在编译产物上复核。布局推断可视化挂在后端 Pass 管线上tilelang/backend/pass_pipeline/pipeline_utils.py 的LayoutVisual通过 JIT 的pass_configs开启。两个开关定义在 tilelang/transform/pass_config.pytilelang.PassConfigKey.TL_LAYOUT_VISUALIZATION_ENABLEtl.layout_visualization_enable默认Falsetilelang.PassConfigKey.TL_LAYOUT_VISUALIZATION_FORMATStl.layout_visualization_formats可取txt/png/pdf/svg/all或逗号组合。最小完整配置如下源自 examples/visual_layout_inference/visual_layout_inference.pyimport tilelang import tilelang.language as T tilelang.jit( pass_configs{ tilelang.PassConfigKey.TL_LAYOUT_VISUALIZATION_ENABLE: True, tilelang.PassConfigKey.TL_LAYOUT_VISUALIZATION_FORMATS: txt,svg, } ) def matmul(A, B, block_M, block_N, block_K, dtypeT.float16, accum_dtypeT.float32): M, N, K T.const(M, N, K) A: T.Tensor((M, K), dtype) B: T.Tensor((K, N), dtype) C T.empty((M, N), dtype) with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads128) as (bx, by): A_shared T.alloc_shared((block_M, block_K), dtype) B_shared T.alloc_shared((block_K, block_N), dtype) C_local T.alloc_fragment((block_M, block_N), accum_dtype) T.clear(C_local) for k in T.Pipelined(T.ceildiv(K, block_K), num_stages3): T.copy(A[by * block_M, k * block_K], A_shared) T.copy(B[k * block_K, bx * block_N], B_shared) T.gemm(A_shared, B_shared, C_local) T.copy(C_local, C[by * block_M, bx * block_N]) return C正常编译或调用 kernel推断 Pass 会遍历算子上的layout_map注解为每个T.Fragment打印文本映射并绘制图像。128 线程、32×32 tile 的C_local输出形如C_local inferred layout: Shape: [32, 32] - [8] Thread: _j // 16 * 64 _i // 16 * 32 _i % 8 * 4 _j % 8 // 2 Index: [_j % 16 // 8 * 4 _i % 16 // 8 * 2 _j % 2] Replicate: 1Thread表达式揭示的是核心信息_j // 16 * 64表示每 16 列跨 64 个线程即 N 方向每 16 元素换一条 lane 组_i // 16 * 32表示 M 方向每 16 行再跨 32 线程_i % 8 * 4 _j % 8 // 2是 8×8 微块内部的 lane 分配——这正是 Hopper MMA 的 32×32 累加器在 128 线程上的标准分布。若你设计时期望的是其他分布对照plot_layout绘制的手工 Fragment 与这里推断出的表达式差异一眼可辨。图像文件按 buffer 名写入./tmp如C_local_layout.svgall格式则同时输出 png/pdf/svg。两个配置细节值得注意一是不设 FORMATS 时只输出纯文本——txt不额外产生图片二是推断 Pass 只对二维输入形状的 Fragment 出图其余形状打印文本后跳过绘图并告警与plot_layout的二维限制保持一致。小结把排错从猜变成看布局对不上的根因绝大多数是手写的映射函数、repeat/replicate组合或后端推断规则三者之一偏离了预期。双路径排错法把这三者全部可视化设计期用plot_layout验证映射函数与片元组装逻辑编译期用推断 Pass 核对编译器在真实 kernel 上选定的布局表达式中间再用view、grid_shape、formats等参数控制观察视角与产物格式。仓库中的 examples/plot_layout 与 examples/visual_layout_inference 两个目录覆盖了从转置、swizzle、MMA/MFMA 片元到 GEMM 推断的全部场景遇到布局说不清时先跑一张图再谈对错。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考