cute 之 Layout

最新推荐文章于 2025-08-23 11:23:17 发布

知识搬运工人

最新推荐文章于 2025-08-23 11:23:17 发布

阅读量210

点赞数

CC 4.0 BY-SA版权

分类专栏： CUDA Architecture GPU 文章标签：人工智能 cute

原文链接：https://zhuanlan.zhihu.com/p/661182311

GPU 同时被 3 个专栏收录

18 篇文章

订阅专栏

CUDA

15 篇文章

订阅专栏

Architecture

2 篇文章

订阅专栏

引自：cute 之 Layout - 知乎

计算机中的内存是一维的线性地址空间，而数学计算问题所要处理的空间经常是高维的。如GEMM（General Matrix Multiplication）问题的数学计算体系是二维计算空间，Deep Learning计算体系是三维以上的计算空间(batch, height, width, channel, etc.)。如何高效的表达高维计算空间，如何高效便捷的将计算所要求的高维空间映射到一维空间变得越来越重要。

历史上对该问题的探究可以分为三个阶段：

第一阶段BLAS的row/col-major + leading dimension描述阶段；
第二阶段Tensor的shape + stride阶段；
第三阶段为Hierarchy Tensor阶段。

第一， BLAS的row/col-major + leading dimension描述阶段：二十世纪七八十年代开发的BLAS（Basic Linear Algebra Subprograms）库中描述矩阵时面对的问题多为二维问题，其通过引入行优先、列优先的概念来刻画一维存储结构和二维逻辑结构的映射关系。为了做某个维度的地址对齐、padding等，引入了leading dimension的概念，其可以描述第二个维度上数据的连续性关系。

第二，Tensor的shape + stride阶段：二十一世纪初Deep Learning在模式识别等领域暂露头角，各种框架应运而生，在Deep Learning计算中为了提升计算密度batch化成了重要手段，这时高维数据的描述成为必须，相关人员在blas中leading dimension基础上进行了扩展，引入了shape和stride描述体系。shape stride描述体系可以很好的刻画一维存储结构和高维逻辑结构的映射关系，并且使得某些操作只需要改变shape stride描述无需对数据实体进行移动，在stride的帮助下通过简单的coordinate和stride点积既可以完成逻辑地址和一维索引的映射。shape stride描述可以表达简单的高维数据，但是其对于具体的维度的轴必须是单调的，在数据的连续性方面要求比较严格。

第三， BLAS的row/col-major + leading dimension描述阶段： 2023年层级的Tensor（Hierarchy Tensor）描述被提出，其在shape stride描述的基础上添加了具体的轴的层级描述（Graphene Tensor IR）[1]，使得数据在某个具体的轴上可以有更为丰富的表达。对于现代GPU计算体系而言，Tensor的分块和复杂映射都可以通过该层级体系（描述+代数）来表达和推导，其构建了复杂逻辑空间和硬件排列映射的一条途径。

Tensor的层次化描述体系在Graphene Tensor IR中进行了详细的描述；在具体实现方面，NVidia开源的cutlass模板中也利用同样的思路进行Tensor描述和计算。为了更好的抽象和利用编译时优化，cutlass利用了C++17实现了cute，它定义了层次化Tensor体系和其之上的代数计算，然后基于该cute描述体系实现了最新hopper架构之上的矩阵运算。Tensor是数据的表达，其表达一个相对独立且有结构的数据体，而Tensor内的数据排布则由Layout来表达。本文将关注这种映射结构，即Layout。本文在行文组织结构上：首先以示例的形式介绍一维向量的表示、二维矩阵的表示，然后介绍层次化的Tensor表示，再后介绍cute中常用的编译时和运行时形状描述，最后本文总结Layout的功能。

简单地讲，引入有层次的描述（Layout）代数来表达计算空间和一维地址空间的映射问题。Layout 是一个数据排列的描述体系，其可以实现将逻辑坐标映射到索引坐标（offset表示）。Layout包含Shape和Stride两部分。其中Shape描述排列的分块层次和结构。Stride描述块内或块间的数据排列连续性。Shape和Stride都是层级的嵌套表示。也就是说Shape可以包含Int数据和Shape。Shape和Stride需要有相同的层次关系。

在介绍有层次的Layout之前，我们先回顾无层次的Layout描述，也就是shape和stride描述的高维tensor。