TileLang 布局可视化实战：用 plot_layout 与编译器推断 Pass 图形化解析 T.Layout 和 T.Fragment 映射

TileLang 布局可视化实战用 plot_layout 与编译器推断 Pass 图形化解析 T.Layout 和 T.Fragment 映射【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelangTileLang 的数据排布Layout / Fragment决定了数据在寄存器、线程与共享内存之间如何搬运是理解其高性能内核生成机制的关键但抽象的索引表达式很难靠阅读代码直接看懂。本文基于仓库文档 docs/tools/layout_visualization.md完整讲解 TileLang 提供的两条布局检查路径设计阶段用tilelang.tools.plot_layout直接绘制T.Layout/T.Fragment以及编译阶段通过pass_configs开启布局推断可视化、查看编译器为每个 fragment 缓冲区选定的真实映射。读完后你能够独立绘制任意自定义排布图、复现 MMA/MFMA 片元布局图并在编译产物中定位编译器推断出的线程映射表达式。两种互补的检查方式设计时绘图与编译期推断可视化TileLang 提供两种相关的布局检查手段对应开发的两个不同阶段tilelang.tools.plot_layout直接绘制你已经拥有的T.Layout或T.Fragment对象。适合在设计数据映射阶段先把自己的排布想法画出来验证布局推断可视化Layout Inference Visualization打印编译器在编译过程中推断出的 fragment 布局并为每个推断出的二维 fragment 保存一张图。适合在检查编译后的内核中编译器最终为某个 buffer 选择了哪种映射。从源码结构看两者是清晰的分工plot_layout 实现 是一个纯 Python 侧的分发函数按对象类型把绘制任务分派给_plot_fragment_layout或_plot_layout_map而编译器侧的可视化则挂载在编译流水线上通过 pass 配置键 控制开关# tilelang/transform/pass_config.py TL_LAYOUT_VISUALIZATION_ENABLE tl.layout_visualization_enable TL_LAYOUT_VISUALIZATION_FORMATS tl.layout_visualization_formats安装可视化依赖PNG、PDF、SVG 等图片输出依赖 Matplotlib。安装 TileLang 的可视化附加依赖即可pip install tilelang[vis]需要注意纯文本形式的编译器输出text-only compiler output不会创建 Matplotlib figure因此只想要文本映射时不强制要求该依赖。直接绘制 T.Layout以 4×4 转置映射为例下面这个例子绘制一个 4×4 的转置映射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, )该调用会写出./tmp/transpose_4x4.png并打印保存路径。理解这张图的关键在于“视图”语义默认viewinput输入视图网格中每个单元代表一个输入位置单元上的数字标签是该输入位置被映射到的扁平化输出位置viewoutput输出视图网格换成输出空间标签改为“映射到该输出位置的源坐标”。从 输入/输出视图的实现 可以看出细节输入视图中N 维输入被压成二维网格的规则是“除最后一维外的所有维合并为行坐标最后一维作为列坐标”见_flatten_to_2d输出视图的网格形状则由映射后实际出现过的最大输出坐标推断。若需要改变输入视图的网格形状可以传grid_shape(rows, cols)其乘积必须等于输入元素总数——源码中对此有显式校验不满足时会抛出带详细说明的ValueError见 校验逻辑。绘制 T.Fragment按线程着色标注 T/L对于T.Fragment描述线程如何持有片元数据的对象绘图的语义完全不同每个网格单元按线程着色并标注两行文字——线程 IDT和线程内的局部寄存器索引Lplot_layout(fragment, namemma_load, formatspdf)Fragment 绘制实现 揭示了两个值得注意的实现细节绘制前会逐元素调用layout.map_forward_thread(index)和layout.map_forward_index(index)枚举全部映射且内部assert要求线程映射与局部索引映射都是单值的——这正是“直接绘制 Fragment 需要二维输入形状且映射为单值”这一限制见下文“局限性”的源码出处着色上有一个 warp 级优化当线程数不小于 32 时同一 warp 内的线程前 32 个改用hsv色轮的 6 倍步进重新取色让 warp 内相邻线程颜色明显区分若线程数少于 32 会打印UserWarning提示建议至少使用一个 warp 的线程以获得最佳观感见 着色逻辑。plot_layoutAPI 详解完整函数签名与各参数含义如下与 源码签名 一一对应plot_layout( layout, save_directory./tmp, namelayout, colormapNone, verboseFalse, formatspdf, viewinput, grid_shapeNone, )参数含义layout一个T.Layout或T.Fragment其他类型会抛出TypeError分发与类型检查。save_directory输出文件目录不存在时自动创建。默认./tmp。name不带扩展名的基础文件名。colormapMatplotlib colormap 名称。不指定时T.Layout默认用SpectralT.Fragment默认用RdPu。verbose为True时在构建绘图过程中逐条打印每个映射。formats字符串pdf、png、svg、all或逗号分隔的组合如png,svg。view仅对T.Layout生效input或output对T.Fragment不使用该参数。grid_shape仅对T.Layout的输入视图生效用(rows, columns)覆盖显示网格其乘积必须等于输入元素总数。几个行为要点plot_layout返回Noneformatsall时会同时写出 PDF、PNG、SVG 三种文件PNG 以 255 DPI 保存见_save_plotformats的解析逻辑很宽松all展开为[pdf, png, svg]逗号分隔串按段切分但只接受字符串传入列表等其他类型会抛出TypeError见_parse_formats高维T.Layout会被扁平化为二维网格除最后一维外的所有维构成行坐标最后一维构成列坐标除plot_layout外该模块还提供了一个 plot_fragment_tv 辅助函数以“线程 × 局部索引”的转置视角TV 视图绘制 Fragment行轴为局部索引I、列轴为线程T可作为进阶诊断工具。示例目录转置、重排、Swizzle 与 MMA/MFMA 片元仓库的 examples/plot_layout 目录提供了一组可直接运行的示例覆盖文档中列出的全部典型场景文件覆盖内容layout_transform.py布局组合与变换4×4/8×8 转置、3D→2D 重排[2,4,8] → (k, i*4j)、行列交织interleavelayout_swizzle.py共享内存 swizzle 映射fragment_mma_load_a.pyNVIDIA CUDA MMA 片元映射fragment_mfma_load_a.pyAMD MFMA 片元映射以 layout_transform.py 为例它展示了如何用T.Layout构造并绘图import tilelang.language as T from tilelang.tools import plot_layout # 4x4 转置 transpose_layout T.Layout([4, 4], lambda i, j: (j, i)) plot_layout(transpose_layout, nametranspose_4x4) # 3D - 2D 重排 转置: (i, j, k) 形状 [2, 4, 8] - (k, i*4j) reshape_layout T.Layout([2, 4, 8], lambda i, j, k: (k, i * 4 j)) plot_layout(reshape_layout, namereshape_3d_to_2d)而 fragment_mma_load_a.py 展示了如何把“基础片元”逐级扩展成 warp 级与 block 级片元并分别绘图先根据 dtype 位宽32/16/8 bit选择对应的shared_*_to_mma_*变换函数通过IndexMap.from_func反推每个(i, j)对应的 lane 与 local id构造出 16×16 的base_layout再用repeat/replicate扩展出 32×16 的warp_layout和 128×32 的block_layout最后对每一级调用plot_layout。这正是验证 MMA 指令数据分布是否符合预期的标准流程。安装依赖后从仓库根目录即可运行任意示例pip install tilelang[vis] python examples/plot_layout/fragment_mma_load_a.py绘图默认写到./tmp目录除非示例中另行指定了输出目录。可视化编译器推断的布局通过 pass_configs 启用第二条工作流面向“检查编译产物”。在 JIT 内核上通过pass_configs启用编译器 pass 即可import 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 kernel(A, block_M, block_N): M, N T.const(M, N) A: T.Tensor((M, N), T.float16) B T.empty((M, N), T.float16) with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads128) as (bx, by): fragment T.alloc_fragment((block_M, block_N), T.float16) T.copy(A[by * block_M, bx * block_N], fragment) T.copy(fragment, B[by * block_M, bx * block_N]) return B照常编译或调用内核即可。对 block 的layout_map注解中出现的每个被推断出的T.Fragmentpass 会打印形如C_local inferred layout: Shape: [32, 32] - [8] Thread: thread-index expression Index: [local-index expression] Replicate: 1图片文件则写入./tmp文件名由 buffer 名派生例如C_local_layout.svg。配置项的具体行为可以对照 pipeline 侧的读取逻辑 与 配置键定义TL_LAYOUT_VISUALIZATION_ENABLE默认为False启用后TL_LAYOUT_VISUALIZATION_FORMATS接受txt、png、pdf、svg、all或逗号分隔的组合省略 formats 设置时退化为仅文本输出只要可视化被启用编译器总会打印文本映射txt只是显式声明文本输出、并不会额外增加图片格式非法的 formats 取值会被检测并报错源码中对非法格式集合有显式校验。仓库中有一个完整的端到端示例 visual_layout_inference.py一个带T.Pipelined(num_stages3)的 32×32 tile matmul 内核开启TL_LAYOUT_VISUALIZATION_FORMATS: svg后运行matmul(a, b, 32, 32, 32)编译器实际打印出的C_local推断布局为C_local inferenced 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]同时该示例还通过torch.testing.assert_close校验了计算正确性——说明可视化 pass 只是编译流程中的观察手段不影响内核语义。局限性与使用建议文档明确列出了四类限制均可在源码中找到对应实现直接绘制T.Fragment要求二维输入形状且线程/局部索引映射为单值。源码中对map_forward_thread和map_forward_index的返回值长度均assert len(...) 1高维或一对多映射会在此处失败编译器生成的图片只覆盖二维输入形状的被推断 fragment。其他形状仍会打印文本映射但会跳过出图并给出警告绘图在 Python 侧逐元素枚举布局。大形状会产生大图并消耗可观的时间与内存建议对超大 tile 先降维或抽样观察输出视图的形状由映射坐标推断。非双射或稀疏的映射可能覆盖某个单元或留下空单元因此输出视图应视为诊断视图diagnostic view而不是布局正确性的校验手段——校验映射是否正确建议以verboseTrue的文本输出或T.Layout的逆映射为准。总体而言plot_layout面向“我设计的映射是否如预期”布局推断 pass 面向“编译器最终选了哪种映射”两者配合使用时你可以先在纸面上用前者推敲 swizzle 与片元分布再在编译期用后者确认生产级内核中每个 fragment buffer 的真实排布从而把 TileLang 内核的性能调优从“猜映射”变成“看图说话”。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考