leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 27 Packages Projects Releases Wiki Activity

CUTLASS 3.0.0 (#786)

Browse Source 
* CUTLASS 3.0.0
This commit is contained in:
Vijay Thakkar
2023-01-23 20:55:28 -05:00
committed by GitHub
parent 66d9cddc83
commit 277bd6e537
377 changed files with 76396 additions and 1186 deletions
Download Patch File Download Diff File Expand all files Collapse all files
media/docs/cute/00_quickstart.md
+75
View File
@@ -0,0 +1,75 @@
# Getting Started With CuTe
CuTe is a collection of C++ CUDA template abstractions for defining and operating on hierarchically multidimensional layouts of threads and data. CuTe provides `Layout` and `Tensor` objects that compactly packages the type, shape, memory space, and layout of data, while performing the complicated indexing for the user. This lets programmers focus on the logical descriptions of their algorithms while CuTe does the mechanical bookkeeping for them. With these tools, we can quickly design, implement, and modify all dense linear algebra operations.
The core abstraction of CuTe are the hierarchically multidimensional layouts which can be composed with data arrays to represent tensors. The representation of layouts is powerful enough to represent nearly everything we need to implement efficient dense linear algebra. Layouts can also be combined and manipulated via functional composition, on which we build a large set of common operations such as tiling and partitioning.
## System Requirements
CuTe shares CUTLASS 3.0's software requirements,
including NVCC with a C++17 host compiler.
## Knowledge prerequisites
CuTe is a CUDA C++ library. It requires C++17
(the revision of the C++ Standard that was released in 2017).
Throughout this tutorial, we assume intermediate C++ experience.
For example, we assume that readers know
how to read and write templated functions and classes, and
how to use the `auto` keyword to deduce a function's return type.
We will be gentle with C++ and explain some things
that you might already know.
We also assume intermediate CUDA experience.
For example, readers must know
the difference between device and host code,
and how to launch kernels.
## Building Tests and Examples
CuTe's tests and examples build and run as part of CUTLASS's normal build process.
CuTe's unit tests live in the [`test/unit/cute`](../../../test/unit/cute) subdirectory.
Its examples live in the [`examples/cute`](../../../examples/cute) subdirectory.
## Library Organization
CuTe is a header-only C++ library, so there is no source code that needs building. Library headers are contained within the top level [`include/cute`](../../../include/cute) directory, with components of the library grouped by directories that represent their semantics.
| Directory | Contents |
|------------------------|------------------------|
| [`include/cute`](../../../include/cute) | Each header in the top level corresponds to one of the fundamental building blocks of CuTe, such as [`Layout`](../../../include/cute/layout.hpp) or [`Tensor`](../../../include/cute/tensor.hpp). |
| [`include/cute/container`](../../../include/cute/container) | Implementations of STL-like container objects, such as tuple, array, aligned array, and array views. |
| [`include/cute/numeric`](../../../include/cute/numeric) | Templates that handle nonstandard floating-point types, unsigned integers, complex numbers, and integer sequence - like fundamental numeric data types. |
| [`include/cute/algorithm`](../../../include/cute/algorithm) | Implementations of utility algorithms such as copy, fill, and clear that automatically leverage architecture-specific features if available. |
| [`include/cute/arch`](../../../include/cute/arch) | Wrappers for architecture-specific matrix-matrix multiply and copy instructions. |
| [`include/cute/atom`](../../../include/cute/atom) | Meta-information for instructions in `arch` and utilities like partitioning and tiling.
## Tutorial
This directory contains a CuTe tutorial in Markdown format.
The file
[`0x_gemm_tutorial.md`](./0x_gemm_tutorial.md)
explains how to implement dense matrix-matrix multiply using CuTe components.
It gives a broad overview of CuTe and thus would be a good place to start.
Other files in this directory discuss specific parts of CuTe.
* [`01_layout.md`](./01_layout.md) describes `Layout`, CuTe's core abstraction.
* [`02_layout_operations.md`](./02_layout_operations.md) describes more advanced `Layout` operations and the CuTe layout algebra.
* [`03_tensor.md`](./03_tensor.md) describes `Tensor`,
a multidimensional array abstraction which composes `Layout`
with an array of data.
* [`04_algorithms.md`](./04_algorithms.md) summarizes CuTe's
generic algorithms that operate on `Tensor`s.
* [`0t_mma_atom.md`](./0t_mma_atom.md) demonstrates CuTe's meta-information and interface to our GPUs'
architecture-specific Matrix Multiply-Accumulate (MMA) instructions.
* [`0x_gemm_tutorial.md`](./0x_gemm_tutorial.md) provides a walkthrough of building a GEMM from scratch using CuTe.
* [`0y_predication.md`](./0y_predication.md) explains what to do
if a tiling doesn't fit evenly into a matrix.
media/docs/cute/01_layout.md
+254
View File
@@ -0,0 +1,254 @@
# CuTe Layouts
## Layout
This document describes `Layout`, CuTe's core abstraction.
A `Layout` maps from (a) logical coordinate space(s)
to a physical index space.
`Layout`s present a common interface to multidimensional array access
that abstracts away the details of how the array's elements are organized in memory.
This lets users write algorithms that access multidimensional arrays generically,
so that layouts can change, without users' code needing to change.
CuTe also provides an "algebra of `Layout`s."
`Layout`s can be combined and manipulated
to construct more complicated layouts
and to partition them across other layouts.
This can help users do things like partition layouts of data over layouts of threads.
## Layouts and Tensors
Any of the `Layout`s discussed in this section can be composed with data -- a pointer or an array -- to create a `Tensor`. The responsibility of the `Layout` is to define valid coordinate space(s) and, therefore, the logical shape of the data and map those into an index space. The index space is precisely the offset that would be used to index into the array of data.
For details on `Tensor`, please refer to the
[`Tensor` section of the tutorial](./03_tensor.md).
## Shapes and Strides
A `Layout` is a pair of `Shape` and `Stride`.
Both `Shape` and `Stride` are `IntTuple` types.
### IntTuple
An `IntTuple` is an integer or a tuple of `IntTuple`s.
This means that `IntTuple`s can be arbitrarily nested.
Operations defined on `IntTuple`s include the following.
* `get<I>(IntTuple)`: The `I`th element of the `IntTuple`. Note that `get<0>` is defined for integer `IntTuples`.
* `rank(IntTuple)`: The number of elements in an `IntTuple`. An int has rank 1, a tuple has rank `tuple_size`.
* `depth(IntTuple)`: The number of hierarchical `IntTuple`s. An int has depth 0, a tuple has depth 1, a tuple that contains a tuple has depth 2, etc.
* `size(IntTuple)`: The product of all elements of the IntTuple.
We write `IntTuple`s with parenthesis to denote the hierarchy. E.g. `6`, `(2)`, `(4,3)`, `(3,(6,2),8)` are all `IntTuple`s.
## Layout
A `Layout` is then a pair of `IntTuple`s. The first defines the abstract *shape* of the layout and the second defines the *strides*, which map from coordinates within the shape to the index space.
As a pair of `IntTuple`s, we can define many similar operations on `Layout`s including
* `get<I>(Layout)`: The `I`th sub-layout of the `Layout`.
* `rank(Layout)`: The number of modes in a `Layout`.
* `depth(Layout)`: The number of hierarchical `Layout`s. An int has depth 0, a tuple has depth 1, a tuple that contains a tuple has depth 2, etc.
* `shape(Layout)`: The shape of the `Layout`.
* `stride(Layout)`: The stride of the `Layout`.
* `size(Layout)`: The logical extent of the `Layout`. Equivalent to `size(shape(Layout))`.
### Hierarchical access functions
`IntTuple`s and thus `Layout`s can be arbitrarily nested.
For convenience, we define versions of some of the above functions
that take a sequence of integers, instead of just one integer.
This makes it possible to access elements
inside of nested `IntTuple` or `Layout`.
For example, we permit `get<I...>(x)`, where `I...` here
and throughout this section is a "C++ parameter pack"
that denotes zero or more (integer) template arguments.
That is, `get<I0,I1,...,IN>(x)` is equivalent to
`get<IN>(` $\dots$ `(get<I1>(get<I0>(x)))` $\dots$ `))`,
where the ellipses are pseudocode and not actual C++ syntax.
These hierarchical access functions include the following.
* `rank<I...>(x) := rank(get<I...>(x))`. The rank of the `I...`th element of `x`.
* `depth<I...>(x) := depth(get<I...>(x))`. The depth of the `I...`th element of `x`.
* `size<I...>(x) := size(get<I...>(x))`. The size of the `I...`th element of `x`.
### Vector examples
Then, we can define a vector as any `Shape` and `Stride` pair with `rank == 1`.
For example, the `Layout`
```
Shape: (8)
Stride: (1)
```
defines a contiguous 8-element vector.
Similarly, with a stride of `(2)`,
the interpretation is that the eight elements
are stored at positions 0, 2, 4, $\dots$.
By the above definition, we *also* interpret
```
Shape: ((4,2))
Stride: ((1,4))
```
as a vector, since its shape is rank 1. The inner shape describes a 4x2 layout of data in column-major order, but the extra pair of parenthesis suggest we can interpret those two modes as a single 1-D 8-element vector instead. Due to the strides, the elements are also contiguous.
### Matrix examples
Generalizing, we define a matrix as any `Shape` and `Stride` pair with rank 2. For example,
```
Shape: (4,2)
Stride: (1,4)
0 4
1 5
2 6
3 7
```
is a 4x2 column-major matrix, and
```
Shape: (4,2)
Stride: (2,1)
0 1
2 3
4 5
6 7
```
is a 4x2 row-major matrix.
Each of the modes of the matrix can also be split into *multi-indices* like the vector example.
This lets us express more layouts beyond just row major and column major. For example,
```
Shape: ((2,2),2)
Stride: ((4,1),2)
0 2
4 6
1 3
5 7
```
is also logically 4x2, with a stride of 2 across the rows but a multi-stride down the columns.
Since this layout is logically 4x2,
like the column-major and row-major examples above,
we can _still_ use 2-D coordinates to index into it.
## Constructing a `Layout`
A `Layout` can be constructed in many different ways.
It can include any combination of compile-time (static) integers
or run-time (dynamic) integers.
```c++
auto layout_8s = make_layout(Int<8>{});
auto layout_8d = make_layout(8);
auto layout_2sx4s = make_layout(make_shape(Int<2>{},Int<4>{}));
auto layout_2sx4d = make_layout(make_shape(Int<2>{},4));
auto layout_2x4 = make_layout(make_shape (2, make_shape (2,2)),
make_stride(4, make_stride(1,2)));
```
## Using a `Layout`
The fundamental use of a `Layout` is to map between logical coordinate space(s) and index space. For example, to print an arbitrary rank-2 layout, we can write the function
```c++
template <class Shape, class Stride>
void print2D(Layout<Shape,Stride> const& layout)
{
for (int m = 0; m < size<0>(layout); ++m) {
for (int n = 0; n < size<1>(layout); ++n) {
printf("%3d ", layout(m,n));
}
printf("\n");
}
}
```
which produces the following output for the above examples.
```
> print2D(layout_2sx4s)
0 2 4 6
1 3 5 7
> print2D(layout_2sx4d)
0 2 4 6
1 3 5 7
> print2D(layout_2x4)
0 2 1 3
4 6 5 7
```
The multi-indices within the `layout_4x4` example are handled as expected and interpreted as a rank-2 layout.
Note that for `layout_1x4`, we're using a 1-D coordinate for a 2-D multi-index in the second mode. In fact, we can generalize this and treat all of the above layouts as 1-D layouts. For instance, the following `print1D` function
```c++
template <class Shape, class Stride>
void print1D(Layout<Shape,Stride> const& layout)
{
for (int i = 0; i < size(layout); ++i) {
printf("%3d ", layout(i));
}
}
```
produces the following output for the above examples.
```
> print1D(layout_8s)
0 1 2 3 4 5 6 7
> print1D(layout_8d)
0 1 2 3 4 5 6 7
> print1D(layout_2sx4s)
0 1 2 3 4 5 6 7
> print1D(layout_2sx4d)
0 1 2 3 4 5 6 7
> print1D(layout_2x4)
0 4 2 6 1 5 3 7
```
This shows explicitly that all of the layouts are simply folded views of an 8-element array.
## Summary
* The `Shape` of a `Layout` defines its coordinate space(s).
* Every `Layout` has a 1-D coordinate space.
This can be used to iterate in a "generalized-column-major" order.
* Every `Layout` has a R-D coordinate space,
where R is the rank of the layout.
These spaces are ordered _colexicographically_
(reading right to left, instead of "lexicographically,"
which reads left to right).
The enumeration of that order
corresponds to the 1-D coordinates above.
* Every `Layout` has an h-D coordinate space where h is "hierarchical." These are ordered colexicographically and the enumeration of that order corresponds to the 1-D coordinates above. An h-D coordinate is congruent to the `Shape` so that each element of the coordinate has a corresponding element of the `Shape`.
* The `Stride` of a `Layout` maps coordinates to indices.
* In general, this could be any function from 1-D coordinates (integers) to indices (integers).
* In `CuTe` we use an inner product of the h-D coordinates with the `Stride` elements.
media/docs/cute/02_layout_operations.md
+710
View File
@@ -0,0 +1,710 @@
# CuTe Layout Operations
CuTe provides an "algebra of `Layout`s."
`Layout`s can be combined and manipulated
to construct more complicated `Layout`s.
This includes tiling and partitioning `Layout`s across other `Layout`s.
In this section, we explain some of these core operations in detail.
## How do I print CuTe objects on host or device?
CuTe comes with different ways to print CuTe objects.
You can print human-readable text,
or you can print LaTeX commands for generating
a beautifully formatted and colored table
describing the CuTe object.
Both of these can be helpful for reasoning about or debugging
layouts, copy atoms, or matrix multiply atoms
(don't worry, we'll explain all of these things in this tutorial).
CuTe's print functions work on either host or device.
Note that on device, printing is expensive.
Even just leaving print code in place on device,
even if it is never called
(e.g., printing in an `if` branch that is not taken at run time),
may generate slower code.
Thus, be sure to remove code that prints on device after debugging.
The following code examples assume that you have a
`using namespace cute;` statement in scope.
### Printing human-readable text
The `cute::print` function has overloads for almost all CuTe types, including Pointers, Layout, Shape, Stride, and Tensors. When in doubt, try calling `print` on it. You might also only want to print on thread 0 of each thread block, or block 0 of the grid. The `thread0()` function returns true only for global thread 0 of the kernel. A typical idiom for printing CuTe objects to print only on thread 0 of block 0.
```c++
if (thread0()) {
print(some_cute_object);
}
```
Some algorithms do different things on different threads or blocks,
so you might sometimes need to print on threads or blocks other than zero.
The header file
[`cute/util/debug.hpp`](../../../include/cute/util/debug.hpp),
among other utilities,
includes the function `bool thread(int tid, int bid)`
that returns `true` if running on thread `tid` and block `bid`.
Some CuTe types have special printing functions that use a different output format.
For example, `print_layout` can display a rank-2 layout in a table
(using plain text formatting).
It has an overload taking a rank-2 matrix layout and a thread layout,
that displays a table with the mapping between threads and values.
Some CuTe types might not have overloads for `print`,
but there are other ways to print their contents.
For example, copy atoms and mma atoms
(see elsewhere in this tutorial)
have a `print_all()` member function.
### Printing LaTeX output
The `cute::print_latex` function works like `cute::print`,
but prints LaTeX commands that you can use
to generate a nicely formatted and colored table.
## Fundamental types
### Layout and its components
This directory includes
[an overview of CuTe's fundamental types for describing layouts](./01_layout.md).
#### Tuple
CuTe starts with a Tuple, which is a finite ordered list of zero or more elements.
In C++, we identify a Tuple with the
[`cute::tuple` class](../../../include/cute/container/tuple.hpp).
`cute::tuple` behaves like `std::tuple`, but it works on device or host,
and it imposes restrictions on its template arguments for performance and simplicity.
#### IntTuple
CuTe then defines an IntTuple as either an integer, or a Tuple of IntTuple.
This recursive definition lets us build arbitrarily nested layouts.
In C++, we identify an IntTuple with [`IntTuple`](../../../include/cute/int_tuple.hpp),
which is just an alias of `cute::tuple`.
Any of the following are thus valid template arguments of IntTuple.
1. "Run-time integers" (or "static integers")
are just ordinary integral types like `int` or `size_t`.
2. "Compile-time integers" include `std::integral_constant`
or subclasses of it that CuTe defines,
such as `Int<Value>` (see below).
These types all have in common
that the value is encoded in the type itself
(as a public `static constexpr value` member).
CuTe defines aliases `_1`, `_2`, `_3` etc.
to the types `Int<1>`, `Int<2>`, `Int<3>` etc.
3. `IntTuple` with any valid template arguments.
CuTe reuses IntTuple for many different things,
including Shape, Stride, Step, and Coord
(see [`include/cute/layout.hpp`](../../../include/cute/layout.hpp)).
In C++, Shape, Stride, Step, and Coord are all aliases for IntTuple.
### Layout
A Layout is a tuple of (Shape, Stride).
Semantically, it implements a mapping from
a "logical" Shape-shaped (multidimensional) index,
to a "physical" 1-D index into an array.
Here is an example of a 2 x 3 array with static strides (3, 1).
```c++
Layout layout = make_layout(make_shape (_2{}, _3{}),
make_stride(_3{}, _1{}));
print_layout(layout);
for (int i = 0; i < size(layout); ++i) {
print(layout(i));
print(", ");
}
print("\n");
print(layout(1, 1));
print("\n");
```
This code produces the following text output.
```text
(_2,_3):(_3,_1)
0 1 2
+---+---+---+
0 | 0 | 1 | 2 |
+---+---+---+
1 | 3 | 4 | 5 |
+---+---+---+
0, 3, 1, 4, 2, 5,
4
```
`print(layout(1, 1))` prints the mapping of
the logical 2-D coordinate (0,1) to 1-D index, which is 4.
You can see that from the table,
which shows the left logical index as the "row,"
and the right logical index as the "column."
### Underscore (`_`)
An Underscore is a special type used for array slices. The underscore punctuation `_` is a constant instance of Underscore. It acts like `:` (the colon punctuation) in Python or Fortran array slices. See [`include/cute/underscore.hpp`](../../../include/cute/underscore.hpp).
### Tile
"A Tile is not a Layout, it's a tuple of Layouts or Tiles or Underscores."
See [`include/cute/tile.hpp`](../../../include/cute/tile.hpp).
The algebraic layout operations discussed below are defined on `Layout`s, but `Tile` allows these operations to recurse and to be applied to sublayouts or particular modes of a given Layout. These are referred to as by-mode operations.
See the section on "Logical Divide" to see an example of using `Tile` to extract portions of a row-mode and portions of a column-mode independently.
## Layout definitions and operations
### Layouts are functions from integers (logical 1-D coordinate) to integers (1-D index)
The `for` loop in the above print example shows how CuTe identifies 1-D coordinates with a column-major layout of logical 2-D coordinates. Iterating from `i = 0` to `size(layout)` (which is 6), and indexing into our layout with the single integer coordinate `i`, traverses the layout in column-major fashion, even though this is a row-major layout. You can see this from the output of the `for` loop (0, 3, 1, 4, 2, 5). CuTe calls this index `i` a "1-D coordinate," versus the "natural coordinate," which would be the logical 2-D coordinate.
If you're familiar with the C++23 feature `mdspan`,
this is an important difference between
`mdspan` layout mappings and CuTe `Layout`s.
`mdspan` layout mappings are *one way*:
they always take a multidimensional logical coordinate,
and they return an integer offset.
Depending on the strides,
the offset may skip over elements of the physical 1-D array.
Thus, `mdspan`'s offset does NOT mean the same thing as
the 1-D logical coordinate `i` in the `for` loop above.
You can iterate correctly over any CuTe `Layout`
by using the 1-D logical coordinate.
`mdspan` doesn't have an idea of a 1-D logical coordinate.
### Rank, depth, size, cosize
*Rank*: the tuple size of the layout's shape.
*Depth*: the depth of the layout's shape. A single integer has depth 0. A tuple has depth 1 + the max depth of its components.
*Size*: Size of the shape; size of the domain of the function. This is the product of all extents in the layout's shape.
*Cosize*: Size of the function's codomain (not necessarily the range); for a layout A, A(size(A) - 1) + 1. (Here, we use size(A) - 1 as a 1-D logical coordinate input.)
### Layout compatibility
We say that layouts A and B are *compatible* if their shapes are compatible. Shape A is compatible with shape B if any natural coordinate of A is also a valid coordinate for B.
### Flatten
The `flatten` operation "un-nests" a potentially nested Layout. For example,
```c++
Layout layout = Layout<Shape <Shape <_4, _3>, _1>,
Stride<Stride<_3, _1>, _0>>{};
Layout flat_layout = flatten(layout);
```
results in `flat_layout` having the following type
```text
Layout<Shape<_4, _3, _1>, Stride<_3, _1, _0>>
```
and
```c++
Layout layout = Layout<Shape <_4, Shape <_4, _2>>,
Stride<_4, Stride<_1, _16>>>{};
Layout flat_layout = flatten(layout);
```
results in `flat_layout` having the following type
```text
Layout<Shape<_4, _4, _2>, Stride<_4, _1, _16>>
```
Hierarchical Layouts and flattening let us reinterpret tensors in place as matrices, matrices as vectors, vectors as matrices, etc. This lets us implement arbitrary tensor contractions as batched matrix multiply, by combining the contraction modes into a single mode, and combining the A, B, C, and "batch" modes as needed to reach the desired form.
### Coalesce
The `coalesce` operation first flattens the layout, then combines all the modes that are possible to combine, starting with mode 0 (the leftmost mode) and moving right. If all the modes can be combined, then this results in a 1-D layout expressing what array elements the original layout accesses.
For example,
```text
layout: (_2,(_1,_6)):(_1,(_6,_2))
coalesce(layout): _12:_1
```
What does it mean to "combine" modes? In the above example, the flattened layout is (2, 1, 6) : (1, 6, 2).
1. If we look at the leftmost two modes, this is just a vector of length 2 and stride 1. The middle mode has extent 1, so the corresponding stride 6 would not be observed anyway. This leaves us with (2, 6) : (1, 2).
2. The intermediate result (2, 6) : (1, 2) is just a 2 x 6 column-major matrix, which can be coalesced into a vector of length 12 and stride 1.
More formally, "combining all the modes" means a left fold, where the binary operation that combines two modes has three cases.
1. If the leftmost layout is s1:d1, and the next layout is 1:d0, then combine into s1:d1. This generalizes Step 1 above. If a mode has extent 1, we can't observe its stride, so we can skip the mode.
2. If the leftmost layout is 1:d1, and the next layout is s0:d0, then combine into s0:d0. Again, if a mode has extent 1, we can't observe its stride, so we can skip the mode.
3. If the leftmost layout is s1:d1, and the next layout is s0 : s1*d1, then combine into s0 * s1 : d1. This generalizes Step 2 above. One can call this "noticing a column-major layout sequence."
That's it! For example, the result of coalescing the row-major layout (2, 2) : (2, 1) is (2, 2) : (2, 1), the same layout, because none of the above three cases applies.
### Complement
#### Definition
The complement B of a layout A with respect to an integer M satisfies the following properties.
1. $A$ and $B$ are *disjoint*: $A(x) \neq B(x)$ for all $x \neq 0$ in the domain of $A$.
2. B is *ordered*: $`B(x-1) < B(x)`$ for all $x$ in $\{0, 1, \dots, size(B) - 1\}$.
3. B is *bounded* by M: $size(B) \geq M / size(A)$, and $cosize(B) \leq floor(M / cosize(A)) * cosize(A)$.
Regarding disjointness: we need to specify $x \neq 0$ because CuTe layouts are linear. That is, if the domain is nonempty, the range always contains zero.
Regarding the ordered property: CuTe layouts are hierarchically strided, so this implies that if size(B) is nonzero, then the strides of B are all positive.
#### Examples
complement(4:1, 24) is 6:4.
1. The result is disjoint of 4:1, so it must have a stride of at least 4 (since it includes 0, but must skip over 1, 2, 3).
2. The size of the result is $\geq 24 / 4 = 6$. (This plus Step (1) means that the cosize is at least 24.)
3. The cosize of the result is $\leq (24 / 4) * 4 = 24$. (This plus Step (2) means that the cosize is exactly 24.)
4. The only (1-D) layout with size 6 and cosize 24 is 6:4.
complement(6:4, 24) is 4:1.
1. 4:1 is disjoint of 6:4, but so is s:d
for any s > 0 and d > 20.
2. The size of the result is $\geq 24 / 6 = 4$.
3. The cosize of the result is $\leq (24 / 21) * 21 = 21$.
4. The stride cannot be greater than 20
(else (2) would contradict (3)),
so it must be less than 4.
5. This leaves 4:1 by elimination.
### Composition
Layouts are functions, so composition of layouts is just composition of functions. The composition $A \circ B$ means "apply the layout B first, then treat the result as a 1-D logical coordinate input to the layout A, and apply A to it." Very often, this composition can be represented as another Layout.
#### Rules for computing composition
Both humans and CuTe compute composition using the following rules.
1. $A \circ B$ has a shape that is compatible with B. In function composition, the rightmost function defines the domain. For `Layout`s this means that any valid coordinate for $B$ can also be used as a coordinate for $A \circ B$.
2. Concatenation: A layout can be expressed as the concatenation of its sublayouts. We denote concatenation with parentheses: $B = (B_0,B_1,...)$. The CuTe function `make_layout`, when given zero or more `Layout`s, concatenates them.
3. Composition is (left-)distributive with concatenation: $A \circ B = A \circ (B0, B1, ...) = (A \circ B0, A \circ B1, ...)$.
4. "Base case": For layouts $A = a : b$ and $B = c : d$ with integral shape and stride, $A \circ B = R = c : (b * d)$.
5. By-mode composition: Let $\langle B, C \rangle$ (angle brackets, not parentheses)
denote a tuple of two layouts B and C, not their concatenation. Let A = (A0, A1).
Then, $A \circ \langle B, C \rangle = (A0, A1) \circ \langle B, C \rangle = (A0 \circ B, A1 \circ C)$.
This allows the application of composition independently to sublayouts of $A$.
#### Examples: Reshape a vector into a matrix
This section gives two composition examples. Both start with a vector with layout $20:2$ (that is, the vector has 20 elements, and the stride between each is 2). They compose this vector with a 4 x 5 matrix layout. This effectively "reshapes" the vector in place into a matrix.
##### Example 1
$20:2 \circ (4,5) : (1,4)$.
This describes interpreting the vector $20:2$
as a 4 x 5 column-major matrix.
The resulting layout has shape $(4,5)$,
because in function composition,
the rightmost function defines the domain.
What are the strides?
1. A layout can be expressed as the concatenation of its sublayouts,
so $(4,5) : (1,4)$ is $(4:1, 5:4)$.
2. Composition is distributive, so
$20:2 \circ (4:1, 5:4)$ is $(20:2 \circ 4:1, 20:2 \circ 5:4)$.
3. $20:2 \circ 4:1$ has shape 4 (rightmost function defines the domain)
and stride $2 = 2 \cdot 1$.
4. $20:2 \circ 5:4$ has shape 5 and stride $8 = 2 \cdot 4$.
5. Result: (4:2, 5:8), which by concatenation is (4,5) : (2,8).
#### Example 2
$20:2 \circ (4,5) : (5,1)$.
This describes interpreting the vector 20:2
as a 4 x 5 row-major matrix.
The resulting layout has shape $(4,5)$, just as before. What are the strides?
1. By deconcatenation, $(4,5) : (5,1)$ is $(4:5, 5:1)$.
2. Composition is distributive, so $20:2 \circ (4:5, 5:1)$ is $(20:2 \circ 4:5, 20:2 \circ 5:1)$.
3. $20:2 \circ 4:5$ has shape $4$ and stride $10 = 2 \cdot 5$.
4. $20:2 \circ 5:1$ has shape $5$ and stride $2 = 2 \cdot 1$.
5. Result: (4:10, 5:2), which by concatenation is (4,5) : (10,2).
### Product
CuTe includes four different kinds of layout products.
1. `logical_product`
2. `blocked_product`
3. `raked_product`
4. `tiled_product`
`logical_product(A, B)` results in a layout where each element of layout B
has been replaced by a "copy" of layout A.
The other three products offer variations of this idea.
#### Example: Tiled matrix
Suppose that I want to make a matrix consisting of 3 x 4 tiles
in a row-major arrangement,
where each tile is a 2 x 2 column-major matrix.
The Layout of each tile (tile) has Shape (2,2) and Stride (1,2).
The Layout of the "matrix of tiles" (`matrix_of_tiles`)
has Shape (3,4) and Stride (4,1).
##### Blocked product: the intuitive tiling
If I were to deduce by hand what the layout of the tiled matrix should be,
it would look like this.
| | (0,0) | (1,0) | (0,1) | (1,1) | (0,2) | (1,2) | (0,3) | (1,3) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (0,0) | 0 | 2 | 4 | 6 | 8 | 10 | 12 | 14 |
| (1,0) | 1 | 3 | 5 | 7 | 9 | 11 | 13 | 15 |
| (0,1) | 16 | 18 | 20 | 22 | 24 | 26 | 28 | 30 |
| (1,1) | 17 | 19 | 21 | 23 | 25 | 27 | 29 | 31 |
| (0,2) | 32 | 34 | 36 | 38 | 40 | 42 | 44 | 46 |
| (1,2) | 33 | 35 | 37 | 39 | 41 | 43 | 45 | 47 |
The row and column labels use the equivalence of 1-D logical coordinates and 2-D column-major coordinates. The left index in each pair is the row resp. column coordinate of the tile, while the right index in each pair is the row resp. column coordinate of the matrix-of-tiles. The resulting layout has Shape ((2, 3), (2, 4)), and Stride ((1, 16), (2, 4)), and the second mode can be coalesced. The Shape ((2, 3), (2, 4)) is hierarchical, but it is still rank-2 and can be drawn in 2D as above. Note how the row mode of the tile remains part of the row mode of the product, and the column mode of the tile remains a column mode of the product.
The above layout is what `blocked_product(tile, matrix_of_tiles)` produces.
A critical use case for blocked product is "tiling" an "atom"
(some tile that relates to a hardware feature) over a matrix.
```c++
Layout tile = Layout<Shape <_2,_2>,
Stride<_1,_2>>{};
Layout matrix_of_tiles = Layout<Shape <_3,_4>,
Stride<_4,_1>>{};
print_layout(blocked_product(tile, matrix_of_tiles));
```
##### Logical product
The logical product `logical_product(tile, matrix_of_tiles)`
results in Shape ((2, 2), (3, 4)) and Stride ((1, 2), (16, 4)).
| | (0,0) | (1,0) | (2,0) | (0,1) | (1,1) | (2,1) | (0,2) | (1,2) | (2,2) | (0,3) | (1,3) | (2,3) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (0,0) | 0 | 16 | 32 | 4 | 20 | 36 | 8 | 24 | 40 | 12 | 28 | 44 |
| (1,0) | 1 | 17 | 33 | 5 | 21 | 37 | 9 | 25 | 41 | 13 | 29 | 45 |
| (0,1) | 2 | 18 | 34 | 6 | 22 | 38 | 10 | 26 | 42 | 14 | 30 | 46 |
| (1,1) | 3 | 19 | 35 | 7 | 23 | 39 | 11 | 27 | 43 | 15 | 31 | 47 |
Note how the tile appears in the leftmost column and is reproduced
in each column in the same order as the matrix-of-tiles. That is,
the tile can be indexed through the first mode of the result and the
matrix-of-tiles can be indexed through the second mode.
```c++
Layout tile = Layout<Shape <_2,_2>,
Stride<_1,_2>>{};
Layout matrix_of_tiles = Layout<Shape <_3,_4>,
Stride<_4,_1>>{};
print_layout(logical_product(tile, matrix_of_tiles));
```
##### Raked product
The raked product `raked_product(tile, matrix_of_tiles)` results in
Shape ((3, 2), (4, 2)) and Stride ((16, 1), (4, 2)).
| | (0,0) | (1,0) | (2,0) | (3,0) | (0,1) | (1,1) | (2,1) | (3,1) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (0,0) | 0 | 4 | 8 | 12 | 2 | 6 | 10 | 14 |
| (1,0) | 16 | 20 | 24 | 28 | 18 | 22 | 26 | 30 |
| (2,0) | 32 | 36 | 40 | 44 | 34 | 38 | 42 | 46 |
| (0,1) | 1 | 5 | 9 | 13 | 3 | 7 | 11 | 15 |
| (1,1) | 17 | 21 | 25 | 29 | 19 | 23 | 27 | 31 |
| (2,1) | 33 | 37 | 41 | 45 | 35 | 39 | 43 | 47 |
The tile is now interleaved or "raked" with the other 3x4 matrix-of-tiles
instead of appearing as blocks. Other references call this is cyclic
distribution.
This might look familiar if you have ever used ScaLAPACK.
It expresses a 2-D block cyclic distribution of a 6 x 8 matrix
over 4 processes in a 2 x 2 "process grid." See
["The Two-dimensional Block-Cyclic Distribution"](https://netlib.org/scalapack/slug/node75.html#sec2dbcd)
and
["Local Storage Scheme and Block-Cyclic Mapping"](https://netlib.org/scalapack/slug/node76.html#seclocalstorage)
in the ScaLAPACK Users' Guide.
In general, `logical_product` and these variations can produce any interleaving,
including blocked, cyclic, by-mode blocked/cyclic, and intermediate interleavings
that don't have common names.
```c++
Layout tile = Layout<Shape <_2,_2>,
Stride<_1,_2>>{};
Layout matrix_of_tiles = Layout<Shape <_3,_4>,
Stride<_4,_1>>{};
print_layout(raked_product(tile, matrix_of_tiles));
```
### Division
The previous section covered layout products,
that reproduce one layout over another.
This section covers layout *division*.
Functions that divide a layout into components are useful
as a basis for tiling and partitioning layouts.
For example, consider folding a vector into a matrix.
We could imagine an operation, called `logical_divide`,
```c++
Layout vec = Layout<_16,_3>{}; // 16 : 3
Layout col = Layout< _4,_1>{}; // 4 : 1
Layout mat = logical_divide(vec, col); // (4,4) : (3,12)
```
that "takes" the first 4 elements of the vector into the first mode
and leaves the "rest" in the second mode. This is a column-major matrix
view of the data in `vec`.
What if we want a row-major matrix view?
```c++
Layout vec = Layout<_16,_3>{}; // 16 : 3
Layout col = Layout< _4,_4>{}; // 4 : 4
Layout mat = logical_divide(vec, col); // (4,4) : (12,3)
```
Now, every fourth element of the vector is in the first mode and
the "rest" are in the second mode.
Multidimensional, hierarchical indices let us extend this operation
to any layout that "divides" the vector.
```c++
Layout vec = Layout<_16,_3>{}; // 16 : 3
Layout col = Layout< _4,_2>{}; // 4 : 2
Layout mat = logical_divide(vec, col); // (4,(2,2)) : (6,(3,24))
```
```c++
Layout vec = Layout<_16,_3>{}; // 16 : 3
Layout col = Layout< _4,_2>{};
Layout col = Layout<Shape <_2,_2>,
Stride<_4,_1>>{}; // (2,2) : (4,1)
Layout mat = logical_divide(vec, col); // ((2,2),(2,2)) : ((12,3),(6,24))
```
All of the above examples produce a 4x4 matrix
that can be indexed and treated like a normal 4x4 matrix,
but each has a different underlying layout.
Thus, our algorithms can be written using logical coordinates,
without needing to address the detailed indexing that each layout requires.
CuTe includes 3 different kinds of layout division operations.
1. `logical_divide`
2. `zipped_divide`
3. `tiled_divide`
We will summarize these in the sections that follow.
#### Logical divide : the intuitive tiling
Suppose I have the 6 x 8 matrix from the Raked Product section
and want to "collect" the `tile`, turning the Raked Product into
the Blocked Product.
To do this, we would like to gather two elements from the column
and leave the rest, then gather two elements from the row and leave the rest.
Thus, we want to apply `logical_divide` independently to the rows and cols
in order to retrieve the appropriate elements.
In code, we copy the Layout from the result of the Raked Product section, then
specify the elements in the rows and cols we would like to gather.
```c++
Layout raked_prod = Layout<Shape <Shape < _3,_2>,Shape <_4,_2>>,
Stride<Stride<_16,_1>,Stride<_4,_2>>>{};
Tile subtile = make_tile(Layout<_2,_3>{}, // Gather elements 2 : 3 from mode 0
Layout<_2,_4>{}); // Gather elements 2 : 4 from mode 1
print_layout(logical_divide(raked_prod, subtile));
```
Indeed, this does produce the result from the Blocked Product section.
| | (0,0) | (1,0) | (0,1) | (1,1) | (0,2) | (1,2) | (0,3) | (1,3) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (0,0) | 0 | 2 | 4 | 6 | 8 | 10 | 12 | 14 |
| (1,0) | 1 | 3 | 5 | 7 | 9 | 11 | 13 | 15 |
| (0,1) | 16 | 18 | 20 | 22 | 24 | 26 | 28 | 30 |
| (1,1) | 17 | 19 | 21 | 23 | 25 | 27 | 29 | 31 |
| (0,2) | 32 | 34 | 36 | 38 | 40 | 42 | 44 | 46 |
| (1,2) | 33 | 35 | 37 | 39 | 41 | 43 | 45 | 47 |
Of course, any other rearrangement of the rows and cols is also valid.
#### Zipped divide
The `zipped_divide` function applies `logical_divide`, and then gathers the
"subtiles" into a single mode and the "rest" into a single mode.
For example, if we apply `zipped_divide` instead of `logical_divide` in the example above,
```c++
Layout raked_prod = Layout<Shape <Shape < _3,_2>,Shape <_4,_2>>,
Stride<Stride<_16,_1>,Stride<_4,_2>>>{};
Tile subtile = make_tile(Layout<_2,_3>{}, // Gather elements 2 : 3 from mode 0
Layout<_2,_4>{}); // Gather elements 2 : 4 from mode 1
print_layout(zipped_divide(raked_prod, subtile));
```
then we get the result
| | (0,0) | (1,0) | (2,0) | (0,1) | (1,1) | (2,1) | (0,2) | (1,2) | (2,2) | (0,3) | (1,3) | (2,3) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (0,0) | 0 | 16 | 32 | 4 | 20 | 36 | 8 | 24 | 40 | 12 | 28 | 44 |
| (1,0) | 1 | 17 | 33 | 5 | 21 | 37 | 9 | 25 | 41 | 13 | 29 | 45 |
| (0,1) | 2 | 18 | 34 | 6 | 22 | 38 | 10 | 26 | 42 | 14 | 30 | 46 |
| (1,1) | 3 | 19 | 35 | 7 | 23 | 39 | 11 | 27 | 43 | 15 | 31 | 47 |
Note that this is the same layout as the result in the Logical Product section.
That is, the first mode is our original tile (and can be interpreted as a 2x2 matrix itself)
and the second mode is its logical layout within the raked layout.
##### More Examples of Divide
For brevity, shapes can be used with `logical_divide` and `tiled_divide` to quickly split and tile modes of a tensor. For example, this C++ code
```c++
Layout layout = Layout<Shape <_12, _32,_6>,
Stride< _1,_128,_0>>{};
Shape tile_shape = make_shape(_4{},_8{});
Layout logical_divided_tile = logical_divide(layout, tile_shape);
Layout zipped_divided_tile = zipped_divide(layout, tile_shape);
print("layout : "); print(layout); print("\n");
print("tile_shape : "); print(tile_shape); print("\n");
print("logical_divided_tile : "); print(logical_divided_tile); print("\n");
print("zipped_divided_tile : "); print(zipped_divided_tile); print("\n\n");
```
produces the following output when we vary `layout`.
```text
full_layout : (_12,_32,_6):(_1,_128,_0)
tile_shape : (_4,_8)
logical_divided_tile : ((_4,_3),(_8,_4),_6):((_1,_4),(_128,_1024),_0)
zipped_divided_tile : ((_4,_8),(_3,_4,_6)):((_1,_128),(_4,_1024,_0))
full_layout : (_12,(_4,_8),_6):(_1,(_32,_512),_0)
tile_shape : (_4,_8)
logical_divided_tile : ((_4,_3),((_4,_2),_4),_6):((_1,_4),((_32,_512),_1024),_0)
zipped_divided_tile : ((_4,(_4,_2)),(_3,_4,_6)):((_1,(_32,_512)),(_4,_1024,_0))
```
This code
```c++
Layout layout = make_layout(Shape<_8,_8>{},
Stride<_8,_1>{});
Layout tile = make_tile(make_layout(Shape<_4>{}),
make_layout(Shape<_2>{}));
print("layout: ");
print_layout(layout);
print("\n");
print("tile: ");
print(tile);
print("\n");
print("logical_divide: ");
print_layout(logical_divide(layout, tile));
print("zipped_divide: ");
print_layout(zipped_divide(layout, tile));
```
results in the following layouts.
<p align="center">
<img src="../../images/cute/logical_divide-and-zipped_divide.png" alt="logical_divide-and-zipped_divide" height="400"/>
</p>
This code
```c++
Layout layout = make_layout(Shape<_8,_8>{},
Stride<_8,_1>{});
Layout tile = make_tile(make_layout(Shape<_2>{}),
make_layout(Shape<_4>{}));
print("layout: ");
print_layout(layout);
print("\n");
print("tile: ");
print(tile);
print("\n");
print("logical_divide: ");
print_layout(logical_divide(layout, tile));
print("zipped_divide: ");
print_layout(zipped_divide(layout, tile));
```
results in the following layouts.
<p align="center">
<img src="../../images/cute/logical_divide-and-zipped_divide-2.png" alt="logical_divide-and-zipped_divide-2" height="400"/>
</p>
#### Tiled divide
The `tiled_divide` function works like `zipped_divide`,
except that it unpacks the second mode. This is useful when you have a `Tile` that describes all of the elements for a particular operation, for example, and want to gather those together but retain the logical shape of those tiles within the original layout. That is,
```text
Layout Shape : (M, N, L, ...)
Tile Shape : <M', N'>
Tiled Result : ((M', N'), m, n, L, ...)
```
where `m` is `M / M'` and `n` is `N / N'`.
We can consider `m` as the "number of `Tile`s in `M`" and `n` as the "number of `Tile`s in `N`". This style of operation is common when applying MMA Atoms and Copy Atoms.
media/docs/cute/03_tensor.md
+262
View File
@@ -0,0 +1,262 @@
# CuTe Tensors
## A Tensor is a multidimensional array
CuTe's `Tensor` class represents a multidimensional array.
The array's elements can live in any kind of memory,
including global memory, shared memory, and register memory.
### Array access
Users access a `Tensor`'s elements in one of three ways:
* `operator()`, taking as many integral arguments as the number of modes,
corresponding to the element's (possibly) multidimensional logical index;
* `operator()`, taking a `Coord` (an `IntTuple` of the logical indices); or
* `operator[]`, taking a `Coord` (an `IntTuple` of the logical indices).
### Slices: Get a Tensor accessing a subset of elements
Users can get a "slice" of a `Tensor`,
that is, a `Tensor` that accesses a subset of elements
of the original `Tensor`.
Slices happen through the same `operator()`
that they use for accessing an individual element.
Passing in `_` (the underscore character, an instance of `Underscore`)
has the same effect as `:` (the colon character) in Fortran or Matlab:
the resulting slice accesses all indices in that mode of the `Tensor`.
### Tensor's behavior determined by its Layout and Engine
A `Tensor`'s behavior is entirely determined by its two components,
which correspond to its two template parameters: `Engine`, and `Layout`.
For a description of `Layout`,
please refer to [the `Layout` section](./01_layout.md)
of this tutorial, or the [GEMM overview](./0x_gemm_tutorial.md).
An `Engine` represents a one-dimensional array of elements.
When users perform array access on a `Tensor`,
the `Tensor` uses its `Layout` to map from a logical coordinate
to a one-dimensional index.
Then, the `Tensor` uses its `Engine`
to map the one-dimensional index
to a reference to the element.
You can see this in `Tensor`'s implementation of array access.
```c++
decltype(auto) operator[](Coord const& coord) {
return engine().begin()[layout()(coord)];
}
```
One could summarize almost all CuTe use cases as follows:
* create `Layout`s,
* create `Tensor`s with those `Layout`s, and
* invoke (either CuTe's, or custom) algorithms on those `Tensor`s.
### Ownership of the elements
`Tensor`s can be owning or nonowning.
"Owning" `Tensor`s behave like `std::array`.
When you copy the `Tensor`, you (deep-)copy its elements,
and the `Tensor`'s destructor deallocates the array of elements.
"Nonowning" `Tensor`'s behave like a (raw) pointer to the elements.
Copying the `Tensor` doesn't copy the elements,
and destroying the `Tensor` doesn't deallocate the array of elements.
Whether a `Tensor` is "owning" or "nonowning" depends entirely on its `Engine`.
This has implications for developers of generic `Tensor` algorithms.
For example, input `Tensor` parameters of a function
should be passed by const reference,
because passing the `Tensor`s by value
might make a deep copy of the `Tensor`'s elements.
It might also *not* make a deep copy of the elements;
there's no way to know without specializing the algorithm
on the `Tensor`'s `Engine` type.
Similarly, output or input/output `Tensor` parameters of a function
should be passed by (nonconst) reference.
Returning a `Tensor` might (or might not)
make a deep copy of the elements.
The various overloads of the `copy_if` algorithm in
[`include/cute/algorithm/copy.hpp`](../../../include/cute/algorithm/copy.hpp)
take their `src` (input, source of the copy) parameter
as `Tensor<SrcEngine, SrcLayout>& const`,
and take their `dst` (output, destination of the copy) parameter
as `Tensor<DstEngine, DstLayout>&`.
Additionally, there are overloads for mutable temporaries like
`Tensor<DstEngine, DstLayout>&&`
so that these algorithms can be applied directly to slices,
as in the following example.
```c++
copy(src_tensor(_,3), dst_tensor(2,_));
```
In C++ terms, each of the expressions
`src_tensor(_,3)`, and `dst_tensor(2,_)`
is in the "prvalue"
[value category](https://en.cppreference.com/w/cpp/language/value_category),
because it is a function call expression
whose return type is nonreference.
(In this case, calling `Tensor::operator()`
with at least one `_` (`Underscore`) argument
returns a `Tensor`.)
The prvalue `dst_tensor(2,_)` won't match
the `copy` overload taking
`Tensor<DstEngine, DstLayout>&`,
because prvalues can't be bound to
nonconst lvalue references (single `&`).
However, it will match the `copy` overload taking
`Tensor<DstEngine, DstLayout>&&`
(note the two `&&` instead of one `&`).
Calling the latter overload binds the reference
to the prvalue `dst_tensor(2,_)`.
This results in
[creation of a temporary](https://en.cppreference.com/w/cpp/language/implicit_conversion#Temporary_materialization)
`Tensor` result to be passed into `copy`.
### CuTe's provided `Engine` types
CuTe comes with three `Engine` types.
* `ArrayEngine<class T, int N>`: an owning `Engine`,
representing an array of `N` elements of type `T`
* `ViewEngine<Iterator>`: a nonowning `Engine`,
where `Iterator` is a random access iterator
(either a pointer to an array, or something that acts like one)
* `ConstViewEngine<Iterator>`: a nonowning `Engine`,
which is the view-of-const-elements version of `ViewEngine`
### "Tags" for different kinds of memory
`ViewEngine` and `ConstViewEngine` wrap pointers to various kinds of memory.
Users can "tag" the memory with its space -- e.g., global or shared --
by calling `make_gmem_ptr(g)` when `g` is a pointer to global memory,
or `make_smem_ptr(s)` when `s` is a pointer to shared memory.
Tagging memory makes it possible for CuTe's `Tensor` algorithms
to use the fastest implementation for the specific kind of memory.
It also avoids incorrect memory access.
For example, some kinds of optimized copy operations require
that the source of the copy be in global memory,
and the destination of the copy be in shared memory.
Tagging makes it possible for CuTe to dispatch
to those optimized copy operations where possible.
CuTe does this by specializing `Tensor` algorithms
on the `Tensor`'s `Engine` type.
### Engine members
In order for a type to be valid for use as an `Engine`,
it must have the following public members.
```c++
using value_type = /* ... the value type ... */;
using iterator = /* ... the iterator type ... */;
iterator begin() /* sometimes const */;
```
## Constructing a Tensor
### Nonowning view of existing memory
A `Tensor` can be a nonowning view of existing memory.
For this use case, users can create the `Tensor` by calling `make_tensor`
with two arguments: a wrapped pointer to the memory to view, and the `Layout`.
Users wrap the pointer by identifying its memory space:
e.g., global memory (via `make_gmem_ptr`) or shared memory (via `make_smem_ptr`).
`Tensor`s that view existing memory can have either static or dynamic `Layout`s.
Here are some examples of creating `Tensor`s
that are nonowning views of existing memory.
```c++
// Global memory (static or dynamic layouts)
Tensor gmem_8s = make_tensor(make_gmem_ptr(A), Int<8>{});
Tensor gmem_8d = make_tensor(make_gmem_ptr(A), 8);
Tensor gmem_8sx16d = make_tensor(make_gmem_ptr(A), make_shape(Int<8>{},16));
Tensor gmem_8dx16s = make_tensor(make_gmem_ptr(A), make_shape ( 8 ,Int<16>{}),
make_stride(Int<16>{},Int< 1>{}));
// Shared memory (static or dynamic layouts)
Shape smem_shape = make_shape(Int<4>{},Int<8>{});
__shared__ T smem[decltype(size(smem_shape))::value]; // (static-only allocation)
Tensor smem_4x8_col = make_tensor(make_smem_ptr(&smem[0]), smem_shape);
Tensor smem_4x8_row = make_tensor(make_smem_ptr(&smem[0]), smem_shape, GenRowMajor{});
```
### Owning array of register memory
A `Tensor` can also be an owning array of register memory.
For this use case, users can create the `Tensor`
by calling `make_tensor<T>(layout)`,
where `T` is the type of each element of the array,
and `layout` is the `Tensor`'s `Layout`.
Owning `Tensor`s must have a static `Layout`,
as CuTe does not perform dynamic memory allocation in `Tensor`s.
Here are some examples of creating owning `Tensor`s.
```c++
// Register memory (static layouts only)
Tensor rmem_4x8_col = make_tensor<float>(make_shape(Int<4>{},Int<8>{}));
Tensor rmem_4x8_row = make_tensor<float>(make_shape(Int<4>{},Int<8>{}), GenRowMajor{});
Tensor rmem_4x8_mix = make_tensor<float>(make_shape (Int<4>{},Int< 8>{}),
make_stride(Int<2>{},Int<32>{}));
Tensor rmem_8 = make_fragment_like(gmem_8sx16d(_,0));
```
The `make_fragment_like` function makes an owning Tensor of register memory,
with the same shape as its input `Tensor` argument.
## Tensor use examples
### Copy rows of a matrix from global memory to registers
The following example copies rows of a matrix (with any `Layout`)
from global memory to register memory,
then executes some algorithm `do_something`
on the row that lives in register memory.
```c++
Tensor gmem = make_tensor(make_gmem_ptr(A), make_shape(Int<8>{}, 16));
Tensor rmem = make_fragment_like(gmem(_, 0));
for (int j = 0; j < size<1>(gmem); ++j) {
copy(gmem(_, j), rmem);
do_something(rmem);
}
```
This code does not need to know anything the `Layout` of `gmem`
other than that it is rank-2 and that the first mode is a compile-time value.
The following code checks both of those conditions at compile time.
```c++
CUTE_STATIC_ASSERT_V(rank(gmem) == Int<2>{});
CUTE_STATIC_ASSERT_V(is_static<decltype(shape<0>(gmem))>{});
```
A `Tensor` encapsulates the data type, data location,
and possibly also the shape and stride of the tensor at compile time.
As a result, `copy` can dispatch, based on the types and Layouts of its arguments,
to use any of various synchronous or asynchronous hardware copy instructions
and can auto-vectorize the copy instructions in many cases as well.
CuTe's `copy` algorithm lives in
[`include/cute/algorithm/copy.hpp`](../../../include/cute/algorithm/copy.hpp).
For more details on the algorithms that CuTe provides,
please refer to [the algorithms section](./04_algorithms.md)
of the tutorial, or the
[CuTe overview in the GEMM tutorial](./0x_gemm_tutorial.md).
media/docs/cute/04_algorithms.md
+223
View File
@@ -0,0 +1,223 @@
# CuTe Tensor algorithms
This section summarizes the interfaces and implementations
of common numerical algorithms performed on `Tensor`s.
The implementation of these algorithms may be found in the
[include/cute/algorithm/](../../../include/cute/algorithm/)
directory.
## `copy`
CuTe's `copy` algorithm copies the elements of a source `Tensor`
into the elements of a destination `Tensor`.
The various overloads of `copy` can be found in
[`include/cute/algorithm/copy.hpp`](../../../include/cute/algorithm/copy.hpp).
### Interface and specialization opportunities
A `Tensor` encapsulates the data type, data location,
and possibly also the shape and stride of the tensor at compile time.
As a result, `copy` can and does dispatch,
based on the types of its arguments,
to use any of various synchronous or asynchronous hardware copy instructions.
The `copy` algorithm has two main overloads.
The first just takes the source `Tensor` and the destination `Tensor`.
```c++
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst);
```
The second takes those two parameters, plus a `Copy_Atom`.
```c++
template <class... CopyArgs,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Atom<CopyArgs...> const& copy_atom,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst);
```
The two-parameter `copy` overload picks a default implementation
based only on the types of the two `Tensor` parameters.
The `Copy_Atom` overload lets callers override that default
by specifying a nondefault `copy` implementation.
### Parallelism and synchronization depend on parameter types
Either the default implementation or
the implementation selected by a `Copy_Atom` overload
may use none or all available parallelism,
and may have a variety of synchronization semantics.
The behavior depends on `copy`'s parameter types.
Users are expected to figure this out based on their knowledge
of the architecture on which they are running.
(Developers often write a custom optimized kernel
for each GPU architecture.)
The `copy` algorithm may be sequential per thread,
or it may be parallel across some collection of threads
(e.g., a block or cluster).
If `copy` is parallel,
then the collection of participating threads
may need synchronization before any thread in the collection
may assume that the copy operation has completed.
For example, if the participating threads form a thread block,
then users must invoke `__syncthreads()`
or the Cooperative Groups equivalent
before they may use the results of `copy`.
The `copy` algorithm may use asynchronous copy instructions,
such as `cp.async`, or its C++ interface `memcpy_async`.
In that case, users will need to perform
the additional synchronization appropriate to that underlying implementation
before they may use the results of the `copy` algorithm.
[The CuTe GEMM tutorial example](../../../examples/cute/tutorial/sgemm_nt_1.cu)
shows one such synchronization method.
More optimized GEMM implementations use pipelining techniques
to overlap asynchronous `copy` operations with other useful work.
### A generic copy implementation
A simple example of a generic `copy` implementation
for any two `Tensor`s looks like this.
```c++
template <class TA, class ALayout,
class TB, class BLayout>
CUTE_HOST_DEVICE
void
copy(Tensor<TA, ALayout> const& src, // Any logical shape
Tensor<TB, BLayout> & dst) // Any logical shape
{
for (int i = 0; i < size(src); ++i) {
dst(i) = src(i);
}
}
```
This generic `copy` algorithm addresses both `Tensor`s
with 1-D logical coordinates, thus traversing both `Tensor`s
in a logical column-major order.
Some reasonable architecture-independent optimizations
would include the following.
1. If the two `Tensor`s have known memory spaces with optimized
access instructions (like `cp.async`), then dispatch to the
custom instruction.
2. The the two `Tensor`s have static layouts and it can be proven
that element vectorization is valid -- for example, four `LDS.32`s
can be combined into a single `LDS.128` -- then vectorize the source
and destinations tensors.
3. If possible, validate that the copy instruction to be used is
appropriate for the source and destination tensors.
CuTe's optimized copy implementations can do all of these.
## `copy_if`
CuTe's `copy_if` algorithm lives in the same header as `copy`,
[`include/cute/algorithm/copy.hpp`](../../../include/cute/algorithm/copy.hpp).
The algorithm takes source and destination `Tensor` parameters like `copy`,
but it also takes a "predication `Tensor`"
with the same shape as the input and output.
Elements of the source `Tensor` are only copied
if the corresponding predication `Tensor` element is nonzero.
For details on why and how to use `copy_if`,
please refer to the
["predication" section of the tutorial](./0y_predication.md).
## `gemm`
### What `gemm` computes
The `gemm` algorithm takes three `Tensor`s, A, B, and C.
What it does depends on the number of modes
that its `Tensor` parameters have.
We express these modes using letters.
* V indicates a "vector," a mode of independent elements.
* M and N indicate the number of rows resp. columns
of the matrix result C of the BLAS's GEMM routine.
* K indicates the "reduction mode" of GEMM,
that is, the mode along which GEMM sums.
Please see the [GEMM tutorial](./0x_gemm_tutorial.md) for details.
We list the modes of the input `Tensor`s A and B,
and the output `Tensor` C,
using a notation `(...) x (...) => (...)`.
The two leftmost `(...)` describe A and B (in that order),
and the `(...)` to the right of the `=>` describes C.
1. `(V) x (V) => (V)`. The element-wise product of vectors: C<sub>v</sub> += A<sub>v</sub> B<sub>v</sub>. Dispatches to FMA or MMA.
2. `(M) x (N) => (M,N)`. The outer product of vectors: C<sub>mn</sub> += A<sub>m</sub> B_<sub>n</sub>. Dispatches to (4) with V=1.
3. `(M,K) x (N,K) => (M,N)`. The product of matrices: C<sub>mn</sub> += A<sub>mk</sub> B<sub>nk</sub>. Dispatches to (2) for each K.
4. `(V,M) x (V,N) => (V,M,N)`. The batched outer product of vectors: C<sub>vmn</sub> += A<sub>vm</sub> B<sub>vn</sub>. Optimizes for register reuse and dispatches to (1) for each M, N.
5. `(V,M,K) x (V,N,K) => (V,M,N)`. The batched product of matrices: C<sub>vmn</sub> += A<sub>vmk</sub> B<sub>vnk</sub>. Dispatches to (4) for each K.
Please refer to the [GEMM tutorial](./0x_gemm_tutorial.md)
for an overview of CuTe's convention for ordering the modes.
For example, if K appears, it always appears rightmost ("outermost").
If V appears, it always appears leftmost ("innermost").
### Dispatch to optimized implementations
Just like with `copy`, CuTe's implementations of `gemm`
uses its `Tensor` arguments' types to dispatch
to an appropriately optimized implementation.
Also like `copy`, `gemm` takes an optional `MMA_Atom` parameter
that lets callers override the default `FMA` instruction
that CuTe would select based on the `Tensor` arguments' types.
For more information on `MMA_Atom` and on specialization of `gemm`
for different architectures, please refer to the
[MMA section of the tutorial](./0t_mma_atom.md).
## `axpby`
The `axpby` algorithm lives in the header file
[`include/cute/algorithm/axpby.hpp`](../../../include/cute/algorithm/axpby.hpp).
It assigns to $y$ the result of $\alpha x + \beta y$,
where $\alpha$ and $\beta$ are scalars and $x$ and $y$ are `Tensor`s.
The name stands for "Alpha times X Plus Beta times Y,"
and is a generalization of the original BLAS "AXPY" routine
("Alpha times X Plus Y").
## `fill`
The `fill` algorithm lives in the header file
[`include/cute/algorithm/fill.hpp`](../../../include/cute/algorithm/fill.hpp).
It overwrites the elements of its `Tensor` output argument
with a given scalar value.
## `clear`
The `clear` algorithm lives in the header file
[`include/cute/algorithm/clear.hpp`](../../../include/cute/algorithm/clear.hpp).
It overwrites the elements of its `Tensor` output argument with zeros.
## Other algorithms
CuTe provides other algorithms.
Their header files can be found in the
[`include/cute/algorithm`](../../../include/cute/algorithm)
directory.
media/docs/cute/0t_mma_atom.md
+434
View File
@@ -0,0 +1,434 @@
# CuTe's support for Matrix Multiply-Accumulate instructions
In this file, we explain in detail how we support our GPUs'
Matrix Multiply-Accumulate (MMA) hardware instructions in CuTe.
MMAs are architecture-specific.
Different generations of GPU architectures
introduce different sets of MMA instructions.
However, CuTe features such as `Layout`
makes it possible to expose MMAs for use in generic CUDA C++ code.
We do this in two steps.
1. We wrap each MMA's PTX instruction in an "Operation" struct.
2. For each Operation struct, we define a "Traits" struct
that defines all of the meta-information needed to use the Operation.
## CuTe MMA Atoms
CuTe exposes each MMA to generic CUDA C++ code as a pair of structs:
an "Operation" struct,
and an `MMA_Traits` struct templated on the Operation struct type.
An "Operation" struct exposes the PTX instruction
for that specific operation.
It defines the arguments and interface it expects.
Operation structs have minimal software dependencies --
it does not use layouts, tensors, or non-standard numeric data types.
Different structs have different names
that describe what the MMA instruction does.
We will explain the naming scheme below.
A corresponding `MMA_Traits` struct specialization
defines meta-information about the Operation,
such as the compute types, the logical shape of the operation,
and the `Layout`s of threads and values within the operation.
The `MMA_Traits` struct takes the Operation as a template parameter.
CuTe specializes `MMA_Traits` for each Operation type that it supports.
Together, these two types comprise an "Atom" that decouples the complexity of thread and data layouts from the call site of of the PTX instruction. The Atom's Traits struct exposes information that is relevant to a single MMA operation, no matter the granularity at which it operates.
CuTe MMA atoms expose the semantics of a single MMA operation.
This is true regardless of the hardware level at which the MMA operates.
CuTe supports MMA atoms that operate at a variety of hardware levels,
including
* a single thread (e.g., fused multiply-add (FMA) instruction);
* a quadpair (Volta);
* a single warp (Ampere); and
* a warpgroup (Hopper).
### Operation structs
#### Location of files
CuTe provides its Operations structs in the
[`include/cute/arch`](../../../include/cute/arch)
directory, in header files starting with `mma`.
#### Operation struct's name
A CuTe Operation struct's name encodes information about
* its first supported architecture,
* the M, N, and K dimensions that it accepts,
* the types that it takes, and
* the expected A and B layouts.
For example, the Volta section below will refer to the
`SM70_8x8x4_F32F16F16F32_NT` Operation struct defined in
[`include/cute/arch/mma_sm70.hpp`](../../../include/cute/arch/mma_sm70.hpp).
* "SM70" refers to Volta.
* "8x8x4" refers to M = 8, N = 8, and K = 4,
the dimensions of the MMA operation that the quadpair performs
(see below).
* "F32F16F16F32" refers to the element types
of the four matrix operands A, B, C, and D.
An MMA computes D = C + A * B,
so we read the types from left to right:
D is F32 (`float`), A is F16 (half),
B is F16 (half), and C is F32 (`float`).
* "NT" means that A is M-major (not transposed)
and B is N-major (transposed).
#### Contents
An Operation struct has the following members.
##### Type aliases
An Operation struct has four public type aliases:
`DRegisters`, `ARegisters`, `BRegisters`, and `CRegisters`.
For example, the `SM70_8x8x4_F32F16F16F32_NT` Operation struct defined in
[`include/cute/arch/mma_sm70.hpp`](../../../include/cute/arch/mma_sm70.hpp)
defines these as follows.
```c++
using DRegisters = float[8];
using ARegisters = uint32_t[2];
using BRegisters = uint32_t[2];
using CRegisters = float[8];
```
This shows how many values each thread will pass into the PTX instruction
for each of the matrices A, B, C, and D. For this Operation,
each thread passes 8 F32 values each for C and D (hence `float[8]`),
and 4 F16 values each for A and B (hence `uint32_t[2]`;
the instruction packs two 16-bit F16 values
in each of the two 32-bit `uint32_t` values).
##### `fma` static member device function
An operation struct defines a public `static void fma` function.
It is marked with the `CUTE_HOST_DEVICE` macro,
which adds the `__host__ __device__` annotations.
Different Operations define `fma` to take different numbers of arguments,
depending on the PTX MMA instruction.
The implementation protects use of the PTX instruction with a macro,
and raises an `assert` if `fma` is called when the macro is not defined.
This ensures that tests and examples that use this Operation in an Atom
can still compile, even if the PTX instruction is not available.
### Traits
#### Location of files
CuTe provides its Traits structs in the
[`include/cute/atom`](../../../include/cute/atom)
directory, in header files starting with `mma_traits`.
#### Contents
An `MMA_Traits` specialization defines the following public type aliases.
* `ElementDVal`: Compute type of the D matrix
* `ElementAVal`: Compute type of the A matrix
* `ElementBVal`: Compute type of the B matrix
* `ElementCVal`: Compute type of the C matrix
* `Shape_MNK`: Logical MxNxK shape of the MMA operation
* `ThrID`: Logical thread mapping within the single MMA operation
(specifying the quadpair, warp, or warpgroup view)
* `ALayout`: Mapping of (thread,value) pairs to the logical MxK A matrix
* `BLayout`: Mapping of (thread,value) pairs to the logical NxK B matrix
* `CLayout`: Mapping of (thread,value) pairs to the logical MxN C matrix
#### Example
The specialization of MMA_Traits for the
`SM70_8x8x4_F32F16F16F32_NT` Operation lives in the header file
[`include/cute/atom/mma_traits_sm70.hpp`](../../../include/cute/atom/mma_traits_sm70.hpp).
It looks like this.
```c++
template <>
struct MMA_Traits<SM70_8x8x4_F32F16F16F32_NT>
{
using ElementDVal = float;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = float;
using Shape_MNK = Shape<_8,_8,_4>;
using ThrID = SM70_QuadPair;
using ALayout = SM70_8x4_Col;
using BLayout = SM70_8x4_Col;
using CLayout = SM70_8x8_32b;
};
```
The next section will explain these type aliases in detail.
## Volta
This and the following sections show examples of how to construct MMA atoms.
We don't try to explain this for all GPU architectures and MMAs.
Instead, we use selected examples to illustrate the process
of developing new atoms.
Volta architecture implements an HMMA instruction where a group of 8 threads called a quadpair (QP) collaborate to share data and perform an 8x8x4 (fp32 or fp16) matrix multiply-accumulate. (since a warp is 32 threads wide, it would perform an MMA across 4 QPs for a tile size of 16x16x4).
We first take a look at how we would take the ISA semantics of thread and data partitioning for the HMMA instruction, and encode it in a Traits struct. The HMMA NT instruction has the thread-data layout:
<p align="center">
<img src="../../images/cute/HMMA.8x8x4.NT.png" alt="HMMA.8x8x4.NT.png" height="400"/>
</p>
### Types
The HMMA NT above uses types:
```cpp
using ElementDVal = float;
using ElementAVal = half_t;
using ElementBVal = half_t;
using ElementCVal = float;
```
The rest of the `MMA_Traits` will be described in units of these types.
### Shape
The HMMA NT above has shape 8x8x4:
```cpp
// Logical shape of the MMA
using Shape_MNK = Shape <_8,_8,_4>;
```
### Thread ID
If the 32 threads in a warp are logically indexed by [0 ... 31], then the above image contains threads [0,1,2,3]U[16,17,18,19]. These threads make up the 0th quadpair. We can write a thread mapping that maps eight logical thread ids [0,1,2,3,4,5,6,7] of the MMA to a quadpair thread index [0,1,2,3]U[16,17,18,19] of a warp. The layout function has 4 elements with a stride of 1 and 2 of those with a stride of 16. With this, we write a layout that represents a quadpair:
```cpp
// Mapping from (logical thread id) -> (thread idx)
using ThrID = Layout<Shape <_4, _2>,
Stride<_1,_16>>;
```
Again, this layout function maps the logical thread id [0,8) of the MMA operation onto the quadpair thread index [0,4)U[16,20) of a warp.
### Accumulator Mapping
Let us look at exactly how the 8 threads within a QP are mapped to the A, B and C matrices. For the C and D matrices, the above image is broken down a bit more below. On the left is shown the whole QP level view, and on the right is shown the values owned by just thread 0.
<p align="center">
<img src="../../images/cute/HMMA.8x8x4.quadpair.C.png" alt="HMMA.8x8x4.quadpair.C.png" height="400"/>
</p>
The metainformation of this single instruction level view is what we want to encode in CuTe. Specifically, the QP level view in this diagram corresponds to the four MMA traits for [SM70_F32F16F16F32](../../../include/cute/arch/mma_sm70.hpp). These structs contain the `Element` types, the `Shape_MNK`, and the `ThrID` mapping we constructed above. Now, let us take a look at the definition of `CLayout`, the thread-data layout of accumulators. The job of `CLayout` is to construct a mapping between the `(logical_thr_id, logical_val_id)` and `(m, n)` coordinate in the C matrix which can then be used to build up more complicated layouts and operations like the 16x16x4 WMMA.
We can start constructing a `CLayout` from the picture above. As with any CuTe layout, it is a pair of `Shape` and corresponding `Stride`. Let us just look at the shape for now. We know that the HMMA uses 8 threads each of which own 8 values. Therefore, the shape of our mapping must have a size of 8 along two modes. With this, we have
```cpp
// (T8,V8) -> (m,n)
using CLayout = Layout<Shape <_8, _8>,
Stride<_?, _?>; // Stride to be filled in below
```
This is not to be confused with the logical 8x8 shape of the C matrix. This is 8-threads by 8-values. We now want to map those to (m,n) coordinates. Since CuTe layouts return indices rather than coordinates, we choose a column-major encoding of the (m,n) coordinates:
```
(logical_thr_id, logical_val_id) -> (m, n) == m + n * M
```
With this in place, we can start thinking about how to construct the strides in `CLayout`. Let's begin by looking at the strides between threads. Note that
* `(T0,V0)` is located at `(m,n) = (0,0) = 0`
* `(T1,V0)` is located at `(m,n) = (1,0) = 1`
* `(T2,V0)` is located at `(m,n) = (0,2) = 16`
* `(T3,V0)` is located at `(m,n) = (1,2) = 17`
* `(T4,V0)` is located at `(m,n) = (4,0) = 4`
* `(T5,V0)` is located at `(m,n) = (5,0) = 5`
* `(T6,V0)` is located at `(m,n) = (4,2) = 20`
* `(T7,V0)` is located at `(m,n) = (5,2) = 21`
where `T4`,`T5`,`T6`,`T7` are the 4th,5th,6th,7th logical thread id of the MMA corresponding to thread indices of 16,17,18,19 of the warp (recorded in the `ThrID` mapping!).
We note that the pattern can be transcribed to a layout. We can find the position of the 8 threads via
```cpp
using CLayout = Layout<Shape <Shape <_2, _2, _2>, _8>,
Stride<Stride<_1, _16, _4>, _?>;
```
With the exact same approach, we can construct the stride along the `logical value id` mode.
* `(T0,V0)` is located at `(m,n) = (0,0) = 0`
* `(T0,V1)` is located at `(m,n) = (0,1) = 8`
* `(T0,V2)` is located at `(m,n) = (2,0) = 2`
* `(T0,V3)` is located at `(m,n) = (2,1) = 10`
* `(T0,V4)` is located at `(m,n) = (0,4) = 32`
* `(T0,V5)` is located at `(m,n) = (0,5) = 40`
* `(T0,V6)` is located at `(m,n) = (2,4) = 34`
* `(T0,V7)` is located at `(m,n) = (2,5) = 42`
We note that this pattern can also be transcribed to a layout. We can find the position of the 8 values via
```cpp
// (T8,V8) -> (m,n)
using CLayout = Layout<Shape <Shape <_2, _2,_2>, Shape <_2,_2, _2>>,
Stride<Stride<_1,_16,_4>, Stride<_8,_2,_32>>>;
```
And that's all! We can verify that each `(tid,vid)` coordinate in this layout is reliably mapped to the correct (encoded) `(m,n)` coordinate.
In the case of F16 accumulators, the layout is way less complex. Each row of accumulators `(m, :)` is held by a single thread, which makes the layout:
```cpp
using CLayout = Layout<Shape <_8,_8>,
Stride<_1,_8>>;
```
### A and B Layout Mapping
A and B matrix layouts depend on whether the sources are transposed or not. The diagram below shows the thread ID to data ownership map for A and B matrices in the case of NT and TN transposes.
<p align="center">
<img src="../../images/cute/HMMA.8x8x4.quadpair.AB.png" alt="HMMA.8x8x4.quadpair.AB.png" height="400"/>
</p>
Let's look at the TN layout for A matrix first (right side in the diagram). Again, there are the same 8 logical threads, but each threads owns only 4 elements this time. The shape of `ALayout` will then be `Shape<_8, _4>`. As for the strides, we again need a similar mapping between `(m, k) == m + k * M`. Looking down the `M` mode, we go from `(T0, V0)` to `(T1, V0)` which is a stride of 1 for all 8 threads. For the `K` mode, as we go across, we go from `(T0, V0)` to `(T0, V1)`, which makes a stride of 8 for all 4 values. Therefore, the A layout is:
```cpp
// (T8,V4) -> (m,k)
using ALayout = Layout<Shape <_8,_4>,
Stride<_1,_8>>;
```
Source B layout is constructed similarly for the TN HMMA, except that we want write it as `(N,K)` rather than `(K,N)` for convenience. For the strides, as we go across the `N` mode, we go from `(T0, V0)` to `(T1, V0)`, making this a stride of 1 for all 8 threads. As we go down the `K` mode, `(T0, V0)` to `(T0, V1)` which is a stride of 8 for all 4 values. So the B layout is the same as A:
```cpp
// (T8,V4) -> (n,k)
using BLayout = Layout<Shape <_8,_4>,
Stride<_1,_8>>;
```
The layouts in the case of NT are a bit more complicated (left side of the diagram). Going down the `M` mode of `A`, we see the four values of `T0` first and then we see the four values of `T4`. This means we first have a stride of 1 for 4 values, followed by a stride of 4 from `T0` to `T4`. So we have two sub-strides along the `M` mode. For the `K` mode, as we go across, we simply increment the `thr_id`, keeping `val_id` the same, making the stride 8 for 4 threads. This makes the A layout:
```cpp
// (T8,V4) -> (m,k)
using ALayout = Layout<Shape <Shape <_4,_2>,_4>,
Stride<Stride<_8,_4>,_1>>;
```
With the `(N,K)` ordering for B, the layout is the same.
```cpp
// (T8,V4) -> (n,k)
using BLayout = Layout<Shape <Shape <_4,_2>,_4>,
Stride<Stride<_8,_4>,_1>>;
```
For the NN and TT transposes, they are simply combinations of the two layouts we have seen for A and B so far.
## Hopper
Now, we are ready to take a look at the much larger GMMA operation (Group MMA) first introduced with Hopper architecture. These MMA instructions operate at the granularity of 128 threads (4 warps), which are collectively referred to as a warpgroup.
### Thread ID
In the case of Hopper GMMAs, the thread IDs are assigned based on the simple 1D contiguous layout, which makes `thrID` trivial:
```cpp
using ThrID = Layout<_128, _1>;
```
### Accumulator Mapping
Accumulators are mapped hierarchically in GMMA, starting from the concept of a core matrix and building up to a layout for the whole C matrix tile. Let's look at this core matrix first. We only consider fp16 accumulators here, but extensions of fp32 accumulators as trivial as we will see later.
Each core matrix has the layout as shown in the diagram below.
<p align="center">
<img src="../../images/cute/gmma_coremat_cd_fp16.png" alt="gmma_coremat_cd_fp16.png" height="600"/>
</p>
As in the Volta examples, the thread IDs are logical only, and which of the four warps they belong to in the warpgroup is not important.
Then GMMA tiles this core matrix first vertically along the M mode, and then repeats that column of core matrices along the N mode to construct the full MxN tile. This tiling is shown in the image below.
<p align="center">
<img src="../../images/cute/gmma_wg_n_slice.png" alt="gmma_wg_n_slice.png" height="600"/>
</p>
With this image, we are again ready to start building the `CLayout` for `SM90_64x128x16_F16F16F16F16_TN` atom. Same as before, we are constructing a mapping between the `(logical_thr_id, logical_val_id) -> (m, n)` coordinate spaces.
To begin, let's follow the first few threads and values. We immediately see that they are arranged along the `N`-mode with pairs of values and four threads. This gives us
```cpp
// (T128,V4) -> (M64,N8)
using CLayout = Layout<Shape <Shape < _4, ...>, Shape < _2, ...>>,
Stride<Stride<_128, ...>, Stride<_64, ...>>>;
```
To complete the first 8x8 core matrix, the four threads repeat eight times down the `M`-mode:
```cpp
// (T128,V4) -> (M64,N8)
using CLayout = Layout<Shape <Shape < _4, _8, ...>, Shape < _2, ...>>,
Stride<Stride<_128, _1, ...>, Stride<_64, ...>>>;
```
Then, as we go to the next core matrix, we wrap back again to `T0`, but this time to `(T0, V2)`.
```cpp
// (T128,V4) -> (M64,N8)
using CLayout = Layout<Shape <Shape < _4, _8, ...>, Shape < _2, _2>>,
Stride<Stride<_128, _1, ...>, Stride<_64, _8>>>;
```
Finally, we get this entire pattern repeating four times, once for each warp, down the `M`-mode starting at `(m,n) = (16,0) = 16`. where two core matrices that belong to the same warp are stacked on top of each other. This makes the size of the final sub-mode of M 4. As for the stride, this time we go to `(T32, V0)`, which makes it a stride of 32.
```cpp
// (T128,V4) -> (M64,N8)
using CLayout = Layout<Shape <Shape < _4, _8, _4>, Shape < _2, _2>>,
Stride<Stride<_128, _1, _16>, Stride<_64, _8>>>;
```
This is the full `CLayout` for 64x8 accumulators. The GMMA instructions include 64xN variants with `N = [16,32,64,128,256]` where this 64x8 pattern is repeated giving each thread additional values. As this starts at `(m,n) = (0,8) = 512`, this is easy to account for in our `CLayout`. For example, the 64x128 `CLayout` is
```cpp
// (T128,V64) -> (M64,N128)
using CLayout = Layout<Shape <Shape < _4, _8, _4>, Shape < _2, _2, _16>>,
Stride<Stride<_128, _1, _16>, Stride<_64, _8, _512>>>;
```
where we see 16 copies of the 64x8 tile.
### A and B Layout Mapping
GMMA atoms that consume A and B sources directly from shared memory are a bit interesting. The GMMA Descriptor is constructed on an entore tile of A and/or B data in shared memory rather than being partitioned by threads. That is, every thread sees the entire tile of data and the tile is not reordered so that the descriptor can be constructed on it. In `ALayout` form, this can be expressed
```cpp
// (T128,V64x8) -> (M64,K16)
using ALayout = Layout<Shape <_128, Shape <_64,_16>>,
Stride< _0, Stride< _1,_64>>>;
```
That is, all threads are mapped the to `(m,k) = (0,0) = 0` element and the values (and shape of the values) remains unchanged. The GMMA Descriptor Constructor can then inspect the `(M,K)` layout of this data and create an appropriate GMMA Descriptor or produce an error message saying the data is in an invalid layout for GMMA.
media/docs/cute/0x_gemm_tutorial.md
+668
View File
@@ -0,0 +1,668 @@
# CuTe dense matrix-matrix multiply tutorial
This section uses the CuTe functionality to write
a dense matrix-matrix multiply implementation.
## A simple dense matrix-matrix multiply example
In this section, we will go through
[this example](../../../examples/cute/tutorial/sgemm_nt_1.cu).
It illustrates a blocked GPU implementation of GEMM
that uses the building blocks of CuTe
to construct global and shared memory layout mappings
and partition threads among them.
This example is closest to the blocked GEMM
that a computer science student might be asked to implement
in a first-year graduate school
or upper-division undergraduate scientific computing course.
Readers who understand this section may also wish to study
CUTLASS's implementation of the stream-K GEMM algorithm,
which uses many features of CuTe.
### Filename and high-level interface
First, let's look at the example's filename `sgemm_nt_1.cu`.
"SGEMM" is the BLAS (Basic Linear Algebra Subroutines) abbreviation
for "Single-precision real, GEneral, Matrix-matrix Multiply."
(If we want to refer to matrix-matrix multiply for all data types,
we say "GEMM.")
The BLAS project started in the 1970s.
You can learn more about its history in Turing Award winner Jack Dongarra's
2004 Oral History interview by SIAM
(the Society for Industrial and Applied Mathematics),
and also in the C++ Standard document [P1417](https://wg21.link/p1417).
The abbreviation SGEMM unpacks as follows.
* "Single-precision" is Fortran-speak for float.
The BLAS supports four different matrix or vector element types:
* S for single precision (`float`),
* D for double precision (`double`),
* C for complex float (like C++'s `std::complex<float>`,
where each of the real and imaginary components has type `float`),
and
* Z for complex double (like C++'s `std::complex<double>`).
* "GEneral" means that the matrix is represented
as a two-dimensional dense array
and not assumed to have any kind of symmetry.
The BLAS supports a variety of matrix representations,
including
* SY: SYmmetric,
* HE: HErmitian,
* TR: TRiangular,
* GB: General Banded,
* SB: Symmetric Banded,
* SP: Symmetric Packed, and
* TP: Triangular Packed.
* MM means "Matrix-matrix multiply," as opposed to other operations,
like MV (Matrix-Vector multiply).
The string "nt" in the filename means that
the first input matrix A is "Not transposed,"
while the second input matrix B is "Transposed."
That is, the function computes `C := beta * C + alpha * A * B^T`,
where the superscript T denotes the transpose of the matrix.
(We never change the input matrix in place or
store its entire transpose explicitly.
Instead, we reinterpret its data in place.)
GEMM's TRANSA and TRANSB arguments lets users specify
the transpose or Hermitian transpose (if complex)
of either or both input matrices A or B.
It turns out that implementations favor this "NT" case,
along with "TN" (A is Transposed, B is Not transposed).
We will explain why below.
As described, the original BLAS GEMM specifies
the dimensions of its matrices
as A is M x K, B is K x N, and C is M x N.
Out of convenience, CuTe interprets A
as M x K, B as N x K, and C as M x N. Instead of row-major or column-major (or Transposed
and Not-Transposed like above), we like to be more specific with M-major, N-major, or K-major.
Regardless, we'll still use the BLAS "NT" notation for high-level descriptions
of kernels when it's appropriate.
Now, let's look at the code.
We'll start with the kernel entry point `gemm_device`
at the top of the file.
```c++
template <class MShape, class NShape, class KShape,
class TA, class AStride, class ABlockLayout, class AThreadLayout,
class TB, class BStride, class BBlockLayout, class BThreadLayout,
class TC, class CStride, class CBlockLayout, class CThreadLayout,
class Alpha, class Beta>
__global__ static
__launch_bounds__(decltype(size(CThreadLayout{}))::value)
void
gemm_device(MShape M, NShape N, KShape K,
TA const* A, AStride dA, ABlockLayout blockA, AThreadLayout tA,
TB const* B, BStride dB, BBlockLayout blockB, BThreadLayout tB,
TC * C, CStride dC, CBlockLayout , CThreadLayout tC,
Alpha alpha, Beta beta);
```
There are many template parameters;
we'll explain them all in due time.
`TA`, `TB`, and `TC` are the element types
of the matrices `A`, `B`, and `C`, respectively.
The two scalar constants `alpha` and `beta`
are part of what GEMM computes: `C = beta * C + alpha * A * B`.
Unlike the (traditional Fortran and C) BLAS,
CuTe lets you mix different matrix element types and/or scalar types.
The compiler will help, but it's somewhat up to you
to use types that are safe and efficient on the GPU.
For example, a custom arbitrary-precision real type
that does dynamic allocation inside may not work on the GPU at all.
Even if it does, it may not perform well.
This leaves five kinds of things to explain:
1. Shapes
2. Strides
3. Block layouts
4. Thread layouts
5. Launch bounds
### Shapes
The original Fortran BLAS GEMM lists the matrices' dimensions
in the order M, N, K. CuTe also uses this convention.
The "MShape" is just M,
the NShape is just N,
and the KShape is just K.
In this example, they are dynamic (run-time) values
defined at the top of the `gemm` host function
that invokes the device kernel.
```c++
// Define shapes (dynamic)
auto M = int(m);
auto N = int(n);
auto K = int(k);
```
Note that the function takes M, N, and K.
It doesn't take the shapes of the three matrices separately,
as (say) three different `Shape<int, int>` objects.
This is because matrix-matrix multiply constrains the shapes.
There's nothing mysterious about `int` here;
it's the usual C++ built-in integral type.
`auto M = int(m)` is a way to say
"convert `m` to an `int` if it's not already an `int`,
and assign it to the freshly declared variable `M`."
CuTe also has a capitalized `Int<Value>` templated type
for representing values as compile-time constants.
For example, `Int<5>` represents a compile-time `int` value 5.
(CuTe implements these as subclasses
of the C++ Standard Library class `std::integral_constant`.)
The above `gemm_device` function is templated on the types
of M, N, and K; this shows that CuTe can represent dimensions
as either run-time or compile-time values.
If you're familiar with the mdspan class going into C++23,
you might notice that CuTe represents shapes
a bit differently from mdspan.
mdspan uses `extents<class IndexType, size_t ... Extents>`
to represent a shape.
The `Extents` are zero or more compile-time values
(see below) representing the dimensions in the shape.
The `Extents...` are "non-type template parameters" (NTTPs) --
that is, they are not types, but compile-time values of type `size_t`.
If you use the special reserved `size_t` value `std::dynamic_extent`
as an extent value,
the resulting dimension is a run-time value
and is stored in the `extents` instance.
Any other extent value is a compile-time value
that is encoded in the extents type itself.
In contrast, CuTe represents a shape as `Shape<class ... Types>`.
The `Types...` are actual types, not NTTPs.
A built-in integral type like `int` or `uint64_t`
denotes a run-time dimension that is stored in the `Shape` instance,
while a compile-time value like `Int<5>`
encodes a compile-time dimension.
For example, the CuTe equivalent of
`extents<int, 3, dynamic_extent, 5>`
is `Shape<Int<3>, int, Int<5>>`.
#### Compile-time-ness of values
C++ values have three levels of "compile-time-ness":
1. dynamic (run-time) values,
2. constexpr values, and
3. static (compile-time) values.
(Rather than saying "C++ has,"
it's more accurate to say "C++17 has."
C++20 introduces `consteval` or "immediate" functions,
which make attempting to evaluate the function at run time
(any call not in an unevaluated context) a compiler error.
We'll ignore those for this tutorial,
since CuTe only requires C++17.)
The `constexpr` keyword was introduced in C++11.
It means something like
"the compiler can evaluate this expression at compile time."
It does NOT mean "the compiler MUST evaluate this at compile time."
If you use a `constexpr` expression in a `static_assert`
or as a non-type template argument,
then the compiler must evaluate the expression at compile time.
However, for `constexpr` occurring in other places,
the compiler may choose to store the value in registers or memory,
and/or do computations with the value at run time.
In some cases, the compiler must do that.
The following example shows that the compiler
might need to store `constexpr` values in memory sometimes.
```c++
// Some function defined in a different compilation unit.
extern int foo(int* x);
int bar()
{
constexpr int value = 42; // a compile-time constant
// Even constexpr variables have a sizeof,
// because we still might need to take their address.
static_assert(sizeof(value) == 4);
// Compiler can't inspect foo to see how it uses the value,
// so it has to store the value in some memory location
// so that we can pass its address to the function.
return foo(&value);
}
```
"Static" is an unfortunately overloaded term in C++. Sometimes it means "the opposite of instance," like a "static function" or "static member" of a class. (Some programming languages, like Java, say "class method" to refer to a "static function of a class.") That's not what we mean here. Instead, we mean "part of a compile-time type." For example, `Int<1>` encodes the value 1 at compile time, as part of the type of a templated class `Int<Value>`. `Int<3>` and `Int<4>` have different types. You can get the value of of the type like this: `Int<3>::value`. (The `value` is a `static constexpr` member of the class, where "static" means "opposite of instance.") As soon as you go from `Int<3>` to `Int<3>::value`, you've gone from (3) above (a compile-time value) to (2) above (a `constexpr` value). In some situations, this may mean that the compiler treats it as a run-time value.
#### Strides
We define a layout using both shapes and strides.
The shape just tells you the dimensions (modes, etc.) of the array.
The strides tell you the mapping from a multidimensional index
into a one-dimensional offset.
Here, we're describing the shapes and strides
of the "global" matrices A, B, and C.
The example defines the global matrices' strides
near the top of the `gemm` function.
```c++
// Define strides (mixed)
auto dA = make_stride(Int<1>{}, ldA); // (dM,dK)
auto dB = make_stride(Int<1>{}, ldB); // (dN,dK)
auto dC = make_stride(Int<1>{}, ldC); // (dM,dN)
```
To evaluate this mapping for a given multidimensional index, take the dot product of the indices with the strides. For example, the offset of `A(index_m, index_k)` is `index_m * 1 + index_k * ldA`. Note the implications for the compile-time-ness of the offset. Any run-time value among either the shape or the strides makes the offset a run-time value. Of course, if a particular stride is a compile-time constant (especially 1), it's easier for the compiler to optimize the arithmetic and result.
Note that in the original source code,
this example is missing the comments after each line.
We've added them in here,
as they stir a brief digression about shapes and modes.
The comment after B says (dN, dK), not (dK, dN).
This means that B is treated as an N x K matrix
instead of a K x N matrix.
As mentioned, CuTe follows the convention
that the meaning of matrix modes is
(M,K) for A, (N,K) for B, and (M,N) for C.
In particular, CuTe's convention is that
"the reduction mode is outermost."
The "reduction mode" of `Shape<M, N, K>` is K.
That's the mode over which we do a reduction,
that is, sum up products of matrix entries.
The K mode disappears in the output C.
"Outermost" here means "rightmost"
(literally, appearing rightmost in the list M, N, K).
Note that the shapes form a kind of Einstein tensor notation.
GEMM does Shape<M, N> = Shape<M, K> * Shape<K, N>.
In Einstein notation, the repeated index indicates
a sum of that term over all values of K.
We say in general that the leftmost mode is the "inner(most)" mode,
and the rightmost mode is the "outer(most)" mode.
This is because,
along with CuTe's convention of thinking of arrays as logically column major,
the leftmost mode is most commonly the mode with the most spatial locality.
It's very often the "most contiguous" mode.
For this reason, it's "the mode that we want in the innermost loop"
(in the nesting of loops that implements GEMM).
This is why we call it the "innermost" mode.
Its contiguity means that also call the innermost mode the "vector mode."
The vector mode also has special meaning:
it contains all of the information needed
to execute the smallest possible computation or communication operations
on hardware, that is, what CuTe calls the "atoms."
Modes are like units conceptually.
For example, you shouldn't mix M-mode indices with K-mode indices.
However, CuTe does nothing to enforce this.
(For example, CuTe does not require use of "tagged" index types.
Indexing works with the usual integer types.)
The previous paragraph relates to shapes, not strides.
Returning to the strides, the above code describes these strides as "mixed."
This means that they include both run-time and compile-time values.
For example, the stride between A(m, k) and A(m+1, k) is `Int<1>`,
a compile-time value 1. The stride between A(m, k) and A(m, k+1),
however, is `ldA`, the "leading dimension of A," a run-time value.
The "leading dimension" of a matrix
refers to the stride between consecutive columns of a column-major matrix
(where the stride between consecutive rows is 1),
or the stride between consecutive rows of a row-major matrix
(where the stride between consecutive columns is 1).
This is a naming convention from the BLAS
and libraries that use it, like LAPACK.
For the purpose of this tutorial, it's just a naming convention
for "the stride that isn't the compile-time constant 1."
#### M-major, N-major, K-major
Note that we haven't uttered the phrases "column-major" or "row-major" here. This is where the experience of a BLAS user diverges from the experience of a BLAS implementer. BLAS users speak of "column-major" and "row-major" layouts. C++23's `mdspan` class encodes these as `layout_left` resp. `layout_right`. However, we don't speak of "column-major" or "row-major" in our GEMM implementations.
We say that a matrix is "M-major" if it is stride 1 in the M-mode, "N-major" if it is stride 1 in the N-mode, or "K-major" if it is stride 1 in the K-mode. In the above code, A has shape (M, K) and strides (1, ldA). Since A has stride 1 in the M mode, we say that A is "M major." B has shape (N, K) and strides (1, ldB), so B is "N-major." Similarly, C has shape (M, N) and strides (1, ldC), so C is "M major."
How do we translate this into the BLAS user's experience?
The following table illustrates for B and C.
(Throughout the table, "Impl" stands for "implementation.")
Note that the implementation reverses the order of B's modes,
and flips B's strides.
Recall that one evaluates a layout
by taking the dot product of the indices and strides.
Thus, reversing the order of both the modes and the strides
does not change this evaluation.
| Matrix | User's shape | User's layout | User's strides | Impl layout | Impl shape | Impl strides |
| --- | --- | --- | --- | --- | --- | --- |
| C | M x N | Column major | (1, LDC) | M-major | (M, N) | (1, LDC) |
| A | M x K | Column major | (1, LDA) | M-major | (M, K) | (1, LDA) |
What about the matrix B? We explained above that B is N-major. How would that translate back into the BLAS user's experience? We take a hint here from the filename including "nt." The "nt" part of the name means that A is not transposed, while B is transposed. The BLAS convention (see e.g., [the documentation for DGEMM](https://netlib.org/lapack/explore-html/d1/d54/group__double__blas__level3_gaeda3cbd99c8fb834a60a6412878226e1.html)) is that if you take the transpose, then the dimensions refer to the transpose ("with op( A ) an m by k matrix, op( B ) a k by n matrix and C an m by n matrix"). Thus, this example actually computes `C = beta * C + alpha * A * B^T`, where `B^T` is an K x N matrix with strides (LDB, 1). The user's "original" matrix B is thus N x K, with strides (1, LDB) -- that's a column-major layout. (Reversing the modes and the strides preserves the layout, since evaluating the layout mapping just takes the dot product of indices and strides.) This lets us expand the above table to include B.
| Matrix | Transposed? | User's shape | User's layout | User's strides | Impl layout | Impl shape | Impl strides |
| --- | --- | --- | --- | --- | --- | --- | --- |
| C | No | M x N | Column major | (1, LDC) | M-major | (M, N) | (1, LDC) |
| A | No | M x K | Column major | (1, LDA) | M-major | (M, K) | (1, LDA) |
| B | Yes | N x K | Column major | (1, LDB) | N-major | (N, K) | (1, LDB) |
CuTe developers say: "In CuTe, you can't tell transposed
apart from non-transposed, MN-major from K-major, etc.
without inspecting the strides."
It's now a bit more clear what that means.
CuTe doesn't see whether A or B are transposed.
Instead, CuTe sees shapes and strides.
A CuTe developer must reason backwards from the shapes and strides
in order to see what the BLAS user sees.
Why does CuTe do this? Consider that matrix multiply performs a reduction in the K-mode. From the user's perspective, it's reducing across rows of the first input matrix, but across columns of the second input matrix. If we instead mentally flip the modes of the first input matrix, then the implementation reduces over columns (the K mode) of both input matrices. This leads to two cases in which the implementation can effectively treat both input matrices in the same way. (If you call it with A and B reversed, it should even give the same results for these cases.)
| Case | User asks for A | User asks for B | Abbreviation |
| --- | --- | --- | --- |
| A is M major, B is N major | Not transposed | Transposed | NT |
| A and B are both K major | Transposed | Not transposed | TN |
This is why an introductory example starts with NT or TN.
For a summary of the four different transpose options for A and B,
and their corresponding implementation layouts,
please see the table below.
| Transpose abbreviation | User sees A transposed? | User sees B transposed? | A's impl layout | B's impl layout |
| --- | --- | --- | --- | --- |
| NT | No | Yes | M major | N major |
| TN | Yes | No | K major | K major |
| NN | No | No | M major | K major |
| TT | Yes | Yes | K major | N major |
#### MN-major and K-major
As we mentioned above, there are two "preferred arrangements," TN and NT. In the TN arrangement, both A and B are K-major. In the NT arrangement, A is M-major and B is N-major. Even though the two stride-1 modes in NT have different names, it's still the leftmost mode for both A and B that has stride 1. Thus, we can think of the NT arrangement as "MN-major," analogous to how the TN arrangement is "K-major."
The two preferred arrangements tend to work themselves into implementations, particularly when they use hardware instructions for accelerating matrix multiplies of blocks. In some cases, the hardware instruction may require NT (MN-major) or TN (K-major). For NN or TT, such instructions would require an intermediate transpose -- for example, when loading from global memory to shared memory.
### Block layouts
Efficient matrix multiply implementations loop over blocks.
For example, a typical GPU implementation strategy
is for each thread block to iterate over some number of blocks.
In the example, this loop occurs near the end of `gemm_device`.
```c++
// TUTORIAL: Example of a very simple compute loop
// Data is read from global to shared memory via the tA|tB partitioning
// gemm(.) operates on the shared memory directly via the tC partitioning
auto k_max = size<2>(tAgA);
for (int k = 0; k < k_max; ++k)
{
// Copy A and B blocks from global memory to shared memory.
copy(tAgA(_,_,k), tAsA);
copy(tBgB(_,_,k), tBsB);
// On some architectures, copy may be asynchronous.
// This may call for extra synchronization instructions
// beyond just __syncthreads().
__syncthreads();
// Compute gemm on shared memory input and register accumulator.
// The "epilogue" after this loop will copy the accumulator
// from the register file into global memory.
gemm(tCsA, tCsB, tCrC);
__syncthreads();
}
```
We will explain the notation in this loop below. The important things to remember are that the coordinate `k` loops over the blocks which the calling thread is supposed to compute, the `copy` functions copy A resp. B blocks from global memory (the first argument) to shared memory (the second argument -- same as C++'s `std::copy`, but the opposite of `memcpy`), and the `gemm` function computes C += A * B on the shared memory blocks.
It turns out that copy takes an optional first argument, the "atom," as in the following.
```c++
copy(atom, source, destination);
```
The "atom" is metadata that explains how to do the copy operation.
There are a few topics to push onto the stack.
The copy function call shows a notation for taking slices of a tensor. A CuTe `Tensor` is a multidimensional array view. It consists of a pointer and a `Layout`. You can learn more about `Tensor`s elsewhere in CuTe's documentation, but for now, please note that `tAgA(_,_,k)` means "create a Tensor that views (i, j, k) for all valid i, all valid j, and a specific value of k." The result has rank one less than the original Tensor. CuTe's underscore means the same thing as a single stand-alone colon in Fortran or Matlab. Note also that CuTe uses the same notation for slices as for tensor indexing. The implementation can distinguish the two cases by checking whether any of the arguments is an underscore. In contrast, the C++23 class mdspan uses a separate function, `submdspan` (not in C++23, and proposed for C++26; see [P2630](https://wg21.link/p2630)), for slicing.
Fully understanding what `copy` and `gemm` do calls for learning about thread layouts as well, so we will wait to explain them completely. For now, note that these functions are implicitly parallel, as they are called collectively by all threads in a thread block.
The block dimensions are defined near the top of the host function `gemm`.
```c++
// Define block sizes (static)
auto bM = Int<128>{};
auto bN = Int<128>{};
auto bK = Int< 8>{};
```
We see that these are fully compile-time dimensions. This is often the case, especially when we use hardware instructions that only work for certain problem dimensions. Three lines of code immediately below these construct the block layouts.
```c++
// Define the block layouts (static)
auto sA = make_layout(make_shape(bM,bK));
auto sB = make_layout(make_shape(bN,bK));
auto sC = make_layout(make_shape(bM,bN));
```
Here, the block layouts just come from the block dimensions. A Layout has two things: a Shape, and Strides. If the caller does not provide Strides, then CuTe computes Strides corresponding to the default "column-major" arrangement of data. This just happens to match the global matrices' layouts, but in general doesn't have to. For example, in the NN or TT cases, we may want to transpose one of the input matrices when copying from global memory to shared memory.
The example "comments out" some code that prints all the layouts on "thread 0" of each thread block. If you enable the printing code and run the example, it will print all the layouts. For example, sA prints as
```
sA
(_128,_8)
(_1,_128)
```
and sB prints as
```
sB
(_128,_8)
(_1,_128)
```
consistently with the definitions above.
If you have looked at other GEMM examples in CuTe, you might be wondering about hardware matrix-matrix multiply instructions. Those instructions tend to require certain values for shapes and strides, that may be a function of the matrix's element type. CuTe knows about these instructions and their required shapes and strides. We will go into more detail about that elsewhere.
The `gemm_device` top-level kernel uses these block layouts to allocate shared memory buffers for A and B tiles.
```c++
// Shared memory buffers
__shared__ TA smemA[cosize_v<ABlockLayout>];
__shared__ TB smemB[cosize_v<BBlockLayout>];
```
Note how the shared memory buffers' sizes depend only on the A resp. B layouts (and element sizes). What's a `cosize_v`? The "`_v`" is a C++ naming convention that specifies a function from one or more template argument(s), to a value. In this case, it's a number of elements. A layout is a function from a set of multidimensional coordinates to a set of one-dimensional array offsets. It's a function, so we can speak of its domain and codomain. The "cosize" of a layout is the size of its codomain. (See e.g., CuTe's implementation of `Layout`.) If we want to allocate a linear array, for which all the offsets produced by a layout are valid, then we can use the cosize of the layout as the length of the array (in terms of number of elements, not in terms of number of bytes).
### Thread layouts
CuTe uses a `Layout` to describe the assignment of threads to work items.
In this example, the host function `gemm` constructs the thread layouts
for A, B, and C.
```c++
// Define the thread layouts (static)
auto tA = make_layout(make_shape(Int<32>{}, Int< 8>{}));
auto tB = make_layout(make_shape(Int<32>{}, Int< 8>{}));
auto tC = make_layout(make_shape(Int<16>{}, Int<16>{}));
```
That is, the thread layout for the A read is M-major 32x8, for the B read is N-major 32x8, and for the C compute/write is M-major 16x16. These thread layouts will partition the data for their respective stages.
#### The example uses compile-time thread and block layouts
Note that the device function `gemm_device` insists that all the thread and block layouts are static -- that is, known at compile time. You can see this from the `CUTE_STATIC_ASSERT` statements near the top of `gemm_device`. `CUTE_STATIC_ASSERT` is a wrapper for `static_assert`, which fails at compile time if its condition is `false`.
```c++
// Preconditions
CUTE_STATIC_ASSERT(is_static<ABlockLayout>::value);
CUTE_STATIC_ASSERT(is_static<BBlockLayout>::value);
CUTE_STATIC_ASSERT(is_static<CBlockLayout>::value);
CUTE_STATIC_ASSERT(is_static<AThreadLayout>::value);
CUTE_STATIC_ASSERT(is_static<BThreadLayout>::value);
CUTE_STATIC_ASSERT(is_static<CThreadLayout>::value);
```
Use of static layouts has two advantages. First, it makes it easier to prove correctness of the algorithm. If the code compiles, it's likely correct. (On the other hand, new CuTe users may find themselves doing more debugging at compile time than they have before.) Second, it makes it easier and faster for CuTe to dispatch to the correct optimized implementations (called "atoms" -- see below) for copying blocks and performing matrix multiplies.
#### The example's block gemm is parallel over elements of C
In the actual device function, `tC` has layout `CThreadLayout`. You might recall that the kernel function `gemm_device` uses `CThreadLayout` to derive the launch bounds, specifically the maximum number of threads per block. The launch bounds show up in the declaration of `gemm_device`.
```c++
template <class MShape, class NShape, class KShape,
class TA, class AStride, class ABlockLayout, class AThreadLayout,
class TB, class BStride, class BBlockLayout, class BThreadLayout,
class TC, class CStride, class CBlockLayout, class CThreadLayout,
class Alpha, class Beta>
__global__ static
__launch_bounds__(decltype(size(CThreadLayout{}))::value)
void
gemm_device(MShape M, NShape N, KShape K,
TA const* A, AStride dA, ABlockLayout blockA, AThreadLayout tA,
TB const* B, BStride dB, BBlockLayout blockB, BThreadLayout tB,
TC * C, CStride dC, CBlockLayout , CThreadLayout tC,
Alpha alpha, Beta beta);
```
The "size" of `CThreadLayout` is the total number of threads, 16 * 16 = 256. (We take `::value` because the size is actually `Int<256>`, a compile-time constant with a `static constexpr int value = 256` member.) This suggests that the block gemm function (in the loop over blocks) parallelizes over elements of the C block. We can see this as well from the kernel launch (at the end of the `gemm` host function), which uses the size of `CThreadLayout` as the block dimension.
```c++
// Define the thread layouts (static)
auto tA = make_layout(make_shape(Int<32>{}, Int< 8>{}));
auto tB = make_layout(make_shape(Int<32>{}, Int< 8>{}));
auto tC = make_layout(make_shape(Int<16>{}, Int<16>{}));
dim3 dimBlock(size(tC));
dim3 dimGrid(ceil_div(size(M), size(bM)),
ceil_div(size(N), size(bN)));
gemm_device
<<< dimGrid, dimBlock, 0, stream >>>
(M, N, K,
A, dA, sA, tA,
B, dB, sB, tB,
C, dC, sC, tC,
alpha, beta);
```
Note that dimBlock is single-dimensional (despite being a dim3), as the size of a layout is a single value. We can see this also because the example only ever uses `threadIdx.x`, not `threadIdx.y`. Yet, C's thread layout has shape (16, 16). What's with that? Recall that a thread layout maps from a "logical" coordinate space (possibly multidimensional tuples of indices) to (one-dimensional) integer indices. In this case, `CThreadLayout` maps from pairs of indices in the Cartesian product space {0, 1, 2, ..., 15} x {0, 1, 2, ..., 15}, to one-dimensional indices 0, 1, 2, ..., 255. The latter, the output of `CThreadLayout`, is the actual thread index `threadIdx.x` in this case. `CThreadLayout` has only a shape (16, 16) and no nondefault strides, so it uses CuTe's default column-major arrangement (with strides (1, 16) in this case).
#### What does `local_tile` do?
The following code near the top of `gemm_device`
operates on the "global" (input and output) matrices A, B, and C
(where mA, mB, and mC are their Tensor representations).
```c++
// Get the appropriate blocks for this thread block --
// potential for thread block locality
auto blk_shape = make_shape(size<0>(sA), size<0>(sB), size<1>(sB)); // (BLK_M,BLK_N,BLK_K)
auto blk_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k)
Tensor gA = local_tile(mA, blk_shape, blk_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k)
Tensor gB = local_tile(mB, blk_shape, blk_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
Tensor gC = local_tile(mC, blk_shape, blk_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N)
```
There are two new features here:
* `make_coord`, which returns a `Coord`, a multidimensional index which can be used as the input of a `Layout`; and
* `local_tile`, which we will explain below.
The `Coord`(inate) `blk_coord` refers to the set of blocks (indexed by k -- the underscore here indicating a free parameter) our thread block will access. (The index k here doesn't mean the K mode; it's the same index as in the loop over blocks that does the computation.)
If we print out the `gA`, `gB`, and `gC` layouts, we get the following.
```
gA
(_128,_8,512)
(_1,5120,40960)
gB
(_128,_8,512)
(_1,5120,40960)
gC
(_128,_128)
(_1,5120)
```
All of these layouts come from the original input or output matrices A, B, and C. Thus, they preserve the original strides, which are all the same in this example (when using default problem dimensions), 5120. This is most easily seen in the gC layout. For the other layouts, there is a clue in 5120 * 8 = 40960. That is, every time we increase k by one, we "skip over 8 columns" of the global matrix, over to the next block of data. This illustrates an important feature of CuTe, that it can view the same data with different modes and/or strides, as a way to identify parallelism or locality.
## Next steps
The above "simple GEMM" example's performance on many problems
is asymptotically optimal
with respect to the GPU's floating-point throughput.
Getting nearly peak performance
relative to the GPU's floating-point throughput,
for a wider variety of problem dimensions,
calls for more advanced techniques.
Please refer to other examples in this repository
to learn more about those techniques.
For example, the
[predication section of the tutorial](./0y_predication.md)
explains what to do if a matrix tiling
doesn't perfectly divide the matrix.
### Implement GEMM as generalized tensor constraction (GETT)
"GETT" here stands for "general(ized) tensor times tensor,"
a tensor contraction.
CuTe permits matrices to have nested `Layout`s.
For example, a matrix A can have a nested `Layout` for its M and N modes.
This means that we can use a "matrix" (`Tensor` with two modes)
to represent any `Tensor`.
This amounts to a "native hierarchical representation."
As a result, we can implement GETT by using
our existing GEMM implementation layers,
with a little bit of fancy custom predication for the K mode.
This is because the stride type of A
and the problem shape itself
are CuTe Shapes and Strides.
This lets us represent the hierarchical modes
of a tensor contraction problem
(which still fundamentally only have 4 modes --
batch mode,
two outer modes (one for A and one for B),
and one reduction mode --
each of which can now have as many nested modes as you want
for the contraction's inputs).
We thus implement GETT as contraction just in one mode -- the K mode.
However, K itself can be hierarchical and can have noncontiguous strides.
We can reorder the modes such that all contraction modes
become a single, possibly hierarchical K mode in the kernel.
This is how we would encode a contraction in multiple modes at once.
media/docs/cute/0y_predication.md
+217
View File
@@ -0,0 +1,217 @@
# Predication: What to do when tiling isn't perfect
The [GEMM tutorial](./0x_gemm_tutorial.md) shows how
we compute a matrix-matrix multiply
by iterating over tiles of the input matrices and output matrix.
The examples all assume that the tiles fit evenly into the matrices,
with no remainder.
What do we do if this is not the case?
For example, we might want to tile a 41 x 55 matrix into 4 x 8 tiles,
but 41 / 4 is 10 remainder 1, and 55 / 8 is 6 remainder 7.
What do we do with those "leftover" parts of the matrix?
Another way to say this, is that `logical_divide`
(CuTe's way of tiling layouts) "rounds up."
For example, if `N` is the layout (1000, 1) and `B` is the layout (128, 1),
then `logical_divide(N, B)` is the layout ((128, 8), (1, 128)).
This effectively rounds up the original shape N = 1000
into an 128 x 8 matrix (as if N = 1024).
What about those last 24 elements,
that aren't part of the original data?
The idiomatic CuTe way to solve this problem is through "predication."
Rather than trying to reason about the "remainder tiles,"
CuTe instead rounds up, but only tries to access data in each tile
that are part of the matrix.
This corresponds well with how our GPUs optimize:
branches without warp divergence are relatively fast.
It also matches the usual CUDA idiom
when dividing N work items in 1-D fashion over B thread blocks:
first test if "my thread" is out of bounds before doing work.
There are a few ways to figure out
which elements need to be predicated.
In-kernel GEMMs like to do this in the following way.
```c++
// Create the predicate tensor
Layout idA = make_layout(shape(A)); // e.g. 1000:1
Layout idAB = logical_divide(idA, B); // e.g. (128,8):(1,128)
Tensor pred = make_tensor<bool>(shape(idAB));
for (int i = 0; i < size(pred); ++i) {
pred(i) = idAB(i) < size(A);
}
// ... intervening code ...
// Use the predicate tensor. c is some coordinate.
// This code would likely live inside some algorithm.
if (pred(c)) { copy(idAB(c), smem(c)); }
```
The general procedure is that we
1. create an "identity" layout (`Layout idA = make_layout(shape(A))`,
in the above example) with the same shape as our original data;
2. repeat the same tiling/partitioning/slicing (possibly rounding up)
on that identity layout (`Layout idAB = logical_divide(idA, B)`);
3. create a "predicate tensor" by comparing the coordinates
of that reference layout with the bounds of the original layout;
and then
4. use the predicate tensor to mask off accesses to out-of-bounds elements.
For example, suppose that we've partitioned A and B tiles
across threads as follows.
```c++
Tensor tAgA = local_partition(gA, tA, thread_idx); // (THR_M,THR_K,k)
Tensor tAsA = local_partition(sA, tA, thread_idx); // (THR_M,THR_K,PIPE)
Tensor tBgB = local_partition(gB, tB, thread_idx); // (THR_N,THR_K,k)
Tensor tBsB = local_partition(sB, tB, thread_idx); // (THR_N,THR_K,PIPE)
```
`tAgA` and `tBgB` partition the global A resp. B matrices over threads,
and `tAsA` and `tBsB` partition the shared memory tiles of A resp. B over threads.
The following code creates predicate tensors
corresponding to `tAgA` and `tBgB`.
They will be computed once in the prologue.
and will be used to mask off instructions in the inner loop.
```c++
Tensor tApA = make_tensor<bool>(make_shape (size<0>(tAgA), size<1>(tAgA)),
make_stride( Int<1>{}, Int<0>{}));
Tensor tBpB = make_tensor<bool>(make_shape (size<0>(tBgB), size<1>(tBgB)),
make_stride( Int<1>{}, Int<0>{}));
```
We're only thread-parallelizing over the leftmost (row) dimension,
so we only need to predicate over the leftmost dimension.
Thus, we can make the rightmost (column) stride zero,
since we will never actually address the rightmost dimension.
The following code creates "two-dimensional identity tensors"
that map coordinates (m,k) -> (m,k)
for the tile of data within the thread block.
```c++
Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) -> (blk_m,blk_k)
Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) -> (blk_n,blk_k)
```
The following lines then tile and partition
the two reference tensors
in exactly the same way the data were tiled and partitioned
into `tAsA` and `tBsB`.
```c++
Tensor tAcA = local_partition(cA, tA, thread_idx);
Tensor tBcB = local_partition(cB, tB, thread_idx);
```
Tiling and partitioning affect the offset and domain,
but not the codomain of the tensors,
so we're left with tensors that map `(thr_m,thr_k) -> (m,k)`
where `(thr_m,thr_k)` is this particular thread's subtensor of the tile
and `(m,k)` is the original codomain: a coordinate into the original tile.
The unrolled loops in the code below then compare
the m- and n-coordinates of those tensors with our known maximums
to mask off elements we are not allowed to access.
```c++
Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) -> (blk_m,blk_k)
Tensor tAcA = local_partition(cA, tA, thread_idx);
Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) -> (blk_n,blk_k)
Tensor tBcB = local_partition(cB, tB, thread_idx);
// Populate
CUTE_UNROLL
for (int m = 0; m < size<0>(tApA); ++m) {
tApA(m,0) = get<0>(tAcA(m,0)) < m_max_coord;
}
CUTE_UNROLL
for (int n = 0; n < size<0>(tBpB); ++n) {
tBpB(n,0) = get<0>(tBcB(n,0)) < n_max_coord;
}
```
Those last `for` loops fill in the two predicate tensors.
In this case, we only need to predicate over the leftmost dimension,
so we only address `(m,0)` resp. `(n,0)`.
We can then use the predicate tensors in `copy_if`
to copy only the elements for which the corresponding
predicate tensor elements are nonzero.
```c++
// Prefetch k_tile=0, gate these on k_residue as well
CUTE_UNROLL
for (int k = 0; k < size<1>(tAsA); ++k) {
if (get<1>(tAcA(0,k)) >= -k_residue) { // some other condition on the column index
copy_if(tApA, tAgA(_,k,0), tAsA(_,k,0));
}
}
CUTE_UNROLL
for (int k = 0; k < size<1>(tBsB); ++k) {
if (get<1>(tBcB(0,k)) >= -k_residue) { // some other condition on the column index
copy_if(tBpB, tBgB(_,k,0), tBsB(_,k,0));
}
}
```
Here are some advantages of this "reference tensor" approach.
1. It doesn't depend on the layout/strides of the tensor
being predicated, just the logical bounds being imposed.
2. The partitioning stage can be anything.
3. It naturally extends to any-dimensional predication.
4. It's a natural generalization of a typical CUDA 1-D
parallel vector access pattern,
which computes an access index `k`
(e.g., as `blockDim.x * blockIdx.x + threadIdx.x`)
and then predicates access to the vector's `k`-th element
on whether `k` is in bounds.
As an example of (3), the epilogue predication does exactly the same thing,
```c++
// Repeat with a tensor of coordinates for predication
Tensor cC = make_identity_tensor(make_shape(size<0>(gC), size<1>(gC)));
Tensor tCcC = thr_mma.partition_C(cC);
const bool isBetaZero = (beta == 0);
CUTE_UNROLL
for (int i = 0; i < size(tCrC); ++i) {
if (elem_less(tCcC(i), make_coord(m_max_coord,n_max_coord))) {
tCgC(i) = isBetaZero ? alpha * tCrC(i) : alpha * tCrC(i) + beta * tCgC(i);
}
}
```
but with the mma responsible for the tiling/partitioning `tCcC`
so that the reference subtensor matches the accumulator's subtensor.
Then, the reference subtensor is predicated against the `if` bounds
(in both m- and n-coordinates) inside the `for` loop.
Another way to explain this is that we don't modify the tiles
to give you the "right" extents so that you never overrun.
Instead, we let you query the original coordinate
to see if that coordinate overruns.
This avoids all branching and variable/dynamic loop bounds
(thus maintaining load balance and synchronicity,
both very important in-kernel) in favor of predication.
It's also general enough to extend to all ranks,
all layouts of threads and data,
and all tiling/partitioning patterns.