AI

A note on A note on the algebra of CuTe Layouts

Posted by w@hidva.com on October 19, 2024

这篇文章纪录了对 A note on the algebra of CuTe Layouts 的学习笔记, 侧重于记录我当时没有看懂的部分, 以及对书中部分知识点的一些扩展. 非常零散, 不成系统; 最好是结合 A note on the algebra of CuTe Layouts 使用. A note on the algebra of CuTe Layouts 是之前在学习 cute 时找到的一篇对 cute layout 深入介绍的文章. 本文章会用到 这里 提到的一些数学常识.

$x \mapsto (x \bmod M_0, \left\lfloor\frac{x}{M_0}\right\rfloor \bmod M_1, \ldots, \left\lfloor\frac{x}{M_0 \cdots M_{\alpha-1}}\right\rfloor \bmod M_\alpha)$ 这个公式便是 1-D coordinate space 到 R-D coordinate space 转换. 以 $(M_0, M_1)$ 为例, 别忘了这里 $x \lt M_0 \times M_1$, 所以 $\left\lfloor\frac{x}{M_0}\right\rfloor \bmod M_1 = \left\lfloor\frac{x}{M_0}\right\rfloor$.

$\widehat{f_L}, f_L$ 区别在于两者定义域不一样了. $f_L$ 定义域是 $[0, M)$, $\widehat{f_L}$ 定义域是 $\mathbb{N}$.

2.6.zy1 B=complement(A, M)’s layout function is strictly increasing.

证明: 对于 $\forall x, y \in [0, \text{size}(B))$, 由于 $x \lt y$ 易证从右往左 $\exists i, x_i \ne y_i$ 且 $x_i \lt y_i$.

\[\begin{align} \iota_B(x) &= (x', x_0, \cdots, x_i, x_{i+1} \cdots, x_{\alpha}) \\ \iota_B(y) &= (y', y_0, \cdots, y_i, y_{i+1} \cdots, y_{\alpha}) \\ x_{i+1} &= y_{i+1}, \cdots x_{\alpha} = y_{\alpha} \end{align}\]

同时易证:

\[\begin{align} f_B(x) &\le (d_0 - 1) 1 + (\frac{d_1}{N_0 d_0} - 1) N_0 d_0 + \cdots + (\frac{d_i}{N_{i-1} d_{i-1}} - 1) N_{i-1} d_{i-1} + (y_i - 1) N_i d_i + G \\ f_B(y) &\ge 0 * 1 + 0 N_0 d_0 + \cdots + 0 N_{i-1} d_{i-1} + y_i N_i d_i + G \\ G &= x_{i+1} N_{i+1} d_{i+1} + \cdots + x_{\alpha} N_{\alpha} d_{\alpha} \\ &= y_{i+1} N_{i+1} d_{i+1} + \cdots + y_{\alpha} N_{\alpha} d_{\alpha} \end{align}\]

所以我们只要证明:

\[\begin{align} &(d_0 - 1) 1 + (\frac{d_1}{N_0 d_0} - 1) N_0 d_0 + \cdots + (\frac{d_i}{N_{i-1} d_{i-1}} - 1) N_{i-1} d_{i-1} + (y_i - 1) N_i d_i \lt y_i N_i d_i \\ \text{LEFT} &= d_0 - 1 + d_1 - N_0 d_0 + \cdots + d_i - N_{i-1} d_{i-1} + y_i N_i d_i - N_i d_i \\ &= (1 - N_0) d_0 - 1 + (1 - N_1) d_1 + \cdots + (1 -N_i) d_i + y_i N_i d_i \\ \end{align}\]

易证 $(1 - N_0) d_0 - 1 + (1 - N_1) d_1 + \cdots + (1 -N_i) d_i \lt 0$, 证明完毕.

Proposition 2.7. 略作补充

  • 首先证明 $f_C$ 值域最小值是 0, 最大值是 M - 1. 即 $f_C$ 是 $\lbrace x \in \mathbb{N} \mid x \in [0, M) \rbrace$ 到 $\lbrace y \in \mathbb{N} \mid 0 \le y \lt M \rbrace$ 的映射. 之后再证明 $f_C$ 是单射, 则 $f_C$ 一定是满射.

  • 话说我本来以为对于任意 layout function 其一定是单射, 但居然想了个 (3, 2) : (2, 4) 的反例.

  • $f_C$ 是双射也进一步加深了 cutlass 文档原文 The complement effectively “repeats” the original layout (displayed in the other colors) such that the codomain size of the result is 24. 这里 “repeats” 的理解.

  • 易得 size(C) == cosize(C)

Corollary 2.8. 略作补充

  • By Proposition 2.7, we have that $f_A(I \cap J) \cap f_B(I \cap J) = \lbrace 0 \rbrace$. 我这里补充下:

