成果:把布局关系做成编译器可用的系统
在 FlagLang / FlagLib 中,我们实现的不只是一个布局记法,而是一套具备类型、运算和可执行查询的布局基础。它让坐标属于哪个空间、变换如何组合、结果采用什么数学表示,都能保留在编译器理解的对象中。
把结构一致性纳入表示
layout_t 将 extent 与 stride 捆绑在同一棵树中,消除两棵独立树可能不同构的问题。layout_mode 同时携带 extent、stride 与 axis,结构关系不再依赖调用方反复维护。
让空间身份与变换关系可以追踪
具名 Axis、IndexSpace、TypedLayout 与 DomainMap 共同描述域和映射;分级偏移保留目标空间,重索引与 coalesce 保留坐标对应。这为组合变换和合法性检查提供明确输入。
让仿射、XOR 与有限映射各有实现
整数仿射布局、F2 位线性布局与 TableLayout 形成多种表示。swizzle 保留 base 与可逆 transform;紧凑表示不能承载结果时,有限表可保存精确映射,而不是伪造 shape/stride。
这些成果的价值在于:布局不再只是后端收到的一组参数,而成为编译过程中可以检查、组合并执行查询的语义对象。下面的原理和例子,解释这套实现为什么值得这样设计。
实现依据:从设计对象到可执行操作
| 已落地能力 | 实现依据 | 源码中的检查重点 |
|---|---|---|
| F2 布局与可逆 swizzle | f2_layout.hpp | 按 XOR 规则构造变换,拒绝奇异矩阵与越界列 |
| 跨表示组合与投影 | table_layout.hpptest_table_layout.cpp | 仿射/F2 组合后保留目标空间与逐点结果 |
| 重索引与域变换 | reindex_tablecoalesce_table | 变换前后保持可执行映射语义,返回精确源域同构 |
表格描述已读取的实现与测试源码;本次文章调整未重新运行 C++ 测试。
这些成果解决什么问题
矩阵乘法的公式只告诉我们算什么。GPU 还需要知道:哪个线程持有哪些元素,数据位于共享内存还是寄存器,计算指令期待怎样的排列。相同的逻辑矩阵,换一种物理映射,就可能改变内存访问、数据搬运和指令选择。
FlagLang 的关键设计,是让这些关系成为编译器能够理解和组合的对象。Layout 连接逻辑域与物理位置;它与计算语义一起进入编译过程,供验证、重索引和目标实现使用。
用一个地址例子解释设计
设一个 4×4 矩阵,行主序地址为 4r+c。如果按 2×2 tile 连续存储,则地址变成 8⌊r/2⌋+4⌊c/2⌋+2(r mod 2)+(c mod 2)。例如同一个坐标 (1,1),两者分别得到 5 和 3。逻辑元素没有变,物理位置变了。
点选一个元素,观察物理位置
教学模型,编号以元素为单位。XOR 模式为 4r+(c XOR r),演示可逆置换;不等同于具体 GPU 的 bank-conflict 测量或完整硬件 swizzle。
这个例子说明,shape 不足以描述一个 tile。编译器必须保存映射本身,才能在改变数据组织时继续追踪“同一个元素”。
具名轴:让坐标携带身份
FlagLang 的 Axis 不只是某一维的序号,它有名义身份。IndexSpace 给轴附上范围;TypedLayout 把映射绑定到源域和目标轴。行轴、归约轴、线程位置与字节地址不会仅因为都用整数表示就自动成为同一种量。
FlagLib 的 graded_offset_t 以 axis/value 分量表达偏移。slice 产生的位移因此可以保留所属空间,而不会被压成一个脱离单位的整数。这让“映射是否合法”成为可检查的关系,而不只是作者脑中的约定。
重索引:变换可以组合,语义可以追溯
设布局 L : D → P 将逻辑域映射到物理位置;将逻辑域拆分或重排后,用 R : D′ → D 说明新坐标对应旧坐标。新布局就是 L′ = L ∘ R。
新坐标 D′ ── R ──→ 原逻辑域 D ── L ──→ 物理位置 P
重索引 原布局
L′(x) = L(R(x))这里拆分 tile 与搬动数据是不同操作:前者可以只改变描述域的方式,后者才改变物理映射。DomainMap 显式记录域之间的关系,让编译器不必从几组 shape 参数反推这两种意图。table_layout.hpp 中的 coalesce 还返回源域同构关系,保留变换前后的对应。
一个 Layout 系统,多种数学表示
FlagLib 没有把所有映射硬塞进 shape/stride。它按运算的数学性质提供不同表示,并保留它们之间的语义联系。
| 表示 | 适合表达什么 | 关键设计 |
|---|---|---|
| 整数仿射 | stride、分块、带偏移的映射 | extent 与 stride 捆绑在同一棵树中 |
| F2 位线性 | XOR 组合与 swizzle | 用 GF(2) 矩阵表达,不把整数加法冒充 XOR |
| 有限表表示 | 有限域上的精确映射 | 当紧凑表示不封闭时保留精确结果 |
尤其是 swizzle:f2_swizzled_layout_t 保留 base 与 transform,构造器检查变换可逆性,拒绝奇异矩阵。变换的数学结构仍在,后续操作可以继续理解它。
有限表表示补足另一端:某个运算的结果无法继续用简洁仿射形式表达,不代表它在语义上没有结果。TableLayout 保存有限映射图,并为组合、乘积和重索引提供精确承载。它可能付出存储与枚举成本;它的价值在于保持语义,而不是宣称任何布局都能廉价执行。
这些设计怎样服务编译
捆绑结构消除 shape/stride 树不一致;具名轴避免混用坐标空间;组合关系保留变换来路;不同数学表示让仿射与 XOR 各自按正确规则运算。查询还区分 defined、oob 和 undef,避免将域外或未定义结果误当成合法地址。
这些能力共同构成 FlagLang 的布局基础:作者写出计算和映射关系,编译器能继续检查并推导实现。Layout 本身并不自动带来某个固定加速比;真正的性能结果要看它如何与存储、Atom、流水线和调度协作。
源码与阅读依据
依据本地 FlagLib 开发快照 d3a671b0e7:include/flaglib/flaglib.hpp、f2_layout.hpp、table_layout.hpp,test/test_table_layout.cpp,以及 FlagLang README。交互图为本专栏原创教学示意。
讲解方式参考 Modern GPU Programming for MLSys:Data Layout,按具体地址、映射公式和编译器设计逐层展开。本文介绍 FlagLang,参考教材使用 TIRx。