Tuple定义

定义

如果\(V\)是一个集合，那么由\(V\)中元素组成的有限长度有序列表就被叫做一个元组(tuple)

\[ X = (x_1, \cdots, x_m) \]

\(x_i \in V, 1\le i \le m\)。

\(X\)的长度为

\[ \text{len}(X) = m \]

我们最感兴趣的是\(V = \mathbb{Z}\)，也就是由整数组成的tuple。当\(X\)是整数组成的tuple的时候，我们定义其size为：

\[ \text{size}(X) = x_1\cdots x_m \]

下面是一些整数组成的tuple的例子

示例

\[ X =(3,2,1), \text{len}(X) = 3, \text{size}(X) = 6 \]
\[ X =(32), \text{len}(X) = 1, \text{size}(X) = 32 \]
\[ X =(), \text{len}(X) = 0, \text{size}(X) = 1 \]

定义

如果\(X = (x_1,\cdots, x_m)\)， \(Y = (y_1, \cdots, y_n)\)都是tuples，它们的连接(concatenation)定义为：

\[ X \star Y = (x_1,\cdots,x_m, y_1, \cdots, y_n) \]

例如 \(X = (2,3), Y=(1,4)\)， 那么

\[ X \star Y = (2,3,1,4) \]

定义

如果\(X, Y\)都是tuple，当存在另外一个tuple \(X'\)，使得

\[ X \star X' = Y \]

我们称\(X\)整除\(Y\)。

注

如果\(n\)是一个正整数，那么我们记

\[ [0, n) = \{0, 1, \cdots, n-1\} \]

如果\(S = (s_1, \cdots, s_m)\)是一个tuple，那么我们记

\[ [0, S) = [0, s_1)\times \cdots \times [0, s_m) \]

也就是\((x_1, \cdots, x_m), 0\le x_i < s_i\)的集合。

示例

如果\(S = (3, 2)\)，那么

\[ [0, S) = \{(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)\} \]

嵌套Tuple

定义

一个profile P满足下面条件之一： - \(P = *\) - \(P = (P_1, \cdots, P_r)\)，其中\(P_1, \cdots, P_r\)都是profile（\(r\geq 0\)，允许空tuple）。

示例

下面是一些profile的示例：

\[ \begin{aligned} P_1 &= () \\ P_2 &= * \\ P_3 &= (*, *) \\ P_4 &= (*, (*, *)) \\ P_5 &= ((*, *), (*, *)) \\ \end{aligned} \]

定义

假设\(P\)是一个profile，我们有

\[ \begin{aligned} \text{rank}(P) &= \begin{cases} 1, & P = * \\ r, & P = (P_1, \cdots, P_r) \end{cases} \\ \text{len}(P) &= \begin{cases} 1, & P = * \\ \sum_{i=1}^r \text{len}(P_i), & P = (P_1, \cdots, P_r) \end{cases} \\ \text{depth}(P) &= \begin{cases} 0, & P = * \\ 1+\max_{1\le i\le r}(\text{depth}(P_i)), & P = (P_1, \cdots, P_r) \end{cases} \end{aligned} \]

示例

\[ \begin{aligned} P &= * & \text{rank}(P) = 1, &&\text{len}(P) = 1, &&\text{depth}(P) = 0 \\ P &= (*, *) &\text{rank}(P) = 2, &&\text{len}(P) = 2, &&\text{depth}(P) = 1 \\ P &= (*, (*, *)) &\text{rank}(P) = 2, &&\text{len}(P) = 3, &&\text{depth}(P) = 2 \\ P &= ((*, *), (*, *)) &\text{rank}(P) = 2, &&\text{len}(P) = 4, &&\text{depth}(P) = 2 \\ \end{aligned} \]

定义

\(P\)是一个rank为\(r\)的profile，ith mode of P(\(1\le i < r\))定义为

\[ \text{mode}_i(P) = \begin{cases} P, & \text{depth}(P) = 0 \\ P_i, & P=(P_1, \cdots, P_r),\ \text{depth}(P) \ge 1 \end{cases} \]

示例

例如\(P = (*, (*, *))\)，我们有

\[ \begin{aligned} \text{mode}_1(P) &= * \\ \text{mode}_2(P) &= (*, *) \end{aligned} \]

定义

\(V\)是一个集合，由\(V\)中的元素组成的nested tuple 的\(X\)是由一个tuple \(X^\text{b}\)和一个profile \(P\)组成 - \(X^{\text{b}} = (x_1, \cdots, x_m)\)， \(x_i \in V, \forall i \in [1, m]\)，是一个tuple，称为flattening of X； - \(\text{prof}(X) = P\)，叫做profile of X。

把\(V\)中元素组成的所有的嵌套tuples记作\(\text{Nest}(V)\)。

示例

\[ \begin{aligned} X &= 3 &X^{\text{b}} &= (3) &\text{prof}(X) &= * \\ X &= (2, 2) &X^{\text{b}} &=(2, 2) &\text{prof}(X) &= (*, *) \\ X &= (2, (3, 4)) &X^{\text{b}} &=(2, 3, 4) &\text{prof}(X) &= (*, (*, *)) \end{aligned} \]

定义

假设\(X\)是由\(V\)中元素组成的nested tuple，那么 - rank of X

\[ \text{rank}(X) = \text{rank}(X^\text{b}) \]
  • length of X
\[ \text{len}(X) = \text{len}(P) = \text{len}(X^{\text{b}}) \]
  • depth of X
\[ \text{depth}(X) = \text{depth}(P) \]
  • 如果\(V = \mathbb{Z}\)， size of X
\[ \text{size}(X) = \text{size}(X^\text{b}) \]

从数学 Tuple 到 CuTe 代码

cute中tuple.hpp包含了Tupe相关的声明，省略构造及访问函数体；eso::ESO_t 是库内部用于紧凑存储元素的基类：

template <class... T>
struct tuple : eso::ESO_t<T...> {
  CUTE_HOST_DEVICE constexpr tuple();
  CUTE_HOST_DEVICE constexpr tuple(T const&... t);
};
template <> struct tuple<> {};

template <class... T>
CUTE_HOST_DEVICE constexpr tuple<T...> make_tuple(T const&... t);

template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr decltype(auto) get(tuple<T...> const& t) noexcept;

template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr decltype(auto) get(tuple<T...>& t) noexcept;

template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr decltype(auto) get(tuple<T...>&& t) noexcept;

rank, depth, size函数在int_tuple.hpp中：

template <int... Is, class IntTuple>
CUTE_HOST_DEVICE constexpr auto rank(IntTuple const& t);

template <int... Is, class IntTuple>
CUTE_HOST_DEVICE constexpr auto depth(IntTuple const& t);

template <int... Is, class IntTuple>
CUTE_HOST_DEVICE constexpr auto size(IntTuple const& a);

flatten函数定义在tupe_algorithms.hpp中，它会把嵌套的Tupe展开

template <class T>
CUTE_HOST_DEVICE constexpr
auto flatten(T const& t)
{
  if constexpr (is_tuple<T>::value) {
    if constexpr (is_flat<T>::value) {      // Shortcut for perf
      return t;
    } else {
      return filter_tuple(t, [](auto const& a) { return flatten_to_tuple(a); });
    }
  } else {
    return t;
  }

  CUTE_GCC_UNREACHABLE;
}

其中传入的T是nested tuple。

真实使用中，我们比较少直接使用tuple，而是使用tupe的别名，最经常使用的有

  • Shape
  • Stride
  • Coord
  • Tile
  • Step

它们定义在layout.hpp文件里面：

template <class... Shapes> using Shape = cute::tuple<Shapes...>;
template <class... Strides> using Stride = cute::tuple<Strides...>;
template <class... Coords> using Coord = cute::tuple<Coords...>;
template <class... Layouts> using Tile = cute::tuple<Layouts...>;
template <class... Strides> using Step = cute::tuple<Strides...>;

同一文件里面还定义了生产它们的工厂函数：

template <class... Ts>
CUTE_HOST_DEVICE constexpr Shape<Ts...> make_shape(Ts const&... t);
template <class... Ts>
CUTE_HOST_DEVICE constexpr Stride<Ts...> make_stride(Ts const&... t);
template <class... Ts>
CUTE_HOST_DEVICE constexpr Coord<Ts...> make_coord(Ts const&... t);
template <class... Ts>
CUTE_HOST_DEVICE constexpr Tile<Ts...> make_tile(Ts const&... t);
template <class... Ts>
CUTE_HOST_DEVICE constexpr Step<Ts...> make_step(Ts const&... t);

这些工厂函数都返回由实参类型组成的 Tuple，用不同名字表达用途：

  • Shape表示一个tensor的形状；

  • Stride表示一个tensor的stride；

  • Coord 表达坐标；

  • Tile 组织各方向的 tiler；

  • Step 表达各方向的选择或顺序。

现在将前面介绍的静态整数与普通整数放在一起：

#include <cute/tensor.hpp>
using namespace cute;

int n = 5;
auto x = make_tuple(_2{}, make_tuple(_3{}, n));
// 数学表达：(2,(3,5))，其中最后一个 5 是运行时值

auto first = get<0>(x);         // Int<2>
auto inner = get<1>(x);         // (Int<3>, int)
int last = get<1>(get<1>(x));    // 5

auto s = make_shape(_2{}, make_shape(_3{}, n));
auto d = make_stride(_1{}, make_stride(_2{}, _6{}));

make_tuple、make_shape、make_stride 都能建立嵌套整数结构；后两者用名字表达“这组数用于形状”或“这组数用于步长”。get<I> 取第 I 个顶层 mode，C++ 下标从 0 开始，它不会自动把嵌套 Tuple 展平，因此取得最里面的 n 需要先进入第二个 mode。

Shape 和 Stride 中，哪些值写成静态整数

同一个数学形状 (4,8) 可以有三种代码表达：

int m = 4, n = 8;
auto dynamic_shape = make_shape(m, n);          // (4,8)
auto static_shape  = make_shape(_4{}, _8{});     // (_4,_8)
auto mixed_shape   = make_shape(_4{}, n);        // (_4,8)

它们当前的数值形状相同，但前者的两个值由对象保存，中间两个值由类型确定，最后一种只把第一个值放进类型。Tuple 的层次和元素个数在这里仍由类型确定；“动态 shape”并不表示可以在运行时随意增加一个维度。

实际使用时，全局矩阵的 M/N/K 常由调用者传入，而一个 kernel 固定使用的 tile 大小以及连续维的步长 1 常写成静态整数。例如，上一章的列优先矩阵可以这样准备形状和步长：

int m = 4, n = 8;
auto shape  = make_shape(m, n);       // (4,8)
auto stride = make_stride(_1{}, m);  // (_1,4)

坐标 (i,j) 对应的元素偏移仍是 i + j*m。第一维每次前进一个元素，这个 1 与矩阵尺寸无关，适合用 _1{} 表达；第二维的步长由运行时传入的 m 决定。 后面将 shape 和 stride 交给 make_layout，就得到完整 Layout。

固定的 tile 大小适合静态表达，需要在不同调用间变化的矩阵尺寸适合动态表达。两种信息可以存在于同一个 Tuple 中；具体硬件操作若要求某些维度或步长静态已知，再按该操作的约束选择类型。