CuTe Layout 代数详解Coalesce、Composition、Complement 与 Product/Divide 运算【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlassCuTe 是 NVIDIA CUTLASS 中用于高性能线性代数的高层张量抽象库其核心思想是把Layout视为整数到整数的函数。本文基于 media/docs/cpp/cute/02_layout_algebra.md 系统讲解 CuTe 的Layout 代数——包括coalesce化简、composition函数复合、complement补集、logical_divide除法分块与logical_product乘法复制等核心运算。读完本文你将掌握如何用这些运算从简单 Layout 组合出复杂的数据/线程布局为理解后续的Tensor、MMA 原子操作与 threadblock 级 tiling 打下基础。所有运算的源码实现集中在 include/cute/layout.hpp配套单元测试位于 test/unit/cute/core/。总览Layout 代数的三大支柱CuTe 提供了一整套Layout 的代数支持以不同方式组合 LayoutLayout 函数复合functional composition把两个 Layout 当作函数复合起来Layout 乘积product按照一个 Layout 去复制另一个 LayoutLayout 除法divide按照一个 Layout 去切分另一个 Layout。用简单 Layout 构建复杂 Layout 的常用工具依赖product把例如数据的Layout 切分到例如线程的其他 Layout 上的分区工具依赖divide而所有这些工具又都建立在 Layout 的函数复合之上。在开始之前请先牢记上一节的核心结论Layout 是整数到整数的函数。接下来我们从最简单的coalesce开始逐步构建这套代数工具。Coalesce对整数函数的化简coalesce是对整数到整数函数的化简simplify操作。如果我们只关心输入整数那么可以在不改变 Layout 作为函数的前提下操纵 Layout 的 shape 与 mode 数量。coalesce唯一不能改变的是 Layout 的size。其单元测试位于 test/unit/cute/core/coalesce.cpp测试中直接断言了三条后置条件// post size(a result) size(a layout) // post depth(a result) 1 // post for all i, 0 i size(a layout), a result(i) a layout(i) Layout coalesce(Layout const layout)例如auto layout LayoutShape _2,Shape _1,_6, Stride_1,Stride_6,_2{}; auto result coalesce(layout); // _12:_1可以看到结果 mode 更少、更简单。当坐标映射与索引映射在运行时动态执行时这种化简能省去若干次运算。化简的直觉列主序column-majorLayout 如(_2,_4):(_1,_2)在一维坐标下与_8:_1行为完全一致size 为静态 1 的 mode其自然坐标恒为静态 0无论 stride 是什么都可以直接忽略。四种基本情况推广到只有两个整数 modes0:d0与s1:d1的情形记合并结果为s0:d0 s1:d1则存在四种情况s0:d0 _1:d1 s0:d0忽略 size 为静态 1 的 mode_1:d0 s1:d1 s1:d1忽略 size 为静态 1 的 modes0:d0 s1:s0*d0 s0*s1:d0若第二个 mode 的 stride 等于第一个 mode 的 size 与 stride 之积则二者可合并s0:d0 s1:d1 (s0,s1):(d0,d1)否则什么都做不了只能分开处理。仅此而已我们可以展平任意 Layout再对相邻 mode 依次应用上述二元运算从而coalesce整个 Layout 的所有 mode。这正是在 include/cute/layout.hpp 中coalesce的实现思路其内部递归地对 shape/stride 的整数元组做合并。按 mode 的 CoalesceBy-mode Coalesce有时我们确实关心 Layout 的形状但仍然想化简。例如手上是一个二维 Layout希望结果保持二维。为此coalesce有一个带额外参数的重载// Apply coalesce at the terminals of trg_profile Layout coalesce(Layout const layout, IntTuple const trg_profile)用法如下auto a LayoutShape _2,Shape _1,_6, Stride_1,Stride_6,_2{}; auto result coalesce(a, Step_1,_1{}); // (_2,_6):(_1,_2) // 等价于 auto same_r make_layout(coalesce(layout0(a)), coalesce(layout1(a)));该函数会递归进入Step_1,_1{}每当遇到整数值本身无关紧要只是标志位就对对应的子 Layout 应用coalesce遇到元组则继续递归。这种先把 Layout 当作一维整数函数定义操作再推广到任意形状的范式在 CuTe 中会反复出现请务必记住。CompositionLayout 的函数复合Layout 的函数复合是 CuTe 的核心几乎所有高层操作都会用到。还是从Layout 是整数到整数的函数出发Layout复合的结果仍然是另一个Layout。先看一个完整例子函数复合R : A o B R(c) : (A o B)(c) : A(B(c)) 例 A (6,2):(8,2) B (4,3):(3,1) R( 0) A(B( 0)) A(B(0,0)) A( 0) A(0,0) 0 R( 1) A(B( 1)) A(B(1,0)) A( 3) A(3,0) 24 R( 2) A(B( 2)) A(B(2,0)) A( 6) A(0,1) 2 R( 3) A(B( 3)) A(B(3,0)) A( 9) A(3,1) 26 R( 4) A(B( 4)) A(B(0,1)) A( 1) A(1,0) 8 R( 5) A(B( 5)) A(B(1,1)) A( 4) A(4,0) 32 R( 6) A(B( 6)) A(B(2,1)) A( 7) A(1,1) 10 R( 7) A(B( 7)) A(B(3,1)) A(10) A(4,1) 34 R( 8) A(B( 8)) A(B(0,2)) A( 2) A(2,0) 16 R( 9) A(B( 9)) A(B(1,2)) A( 5) A(5,0) 40 R(10) A(B(10)) A(B(2,2)) A( 8) A(2,1) 18 R(11) A(B(11)) A(B(3,2)) A(11) A(5,1) 42最令人惊叹的观察是上面定义的函数R(c) k可以写成另一个 LayoutR ((2,2),3):((24,2),8)并且满足compatible(B, R)即B的每个坐标都可以当作R的坐标使用。这正是函数复合的应有之义因为B定义了R的定义域。test/unit/cute/core/composition.cpp 中包含大量示例与经过检查的后置条件与我们上面所述完全一致// post compatible(a layout_b, a result) // post for all i, 0 i size(a layout_b), a result(i) a layout_a(a layout_b(i))) Layout composition(LayoutA const layout_a, LayoutB const layout_b)如何计算 Composition首先给出两条观察B (B_0, B_1, ...)一个 Layout 可以表达为其子 Layout 的拼接concatenationA o B A o (B_0, B_1, ...) (A o B_0, A o B_1, ...)当B是单射injective时复合对拼接是左分配的。据此不失一般性地假设B s:d是形状和步长为整数的 Layout并假设A是已展平、已 coalesce 的 Layout。当A是整数的A a:b结果很平凡R A o B a:b o s:d s:(b*d)。复合结果R就是A的以d为步长的前s个元素。当A是多 mode 的A o B A o s:ds、d为整数意味着我们想做两件事找出一个能产生A的每第d个元素的 Layout。该中间 Layout 的 shape 可以通过从左侧开始把d从A的 shape 中逐步除出去来计算。例如(6,2) / 2 (3,2)(6,2) / 3 (2,2)(6,2) / 6 (1,2)(6,2) / 12 (1,1)(3,6,2,8) / 3 (1,6,2,8)(3,6,2,8) / 6 (1,3,2,8)(3,6,2,8) / 9 (1,2,2,8)(3,6,2,8) / 72 (1,1,1,4)带 stride 的 Layout 的 stride 则由上述操作的余数去缩放A的 stride。例如最后一个例子(3,6,2,8):(w,x,y,z) / 72其步长 Layout 的 stride 为(72*w,24*x,4*x,2*z)。你可能已经注意到只能对某些特定的值做 shape 除法才能得到有意义的结果这被称为stride 整除条件stride divisibility conditionCuTe 会在编译期尽可能静态检查。保留新 stridedA的前s个元素使结果与B形状兼容。这可以通过从左侧开始把前s个元素从A的 shape 中模除出去计算。例如(6,2) % 2 (2,1)(6,2) % 3 (3,1)(6,2) % 6 (6,1)(6,2) % 12 (6,2)(3,6,2,8) % 6 (3,2,1,1)(3,6,2,8) % 9 (3,3,1,1)(1,2,2,8) % 2 (1,2,1,1)(1,2,2,8) % 16 (1,2,2,4)该操作使结果 shape 与B兼容同样需要满足shape 整除条件CuTe 在可能时也会静态检查。由上述示例可以构造出复合(3,6,2,8):(w,x,y,z) o 16:9 (1,2,2,4):(9*w,3*x,y,z)。示例 1 —— 计算复合的完整工作示例下面给出一个更复杂的复合例子两个操作数 Layout 都是多 mode 的用于演示上述概念函数复合R : A o B R(c) : (A o B)(c) : A(B(c)) 例 A (6,2):(8,2) B (4,3):(3,1) 1. 利用左分配律与拼接性质把复合写成 R A o B (6,2):(8,2) o (4,3):(3,1) ((6,2):(8,2) o 4:3, (6,2):(8,2) o 3:1)计算(6,2):(8,2) o 4:3先计算 strided layout(6,2):(8,2) / 3 (6/3,2):(8*3,2) (2,2):(24,2)再保持 shape 兼容(2,2):(24,2) % 4 (2,2):(24,2)计算(6,2):(8,2) o 3:1先计算 strided layout(6,2):(8,2) / 1 (6,2):(8,2)再保持 shape 兼容(6,2):(8,2) % 3 (3,1):(8,2)把两部分拼接并对每个 mode 做 coalesce得到最终结果R A o B ((2, 2), 3): ((24, 2), 8)示例 2 —— 把 Layout 重塑为矩阵20:2 o (5,4):(4,1)。用复合表述把20:2解释为一个按行主序排列的 5x4 矩阵。 20:2 o (5:4,4:1)把(5,4):(4,1)写成子 Layout 的拼接 (20:2 o 5:4, 20:2 o 4:1)左分配律。20:2 o 5:4 5:8平凡情形20:2 o 4:1 4:2平凡情形 (5:8, 4:2)复合结果表示为子 Layout 的拼接 (5,4):(8,2)最终复合 Layout。示例 3 —— 把 Layout 重塑为矩阵(10,2):(16,4) o (5,4):(1,5)把(10,2):(16,4)解释为一个按列主序排列的 5x4 矩阵。 (10,2):(16,4) o (5:1,4:5)把(5,4):(1,5)写成子 Layout 的拼接 ((10,2):(16,4) o 5:1, (10,2):(16,4) o 4:5)左分配律。(10,2):(16,4) o 5:1 (5,1):(16,4)对 shape 5 做模除(10,2):(16,4) o 4:5 (2,2):(80,4)对 stride 5 做整除 ((5,1):(16,4), (2,2):(80,4))复合结果表示为子 Layout 的拼接 (5:16, (2,2):(80,4))按 mode coalesce (5,(2,2))):(16,(80,4))最终复合 Layout。使用编译期 shape 与 stride 时CuTe 会精确得到这个结果。下面这段 C 代码打印(_5,(_2,_2)):(_16,(_80,_4))Layout a make_layout(make_shape (Int10{}, Int2{}), make_stride(Int16{}, Int4{})); Layout b make_layout(make_shape (Int 5{}, Int4{}), make_stride(Int 1{}, Int5{})); Layout c composition(a, b); print(c);如果改用动态整数下面的代码打印((5,1),(2,2)):((16,4),(80,4))Layout a make_layout(make_shape (10, 2), make_stride(16, 4)); Layout b make_layout(make_shape ( 5, 4), make_stride( 1, 5)); Layout c composition(a, b); print(c);两种结果看起来不同但数学上等价shape 中的 1 不影响 Layout 作为一维坐标到整数或二维坐标到整数的函数。在动态情形下由于包含这些动态 size-1 mode 的元组的静态 rank 和类型固定CuTe 无法通过 coalesce 把它们化简掉。按 mode 的 CompositionBy-mode Composition与按 mode 的coalesce类似并且是通往通用 tiling 操作的一步有时我们关心A的形状希望只对个别 mode 应用composition。例如一个二维 Layout想沿列取某个子 Layout沿行取另一个子 Layout。为此composition的第二个参数B也支持Tiler。一般地tiler 是一个 Layout 或 Layout 的元组注意这是对IntTuple的推广用法如下// (12,(4,8)):(59,(13,1)) auto a make_layout(make_shape (12,make_shape ( 4,8)), make_stride(59,make_stride(13,1))); // 3:4, 8:2 auto tiler make_tile(Layout_3,_4{}, // 对 mode-0 应用 3:4 Layout_8,_2{}); // 对 mode-1 应用 8:2 // (_3,(2,4)):(236,(26,1)) auto result composition(a, tiler); // 等价于 auto same_r make_layout(composition(layout0(a), get0(tiler)), composition(layout1(a), get1(tiler)));我们常用LayoutA, LayoutB, ...这种尖括号记法来区分Tiler与之前用过的拼接记法(LayoutA, LayoutB, ...)。上面的result可以形象地理解为原 Layout 中高亮出的 3x8 子 Layout为了方便CuTe 也把Shape解释为一种 tiler——即stride 全为 1 的 Layout 元组// (12,(4,8)):(59,(13,1)) auto a make_layout(make_shape (12,make_shape ( 4,8)), make_stride(59,make_stride(13,1))); // (3, 8) auto tiler make_shape(Int3{}, Int8{}); // 等价于 3:1, 8:1 // auto tiler make_tile(Layout_3,_1{}, // 对 mode-0 应用 3:1 // Layout_8,_1{}); // 对 mode-1 应用 8:1 // (_3,(4,2)):(59,(13,1)) auto result composition(a, tiler);这里的result同样可以理解为原 Layout 中高亮出的 3x8 子 LayoutComposition Tilers总结归纳起来Tiler是以下三种对象之一一个Layout一个Tiler的元组一个Shape将被解释为 stride-1 的Layout构成的 tiler。以上任何对象都可以作为composition的第二个参数。对于 (1)无论两个 Layout 的 rank 如何我们都把composition视为两个整数到整数函数之间的复合对于 (2) 和 (3)composition会按 mode 在A与B的对应 mode 上执行直到遇到情形 (1)。这种设计让composition既能按 mode取出张量指定 mode 的任意子 Layout给我这个 MxNxL 张量的 3x5x8 子块也能把整块数据当作一维向量进行重塑与重排用这个奇怪的元素顺序把 8x16 数据块重排成 32x4 块。按 mode 的情形会在后续 threadblock tiling 示例中频繁出现一维重塑与重排则用于 MMA 中对线程与数值应用任意分区模式。Complement找出剩下的布局在进入 product 与 divide 之前还差最后一个操作。可以把composition理解为 LayoutB从另一个 LayoutA中挑选某些坐标。但那些没被挑选的坐标怎么办为了实现通用 tiling我们希望既能选出任意元素tile也能描述这些 tile 的布局——即剩下的部分或者说其余。complement正是尝试找出表示其余即未被该 Layout 触及的元素的另一个 Layout。test/unit/cute/core/complement.cpp 中包含大量示例与检查过的后置条件// post cosize(make_layout(a layout_a, a result))) size(a cotarget) // post cosize(a result) round_up(size(a cotarget), cosize(a layout_a)) // post for all i, 1 i size(a result), // a result(i-1) a result(i) // post for all i, 1 i size(a result), // for all j, 0 j size(a layout_a), // a result(i) ! a layout_a(j) Layout complement(LayoutA const layout_a, Shape const cotarget)即LayoutA关于 ShapeIntTupleM的补集R满足如下性质有界R的 size以及 cosize以size(M)为上界有序R的 stride 为正且递增因此R是唯一的定义域不相交A与R的 codomain 不相交R试图补全A的 codomain。上面的cotarget参数最常见的是整数——可以看到我们只用到了size(cotarget)。不过有时指定一个带静态属性的整数很有用。例如28是动态整数而(_4,7)是 size 为 28 且静态已知可被_4整除的 shape。两者数学上会产生相同的complement但额外信息能让complement尽可能保留结果的静态性。Complement 示例complement在静态 shape 与 stride 上最有效因此下面把整数都视为静态。动态 shape/stride 以及 IntTuplecotarget的类似示例见 test/unit/cute/core/complement.cpp。complement(4:1, 24)得到6:4。注意(4,6):(1,4)的 cosize 为 24。Layout4:1实际上被6:4重复了 6 次。complement(6:4, 24)得到4:1。注意(6,4):(4,1)的 cosize 为 24。6:4中的洞被4:1填满。complement((4,6):(1,4), 24)得到1:0。无需追加任何东西。complement(4:2, 24)得到(2,3):(1,8)。注意(4,(2,3)):(2,(1,8))的 cosize 为 24。4:2中的洞先被2:1填满然后整体被3:8重复 3 次。complement((2,4):(1,6), 24)得到3:2。注意((2,4),3):((1,6),2)的 cosize 为 24 且产生唯一索引。complement((2,2):(1,6), 24)得到(3,2):(2,12)。注意((2,2),(3,2)):((1,6),(2,12))的 cosize 为 24 且产生唯一索引。作为可视化上图描绘了最后一个例子的 codomain原 Layout(2,2):(1,6)的像染成灰色complement有效地重复原布局以其他颜色显示使结果 codomain 的大小为 24。补集(3,2):(2,12)可以看作重复的布局。DivisionTiling用另一个 Layout 除一个 Layout最后我们可以定义 Layout 之间的除法。把 Layout 划分为多个组成部分的函数是 tiling 与分区的基础。本节先定义logical_divide(Layout, Layout)——它同样把所有 Layout 视为整数到整数的一维函数——再用该定义构造多维 Layout 除法。非正式地说logical_divide(A, B)把 LayoutA拆成两个 mode第一个 mode 里是B指向的所有元素第二个 mode 里是B没有指向的所有元素。形式化地可写作$$A \oslash B : A \circ (B,B^*)$$其实现为见 include/cute/layout.hpptemplate class LShape, class LStride, class TShape, class TStride auto logical_divide(LayoutLShape,LStride const layout, LayoutTShape,TStride const tiler) { return composition(layout, make_layout(tiler, complement(tiler, size(layout)))); }注意它仅由拼接concatenation、复合composition与补集complement定义。那这到底是什么第一个 mode 里是B指向的所有元素这显然就是复合A o B。第二个 mode 里是B没有指向的所有元素B没有指向的元素听起来就像补集B*以A的 size 为界。正如上面 complement 一节所见补集可以描述为B的重复布局。如果B是tiler那么B*就是 tile 的布局。一维 Logical Divide 示例考虑用 tilerB 4:2对一维 LayoutA (4,2,3):(2,1,8)做 tiling。非正式地说我们有一个由A定义的、24 个元素的存储顺序想要提取出以步长 2 排列的 4 元素 tile。按上述实现的三步计算B 4:2在size(A) 24下的补集为B* (2,3):(1,8)拼接(B,B*) (4,(2,3)):(2,(1,8))A (4,2,3):(2,1,8)与(B,B*)的复合为((2,2),(2,3)):((4,1),(2,8))。上图把A描绘成一维 Layout其中B指向的元素用灰色高亮。LayoutB描述我们的数据tileA中共有六个这样的 tile每种颜色一个。除法之后结果的第一个 mode 是数据 tile第二个 mode 遍历每个 tile。二维 Logical Divide 示例利用上面定义的Tiler概念这可以立即推广到多维 tiling。下面的例子就是用Tiler对一个二维 Layout 的行和列按 mode 应用logical_divide。与上面的二维 composite 示例类似考虑二维 LayoutA (9,(4,8)):(59,(13,1))希望在列mode-0上应用3:3在行mode-1上应用(2,4):(1,8)。于是 tiler 可以写成B 3:3, (2,4):(1,8)。上图把A描绘成二维 LayoutB指向的元素用灰色高亮。LayoutB描述我们的数据tileA中共有十二个这样的 tile每种颜色一个。除法后结果每个 mode 的第一个 mode 是数据 tile第二个 mode 遍历每个 tile。在这个意义上该操作可以看作一种gather操作或对行和列的排列。注意结果每个 mode 的第一个 mode 是子 Layout(3,(2,4)):(177,(13,2))这正是如果改用composition而非logical_divide会得到的结果。Zipped、Tiled、Flat 除法在图中高亮 tile 很容易但实际操作起来仍然不便如何切出第3个 tile、第7个 tile 或第(1,2)个 tile以便继续处理为此logical_divide提供了一些便捷变体。假设我们有一个Layout和一个某种形状的Tiler则每个操作都会应用logical_divide但可能把 mode 重排成更方便的形式Layout Shape : (M, N, L, ...) Tiler Shape : TileM, TileN logical_divide : ((TileM,RestM), (TileN,RestN), L, ...) zipped_divide : ((TileM,TileN), (RestM,RestN,L,...)) tiled_divide : ((TileM,TileN), RestM, RestN, L, ...) flat_divide : (TileM, TileN, RestM, RestN, L, ...)例如zipped_divide先应用logical_divide然后把所有子 tile收集到一个 mode、把所有rest收集到另一个 mode// A: shape 为 (9,32) auto layout_a make_layout(make_shape (Int 9{}, make_shape (Int 4{}, Int8{})), make_stride(Int59{}, make_stride(Int13{}, Int1{}))); // B: shape 为 (3,8) auto tiler make_tile(Layout_3,_3{}, // 对 mode-0 应用 3:3 LayoutShape _2,_4, // 对 mode-1 应用 (2,4):(1,8) Stride_1,_8{}); // ((TileM,RestM), (TileN,RestN))shape 为 ((3,3), (8,4)) auto ld logical_divide(layout_a, tiler); // ((TileM,TileN), (RestM,RestN))shape 为 ((3,8), (3,4)) auto zd zipped_divide(layout_a, tiler);于是第3个 tile 的偏移是zd(0,3)第7个 tile 的偏移是zd(0,7)第(1,2)个 tile 的偏移是zd(0,make_coord(1,2))。tile 本身的 Layout 永远是layout0(zd)。事实上恒有layout0(zipped_divide(a, b)) composition(a, b)。需要指出logical_divide在置换 mode 内元素的同时保留了 mode 的语义——A的Mmode 仍是结果的MmodeNmode 仍是结果的Nmode。zipped_divide则不是这样其结果 mode-0 是Tile本身无论Tiler的 rank 是多少mode-1 是这些 tile 的布局。把它们画成二维 Layout 并不总是有意义因为Mmode 更应称为tile-modeNmode 更应称为rest-mode。尽管如此我们仍然可以像下面这样把结果画成二维为了清晰前面各图保留了每个 tile 的颜色。显然跨 tile 迭代现在等价于沿这个布局的一行迭代而 tile 内元素迭代等价于沿一列迭代。正如我们将在Tensor一节看到的这在 tile 内部或跨 tile 的数据分区中可以发挥巨大作用。ProductTiling用另一个 Layout 乘一个 Layout最后我们可以定义 Layout 之间的乘法。本节先定义logical_product(Layout, Layout)——同样把所有 Layout 视为整数到整数的一维函数——再用该定义构造多维 Layout 乘积。非正式地说logical_product(A, B)产生一个二 mode 的 Layout第一个 mode 是 LayoutA第二个 mode 是 LayoutB但B的每个元素都被替换为A的唯一复制。形式化地可写作$$A \otimes B : (A, A^* \circ B)$$CuTe 中的实现为见 include/cute/layout.hpptemplate class LShape, class LStride, class TShape, class TStride auto logical_product(LayoutLShape,LStride const layout, LayoutTShape,TStride const tiler) { return make_layout(layout, composition(complement(layout, size(layout)*cosize(tiler)), tiler)); }注意它同样仅由拼接、复合与补集定义。那这又是什么第一个 mode 是 LayoutA这显然就是A的一份拷贝。第二个 mode 是 LayoutB但每个元素都被替换为A的唯一复制A的唯一复制听起来就像补集A*以B的 cosize 为界。正如 complement 一节所见补集可以描述为A的重复布局。如果A是tile那么A*就是可供B使用的重复布局。一维 Logical Product 示例考虑按照B 6:1复制一维 LayoutA (2,2):(4,1)。非正式地说我们有一个由A定义的 4 元素一维 Layout想把它复制 6 次。按上述实现的三步计算A (2,2):(4,1)在6*4 24下的补集为A* (2,3):(2,8)A* (2,3):(2,8)与B 6:1的复合为(2,3):(2,8)拼接(A,A* o B) ((2,2),(2,3)):((4,1),(2,8))。上图把A和B描绘成一维 Layout。LayoutB描述A复制的数量与顺序为清晰起见用不同颜色标识。乘积之后结果的第一个 mode 是数据 tile第二个 mode 遍历每个 tile。注意该结果与一维 Logical Divide 示例的结果完全一致。当然改变B就可以改变乘积中 tile 的数量与顺序。例如上图中B (4,2):(2,1)共有 8 个重复 tile而非 6 个且顺序不同。二维 Logical Product 示例利用前面发展出的按 modetiler策略也可以写出多维乘积。上图演示了用tiler按 mode 应用logical_product。尽管这不是推荐做法结果仍是一个 rank-2 的 Layout一个 2x5 行主序块被平铺到一个 3x4 列主序的排列中。不推荐的原因在于上式中的tiler B极不直观——要构造它你需要完全知道A的 shape 与 stride。我们希望能以更直观、让A和B相互独立的方式来表达按 LayoutB平铺 LayoutA。Blocked 与 Raked 乘积blocked_product(LayoutA, LayoutB)与raked_product(LayoutA, LayoutB)是在一维logical_product之上的对 rank 敏感的变换让我们能表达最常需要、也更直观的Layout乘积。实现这些函数的关键观察是logical_product的兼容性后置条件// post rank(result) 2 // post compatible(layout_a, layout0(result)) // post compatible(layout_b, layout1(result))因为A总是与结果 mode-0 兼容、B总是与结果 mode-1 兼容所以只要让A和B同 rank就能在乘积之后重新结合相同语义的 modeA的列 mode 可与B的列 mode 结合A的行 mode 可与B的行 mode 结合等等。这正是blocked_product与raked_product所做的也是它们被称为rank-sensitive的原因。与 CuTe 中其他接受Layout参数的函数不同它们关心参数的顶层 rank以便在logical_product之后重新结合每个 mode。上图展示了与tiler方法相同的结果但参数直观得多一个 2x5 行主序 Layout 作为 tile排布在 3x4 列主序排列中。同时注意blocked_product还帮我们把 mode-0 做了coalesce。类似地raked_product以略微不同的方式结合 mode结果列 mode 不再由A的列 mode 后接B的列 mode 构成而是由B的列 mode 后接A的列 mode 构成。这导致tileA不再以块的形式出现而是与tile 的布局B交织或耙在一起。其他文献称之为循环分布cyclic distribution。Zipped 与 Tiled 乘积与zipped_divide、tiled_divide类似zipped_product与tiled_product只是把按 modelogical_product的结果 mode 重新排列Layout Shape : (M, N, L, ...) Tiler Shape : TileM, TileN logical_product : ((M,TileM), (N,TileN), L, ...) zipped_product : ((M,N), (TileM,TileN,L,...)) tiled_product : ((M,N), TileM, TileN, L, ...) flat_product : (M, N, TileM, TileN, L, ...)小结与进阶指引至此CuTe Layout 代数的核心算子已全部介绍完毕coalesce负责化简、composition负责函数复合、complement负责找出剩余而logical_divide/logical_product及其各种变体zipped_divide、tiled_divide、flat_divide、blocked_product、raked_product、zipped_product、tiled_product等则分别负责按 tiler 切分或复制布局。可以验证logical_divide与logical_product都只由拼接、复合与补集三个原语定义这正是 CuTe 代数体系自洽性的体现。想要深入验证本文结论可以直接运行仓库中的单元测试test/unit/cute/core/coalesce.cpp验证coalesce的 size 保持、深度 1 与逐坐标一致test/unit/cute/core/composition.cpp验证composition的兼容性与逐坐标求值test/unit/cute/core/complement.cpp验证complement的有界、有序、不相交三性质test/unit/cute/core/logical_divide.cpp 与 test/unit/cute/core/logical_product.cpp验证 divide/product 及其各种变体。这些运算在 CUTLASS 中无处不在threadblock 级别的数据平铺、MMA 原子操作中的线程与数值分区、Tensor的切片与重排全都建立在这套 Layout 代数之上。掌握了本文的代数工具再阅读后续的Tensor与 MMA 相关章节见 media/docs/cpp/cute/ 目录下的其他文档就会水到渠成。CopyrightCopyright (c) 2017 - 2026 NVIDIA CORPORATION AFFILIATES. All rights reserved. SPDX-License-Identifier: BSD-3-Clause【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考