TileLang Layout 系统深度指南Layout/Fragment 核心类型、CuTe 代数与布局推断【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang本文以 TileLang 仓库中的布局系统开发文档.agents/skills/tilelang-layout/SKILL.md为骨架结合 src/layout/layout.h、tilelang/layout/ 与 maint/layout_inference/ 等源码与测试工具系统讲解tl::Layout/tl::Fragment的数据模型、复刻replication语义、CuTe 布局代数及其 TileLang↔CuTe 转换、布局推断layout inference的三种层级与可插拔代价模型并给出调试、验证与 Python 侧检查的完整实践路径。读完本文你将能够理解 TileLang 中逻辑坐标→物理坐标的映射机制读懂 Fragment 的打印输出与逆向打包规则掌握tl.layout_cost_model两个策略的行为差异并会用maint/layout_inference工具链做布局回归验证。一、核心类型Layout 与 Fragment1.1tl::Layout逻辑坐标到物理坐标的映射tl::Layout的核心职责是把逻辑坐标映射为物理坐标其数据模型定义在 src/layout/layout.h 的LayoutNode中由两部分组成input_size_ffi::ArrayPrimExpr逻辑形状logical shape即输入各维的 extentforward_index_ffi::ArrayPrimExpr前向索引每个输出轴对应一个PrimExpr表达式。表达式使用进程级process-wide占位变量书写InputPlaceholder(i)打印时表现为_i、_j等复刻占位变量为ReplicationPlaceholder()打印为_rep见 src/layout/layout.h。关键 API 一览均可从 Python 侧经 FFI 调用见 tilelang/layout/layout.pyAPI语义InputShape()/OutputShape()输入形状 / 输出形状。注意OutputShape()不是存储的字段而是由分析器Analyzer根据各表达式的取值范围推导出来的Python 侧对应get_input_shape()/get_output_shape()GetForwardIndex()取出前向索引表达式数组Pythonget_forward_index()/indexForward(vars)给定输入坐标计算物理输出坐标Inverse()求逆映射InverseWithLevel()带迭代映射层级信息的求逆见下文 Fragment 的规范打包Reshape()重设逻辑形状支持rescale_num/rescale_den处理不同 dtype 别名视图时的元素尺寸换算DetectInjective()校验不同逻辑坐标 → 不同物理坐标的单射性DebugOutput()结构化的调试打印Pythonrepr直接使用它从 Python 构造一个Layout非常直观Layout(shape, forward_fn)其中forward_fn接收每个维度的变量并返回前向索引表达式详见 tilelang/layout/layout.py。__call__与map_forward_index则用 TVM 的IndexMap完成具体坐标的求值。1.2tl::Fragment带线程维度的寄存器缓冲布局tl::Fragment继承Layout额外增加线程维度用于描述寄存器缓冲register buffer中数据的物理分布定义见 src/layout/layout.hPython 包装见 tilelang/layout/fragment.py。核心成员forward_thread_一个PrimExpr描述逻辑点被哪个线程持有可能引用ReplicationPlaceholder()_repreplicate_size_复刻数thread_range_通过BindThreadRange设置用于 warp specialization 时偏移线程空间前向映射本身始终保持在归一化的[0, T)范围内。两个推导而非存储的量ThreadExtent() 在有界定义域下max(forward_thread) 1OutputShape()同样由分析器推导。复刻replication语义replicate_size表示每个逻辑点由多少个线程共同持有例如广播读场景下每个线程都持有整份数据即FullyReplicated。配套操作Replicate / DeReplicate / CondenseReplicateVar用于操控复刻轴。Python 侧可分别通过replicate()、condense_rep_var()调用见 tilelang/layout/fragment.py。1.3 已经踩过的三个坑Gotchas文档明确列出了几个经过验证、极易踩中的点FragmentNode::GetForwardVars()会前置复刻占位变量。当replicate_size 1时它把ReplicationPlaceholder()前插到变量列表头部src/layout/layout.cc而输入占位变量始终是尾部的InputDim()条目。因此Python 代码若把get_forward_vars()与形状逐位zip对复刻 Fragment 而言是错误的。规范打包顺序(thread, slot, rep)的打包逻辑位于FragmentNode::InverseWithLevelsrc/layout/layout.cc。它将 rep 变为尾部的一个普通输入维度extent 为ReplicateExtent()替换占位变量thread 作为尾部输出。于是逆映射(slot..., thread) - (coords..., rep)rep 永远在最后。loop_partition.cc正是消费这一顺序并在存储时生成 replica-zero 保护。需要把 Fragment 当普通多输出布局用时请严格镜像这一打包方式。结构化打印通过DebugOutput()/ Pythonrepr输出例如Fragment((2,) - (2,), replicate: 256, thread: _rep, index: (_i,), thread_range: I.Range(0, 256))结构化字段从 Python 侧可直接访问replicate_size、get_thread_size()、forward_thread、forward_index、thread_range、get_input_shape()方便在调试与测试中断言字段值而非匹配打印字符串。1.4 Swizzle共享内存的 XOR 布局Swizzle见 tilelang/layout/swizzle.py、src/layout/swizzle_mode.h是共享内存中基于 XOR 的布局用于缓解 bank conflict。它不能用 strided 布局表达在 CuTe 侧被建模为独立的Swizzle函子ComposedLayout而永远不是(shape, stride)模式。Swizzle 模式枚举定义在 src/layout/swizzle_mode.hNONE、SWIZZLE_32B、SWIZZLE_64B、SWIZZLE_128B其字节宽度BBits分别对应 0、1、2、3。此外 src/layout/layout.h 提供了MakeSwizzledLayout、MakeWgmmaSwizzledLayout、MakeTcgen05MmaSwizzledLayout、MakeFullBankSwizzleLayout等构造器以及DetectSwizzleMode与MergeSwizzleLayouts按更小粒度合并两个 swizzle 布局。二、CuTe 布局代数TileLang 内建的全套 CuTe 实现TileLang 为 MMA/TMA 后端维护了一套完整的 CuTe 布局代数实现C 侧在namespace tvm::tl::cutesrc/layout/cute_layout.h、src/layout/cute_layout.ccPython 侧为tilelang.layout.cutetilelang/layout/cute.py。2.1 基本约定布局表示为(shape, stride)的IntTuple树列主序column-major第一维最快——与 TileLang 行主序的直觉相反是转换时最容易出错的地方多输出陪域使用ScaledBasis步长vaxis、Ei记法。Python 侧E(mode)返回单位基向量ScaledBasis(value, mode)表示value沿mode方向的缩放基。2.2 代数操作全集Python 侧全部以函数形式暴露见 tilelang/layout/cute.py操作Python API说明合并连续模式coalesce(layout, max_extentNone)支持max_extent上限右逆right_inverse(layout)部分函数见下方约定左逆left_inverse(layout)复合composition(lhs, rhs)补complement(layout, cotarget)大小cosize/size/coshape陪域大小划分logical_divide/logical_product/tiled_product/blocked_product过滤/限制filter/restrict变形with_shape、make_layout、make_column_major_layout、make_row_major_layout、make_identity_layoutmake_layout在只传 shape 时默认生成列主序布局解析/打印parse/Print精确的 CuTe 拼写2.3 转换器probe-then-prove先探测、后证明三个*FromTileLang转换器tilelang/layout/cute.py 与Layout.from_tilelang、Layout.from_tilelang_hierarchical、ComposedLayout.from_tilelangLayoutFromTileLang针对单一扁平输出。多输出的tl::Layout按行主序序列化因此输出[thread, slot]会得到thread * slots slot的扁平布局LayoutFromTileLangHierarchical按轴per-axis恢复基准轴上的映射ComposedLayoutFromTileLang恢复 swizzle。三者都采用probe-then-prove先在 one-hot 点上数值探测 stride再符号化证明等价性——错误的恢复不可能漏网失败时返回None/nullopt而不是抛出异常。2.4 失败约定必须遵守只有上述三个*FromTileLang转换器返回 Optional其余所有代数操作在其前置条件不满足时直接 ICHECK 崩溃包括composition 的整除性、probe 中出现非常数 extent、complement 非单射、restrict 的秩不匹配等。因此输入不可信时务必包一层调用做保护RightInverse是部分函数它静默丢弃 stride-0 与非 const-stride 的模式只对最大连续链求逆。请通过size(right_inverse(F)) size(F)校验双射性。2.5 转换注意事项转换器只读取GetForwardIndex()Fragment 的forward_thread与复刻会被静默忽略除非你先把它们打包进一个普通的多输出 Layout使用上文 1.3 的规范打包方式用尾部输入变量替换ReplicationPlaceholder。如果复刻变量泄漏进探测过程转换会优雅地返回None。三、布局推断Layout Inference三级策略与代价模型3.1 整体流程tl.LayoutInference实现于 src/transform/layout_inference/layout_inference.cc为每个 fragment 缓冲与并行循环嵌套分配布局分三个层级strict严格来自注解与 MMA 指令强加的布局直接采用common公共通过共享缓冲做 BFS 传播free自由对每个连通分量connected component尝试把每个成员作为推断根保留代价最低的方案。推断结果以 IR 注解的形式落盘layout_mapBuffer → Layout挂在 SBlock 上parallel_loop_layoutFragment挂在最外层并行 For 上配套注解还有parallel_loop_predicate、parallel_loop_requires_padding_guard、coalesced_width见 src/layout/layout.h。ParallelLoopLayoutValidatorsrc/transform/layout_inference/parallel_loop_layout_validator.h负责强制注解契约。3.2 可插拔的最廉价策略最廉价由可插拔策略决定实现于 src/transform/layout_inference/layout_cost_model.h 与 layout_cost_model.cc通过环境变量tl.layout_cost_model选择register-count默认仅按寄存器槽位总数排序RegisterCountCostModelio-aware可选开启IOAwareCostModel遍历分量内所有触碰全局内存的语句fragment↔global 拷贝、直接访问 global 的并行循环在尝试布局下按max(bandwidth bytes, issue bytes)计费并在 CuTe 代数上做符号化评分pack →LayoutFromTileLang→RightInverse→Composition从合并后的模式读出向量宽度按 warp/step 粒度统计 segment 数。寄存器作为字典序 tiebreak。模型之外的语句按保守最坏情况计费——一次尝试绝不能从不透明中获利。评分结构为AttemptCost{mem, regs}比较规则mem优先regs次之见 layout_cost_model.h其中mem还包含寄存器数组因线程相关索引而溢出到 local memory 的流量估算。3.3 硬件几何参数化通道宽度lane widthMaxVectorLoadBitssrc/transform/loop_vectorize.h与向量化器共享保证模型对宽度的信念与代码生成一致warp 大小取自 target 的thread_warp_sizesegment 粒度128B见BindMemoryGeometry。3.4 评分公式的 Python 奇偶校验评分公式由 Python 奇偶校验守护maint/layout_inference/run.py --cute将符号化评分器与独立的 NumPy 精确枚举 oracle 对比。修改layout_cost_model.cc时必须同步更新 maint/layout_inference/cute_model.py保持二者一致。四、下游消费者布局如何影响代码生成布局注解最终被以下关键 pass 消费src/transform/loop_partition.cc通过 fragment 逆映射按线程划分并行循环并为复刻存储生成 replica-zero 保护src/transform/loop_vectorize.cc规划向量宽度GetVectorizeSize、IndicesCanVectorize代价模型镜像其判断TMA/MMA 降低src/cuda/op/tma_layout.cc 与producer_consumer_ws.cc通过ComposedLayoutFromTileLang恢复共享缓冲的 swizzletcgen05/wgmma 宏生成器tilelang/cuda/intrinsics/macro/则通过to_tilelang/from_tilelang_hierarchical往返 TMEM 布局。五、调试与验证工具5.1maint/layout_inference/验证框架该目录是布局验证工具链详见 maint/layout_inference/README.md每个cases/*.py构造一个已知正确答案的PrimFunc驱动run.py在两种代价策略下分别推断并与expected/*.json中的黄金布局快照对比。核心用法python run.py # 校验所有 case 与 expected/ 一致 python run.py --case NAME # 按子串过滤单个 case python run.py --show # 同时打印推断出的布局 python run.py --record # 用当前行为重写黄金快照录完必须人工审查 diff python run.py --anchor # 完整 lower 后对照 VECTOR_ANCHOR 检查设备端 TIR 的每缓冲向量宽度 python run.py --cute # 符号评分器 vs 独立精确枚举 oracle 的奇偶校验各选项的含义来自 maint/layout_inference/run.py 的 docstring 与 README--anchor闭环验证模型假设的向量宽度与向量化器实际发出的宽度是否一致不一致说明模型的宽度信念与代码生成发生了漂移——这正是共享MaxVectorLoadBits策略要防止的问题--cute中cute_model.py按(coords..., rep) - [thread, slot]的规范打包FragmentNode::InverseWithLevel把每个 fragment 经cute.Layout.from_tilelang转成单一strided 布局用right_inversecomposition推导字节地址布局从合并的 slot 模式读出向量宽度oracle.py则是独立的 numpy 实现。每个黄金布局都以 load 和 store 两种方式双路径评分含复刻门控(V, issue, bw, segments)必须完全一致。README 记录的现状是88/88 条语句匹配、100% 转换命中率--record只是录制而非批准录制后必须人工阅读expected/下的 diff确认每个变更的布局都是预期行为即使录制模式下结构不变量case 的check回调依然强制执行。每个 case 文件定义VARIANTS名字 → 返回新PrimFunc的可调用对象与可选的check(variant, model, result)断言。result中每个布局都是结构化 dictcommon.layout_to_dict例如{kind: Fragment, input_shape: [2], output_shape: [2], forward_index: [_i], replicate: 256, threads: 256, forward_thread: _rep, thread_range: [0, 256]}因此黄金 diff 能精确定位哪个字段发生了移动如replicate: expected 1, got 128check也能断言字段值而非匹配打印格式。当前 8 个 case 覆盖的场景见下表摘自 maint/layout_inference/README.mdcase固定了什么行为elementwise_copy基线两种模型必须对合并、向量化的往返布局达成一致主等分锚点fp8_copy1 字节 dtype共享宽度策略宽端上的 16 元素向量宽度broadcast_readIssue #1729两模型有意分歧——register-count 保留线程折叠的遗留病态布局黄金记录之io-aware 必须选择全复刻 非复刻合并循环由check强制transposed_store加载与存储把布局拉向相反方向fp32 变体两模型选择不同布局值得基准测试mixed_dtype_chainfp16/fp32 fragment 对共处一个分量向量化按冲突 dtype 定宽reduce_broadcastSoftmax 形行归约 广播消费最常用的真实 kernel 分量两模型一致offset_region_copy区域偏移携带块索引外部变量的多块平铺拷贝偏移区域必须与零偏移区域排序一致shared_stagingglobal→shared→fragment→global 链shared 侧拷贝在 io 模型之外fragment 仅由拷出决定扩展规范每次修改推断、代价模型或转换器时都应为该变更新增一个 case黄金快照先--record录制再由人工审查。5.2 其他诊断手段DLOG 日志推断与代价模型 pass 在 debug 构建下通过 DLOG 记录决策布局可视化tl.layout_visualization_enable可渲染布局Python 侧检查在模块上运行tl.transform.LayoutInference()并读取注解提取惯用法见 maint/layout_inference/common.py或调用cute.Layout.from_tilelang(...)查看布局的(shape, stride)正规形式。六、文件地图速查领域文件核心类型src/layout/layout.h、src/layout/layout.cc、tilelang/layout/layout.py、tilelang/layout/fragment.pySwizzlesrc/layout/swizzle_mode.h、tilelang/layout/swizzle.pyCuTe 代数src/layout/cute_layout.h、src/layout/cute_layout.cc、tilelang/layout/cute.py推断src/transform/layout_inference/layout_inference.cc代价模型src/transform/layout_inference/layout_cost_model.h、layout_cost_model.ccMMA 布局src/layout/gemm_layouts.cc、src/layout/tcgen05_layout.h验证工具链maint/layout_inference/含 README.md、run.py、cute_model.py、oracle.py七、总结与最佳实践清单坐标系直觉要切换TileLang 布局是逻辑→物理映射输出形状是推导值CuTe 布局是列主序(shape, stride)与行主序直觉相反处理 Fragment 时牢记三个 GotchaGetForwardVars()前置 rep 变量规范打包为(slot..., thread) - (coords..., rep)且 rep 在尾部打印输出中的结构化字段可直接在 Python 断言转换器与代数操作的失败语义不同只有*FromTileLang返回 Optional其余操作前置条件失败即 ICHECK 崩溃不可信输入务必包一层调用RightInverse是部分函数用size比对校验双射自由模式搜索的代价策略默认register-countio-aware按全局内存流量带宽字节 vs issue 字节取最大符号化计费并共享向量化器的宽度策略修改评分公式必须同步cute_model.py并保持--cute奇偶校验通过回归纪律任何对推断、代价模型或转换器的改动都在maint/layout_inference下新增 case、录制黄金、人工审查后再提交。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考