On CuTe layouts

[work in progress. many citations missing, and many writings here could be wrong.]

[Hieu’s personal blog index]

Introduction

CuTe is a library in NVIDIA’s CUTLASS. CuTe makes GPU programming flexible and generalizable across different GPU architectures, such as Volta, Ampere, and Hopper. Thanks to these benefits, CuTe plays the central role in the development of many efficient deep learning libraries, including but not limited to FlashAttention, FasterTransformer, and xFormers.

CuTe’s benefits stem from the ingenious design of its central concept: layout. A layout represents a map from a $D$-dimensional coordinate to an integer. Despite this very basic intuition, layouts accompanied by the operations on them have significantly improved the experience of writing performant GPU programs.

Motivations

Despite the benefits of CuTe, we find CuTe’s original documentation somewhat inadequate, especially at building a foundational understanding of CuTe’s concepts. For instance, the complement operation is a fundamental one with CuTe layout, but the doc does not explain how to compute a complement. While it is okay to treat these low-level operations as black boxes and just move on writing good GPU programs with CuTe, we do feel the need to have a rigorous foundation of the concepts when writing our programs.

Goal. Building such rigurous foundation is the goal of this blog post. In particular, we attempt to formalize the definition of CuTe layout and a few of its accompanying operations, most importantly complementation, composition, and logical division.

Our source of inspiration. This blog post is not the first attempt at building foundational understanding of CuTe. We are inspired by Jay Shah’s note on the same topic. Jay’s note elicits the concepts around layouts very well, perhaps better than what we are doing, and we recommend the note to anyone who wants to deepen their undersatnding of CuTe.

In this post, we provide a different treatment of CuTe’s layout. Jay’s note states concise standards to test which layouts are “admissable for complementation” or “composition”. In order to derive such standards, they introduced some new concepts, such as sorting the modes of layouts. A consequence of their new standards anc concepts is that some layouts are not admissable for complement or composition. Ruling out such layouts makes sense practically, since we rarely encounter them while programming, but doing so means that the theory constructed around layouts does not faithfully reflect the behaviors of layouts in the CuTe library. We do not say such is a bad consequence. Instead, augmented by the intuitions harvested from Jay’s note, here we provide an algorithmic treatment of layouts, which we hope more closely reflects their behaviors in the CuTe library.

Our roadmap. We treat a layout as a way to represent function $f: \mathbb{N} \to \mathbb{N}$ that maps an integer to a multilinear function’s output defined by the layout. We explain that obtaining such a function from a layout is straightforward, and develop an algorithm to find a layout corresponding a certain function (it’s a deterministic algorithm, but it does not elicit a concise standard like Jay’s treatment). Then, using the correspondence between layouts and these functions via an algorithm, and then use this correspondence to construct layout operations such as complementation, composition, and logical division.

Other layout operations, such as coalescing, left inverse, and right inverse, are relatively easy to define via these layout functions, so we do not cover them here.

Layout

There are some terminologies associated with the definition of layout:

• Each element in the size tuple is sometimes called an extent of $L$.

• $\text{size}(L)$ is simply the product of all extents of $L$: $\text{size}(L) = n_0 n_1 \cdots n_{D-1}$.

• The tuple $S$ is called the layout’s stride.

• The maximum offset that the layout can represent, i.e., $1 + \sum_{i=0}^{D-1}s_i (n_i - 1)$, is the layout’s cosize and is denoted by $\text{cosize}(L)$.

• The pair $(n_i, s_i)$, sometimes written $(n_i) : (s_i)$, is called the $i^\text{th}$ mode of $L$.

In the next section, we will study the canonical functions of layouts, which represent how to map coordinates to offsets in layouts.

Canonical functions

For brevity, when clear from context, we drop the word “canonical”, and might also use the overloaded notation:

Other than the canonical multivariate function, we are also interested in the canonical singlevariate function of a layout. The singlevariate function is constructed from the layout’s multivariate function via the natural isomorphism between $[0, n_0 n_1 \cdots n_{D-1})$ and $[0, n_0) \times [0, n_1) \times \cdots \times [0, n_{D - 1})$. Intuitively, this isomorphism is the enumerating of the points on the integral lattice $[0, n_0) \times [0, n_1) \times \cdots \times [0, n_{D - 1})$. Here comes the tricky piece: this isomorphism is itself a layout:

