https://docs.nvidia.com/cutlass/media/docs/cpp/cute/01_layout.html
CuTe 布局 (CuTe Layouts)
本文档描述了 Layout(布局),它是 CuTe 的核心抽象概念。从根本上说,一个 Layout 将一个或多个坐标空间 (coordinate space) 映射到一个索引空间 (index space)。
布局为多维数组访问提供了一个通用接口,它抽象了数组元素在内存中组织方式的细节。这使得用户可以编写通用地访问多维数组的算法,这样即使布局发生变化,用户的代码也无需更改。例如,一个行优先的 MxN 布局和一个列优先的 MxN 布局在软件中可以以相同的方式处理。
CuTe 还提供了一套“布局代数 (algebra of Layouts)”。布局可以被组合和操作,以构建更复杂的布局,并可以将布局在其他布局之上进行平铺 (tiling)。这可以帮助用户完成诸如将数据的布局划分到线程的布局上等任务。
基本类型和概念 (Fundamental Types and Concepts)
整数 (Integers)
CuTe 大量使用了动态整数(仅在运行时已知)和静态整数(在编译时已知)。
-
动态整数(或“运行时整数”)就是普通的整型类型,如
int
、size_t
或uint16_t
。任何能被std::is_integral<T>
接受的类型在 CuTe 中都被视为动态整数。 -
静态整数(或“编译时整数”)是像
std::integral_constant<Value>
这样的类型的实例。这些类型将值编码为静态constexpr
成员。它们也支持转换为其底层动态类型,因此可以在包含动态整数的表达式中使用。CuTe 定义了自己兼容 CUDA 的静态整数类型cute::C<Value>
,以及重载的数学运算符,使得对静态整数进行数学运算的结果仍然是静态整数。CuTe 定义了快捷别名Int<1>
,Int<2>
,Int<3>
和_1
,_2
,_3
以方便使用,你在示例中会经常看到它们。
CuTe 试图以相同的方式处理静态和动态整数。在接下来的示例中,所有动态整数都可以替换为静态整数,反之亦然。当我们在 CuTe 中说“整数”时,我们几乎总是指静态或动态整数。
CuTe 提供了许多特征 (traits) 来处理整数:
-
cute::is_integral<T>
:检查 T 是静态还是动态整数类型。 -
cute::is_std_integral<T>
:检查 T 是否是动态整数类型。等同于std::is_integral<T>
。 -
cute::is_static<T>
:检查 T 是否是一个空类型(因此实例化不能依赖于任何动态信息)。等同于std::is_empty
。 -
cute::is_constant<N,T>
:检查 T 是静态整数并且其值等于 N。
有关更多信息,请参阅 integral_constant
的实现。
元组 (Tuple)
元组是一个包含零个或多个元素的有限有序列表。cute::tuple
类的行为类似于 std::tuple
,但可在设备和主机上工作。它对其模板参数施加了限制,并简化了实现以提高性能和简洁性。
整数元组 (IntTuple)
CuTe 定义了 IntTuple 概念,它要么是一个整数,要么是一个由 IntTuple 组成的元组。注意这个递归定义。在 C++ 中,我们定义了针对 IntTuple 的操作。
IntTuple 的示例包括:
-
int{2}
,动态整数 2。 -
Int<3>{}
,静态整数 3。 -
make_tuple(int{2}, Int<3>{})
,包含动态-2 和静态-3 的元组。 -
make_tuple(uint16_t{42}, make_tuple(Int<1>{}, int32_t{3}), Int<17>{})
,包含动态-42、(静态-1 和动态-3 的元组)和静态-17 的元组。
CuTe 重用了 IntTuple 概念来表示许多不同的东西,包括 Shape(形状)、Stride(步长)、Step(步骤)和 Coord(坐标)(参见 include/cute/layout.hpp
)。
在 IntTuple 上定义的操作包括:
-
rank(IntTuple)
:IntTuple 中的元素数量。单个整数的秩为 1,元组的秩为tuple_size
。 -
get<I>(IntTuple)
:IntTuple 的第 I 个元素,其中 I < rank。对于单个整数,get<0>
就是该整数本身。 -
depth(IntTuple)
:分层 IntTuple 的数量。单个整数的深度为 0,整数元组的深度为 1,包含整数元组的元组的深度为 2,依此类推。 -
size(IntTuple)
:IntTuple 所有元素的乘积。
我们用括号表示 IntTuple 的层次结构。例如,6
, (2)
, (4,3)
, 和 (3,(6,2),8)
都是 IntTuple。
形状和步长 (Shapes and Strides)
Shape 和 Stride 都是 IntTuple 概念。
布局 (Layout)
一个 Layout 是一个由 (Shape, Stride) 组成的元组。从语义上讲,它通过 Stride 实现了从 Shape 内的任何坐标到索引的映射。
张量 (Tensor)
一个 Layout 可以与数据(例如,指针或数组)组合以创建一个 Tensor。Layout 生成的索引用于下标一个迭代器以检索相应的数据。有关 Tensor 的详细信息,请参阅教程的 Tensor 部分。
布局的创建和使用 (Layout Creation and Use)
一个 Layout 是一对 IntTuple:Shape 和 Stride。第一个元素定义了 Layout 的抽象形状,第二个元素定义了步长,它将形状内的坐标映射到索引空间。
我们在 Layout 上定义了许多操作,类似于在 IntTuple 上定义的操作:
-
rank(Layout)
:Layout 中的模式数量 (number of modes)。等同于 Layout 形状的元组大小。 -
get<I>(Layout)
:Layout 的第 I 个子布局,其中 I < rank。 -
depth(Layout)
:Layout 形状的深度。单个整数的深度为 0,整数元组的深度为 1,整数元组的元组的深度为 2,依此类推。 -
shape(Layout)
:Layout 的形状。 -
stride(Layout)
:Layout 的步长。 -
size(Layout)
:Layout 函数定义域的大小。等同于size(shape(Layout))
。 -
cosize(Layout)
:Layout 函数值域的大小(不一定是范围)。等同于stride * (size - 1) + 1
。(译者注:原文公式A(size(A) - 1) + 1
可能不准确,这里根据上下文和概念意译)
分层访问函数 (Hierarchical access functions)
IntTuple 和 Layout 可以任意嵌套。为了方便起见,我们定义了上述一些函数的版本,这些版本接受一个整数序列,而不仅仅是一个整数。这使得可以更轻松地访问嵌套的 IntTuple 或 Layout 内部的元素。例如,我们允许 get<I...>(x)
,其中 I...
是一个“C++ 参数包”,表示零个或多个(整数)模板参数。这些分层访问函数包括:
-
get<I0,I1,...,IN>(x)
:=get<IN>(...(get<I1>(get<I0>(x)))...)
。提取 x 的第 I0 个元素的第 I1 个元素的 ... 的第 IN 个元素。 -
rank<I...>(x)
:=rank(get<I...>(x))
。x 的第 I... 个元素的秩。 -
depth<I...>(x)
:=depth(get<I...>(x))
。x 的第 I... 个元素的深度。 -
shape<I...>(x)
:=shape(get<I...>(x))
。x 的第 I... 个元素的形状。 -
size<I...>(x)
:=size(get<I...>(x))
。x 的第 I... 个元素的大小。
在下面的示例中,你会看到使用 size<0>
和 size<1>
来确定布局或张量的第 0 和第 1 模式的循环边界。
构建一个布局 (Constructing a Layout)
可以通过多种方式构建一个 Layout。它可以包含编译时(静态)整数或运行时(动态)整数的任意组合。
// 1D 布局,静态大小 8,步长 1
Layout s8 = make_layout(Int<8>{});
// 1D 布局,动态大小 8,步长 1
Layout d8 = make_layout(8);
// 2D 布局,形状 (静态2, 静态4),默认生成列优先步长 (1, 2)
Layout s2xs4 = make_layout(make_shape(Int<2>{},Int<4>{}));
// 2D 布局,形状 (静态2, 动态4),默认生成列优先步长 (1, 2)
Layout s2xd4 = make_layout(make_shape(Int<2>{},4));
// 2D 布局,形状 (静态2, 动态4),显式指定步长 (12, 1)
Layout s2xd4_a = make_layout(make_shape (Int< 2>{},4),
make_stride(Int<12>{},Int<1>{}));
// 2D 布局,形状 (静态2, 动态4),使用 LayoutLeft 标签生成列优先步长 (1, 2)
Layout s2xd4_col = make_layout(make_shape(Int<2>{},4),
LayoutLeft{});
// 2D 布局,形状 (静态2, 动态4),使用 LayoutRight 标签生成行优先步长 (4, 1)
Layout s2xd4_row = make_layout(make_shape(Int<2>{},4),
LayoutRight{});
// 分层 2D 布局,形状 (2, (2,2)),显式指定分层步长 (4, (2,1))
Layout s2xh4 = make_layout(make_shape (2,make_shape (2,2)),
make_stride(4,make_stride(2,1)));
// 分层 2D 布局,形状同上,使用 LayoutLeft 标签生成步长 (1, (2,4))
Layout s2xh4_col = make_layout(shape(s2xh4),
LayoutLeft{});
make_layout
函数返回一个 Layout。它推导函数参数的类型并返回一个具有适当模板参数的 Layout。类似地,make_shape
和 make_stride
函数分别返回一个 Shape 和 Stride。由于构造函数模板参数推导 (CTAD) 的限制以及为了避免重复静态或动态整数类型,CuTe 经常使用这些 make_*
函数。
当省略 Stride 参数时,它会根据提供的 Shape 默认使用 LayoutLeft
生成。LayoutLeft
标签从左到右将 Shape 进行排他性前缀乘积 (exclusive prefix product) 来构造步长,而不考虑 Shape 的层次结构。这可以被认为是“广义的列优先步长生成”。LayoutRight
标签从右到左将 Shape 进行排他性前缀乘积来构造步长,同样不考虑 Shape 的层次结构。对于深度为 1 的形状,这可以被认为是“行优先步长生成”,但对于分层形状,产生的步长可能令人惊讶。例如,上面的 s2xh4
的步长可以用 LayoutRight
生成。
在每个上述布局上调用 print
会产生以下结果:
s8 : _8:_1 // 静态8:静态1
d8 : 8:_1 // 动态8:静态1
s2xs4 : (_2,_4):(_1,_2) // 形状(静态2,静态4):步长(静态1,静态2) 列优先
s2xd4 : (_2,4):(_1,_2) // 形状(静态2,动态4):步长(静态1,静态2) 列优先
s2xd4_a : (_2,4):(_12,_1) // 形状(静态2,动态4):步长(静态12,静态1) 自定义
s2xd4_col : (_2,4):(_1,_2) // 形状(静态2,动态4):步长(静态1,静态2) 列优先
s2xd4_row : (_2,4):(4,_1) // 形状(静态2,动态4):步长(动态4,静态1) 行优先
s2xh4 : (2,(2,2)):(4,(2,1)) // 形状(动态2, (静态2,静态2)):步长(动态4, (静态2,静态1)) 分层自定义
s2xh4_col : (2,(2,2)):(_1,(2,4)) // 形状(动态2, (静态2,静态2)):步长(静态1, (静态2,静态4)) LayoutLeft生成的分层
使用了 Shape:Stride
表示法来表示 Layout。_N
表示法是静态整数的简写,而其他整数是动态整数。注意,Shape 和 Stride 都可以由静态和动态整数组成。
还要注意,Shape 和 Stride 被假定为是一致的 (congruent)。也就是说,Shape 和 Stride 具有相同的元组结构。对于 Shape 中的每个整数,Stride 中都有一个对应的整数。这可以用以下方式断言:
static_assert(congruent(my_shape, my_stride));
使用布局 (Using a Layout)
Layout 的基本用途是在由 Shape 定义的坐标空间和由 Stride 定义的索引空间之间进行映射。例如,要在一个 2-D 表中打印任意的 rank-2 布局,我们可以编写函数:
template <class Shape, class Stride>
void print2D(Layout<Shape,Stride> const& layout)
{
for (int m = 0; m < size<0>(layout); ++m) { // 遍历第0维 (行)
for (int n = 0; n < size<1>(layout); ++n) { // 遍历第1维 (列)
printf("%3d ", layout(m,n)); // 使用 2D 坐标 (m, n) 索引布局
}
printf("\n");
}
}
对上面的例子产生以下输出:
> print2D(s2xs4) // 标准列优先
0 2 4 6
1 3 5 7
> print2D(s2xd4_a) // 自定义步长,看起来像小矩阵在大矩阵中的偏移
0 1 2 3
12 13 14 15
> print2D(s2xh4_col)// LayoutLeft生成的分层布局,表现和标准列优先一样
0 2 4 6
1 3 5 7
> print2D(s2xh4) // 自定义分层布局,访问模式不同
0 2 1 3
4 6 5 7
我们可以看到静态、动态、行优先、列优先和分层布局的打印结果。layout(m,n)
语句提供了逻辑 2-D 坐标 (m,n)
到 1-D 索引的映射。
有趣的是,s2xh4
示例既不是行优先也不是列优先。此外,它有三个模式,但仍然被解释为 rank-2,并且我们使用 2-D 坐标。具体来说,s2xh4
在第二个模式中有一个 2-D 多模式,但我们仍然能够对该模式使用 1-D 坐标。更多内容在下一节,但首先我们可以再概括一步。让我们使用 1-D 坐标,并将每个布局的所有模式视为单个多模式。例如,下面的 print1D
函数:
template <class Shape, class Stride>
void print1D(Layout<Shape,Stride> const& layout)
{
for (int i = 0; i < size(layout); ++i) { // 遍历 1D 坐标
printf("%3d ", layout(i)); // 使用 1D 坐标 i 索引布局
}
}
对上面的例子产生以下输出:
> print1D(s2xs4) // 标准列优先的线性化
0 1 2 3 4 5 6 7
> print2D(s2xd4_a) // 自定义步长的线性化,注意顺序
0 12 1 13 2 14 3 15
> print2D(s2xh4_col)// LayoutLeft分层布局的线性化,表现和标准列优先一样
0 1 2 3 4 5 6 7
> print2D(s2xh4) // 自定义分层布局的线性化
0 4 2 6 1 5 3 7
布局的任何多模式,包括整个布局本身,都可以接受 1-D 坐标。更多内容在以下部分。
CuTe 提供了更多的打印实用程序来可视化布局。print_layout
函数生成布局映射的格式化 2-D 表。
> print_layout(s2xh4)
(2,(2,2)):(4,(2,1))
0 1 2 3
+---+---+---+---+
0 | 0 | 2 | 1 | 3 |
+---+---+---+---+
1 | 4 | 6 | 5 | 7 |
+---+---+---+---+
print_latex
函数生成 LaTeX,可以用 pdflatex
编译成相同 2-D 表的彩色矢量图形图像。
向量布局 (Vector Layouts)
我们将任何 rank == 1
的 Layout 定义为一个向量。例如,布局 8:1
可以解释为一个 8 元素向量,其索引是连续的。
Layout: 8:1
Coord : 0 1 2 3 4 5 6 7 <-- 1D 坐标
Index : 0 1 2 3 4 5 6 7 <-- 映射得到的索引
类似地,布局 8:2
可以解释为一个 8 元素向量,其中元素的索引跨度为 2。
Layout: 8:2
Coord : 0 1 2 3 4 5 6 7
Index : 0 2 4 6 8 10 12 14
根据上述 rank-1 的定义,我们也把布局 ((4,2)):((2,1))
解释为一个向量,因为它的形状是 rank-1。内部形状看起来像一个 4x2 的行优先矩阵,但额外的括号表明我们可以将这两个模式解释为一个 1-D 的 8 元素向量。步长告诉我们前 4 个元素的跨度为 2,然后有 2 个这样的第一组元素,跨度为 1。
Layout: ((4,2)):((2,1))
Coord : 0 1 2 3 4 5 6 7
Index : 0 2 4 6 1 3 5 7
我们可以看到第二组 4 个元素是第一组 4 个元素的副本,并额外增加了 1 的跨度。
考虑布局 ((4,2)):((1,4))
。同样是 4 个元素,跨度为 1,然后是 2 个这样的第一组元素,跨度为 4。
Layout: ((4,2)):((1,4))
Coord : 0 1 2 3 4 5 6 7
Index : 0 1 2 3 4 5 6 7
作为一个从整数到整数的函数,它与 8:1
是相同的。它是恒等函数。
矩阵示例 (Matrix examples)
概括来说,我们将任何 rank-2 的 Layout 定义为一个矩阵。例如,
Shape : (4,2)
Stride: (1,4)
0 4
1 5
2 6
3 7
是一个 4x2 列优先布局,沿着列的步长为 1,沿着行的步长为 4。
而
Shape : (4,2)
Stride: (2,1)
0 1
2 3
4 5
6 7
是一个 4x2 行优先布局,沿着列的步长为 2,沿着行的步长为 1。主序 (Majorness) 简单地就是哪个模式具有步长 1。
就像向量布局一样,矩阵的每个模式也可以拆分为多模式。这使我们能够表达超越行优先和列优先的更多布局。例如,
Shape: ((2,2),2)
Stride: ((4,1),2)
0 2
4 6
1 3
5 7
在逻辑上也是 4x2,沿着行的步长为 2,但沿着列有一个多步长。沿着列的前 2 个元素的步长为 4,然后是这些元素的副本,步长为 1。由于此布局在逻辑上是 4x2,与上面的列优先和行优先示例一样,我们仍然可以使用 2-D 坐标来索引它。
布局概念 (Layout Concepts)
在本节中,我们将介绍 Layout 接受的坐标集,以及坐标映射和索引映射是如何计算的。
布局兼容性 (Layout compatibility)
如果布局 A 的 shape 与布局 B 的 shape 兼容,我们就说布局 A 与布局 B 兼容。Shape A 与 Shape B 兼容,如果:
-
A 的 size 等于 B 的 size,并且
-
A 内的所有坐标在 B 内都是有效的坐标。
例如:
-
Shape 24 不兼容 Shape 32。(大小不同)
-
Shape 24 兼容 Shape (4,6)。(1D 坐标 0-23 映射到 2D 坐标 (0,0)-(3,5))
-
Shape (4,6) 兼容 Shape ((2,2),6)。(2D 坐标 (m,n) 映射到分层坐标 (m/2, m%2, n),注意整除关系)
-
Shape ((2,2),6) 兼容 Shape ((2,2),(3,2))。(分层坐标 (i,(j,k)), n 映射到更深的分层坐标 (i, (j,k), (n/3, n%3)),注意整除关系)
-
Shape 24 兼容 Shape ((2,2),(3,2))。(通过上述传递性)
-
Shape 24 兼容 Shape ((2,3),4)。(1D 坐标 0-23 映射到 2D 坐标 (i,j) 再映射到分层坐标 (i/2, i%2, j/3, j%3),注意整除关系)
-
Shape ((2,3),4) 不兼容 Shape ((2,2),(3,2))。((2,3) 和 (2,2) 不兼容,无法建立自然的一一对应坐标映射)
-
Shape ((2,2),(3,2)) 不兼容 Shape ((2,3),4)。(同上)
-
Shape 24 兼容 Shape (24)。(1D 坐标相同)
-
Shape (24) 不兼容 Shape 24。(Shape (24) 是 rank-1,需要 1D 坐标;Shape 24 (作为IntTuple) 是 rank-0(?) 或视为标量,接受方式不同?此例可能旨在说明类型差异)
-
Shape (24) 不兼容 Shape (4,6)。(Shape (24) 是 rank-1,需要 1D 坐标;Shape (4,6) 是 rank-2,需要 2D 坐标。虽然大小相同,但坐标空间维度不同,直接索引方式不兼容)
也就是说,兼容性是 Shapes 上的一个弱偏序关系,因为它是自反的、反对称的和传递的。
布局坐标 (Layouts Coordinates)
有了上面的兼容性概念,我们强调每个 Layout 接受多种类型的坐标。每个 Layout 接受与其兼容的任何 Shape 的坐标。CuTe 通过逆字典序 (colexicographical order)(从右向左读,而不是“字典序lexicographical”从左向右读)在这些坐标集之间提供映射。
因此,所有 Layout 都提供两个基本映射:
-
通过 Shape 将输入坐标映射到相应的自然坐标。
-
通过 Stride 将自然坐标映射到索引。
坐标映射 (Coordinate Mapping)
从输入坐标到自然坐标的映射是在 Shape 内应用逆字典序。
以形状 (3,(2,3))
为例。这个形状有三个坐标集:1-D 坐标、2-D 坐标和自然(h-D)坐标。
1-D | 2-D | Natural | 1-D | 2-D | Natural |
---|---|---|---|---|---|
0 | (0,0) | (0,(0,0)) | 9 | (0,3) | (0,(1,1)) |
1 | (1,0) | (1,(0,0)) | 10 | (1,3) | (1,(1,1)) |
2 | (2,0) | (2,(0,0)) | 11 | (2,3) | (2,(1,1)) |
3 | (0,1) | (0,(1,0)) | 12 | (0,4) | (0,(0,2)) |
4 | (1,1) | (1,(1,0)) | 13 | (1,4) | (1,(0,2)) |
5 | (2,1) | (2,(1,0)) | 14 | (2,4) | (2,(0,2)) |
6 | (0,2) | (0,(0,1)) | 15 | (0,5) | (0,(1,2)) |
7 | (1,2) | (1,(0,1)) | 16 | (1,5) | (1,(1,2)) |
8 | (2,2) | (2,(0,1)) | 17 | (2,5) | (2,(1,2)) |
输入到形状 (3,(2,3))
的每个坐标都有两个等效坐标,并且所有等效坐标都映射到相同的自然坐标。再次强调,因为所有上述坐标都是有效的输入,所以一个具有 Shape (3,(2,3))
的 Layout 可以通过使用 1-D 坐标将其用作 18 个元素的 1-D 数组,通过使用 2-D 坐标将其用作 3x6 元素的 2-D 矩阵,或者通过使用 h-D(自然)坐标将其用作 3x(2x3) 元素的 h-D 张量。
前面的 1-D 打印演示了 CuTe 如何通过 2-D 坐标的逆字典序来识别 1-D 坐标。从 i = 0
迭代到 size(layout)
并使用单整数坐标 i
索引我们的布局,即使布局以行优先或更复杂的方式将坐标映射到索引,也会以这种“广义列优先”顺序遍历 2-D 坐标。
函数 cute::idx2crd(idx, shape)
负责坐标映射。它将获取形状内的任何坐标并计算该形状的等效自然坐标。(示例代码略)
索引映射 (Index Mapping)
从自然坐标到索引的映射是通过取自然坐标与 Layout 的 Stride 的内积来执行的。
以布局 (3,(2,3)):(3,(12,1))
为例。那么一个自然坐标 (i,(j,k))
将产生索引 i*3 + j*12 + k*1
。此布局计算的索引显示在下面的 2-D 表中,其中 i
用作行坐标,(j,k)
用作列坐标。
0 1 2 3 4 5 <== 1-D 列坐标
(0,0) (1,0) (0,1) (1,1) (0,2) (1,2) <== 2-D 列坐标 (j,k)
+-----+-----+-----+-----+-----+-----+
0 | 0 | 12 | 1 | 13 | 2 | 14 |
+-----+-----+-----+-----+-----+-----+
1 | 3 | 15 | 4 | 16 | 5 | 17 |
+-----+-----+-----+-----+-----+-----+
2 | 6 | 18 | 7 | 19 | 8 | 20 |
+-----+-----+-----+-----+-----+-----+
函数 cute::crd2idx(c, shape, stride)
负责索引映射。它将获取形状内的任何坐标,计算该形状的等效自然坐标(如果还不是),并计算与步长的内积。(示例代码略)
布局操作 (Layout Manipulation)
子布局 (Sublayouts)
可以使用 layout<I...>
检索子布局
Layout a = Layout<Shape<_4,Shape<_3,_6>>>{}; // (4,(3,6)):(1,(4,12))
Layout a0 = layout<0>(a); // 4:1 (取第0模式)
Layout a1 = layout<1>(a); // (3,6):(4,12) (取第1模式)
Layout a10 = layout<1,0>(a); // 3:4 (取第1模式的第0子模式)
Layout a11 = layout<1,1>(a); // 6:12 (取第1模式的第1子模式)
或者使用 select<I...>
Layout a = Layout<Shape<_2,_3,_5,_7>>{}; // (2,3,5,7):(1,2,6,30)
Layout a13 = select<1,3>(a); // (3,7):(2,30) (选择第1和第3模式)
Layout a01 = select<0,1,3>(a); // (2,3,7):(1,2,30) (选择第0,1,3模式)
Layout a2 = select<2>(a); // (5):(6) (选择第2模式)
或者使用 take<ModeBegin, ModeEnd>
Layout a = Layout<Shape<_2,_3,_5,_7>>{}; // (2,3,5,7):(1,2,6,30)
Layout a13 = take<1,3>(a); // (3,5):(2,6) (取模式1到2(不包括3))
Layout a14 = take<1,4>(a); // (3,5,7):(2,6,30) (取模式1到3)
// take<1,1> 不允许。不允许空布局。
连接 (Concatenation)
可以向 make_layout
提供 Layout 以进行包装和连接
Layout a = Layout<_3,_1>{}; // 3:1
Layout b = Layout<_4,_3>{}; // 4:3
Layout row = make_layout(a, b); // (3,4):(1,3) (连接成行)
Layout col = make_layout(b, a); // (4,3):(3,1) (连接成列)
Layout q = make_layout(row, col); // ((3,4),(4,3)):((1,3),(3,1)) (连接成更高维)
Layout aa = make_layout(a); // (3):(1) (包装)
Layout aaa = make_layout(aa); // ((3)):((1)) (再次包装)
Layout d = make_layout(a, make_layout(a), a); // (3,(3),3):(1,(1),1) (混合连接和包装)
或者可以使用 append
、prepend
或 replace
进行组合。(示例代码略)
分组和展平 (Grouping and flattening)
可以使用 group<ModeBegin, ModeEnd>
对布局模式进行分组,并使用 flatten
进行展平。
Layout a = Layout<Shape<_2,_3,_5,_7>>{}; // (_2,_3,_5,_7):(_1,_2,_6,_30)
Layout b = group<0,2>(a); // ((_2,_3),_5,_7):((_1,_2),_6,_30) (将模式0-1分组)
Layout c = group<1,3>(b); // ((_2,_3),(_5,_7)):((_1,_2),(_6,_30)) (将模式1-2分组)
Layout f = flatten(b); // (_2,_3,_5,_7):(_1,_2,_6,_30) (展平分组b)
Layout e = flatten(c); // (_2,_3,_5,_7):(_1,_2,_6,_30) (展平分组c)
分组、展平和重新排序模式允许就地重新解释张量,例如将张量视为矩阵,将矩阵视为向量,将向量视为矩阵等。
切片 (Slicing)
可以对布局进行切片,但切片更适合在张量上执行。有关切片详细信息,请参阅张量部分。
总结 (Summary)
-
Layout 的 Shape 定义了它的坐标空间。
-
每个 Layout 都有一个 1-D 坐标空间。这可用于以逆字典序迭代坐标空间。
-
每个 Layout 都有一个 R-D 坐标空间,其中 R 是布局的秩。R-D 坐标的逆字典序枚举对应于上面的 1-D 坐标。
-
每个 Layout 都有一个 h-D(自然)坐标空间,其中 h 是“分层的”。这些按逆字典序排序,并且该顺序的枚举对应于上面的 1-D 坐标。自然坐标与 Shape 一致,因此坐标的每个元素都对应 Shape 的一个元素。
-
Layout 的 Stride 将坐标映射到索引。
-
自然坐标的元素与 Stride 的元素的内积产生最终索引。
-
对于每个 Layout,都存在一个与该 Layout 兼容的整数 Shape。即,该整数 shape 是
size(layout)
。然后我们可以观察到:-
Layout 是从整数到整数的函数。
-
如果你熟悉 C++23 的特性 mdspan
,这是 mdspan
布局映射和 CuTe Layouts 之间的一个重要区别。在 CuTe 中,Layout 是一等公民,本质上是分层的,可以自然地表示超越行优先和列优先的函数,并且同样可以用层次结构坐标进行索引。(mdspan
布局映射也可以表示分层函数,但这需要定义自定义布局。)mdspan
的输入坐标必须与 mdspan
具有相同的形状;多维 mdspan
不接受 1-D 坐标。