证明: 假设 $\exists k \in f_A(I \cap J) \cap f_B(I \cap J), k \ne 0$. 即 $\exists k_0, k_1 \in I \cap J, f_A(k_0) = f_B(k_1) = k$. 则 $f_C((k_0, 0)) = f_A(k_0) = f_B(k_1) = f_C((0, k_1))$, 这与 $f_C$ 是单射矛盾.

  • 我在 cutlass layout 文档中对 B = complement(A, M), C=(A, B) 的理解就是 $f_B()$ 定义了在 C 中这些偏移处重复应用 A layout. 这里一个前提就是 $f_A(I) \cap \widehat{f_B}(I) = \lbrace 0 \rbrace$, 正好本文这里证明了这一前提.

Proposition 2.14. 略作补充

  • Definition 2.13 中 $j, c’$ 是什么? j 是 $\widehat{M’}, N$ 的 division index. $j \lt \alpha$ 时有 $c’ \lt M_j$
\[\begin{align} M &= M_0 \cdots M_{i-1} M_i M_{i+1} \cdots M_{\alpha} \\ M &= r \frac{M_i}{c} M_{i+1} \cdots M_{\alpha - 1} M_{\alpha} \\ \widehat{M'} &= \frac{M_i}{c} M_{i+1} \cdots M_{\alpha - 1} \infty \\ N &= \frac{M_i}{c} M_{i+1} \cdots M_{j-1} c' \end{align}\]
  • $r$ is sent to $c \cdot \delta_i$ 的意思是 $\hat{\iota}(r) = c \cdot \delta_i$, 这是因为 $r = M_0 \cdots M_{i-1} c$ 且当 $ i \lt \alpha$ 时有 $c \lt M_i$. 所以 $\widehat{f_A} (f_B(1)) = \widehat{f_A}(r) = cd_i$. 而容易验证 $f_{A \circ B}(1)$ 在几种情况下均为 $cd_i$. 额外补充下, 在 $A \circ B = ( M_i / c, M_{i+2}, \cdots )$ 情况中, $c \lt M_i, M_i / c \gt 1$, 所以 1 此时对应坐标是 $(1, 0, 0, \cdots, 0)$ 而不会是 $(0, 1, 0, \cdots, 0)$ 等.

  • In the case of $i \lt \alpha$ and $N \le M_i /c$ or $i = \alpha$, this already suffices to show $f_{A \circ B} = \widehat{f_A} \circ f_B$. 我这里补充下.

证明:

\[A \circ B = \begin{cases} (N) : (c d_{\alpha}) \quad \mbox{if } i = \alpha; (1) \\ (N) : (c d_{i}) \quad \mbox{if } i \lt \alpha \mbox{ and } N \le M_i / c; (2) \\ (M_i / c, M_{i+1}, \cdots, M_{j-1}, c'):(cd_i, d_{i+1}, \cdots, d_{j-1}, d_j) \quad \mbox{if } c' \gt 1 ; (3) \\ (M_i / c, M_{i+1}, \cdots, M_{j-1}):(cd_i, d_{i+1}, \cdots, d_{j-1}) \quad \mbox{if } c' = 1 ; (4) \end{cases}\]

对于情况 (1) 此时 $\widehat{f_A}( f_B(x) ) = \widehat{f_A}(xr)$, 这里 $r = M_0 \cdots M_{\alpha-1} c$, 所以 $xr$ 对应坐标是 $(0, 0, \cdots, 0, cx)$ 所以 $\widehat{f_A}(xr) = cx d_{\alpha} = f_{A \circ B}(x)$.

对于情况 (2) 此时 $xr = M_0 \cdots M_{i-1} c x, x \lt N \le M_i / c$ 即 $cx \lt M_i$, 所以 $xr$ 对应坐标是 $(0, \cdots, 0, \underbrace{cx}i, 0, \cdots, 0)$, $\widehat{f_A}(xr) = cx d_i = f{A \circ B}(x)$.

  • In view of the multi-linearity properties of both functions, this implies that $f_{A \circ B} = \widehat{f_A} \circ f_B$. 我这里补充下.

证明: 这里仅考虑 $c’ \gt 1$ 情况, $c’=1$ 时类似:

x 1 $\frac{M_i}{c}$ $M_{i+1} \frac{M_i}{c}$ $M_{j-1} \cdots M_{i+1}\frac{M_i}{c}$
B(x) r $\frac{M_i}{c}r$ $M_{i+1} \frac{M_i}{c}r$ $M_{j-1} \cdots M_{i+1}\frac{M_i}{c}r$
$\iota_C(x)$ $\delta_0$ $\delta_1$ $\delta_2$ $\delta_{j-i}$
$\hat{\iota}(B(x))$ $c \delta_i$ $\delta_{i+1}$ $\delta_{i+2}$ $\delta_j$

对于 $\forall x, \iota_C(x) = (x_i, x_{i+1}, \cdots, x_{j-1}, x_j)$, 则此时:

\[\begin{align} x &= x_j M_{j-1} \cdots M_{i+1}\frac{M_i}{c} + x_{j-1} M_{j-2} \cdots M_{i+1}\frac{M_i}{c} + \cdots + x_{i+1} \frac{M_i}{c} + x_i \\ xr &= x_j M_{j-1} \cdots M_{i+1}\frac{M_i}{c} r + x_{j-1} M_{j-2} \cdots M_{i+1}\frac{M_i}{c} r + \cdots + x_{i+1} \frac{M_i}{c} r + x_i r \\ &= x_j M_{j-1} \cdots M_{i+1} M_i M_{i-1} \cdots M_0 + x_{j-1} M_{j-2} \cdots M_0 + \cdots + x_{i+1} M_i \cdots M_0 + c x_i M_{i-1} \cdots M_0 \end{align}\]

这里 $x_i \lt \frac{M_i}{c}$, 所以 $c x_i \lt M_i$, 所以:

\[\hat{\iota}(xr) = (0, \cdots, 0, \underbrace{c x_i}_i, x_{i+1}, \cdots, x_{j-1}, x_j)\]

由 $\iota_C(x), \hat{\iota}(xr)$ 易得结论.

2.19.zy1, 求证 $A / (BC) = (A / B) / C$, 这里 $A/B$ 是 C/C++ 中的整除, 比如 $5/2 = 2$.

证明: 设 $A/(BC) = K_1, A / B = K_2, K_2 / C = K_3$. 则

\[\begin{align} A &= BC K_1 + r_1, r_1 \lt BC \\ A &= B K_2 + r_2, r_2 \lt B \\ K_2 &= C K_3 + r_3, r_3 \lt C \\ A &= BC K_3 + B r_3 + r_2 \end{align}\]

这里 $r_2 \le B - 1, r_3 \le C - 1, B r_3 + r_2 \le BC - 1 \lt BC$. 由 代数学基础定理 3.3 带余除法 可得 $K_1 = K_3, r_1 = B r_3 + r_2$.

Lemma 2.19 略作补充

  • Now observe that the composite map $(\iota_0, \cdots, \iota_{\gamma}) \circ \iota$ is also the usual isomorphism with respect to the maximal decomposition of $C$. 我一眼没有 observe 出来, 所以这里补充下.

证明: 如下所示这里 $\mbox{size}(C)$ 被分解为 $\mbox{size}(C_0), \mbox{size}(C_1), \cdots, \mbox{size}(C_{\gamma})$, 其中 $\mbox{size}(C_0)$ 进一步被分解为 $\mbox{size}(C_{00}), \cdots, \mbox{size}(C_{0 n_0})$ 这里 $C_{00}$ 即 $C_0’$.

我们现在证明对于任意 $V, (V / C_{00}) \mod C_{01} = ((V \mod C_0) / C_{00}) \mod C_{01}$. 其他情况类似.

\[\begin{align} &V_1 = V \mod C_0, V_2 = V_1 / C_{00}, V_3 = V_2 \mod C_{01} \\ &V_4 = V / C_{00}, V_5 = V_4 \mod C_{01} \\ &V = C_0 K_1 + V_1, V_1 \lt C_0; V_1 = V_2 C_{00} + r_2, r_2 \lt C_{00}; V_2 = C_{01} K_2 + V_3, V_3 \lt C_{01} \\ &V = C_0 K_1 + C_{00} C_{01} K_2 + C_{00} V_3 + r_2, r_2 = V \mod C_{00} \\ &V = V_4 C_{00} + r_2, V_4 = C_{01} K_3 + V_5, V_5 \lt C_{01} \\ &V = C_{00} C_{01} K_3 + V_5 C_{00} + r_2 \\ &V_3 \le C_{01} - 1, C_{00}V_3 \le C_{00} C_{01} - C_{00}, C_{00}V_3 + r_2 \lt C_{00} C_{01} \\ &V_5 C_{00} + r_2 \lt C_{00} C_{01} \end{align}\]

即 $C_{00}V_3 + r_2 = V \mod C_{00} C_{01} = C_{00} V_5 + r_2$. 所以 $V_3 = V_5$.

Theorem 2.18 略作补充

  • the coordinates of subsequent values of $f_{B_k}$ and $f_{B_l}$ will increment by multiples of $w_{i_0}$ and $z_{j_0}$ in indices $i_0$ and $j_0$, by increments of 1 for indices greater than $i_0$ and $j_0$. 具体是个什么情况.

解: 由已知条件可得 $r_k = f_{B_k}(1), r_k = M_0 \cdots M_{i_0 - 1} w_{i_0}$, 当 $i_0 \lt \alpha$ 时 $w_{i_0} \lt M_{i_0}, M_{i_0} = K w_{i_0}$, 这里 K 为正整数. 所以此时:

$x$ $f_{B_k}(x)$ $\iota(f_{B_k}(x))$
1 $r_k$ $(0, \cdots, 0, w_{i_0}, 0, \cdots, 0)$
2 $2 r_k$ $(0, \cdots, 0, 2 w_{i_0}, 0, \cdots, 0)$
K-1 $(K-1) r_k$ $(0, \cdots, 0, (K-1) w_{i_0}, 0, \cdots, 0)$
K $ K r_k$ $(0, \cdots, 0, 0 (i_0), 1, 0, \cdots, 0)$
K+1 $ (K+1) r_k$ $(0, \cdots, 0, w_{i_0}, 1, 0, \cdots, 0)$

P.S. It then suffices to see that the analogous diagram with $\widehat{f_A} \circ f_B$ commutes, 这句话是指 “接下来我们只要证明 $\widehat{f_A} \circ f_B$ commutes 就可以了”, 我一开始理解这句话是指 “易得下图 $\widehat{f_A} \circ f_B$ commutes”, 想了半天没想出来怎么易得的.

  • 假设 $f_{B_l}(1) \gt f_{B_k}(N_k - 1)$, 为何这里坐标相加不会溢出?

解: 如下所示, 最可能会导致溢出的一种情况是 $v_2, z_{j_0}$ 都位于坐标第 $j_0$ 分量处.

\[\begin{align} \iota(f_{B_k}(N_k - 1)) &= (0, \cdots, v_1, v_2, 0, \cdots, 0) \\ \iota(r_l) &= (0, \cdots, 0, z_{j_0}, 0, \cdots, 0) \end{align}\]

此时 $v_2 \lt z_{j_0}, z_{j_0} \le \frac{M_{j_0}}{2}$, 所以 $v + z_{j_0} \lt M_{j_0}$

P.S. 在 Definition 2.16 中, $M’ = M_0 \cdots M_{\alpha-1}$, 若 $r_l = r_k(N_k - 1) = M’$, 此时仍然是符合 Definition 2.17(2) disjoint 要求的, 此时相加会不会溢出? 不会, 因为 $\widehat{f_A}$ 使用 $\widehat{\iota}$ 来计算坐标, 而 $\widehat{\iota}$ 在 $\alpha$ 分量处没有 $\mod M_{\alpha}$ 操作.

Lemma 2.23. 略作补充

  • the additional divisibility condition (3) in Definition 2.11 needed to promote weak left divisibility to left divisibility is necessarily satisfied for all the $N_{\varphi(k)}$ terms. 这句话啥意思啊?

解: 由已知条件可得 $r_{\varphi(\beta)} = M_0 M_1 \cdots M_{i - 1} c$, 这里我们先只考虑 $i \lt \alpha$, 此时 $M_i = cK, K \in \mathbb{N}$.

\[\begin{align} \widehat{M'} &= \frac{M_i}{c} M_{i+1} \cdots M_{\alpha - 1} \infty \\ N_{\varphi(\beta)} &= \frac{M_i}{c} M_{i+1} \cdots M_{j - 1} c' \end{align}\]

这里先只考虑 $j \lt \alpha$, 此时由 weakly left divisible 条件只知道 $c’ \lt M_j$, 并不能确保 $c’$ 可以整除 $M_j$. 现在要证明 M is left divisible by $N_{\varphi(\beta)} r_{\varphi(\beta)}$, 如下所示需要证明 $c’ \mid M_j$

\[M_0 \cdots M_{i-1} M_i M_{i+1} \cdots M_{j-1} c' = N_{\varphi(\beta)} r_{\varphi(\beta)}\]

但没想出咋整, 虽然由 $N_{\varphi(\beta)} r_{\varphi(\beta)} \mid M$ 可以推出 $c’ \mid M_j \cdots M_{\alpha}$. 所以回到 Definition 2.12, 这里 $M$ is left divisible by $r$ 一定是指 $M = M_0 \cdots M_{\alpha}$ 这种因式分解下么? 如果定义:

\[\begin{align} M_x &= M_j \cdots M_{\alpha} \\ M &= M_0 M_1 \cdots M_{j-1} M_x \end{align}\]

那么倒是可以能证明在这种因式分解下 M is left divisible by $N_{\varphi(\beta)} r_{\varphi(\beta)}$. 不对不对, 以 $A=(2, 3, 2, 3):(d_0, d_1, d_2, d_3), B=6:4$ 为例, 按照 Definition 2.13 可以推出 $A \circ B = (3, 2) : (?, ?)$, stride 推不出来…

参考

A note on the algebra of CuTe Layouts

—– 2024-10-16, A note on the algebra of CuTe Layouts, Date: January 16, 2024