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

Updates for 3.1 (#932)

Browse Source
This commit is contained in:
ANIKET SHIVAM
2023-04-29 09:34:27 -04:00
committed by GitHub
parent 6f8596ce3f
commit 7c04f95415
51 changed files with 1796 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files
media/docs/cute/01_layout.md
+28 -16
View File
@@ -3,8 +3,8 @@
## Layout
This document describes `Layout`, CuTe's core abstraction.
A `Layout` maps from (a) logical coordinate space(s)
to a physical index space.
A `Layout` maps from a logical coordinate space
to an 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.
@@ -19,7 +19,11 @@ This can help users do things like partition layouts of data over layouts of thr
## 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.
Any of the `Layout`s discussed in this section can be composed with data -- e.g., a pointer or an array -- to create a `Tensor`.
The `Layout`'s logical coordinate space represents the logical "shape" of the data,
e.g., the modes of the `Tensor` and their extents.
The `Layout` maps a logical coordinate into an index,
which is an offset to 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).
@@ -31,31 +35,31 @@ Both `Shape` and `Stride` are `IntTuple` types.
### IntTuple
An `IntTuple` is an integer or a tuple of `IntTuple`s.
An `IntTuple` is defined recursively as either a single 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`.
* `get<I>(IntTuple)`: The `I`th element of the `IntTuple`. For an `IntTuple` consisting of a single integer, `get<0>` is just that integer.
* `rank(IntTuple)`: The number of elements in an `IntTuple`. An int has rank 1, a tuple has rank `tuple_size`.
* `rank(IntTuple)`: The number of elements in an `IntTuple`. A single integer has rank 1, and 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.
* `depth(IntTuple)`: The number of hierarchical `IntTuple`s. A single integer has depth 0, a tuple of integers has depth 1, a tuple that contains a tuple of integers has depth 2, etc.
* `size(IntTuple)`: The product of all elements of the IntTuple.
* `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.
We write `IntTuple`s with parenthesis to denote the hierarchy. For example, `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.
A `Layout` is then a pair of `IntTuple`s. The first element defines the abstract *shape* of the `Layout`, and the second element 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
Since a `Layout` is just a pair of `IntTuple`s, we can define operations on `Layout`s analogous to those defined on `IntTuple`.
* `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.
* `depth(Layout)`: The number of hierarchical `Layout`s. A single integer has depth 0, a tuple of integers has depth 1, a tuple that contains a tuple of integers has depth 2, etc.
* `shape(Layout)`: The shape of the `Layout`.
@@ -86,7 +90,7 @@ These hierarchical access functions include the following.
### Vector examples
Then, we can define a vector as any `Shape` and `Stride` pair with `rank == 1`.
We define a vector as any `Shape` and `Stride` pair with `rank == 1`.
For example, the `Layout`
```
@@ -95,9 +99,9 @@ Stride: (1)
```
defines a contiguous 8-element vector.
Similarly, with a stride of `(2)`,
For a vector with the same Shape but a Stride of `(2)`,
the interpretation is that the eight elements
are stored at positions 0, 2, 4, $\dots$.
are stored at positions 0, 2, 4, $\dots$, 14.
By the above definition, we *also* interpret
@@ -168,9 +172,17 @@ auto layout_2x4 = make_layout(make_shape (2, make_shape (2,2)),
make_stride(4, make_stride(2,1)));
```
The `make_layout` function returns a `Layout`.
It deduces the returned `Layout`'s template arguments from the function's arguments.
Similarly, the `make_shape` and `make_stride` functions
return a `Shape` resp. `Stride`.
CuTe often uses these `make_*` functions,
because constructor template argument deduction (CTAD)
does not work for `cute::tuple` as it works for `std::tuple`.
## 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
The fundamental use of a `Layout` is to map between logical coordinate space(s) and an index space. For example, to print an arbitrary rank-2 layout, we can write the function
```c++
template <class Shape, class Stride>
media/docs/cute/02_layout_operations.md
+139 -9
View File
@@ -73,6 +73,7 @@ In C++, we identify a Tuple with the
`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.
@@ -136,7 +137,7 @@ This code produces the following text output.
```
`print(layout(1, 1))` prints the mapping of
the logical 2-D coordinate (1,1) to 1-D index, which is 4.
the logical 2-D coordinate (1,1) to the 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."
@@ -302,13 +303,13 @@ Both humans and CuTe compute composition using the following rules.
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, ...)$.
3. Composition is (left-)distributive with concatenation: $A \circ B = A \circ (B_0, B_1, ...) = (A \circ B_0, A \circ B_1, ...)$.
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)$.
denote a tuple of two layouts B and C, not their concatenation. Let $A = (A_0, A_1)$.
Then, $A \circ \langle B, C \rangle = (A_0, A_1) \circ \langle B, C \rangle = (A_0 \circ B, A_1 \circ C)$.
This allows the application of composition independently to sublayouts of $A$.
#### Examples: Reshape a vector into a matrix
@@ -359,6 +360,55 @@ The resulting layout has shape $(4,5)$, just as before. What are the strides?
5. Result: (4:10, 5:2), which by concatenation is (4,5) : (10,2).
#### Example: Reshape a matrix into another matrix
The composition $((20,2):(16,4) \circ (4,5):(1,4))$
expresses reshaping the matrix with layout (20,2):(16:4),
into a 4 x 5 matrix in a column-major way.
1. By deconcatenation, $(4,5) : (1,4)$ is $(4:1, 5:4)$.
2. Composition is distributive, so $(20,2):(16,4) \circ (4:1, 5:4)$ is $((20,2):(16,4) \circ 4:1, (20,2):(16,4) \circ 5:4)$.
3. $(20,2):(16,4) \circ 4:1$ has shape $4$ and stride $16$. (4:1 expresses picking the first 4 consecutive elements of (20,2):(16,4). These elements run down the 0th column (leftmost mode) of the layout, whose stride is 16.)
4. $(20,2):(16,4) \circ 5:4$ has shape $5$ and stride $64 = 4 \cdot 16$.
5. Result: $(4:16, 5:64)$, which by concatenation is $(4,5) : (16,64)$.
We get exactly this result with CuTe
if we use compile-time shapes and strides.
The following C++ code prints `(_4,_5):(_16,_64).`
```c++
using namespace cute;
auto a = make_layout(make_shape(Int<20>{}, _2{}), make_stride(_16{}, _4{}));
auto b = make_layout(make_shape( _4{}, _5{}), make_stride( _1{}, _4{}));
auto c = composition(a, b);
printf("\n");
print(c);
```
Results may _look_ different (but are the same mathematically)
if we use run-time integers.
The following C++ code prints `((4,1),(5,1)):((16,4),(64,4)).`
```c++
using namespace cute;
auto a = make_layout(make_shape(20, 2), make_stride(16, 4));
auto b = make_layout(make_shape( 4, 5), make_stride( 1, 4));
auto c = composition(a, b);
printf("\n");
print(c);
```
((4,1),(5,1)):((16,4),(64,4)) is effectively the same layout
as (4,5) : (16,64), because the 1s in the shape don't affect the layout
(as a mathematical function from one integer to one integer).
CuTe chooses not to simplify layout computations
with run-time values in them as much as it could,
because simplifications involving run-time values have a run-time cost.
### Product
CuTe includes four different kinds of layout products.
@@ -428,7 +478,7 @@ results in Shape ((2, 2), (3, 4)) and Stride ((1, 2), (16, 4)).
| (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,
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.
@@ -456,8 +506,8 @@ Shape ((3, 2), (4, 2)) and Stride ((16, 1), (4, 2)).
| (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 cyclic
distribution.
instead of appearing as blocks. Other references call this a "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
@@ -542,7 +592,87 @@ CuTe includes 3 different kinds of layout division operations.
We will summarize these in the sections that follow.
#### Logical divide : the intuitive tiling
#### Logical divide
##### Example worked in detail
This section will work the following logical divide example in detail.
```c++
Layout a = make_layout(24, 2);
Layout b = make_layout( 4, 2);
Layout c = logical_divide(a, b);
```
Logical divide produces a rank-2 `Layout`,
where mode 0 (the leftmost mode) corresponds to the divisor `b`,
and mode 1 (the rightmost mode) corresponds to the "remainder."
Intuitively, the remainder of 24 divided by 4 is 6,
so we know that mode 1 has 6 elements.
We just don't know its shape yet.
CuTe defines `logical_divide(a, b)` as
`composition(a, make_layout(b, complement(b, size(a))))`.
Here, `size(a)` is 24.
What is `complement(b, 24)`?
Intuitively, it means "the remainder,"
what's left over after applying `b` to 0, 1, 2, $\dots$, 23.
The layout 4:2 means "take 4 elements at even-numbered indices."
The following table overlays the range of 4:2
atop the complement's codomain 0, 1, $\dots$, 23.
| Range of 4:2 | 0 | | 2 | | 4 | | 6 | | | | | |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Codomain | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | $\dots$ | 23 |
Layouts are linear, so their range must include zero.
The complement of 4:2 with respect to 24 is thus a layout whose range
* includes zero;
* does not include any other elements of the range of 4:2
(i.e., satisfies the disjoint property; see above); and
* includes as much of 0, 1, $\dots$, 23 as possible
(so that it forms the "remainder" of 4:2 with respect to 24).
Intuitively, the range of the complement must look like this:
0, 1, 8, 9, 16, 17.
The resulting layout is ordered.
It has size 6 and cosize 18,
so it satisfies the bounded property (see above).
This is the layout (2, 3) : (1, 8).
(Going from this intuitive sense of the complement
to knowing how to compute it directly
is out of scope for this part of the tutorial.)
The following table shows 4:2 with its complement (2, 3) : (1, 8).
| Range of 4:2 | 0 | | 2 | | 4 | | 6 | | | | | | | | | | | | | |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Codomain | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | $\dots$ | 23 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Range of complement | 0 | 1 | | | | | | | 8 | 9 | | | | | | | 16 | 17 | | |
Now we know that `logical_divide`(24:2, 4:2) is
`composition`(24:2, `make_layout`(4:2, (2,3):(1,8))).
The composition of two layouts has the shape of the second (rightmost) layout,
so the resulting shape is (4, (2, 3)).
We see that the leftmost mode 4 corresponds to the divisor 4:2,
and the rightmost mode (2, 3) describes what's "left over"
from the original shape 24.
What are the strides?
We can start from the leftmost mode.
4:2 takes every other element (the even-numbered elements) of 24:2.
That's a stride-2 thing, striding over a stride-2 thing.
The resulting stride is 4.
Similarly, the stride 2 of 24:2
doubles the two strides of the rightmost mode.
The resulting layout is (4, (2, 3)) : (4, (2, 16)).
##### Tiling example
Suppose I have the 6 x 8 matrix from the Raked Product section
and want to "collect" the `tile`, turning the Raked Product into
@@ -607,7 +737,7 @@ 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
#### 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