As $\text{MultiToSingle}$ is an isomorphism, its invert $\text{SingleToMulti}$ is mathematically well-defined. Following is a self-contained formula for $\text{MultiToSingle}$ and $\text{SingleToMulti}$:

Using the $\text{SingleToMulti}$ function above, we can define a layout’s canonical singlevariate function:

Similar to our treatment of canonical multivariate functions, we also drop the terms “canonical” when clear from context. We might also write $L(x)$ instead of $f_L(x)$.

Given any layout $L$, we can easily construct its singlevariate function. The reverse question is much less trivial: for which functions $f: \mathbb{N} \to \mathbb{N}$ there is a layout whose singlevariate function is $f$?

In the next section, we will discuss a general question regarding the canonical functions. For a function $f: [0, M) \to \mathbb{N}$, we say that $L$ admits $f$ if $L(x) = f(x)$ for all $x \in [0, M)$. Under that definition, which function $f: \mathbb{N} \to \mathbb{N}$ is admitted by a layout?

It can be seen from the definition that mode-coalescing does not change the layout’s singlevariate function.

Later in this note, we will discuss an algorithm to determine which functions can be expressed by a layout.

Layout extensions

Now that we have discussed the canonical functions, we can discuss a layout’s extension. This concept is introduced in Jay’s note. While the concept is not presented in CuTe’s original definition, we find it very natural once we have understood the basics of layouts. More importantly, the concept of extension serves to make the mathemetical proofs later in this blog post consistent, and enhances our understanding of layouts.

The multivariate function of $L_\text{ext}$ is defined naturally based on that of $L$, except that the last coordinate can take any value, instead of being restricted to $[0, n_{D-1})$:

The singlevariate function $L_\text{ext}$ is therefore:

What function $f: \mathbb{N} \to \mathbb{N}$ can be admitted by a layout?

            0 = f(0)     | f(n_0)     | f(2*n_0)   | ...
          s_0 = f(1)     | f(n_0+1)   | f(2*n_0+1) | ...
          ...            | ...        | ...        | ...
(n_0-1) * s_0 = f(n_0-1) | f(2*n_0-1) | f(3*n_0-1) | ...

Concatenation

With that definition, it is relatively easy to check that we have a closed formula of layout concatentation. In particular, if:

Then their concatenation is simply obtained by concatenating their sizes and strides:

We also have the following summary regarding the size and cosize of layout concatenations:

As simple as the concatenation operation is, it plays an important role in defining two much more complex yet crucial operations on layouts: complementation and composition. We discuss them in the next sections.

Complement

Note that CuTe’s original definition of complementation specifies the following conditions:

1. $\text{cosize}((A, B)) \leq M$.

2. $\text{size}(B) \geq \left\lceil \dfrac{M}{\text{size}(A)} \right\rceil$.

3. $B(i) \neq A(j)$ for all $i \in [1, \text{size}(B))$ and $j \in [0, \text{size}(A))$.

Not all layouts have a complement. In particular, the bijection requirement in our definition rules out all layout $A$ whose singlevariate function is not injective. Jay’s note offers a method to check whether a layout has a complement layout.

The function-to-layout Algorithm offers a procedure to find $\text{Complement}(A, M)$ for any layout $A$ and positive integer $M$, or to tell that such complement does not exist.

Indeed, the idea is to determine $B$’s singlevariate function based on the given conditions, and then use function-to-layout Algorithm to find $B$, or to tell if there is no such $B$.

Composition

As for complementation, not all pairs of layouts can be composed. To determine two layouts’ composition, we first determine the composite of their single variate function, and then use the function-to-layout Algorithm to find the composition or to conclude that such composition does not exist.

It turns out that by using only arguments with function composition, we can derive the left distribution of composition with respect to concatenation. This is a different approach compared to CuTE’s original documentation, which uses the left distribution property to define composition.

Logical division

We find this way of defining logical division to be very useful, as it offers very assuring to reason about the logical division via $A$ and $B$’s singlevariate functions.

In particular, since $A$ and $B$ are layouts:

• Their singlevariate functions can be easily defined.

• Then, $C = \text{Complement}(B, \text{size}(A))$ is either well-defined or does not exist.

• Then, $\text{Concat}(B, C)$ is well-defined.

• Finally, $A \circ \text{Concat}(B, C)$ is either well-defined or does not exist.

