@fedebruzzone7: 从范畴论的角度：CuTe 布局的范畴基础 https://arxiv.org/abs/2601.05972

从范畴论的视角看：CuTe 布局的范畴基础 🔗https://t.co/BuL7RnPHkH https://t.co/6UcrRiH2m8 — # CuTe 布局的范畴基础 来源：https://arxiv.org/html/2601.05972 Colfax [email protected]年1月 ###### 摘要 NVIDIA的CUTLASS库提供了一套健壮且富有表现力的方法，用于描述和操作GPU上的多维张量数据。这些方法的概念根植于CuTe布局的抽象概念以及此类布局的丰富代数体系，包括复合、逻辑积和逻辑除等运算。在本文中，我们提出一个范畴框架来理解这种布局代数，重点聚焦于一类自然出现的可处理布局。为此，我们定义了两个范畴Tuple{\boldsymbol{\mathsf{Tuple}}}和Nest{\boldsymbol{\mathsf{Nest}}}，其中的态射引出布局。我们定义了这些范畴上态射的一套运算，并证明了它们与相应布局运算的兼容性。此外，我们对由我们构造产生的布局给出了完整的刻画。最后，我们提供了我们范畴构造的Python实现，以及展示与CUTLASS行为对齐的测试。此实现可在我们的git仓库https://github.com/ColfaxResearch/layout-categories找到。 ###### 目录 1. 1引言 (https://arxiv.org/html/2601.05972#Ch1)1. 1.1主要结果摘要 (https://arxiv.org/html/2601.05972#Ch1.S1) 2. 1.2论文组织 (https://arxiv.org/html/2601.05972#Ch1.S2) 3. 1.3相关工作 (https://arxiv.org/html/2601.05972#Ch1.S3) 4. 1.4实现 (https://arxiv.org/html/2601.05972#Ch1.S4) 5. 1.5符号说明 (https://arxiv.org/html/2601.05972#Ch1.S5) 2. 2布局及其代数 (https://arxiv.org/html/2601.05972#Ch2)1. 2.1平面布局 (https://arxiv.org/html/2601.05972#Ch2.S1)1. 2.1.1元组 (https://arxiv.org/html/2601.05972#Ch2.S1.SS1) 2. 2.1.2基本定义 (https://arxiv.org/html/2601.05972#Ch2.S1.SS2) 3. 2.1.3基本运算 (https://arxiv.org/html/2601.05972#Ch2.S1.SS3) 4. 2.1.4平面合并 (https://arxiv.org/html/2601.05972#Ch2.S1.SS4) 5. 2.1.5紧凑平面布局 (https://arxiv.org/html/2601.05972#Ch2.S1.SS5) 6. 2.1.6补元 (https://arxiv.org/html/2601.05972#Ch2.S1.SS6) 7. 2.1.7进一步运算 (https://arxiv.org/html/2601.05972#Ch2.S1.SS7) 8. 2.1.8可处理平面布局 (https://arxiv.org/html/2601.05972#Ch2.S1.SS8) 2. 2.2嵌套元组 (https://arxiv.org/html/2601.05972#Ch2.S2)1. 2.2.1配置 (https://arxiv.org/html/2601.05972#Ch2.S2.SS1) 2. 2.2.2基本定义 (https://arxiv.org/html/2601.05972#Ch2.S2.SS2) 3. 2.2.3替换 (https://arxiv.org/html/2601.05972#Ch2.S2.SS3) 4. 2.2.4细化 (https://arxiv.org/html/2601.05972#Ch2.S2.SS4) 3. 2.3布局 (https://arxiv.org/html/2601.05972#Ch2.S3)1. 2.3.1基本定义 (https://arxiv.org/html/2601.05972#Ch2.S3.SS1) 2. 2.3.2基本运算 (https://arxiv.org/html/2601.05972#Ch2.S3.SS2) 3. 2.3.3合并 (https://arxiv.org/html/2601.05972#Ch2.S3.SS3) 4. 2.3.4相对合并 (https://arxiv.org/html/2601.05972#Ch2.S3.SS4) 5. 2.3.5紧凑布局 (https://arxiv.org/html/2601.05972#Ch2.S3.SS5) 6. 2.3.6补元 (https://arxiv.org/html/2601.05972#Ch2.S3.SS6) 7. 2.3.7复合 (https://arxiv.org/html/2601.05972#Ch2.S3.SS7) 8. 2.3.8逻辑除法 (https://arxiv.org/html/2601.05972#Ch2.S3.SS8) 9. 2.3.9逻辑积 (https://arxiv.org/html/2601.05972#Ch2.S3.SS9) 10. 2.3.10可处理布局 (https://arxiv.org/html/2601.05972#Ch2.S3.SS10) 3. 3布局的范畴 (https://arxiv.org/html/2601.05972#Ch3)1. 3.1范畴Tuple{\boldsymbol{\mathsf{Tuple}}} (https://arxiv.org/html/2601.05972#Ch3.S1)1. 3.1.1基本定义 (https://arxiv.org/html/2601.05972#Ch3.S1.SS1) 2. 3.1.2从元组态射到平面布局 (https://arxiv.org/html/2601.05972#Ch3.S1.SS2) 3. 3.1.3示例 (https://arxiv.org/html/2601.05972#Ch3.S1.SS3) 4. 3.1.4元组态射的实现 (https://arxiv.org/html/2601.05972#Ch3.S1.SS4) 5. 3.1.5元组态射上的运算 (https://arxiv.org/html/2601.05972#Ch3.S1.SS5) 2. 3.2范畴Nest{\boldsymbol{\mathsf{Nest}}} (https://arxiv.org/html/2601.05972#Ch3.S2)1. 3.2.1基本定义 (https://arxiv.org/html/2601.05972#Ch3.S2.SS1) 2. 3.2.2从嵌套元组态射到布局 (https://arxiv.org/html/2601.05972#Ch3.S2.SS2) 3. 3.2.3示例 (https://arxiv.org/html/2601.05972#Ch3.S2.SS3) 4. 3.2.4嵌套元组态射的实现 (https://arxiv.org/html/2601.05972#Ch3.S2.SS4) 5. 3.2.5细化 (https://arxiv.org/html/2601.05972#Ch3.S2.SS5) 6. 3.2.6嵌套元组态射上的运算 (https://arxiv.org/html/2601.05972#Ch3.S2.SS6) 4. 4计算 (https://arxiv.org/html/2601.05972#Ch4)1. 4.1可处理布局的复合 (https://arxiv.org/html/2601.05972#Ch4.S1)1. 4.1.1互细化 (https://arxiv.org/html/2601.05972#Ch4.S1.SS1) 2. 4.1.2从互细化到可复合态射 (https://arxiv.org/html/2601.05972#Ch4.S1.SS2) 3. 4.1.3复合算法 (https://arxiv.org/html/2601.05972#Ch4.S1.SS3) 4. 4.1.4示例 (https://arxiv.org/html/2601.05972#Ch4.S1.SS4) 5. 4.1.5更一般的复合 (https://arxiv.org/html/2601.05972#Ch4.S1.SS5) 6. 4.1.6复合的容许性 (https://arxiv.org/html/2601.05972#Ch4.S1.SS6) 2. 4.2逻辑除法和逻辑积 (https://arxiv.org/html/2601.05972#Ch4.S2)1. 4.2.1逻辑除法示例 (https://arxiv.org/html/2601.05972#Ch4.S2.SS1) 2. 4.2.2逻辑积示例 (https://arxiv.org/html/2601.05972#Ch4.S2.SS2) 5. A范畴论入门 (https://arxiv.org/html/2601.05972#A1)1. A.1什么是范畴？ (https://arxiv.org/html/2601.05972#A1.S1) 2. A.2什么是函子？ (https://arxiv.org/html/2601.05972#A1.S2) 6. 参考文献 (https://arxiv.org/html/2601.05972#bib) ## 第一章引言 在现代计算中，特别是在GPU编程中，性能关键取决于多维数据在内存中的存储和访问方式。尽管我们关心的大多数数据——如图像、视频和机器学习中的张量——本质上是多维的，但计算机的内存本质上是一维的。这意味着当我们想要加载、存储或以其他方式操作数据时，需要将其多维逻辑坐标映射到一维物理坐标。这种映射被称为布局，对于正确高效地读写内存至关重要。此外，针对GPU的SIMT执行模型，布局用于描述和操作线程在数据上的分区。这对于确保优化的内存访问模式和正确调用专用硬件指令（例如用于张量核心的指令）非常重要。 作为一个激励性示例，假设我们想在内存中存储4×8矩阵 A=[12.4787.2134.0856.9345.659.1773.0221.3964.8830.411.7288.0492.5517.0650.9168.773.3377.1961.5829.4615.8280.7544.6239.2891.4026.126.9753.0358.6633.7911.2070.55]A=\begin{bmatrix}12.47&87.21&34.08&56.93&45.65&9.17&73.02&21.39\\ 64.88&30.41&1.72&88.04&92.55&17.06&50.91&68.77\\ 3.33&77.19&61.58&29.46&15.82&80.75&44.62&39.28\\ 91.40&26.12&6.97&53.03&58.66&33.79&11.20&70.55\end{bmatrix}。为此，我们需要为A的每个条目指定一个内存地址。我们通过选择A的第(0,0)个条目的某个地址，并为A的每个其他条目指定一个偏移量来实现。一种常见的选择是行主序布局 012345678910111213141516171819202122232425262728293031Lrow=(4,8):(8,1)=L^{\mathsf{row}}=(4,8):(8,1)=，其中的符号Lrow=(4,8):(8,1)L^{\mathsf{row}}=(4,8):(8,1)表示矩阵的第(i,j)个条目的偏移量是 (i,j)⋅(8,1)=8i+j。另一种常见选择是列主序布局 048121620242815913172125292610141822263037111519232731Lcol=(4,8):(1,4)=L^{\mathsf{col}}=(4,8):(1,4)=。同样，符号Lcol=(4,8):(1,4)L^\{\mathsf\{col\}\}=\(4,8\):\(1,4\)表示矩阵的第(i,j)个条目的偏移量是 (i,j)⋅(1,4)=i+4j。 这些布局非常有用，但不足以满足所有目的。例如，在高性能计算中，通常通过以下步骤计算矩阵乘积AB：1. 将操作数矩阵A和B划分为块（tile），2. 计算各种块的矩阵乘积，3. 组合这些部分结果以获得完整结果AB。例如，我们可以将4×8矩阵A划分为2×2的块，如下所示。 A=[[12.4787.2164.8830.41][34.0856.931.7288.04][45.659.1792.5517.06][73.0221.3950.9168.77][3.3377.1991.4026.12][61.5829.466.9753.03][15.8280.7558.6633.79][44.6239.2811.2070.55]]A=\left[\begin{array}[]{cccc}\begin{bmatrix}12.47&87.21\\ 64.88&30.41\end{bmatrix}&\begin{bmatrix}34.08&56.93\\ 1.72&88.04\end{bmatrix}&\begin{bmatrix}45.65&9.17\\ 92.55&17.06\end{bmatrix}&\begin{bmatrix}73.02&21.39\\ 50.91&68.77\end{bmatrix}\\[8.61108pt] \begin{bmatrix}3.33&77.19\\ 91.40&26.12\end{bmatrix}&\begin{bmatrix}61.58&29.46\\ 6.97&53.03\end{bmatrix}&\begin{bmatrix}15.82&80.75\\ 58.66&33.79\end{bmatrix}&\begin{bmatrix}44.62&39.28\\ 11.20&70.55\end{bmatrix}\end{array}\right]。 假设现在我们想切出A的各个块，并假设它在内存中以列主序格式排列。为此，可以手动计算偏移量：对于第(i,j)个块，索引到该块左上角条目的偏移量为2i+8j。另一方面，为了更好地组织此计算，我们可以使用块的交错布局 028101618242613911171925274612142022283057131521232931Ltiled=L^{\mathsf{tiled}}=，其中列由A的块给出，行由块形状内的坐标给出。这里，我们使用反字典序线性枚举块和块内的坐标，因此布局LtiledL^\{\mathsf\{tiled\}\}的顶层形状为(4,8)。然而，请注意LtiledL^\{\mathsf\{tiled\}\}所示的交错模式意味着它不能被表示为任何步幅a,b的布局(4,8):(a,b)。相反，我们可以将形状(4,8)的模进行分解，并定义 Ltiled=((2,2),(2,4)):((1,4),(2,8))。L^{\mathsf{tiled}}=((2,2),(2,4)):((1,4),(2,8)))。先前的偏移量计算2i+8j随后通过将LtiledL^\{\mathsf\{tiled\}\}应用于坐标(0,(i,j))而出现，而块布局本身由第一个模给出。因此，在将布局LtiledL^\{\mathsf\{tiled\}\}赋予A以形成AtiledA^\{\mathsf\{tiled\}\}后，我们可以获得A的第(i,j)个块为切片 Ai,j=Atiled(,(i,j))。A_{i,j}=A^{\mathsf{tiled}}(;\rule{6.99997pt}{0.4pt};,(i,j))。 CUTLASS中发展的一个关键思想是，有用但更复杂的辅助布局（如LtiledL^\{\mathsf\{tiled\}\}）可以通过某些基本运算从更简单的布局系统地推导出来。对于LtiledL^\{\mathsf\{tiled\}\}，所讨论的运算称为逻辑除法。如果我们用 0145T=(2,2):(1,4)=T=(2,2):(1,4)= 表示块布局，那么LtiledL^\{\mathsf\{tiled\}\}是逻辑除法 Ltiled=Lcol⊘TL^{\mathsf{tiled}}=L^{\mathsf{col}}\oslash T，如下所示。 048121620242815913172125292610141822263037111519232731Lcol=L^{\mathsf{col}}=T=T=0145 028101618242613911171925274612142022283057131521232931Lcol⊘T=L^{\mathsf{col}}\oslash T= 除了逻辑除法，其他基本布局运算包括逻辑积、补元，以及最重要的复合。这些布局运算是CUTLASS的基石，深入理解它们的行为有助于编写正确且高性能的代码。然而，这些运算的定义和构造相当微妙。例如，布局A和B的复合B∘AB\circ A仅在A和B满足某些整除约束时才有定义，CUTLASS会在后台检查这些约束。特别是，两个布局何时可复合，或者如何解释它们的复合，并不总是显而易见的。 ### 1.1主要结果摘要 本文的主要思想是，通过将注意力限制在可处理布局上，我们可以为处理布局开发一个直观且强大的数学框架，其条目满足一个简单的整除条件（参见定义2.3.10.1 (https://arxiv.org/html/2601.05972#Ch2.S3.SS10.Thmtheorem1)）。可处理布局包括实践中遇到的几乎所有布局，例如 - •行主序和列主序布局，它们无处不在， - •紧凑布局，将数据存储在连续的内存地址中， - •投影，广播数据的多个副本，和 - •膨胀，启用填充加载和存储。 如果L是一个可处理布局，我们可以用一个图来表示L。例如，布局LrowL^\{\mathsf\{row\}\}、LcolL^\{\mathsf\{col\}\}和LtiledL^\{\mathsf\{tiled\}\}由以下图表示。 8{\lx@inpgf@ignorespaces 8}4{\lx@inpgf@ignorespaces 4}(4,8):(8,1){\lx@inpgf@ignorespaces(4,8):(8,1)}↭{\lx@inpgf@ignorespaces\leftrightsquigarrow}4{\lx@inpgf@ignorespaces 4}8{\lx@inpgf@ignorespaces 8}8{\lx@inpgf@ignorespaces 8}8{\lx@inpgf@ignorespaces 8}(4,8):(1,4){\lx@inpgf@ignorespaces(4,8):(1,4)}↭{\lx@inpgf@ignorespaces\leftrightsquigarrow}4{\lx@inpgf@ignorespaces 4}4{\lx@inpgf@ignorespaces 4}4{\lx@inpgf@ignorespaces 4}4{\lx@inpgf@ignorespaces 4}8{\lx@inpgf@ignorespaces 8}2{\lx@inpgf@ignorespaces 2}2{\lx@inpgf@ignorespaces 2}((2,2),(2,4)):((1,4),(2,8)){\lx@inpgf@ignorespaces((2,2),(2,4)):((1,4),(2,8))}↭{\lx@inpgf@ignorespaces\leftrightsquigarrow}2{\lx@inpgf@ignorespaces 2}2{\lx@inpgf@ignorespaces 2}4{\lx@inpgf@ignorespaces 4}2{\lx@inpgf@ignorespaces 2}2{\lx@inpgf@ignorespaces 2}。 这些图可以被解释为范畴中的态射。这使我们能够利用范畴论的力量来描述布局及其运算。111对于不熟悉该主题的读者，我们在附录A (https://arxiv.org/html/2601.05972#A1)中提供了范畴论入门。更准确地说，我们定义一个范畴Nest{\boldsymbol{\mathsf{Nest}}}，其对象是正整数的嵌套元组，其态射f:S→Tf:S\to T对应于上述图（详见定义3.1.1.13 (https://arxiv.org/html/2601.05972#Ch3.S1.SS1.Thmtheorem13)和定义3.2.1.1 (https://arxiv.org/html/2601.05972#Ch3.S2.SS1.Thmtheorem1)）。如果L是一个非退化可处理布局（参见定义2.3.1.24 (https://arxiv.org/html/2601.05972#Ch2.S3.SS1.Thmtheorem24)），则存在一个本质上唯一的Nest{\boldsymbol{\mathsf{Nest}}}-态射f编码L，如下对应定理所示。 ###### 定理A.（参见3.2.2.15 (https://arxiv.org/html/2601.05972#Ch3.S2.SS2.Thmtheorem15)） 存在一一对应关系 {非退化可处理布局}{\lx@inpgf@ignorespaces\begin{Bmatrix}\text{非退化}\\ \text{可处理布局}\end{Bmatrix}}{标准形式的非退化Nest-态射}{\lx@inpgf@ignorespaces\begin{Bmatrix}\text{非退化}\\ {\boldsymbol{\mathsf{Nest}}}\text{-态射}\\ \text{标准形式}\end{Bmatrix}} 布局运算，如复合、逻辑除法