v3.9 update (#2203)
* v3.9 update * voidD --------- Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
Yujia Zhai
@@ -0,0 +1,255 @@
|
||||
# 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/](https://github.com/NVIDIA/cutlass/tree/main/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`](https://github.com/NVIDIA/cutlass/tree/main/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](https://github.com/NVIDIA/cutlass/tree/main/examples/cute/tutorial/)
|
||||
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(dst); ++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 two `Tensor`s have static layouts and it can be proven
|
||||
that element vectorization is valid -- for example, four `ld.global.b32`s
|
||||
can be combined into a single `ld.global.b128` -- 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`](https://github.com/NVIDIA/cutlass/tree/main/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`](https://github.com/NVIDIA/cutlass/tree/main/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`](https://github.com/NVIDIA/cutlass/tree/main/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`](https://github.com/NVIDIA/cutlass/tree/main/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`](https://github.com/NVIDIA/cutlass/tree/main/include/cute/algorithm)
|
||||
directory.
|
||||
|
||||
## Copyright
|
||||
|
||||
Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
```
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright notice, this
|
||||
list of conditions and the following disclaimer.
|
||||
|
||||
2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
this list of conditions and the following disclaimer in the documentation
|
||||
and/or other materials provided with the distribution.
|
||||
|
||||
3. Neither the name of the copyright holder nor the names of its
|
||||
contributors may be used to endorse or promote products derived from
|
||||
this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
```
|
||||
Reference in New Issue
Block a user