This way of reasoning does not only offer an assuring mathemtical definition of logical division, but also gives an algorithm to evaluate two layout’s logical division. While we cannot assert that this is the way layouts are implemented in CuTe, we implement logical division operation in Python to help with studying the concepts.

Python implementation

Here, we provide the Python implementation for the basic layout operations. We have only tested the code on very basic cases, so if you use it, please proceed carefully.

r"""Simple experiments with CuTe layout."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import numpy.typing as npt


@dataclass
class Layout:
    r"""A CuTe layout."""

    N: npt.ArrayLike
    S: npt.ArrayLike

    def __post_init__(self):
        self._cum_prod = np.concatenate([[1], np.cumprod(self.N[:-1])])

    def __repr__(self) -> str:
        n_str = ",".join([str(n) for n in self.N])
        s_str = ",".join([str(s) for s in self.S])
        return f"[layout] ({n_str}) : ({s_str})"

    def size(self) -> int:
        r"""The layout's size is the product of all its dimensions."""
        return np.prod(self.N)

    def single2multi(self, x: int) -> tuple[int, ...]:
        r"""1-D coordinate to multi coordinate."""
        return tuple(((x // self._cum_prod) % self.N).tolist())

    def multi2single(self, *args) -> int:
        r"""Multi coordinate to 1-D coordinate."""
        return np.dot(self._cum_prod, np.array(args))

    def f_single(self, x: int) -> int:
        r"""Evaluate the layout's singlevariate function."""
        return np.array(self.single2multi(x)).dot(self.S)


def find_layout(f: npt.ArrayLike) -> Layout | None:
    r"""Returns a layout admitting `f` or `None` if no such layout exists."""

    if f[0] != 0:
        return None

    m = np.size(f)
    s_0 = f[1]

    # this is the layout (m) : (s_0)
    m_s_0 = np.arange(0, m * s_0, s_0)

    # check if (n_0) : (s_0) is a solution. if yes, return
    if np.all(m_s_0 == f):
        return Layout(N=np.array([m]), S=np.array([s_0]))

    # unique index t s.t.: f[i] = i * s_0 for i = 0, ..., t-1 but f[t] != i * t
    t = np.where(m_s_0 != f)[0][0]

    for n_0 in range(t, 0, -1):
        # check if f is consistent with (n_0) : (s_0) as the first mode
        tgt_sz = (m + n_0 - 1) // n_0 * n_0
        pad_sz = tgt_sz - m
        pad_f = np.pad(f, [(0, pad_sz)]).reshape(-1, n_0).transpose()

        # row diff for all columns, except for the last one
        row_d = pad_f[1:, :-1] - pad_f[:-1, :-1]

        # row diff in the last column. need to remove the padded zeros
        last_f = pad_f[: m % n_0, -1]
        last_d = last_f[1:] - last_f[:-1]

        n_0_consistent = np.all(row_d[:, :-1] == s_0) and np.all(last_d == s_0)
        if n_0_consistent:
            r = find_layout(pad_f[0, :])
            if r is None:
                return None
            return Layout(N=np.concatenate([[n_0], r.N]),
                          S=np.concatenate([[s_0], r.S]))

    return None


def test_find_layout():
    """Ensure that find_layout finds the correct layout."""

    def _test_case(n: npt.ArrayLike, s: npt.ArrayLike):
        layout_1 = Layout(N=n, S=s)
        f = [layout_1.f_single(x) for x in range(layout_1.size())]
        layout_2 = find_layout(f)
        assert np.all(layout_1.N == layout_2.N), (
            f"{layout_1.N=} but {layout_2.N=}.")
        assert np.all(layout_1.S == layout_2.S), (
            f"{layout_1.S=} but {layout_2.S=}.")

    _test_case([3, 5, 7], [4, 9, 8])
    _test_case([3, 2], [10, 13])
    _test_case([9], [4])


def test_layout_function():
    """Simple opreations."""
    layout = Layout(N=[3, 4], S=[1, 3])

    assert layout.single2multi(5) == (2, 1)
    assert layout.multi2single(2, 1) == 5
    assert layout.multi2single(2, 3) == 11

    for x in range(layout.size()):
        assert x == layout.multi2single(*layout.single2multi(x))
        assert x == layout.f_single(x)


test_find_layout()
test_layout_function()