Data Layout in GPU Programming

GPU 编程里的数据布局 (Data Layout) 指的数据在显存中的组织方式, 以及张量块 (Tensor Block) 中的元素和 GPU 资源 (线程, 寄存器, 共享内存) 之间的映射关系. 简单说它就是回答 "张量中的第 (i,j) 个元素储存在哪个线程的哪个寄存器里, 或者在共享内存的哪个位置" 这个问题.

Triton 中几种常见的数据布局类型是:

• 分块布局 (Blocked Layout)

布局传播

Triton 中有一个专门的 Pass 用来消除多余的布局转换操作. Pass 的工作流程是:

• 标记锚点 (Anchor) 张量: 锚点张量指的是布局不能被修改的张量, 包括函数参数, Dot 类操作 (因为 GPU 硬件对 Dot 操作通常有特定的布局要求) 的结果, 原子操作, 昂贵的 Load/Store 等. 同时如果一个 reshape 操作被显式地指定了不允许重排布局, 那么它的输出张量也会被标记为锚点张量.
• 传播布局: 从锚点张量开始, 向其所有使用者传播布局信息. 这一步不实际修改布局信息, 只是将某个张量用到的各种布局记录下来 (用 Worklist 实现)
• 解决布局冲突: 如果某个张量