Release v4.0.0 (#2294)
This commit is contained in:
Kihiro Bando
@@ -0,0 +1,154 @@
|
||||
.. _autotuning_gemm:
|
||||
|
||||
Guidance for Auto-Tuning
|
||||
=============================
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
Numerous GEMM kernel code examples are offered within our codebase.
|
||||
When integrating these kernels into frameworks, auto-tuning becomes essential
|
||||
for achieving optimal performance. This involves selecting the appropriate
|
||||
kernel parameters based on the inputs of real applications.
|
||||
Next, we'll briefly introduce some tips on how to perform auto-tuning.
|
||||
|
||||
The auto-tuning process typically involves the following steps:
|
||||
|
||||
1. Define search space
|
||||
2. Benchmark each configuration and select the kernel with the best performance
|
||||
3. Enable caching to reduce the tuning cost
|
||||
|
||||
The search space defines the valid combinations of kernel parameters that can be used to run the kernels.
|
||||
Different inputs (shapes, data types, etc.) typically require different kernel parameters to achieve optimal performance.
|
||||
The search space is related to the kernel. We take the Blackwell GEMM persistent kernel as an example.
|
||||
The search space is as follows:
|
||||
|
||||
- ``mma_tiler_mn``: Defines the dimensions of the matrix tile that each Matrix Multiply-Accumulate (MMA) instruction processes in a single operation.
|
||||
- ``cluster_shape_mn``: Specifies the number of CTAs along each dimension within a cluster. Refer `Parallel Thread Execution ISA documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#tensorcore-5th-generation-family-instructions>`_ for the possible mma tiler size and cluster shape for different tensor data types.
|
||||
- ``use_2cta_instrs``: Whether to utilize Blackwell's 2 CTA instructions for MMA/Copy.
|
||||
- ``use_tma_store``: Whether to use Tensor Memory Access (TMA) instructions to store the result back to global memory.
|
||||
|
||||
After defining the search space, we could traverse all parameter combinations to find the optimal kernel.
|
||||
The ``autotune_gemm`` function below demonstrates a simple exhaustive search approach - it iterates
|
||||
through configurations, compiles and benchmarks each kernel, and returns the best performing one.
|
||||
Since kernel compilation incurs overhead, it's important to cache and reuse compiled kernels
|
||||
to minimize host launch latency. CuTe DSL facilitates this through its separate compilation
|
||||
and execution workflow. More details can be found in :ref:`JIT_Caching`.
|
||||
As demonstrated in the ``autotune_gemm`` function
|
||||
(between the ``begin of cache the compiled GEMM kernel`` and ``end of cache the compiled GEMM kernel`` comments),
|
||||
we can use ``cute.compile()`` to compile a kernel once, cache the compiled result, and reuse the cached JIT executor for multiple kernel
|
||||
executions. We could maintain a global configuration-to-kernel dictionary (``config_kernel_dict``) to cache the compiled GEMM kernels,
|
||||
where each key (``kernel_cache_key``) uniquely identifies a kernel based on its characteristics.
|
||||
Usually we could use the {dtype + kernel configs} as the cached key for GEMM compilation. For example,
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
kernel_cache_key = f"{ab_dtype}x{c_dtype}x{acc_dtype}x{use_2cta_instrs}x{mma_tiler}x{cluster_shape_mn}x{use_tma_store}"
|
||||
|
||||
If the input tensor's layout is static, we should add the shape in the cached key too.
|
||||
Users can customize the ``benchmark`` function to measure kernel execution time.
|
||||
For stable and reliable performance measurements:
|
||||
|
||||
1. Run a few warmup iterations (e.g., 5-10) to stabilize GPU temperature
|
||||
2. Execute multiple timed iterations (e.g., 100-1000) for statistical significance
|
||||
3. Use CUDA events and synchronization for precise timing
|
||||
4. Lock GPU frequencies (SM and memory frequencies) with nvidia-smi
|
||||
5. Process results by removing outliers and using min/avg statistics as measurements.
|
||||
|
||||
This ensures reliable kernel selection through proper benchmarking.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# get the best GEMM kernel for given input tensors
|
||||
def autotune_gemm(
|
||||
a: cute.Tensor,
|
||||
b: cute.Tensor,
|
||||
c: cute.Tensor,
|
||||
stream: cuda.CUstream,
|
||||
use_2cta_instrs_list: List[bool] = [True],
|
||||
use_tma_store_list: List[bool] = [True],
|
||||
mma_tiler_m_list: List[int] = [256],
|
||||
mma_tiler_n_list: List[int] = [256],
|
||||
cluster_shape_m_list: List[int] = [2],
|
||||
cluster_shape_n_list: List[int] = [1],
|
||||
):
|
||||
best_kernel = None
|
||||
min_time = float("inf")
|
||||
# traverse the search space
|
||||
for use_2cta_instrs in use_2cta_instrs_list:
|
||||
for use_tma_store in use_tma_store_list:
|
||||
for mma_tiler_mn in product(mma_tiler_m_list, mma_tiler_n_list):
|
||||
for cluster_shape_mn in product(cluster_shape_m_list, cluster_shape_n_list):
|
||||
acc_dtype = cutlass.Float32
|
||||
hardware_info = cutlass.utils.HardwareInfo()
|
||||
max_active_clusters = hardware_info.get_max_active_clusters(
|
||||
cluster_shape_mn[0] * cluster_shape_mn[1]
|
||||
)
|
||||
# instance a GEMM kernel
|
||||
gemm = PersistentDenseGemmKernel(
|
||||
acc_dtype,
|
||||
use_2cta_instrs,
|
||||
mma_tiler_mn,
|
||||
cluster_shape_mn,
|
||||
use_tma_store,
|
||||
)
|
||||
# begin of cache the compiled GEMM kernel
|
||||
if kernel_cache_key not in config_kernel_dict:
|
||||
# compile gemm kernel
|
||||
compiled_gemm = cute.compile(
|
||||
gemm,
|
||||
a,
|
||||
b,
|
||||
c,
|
||||
max_active_clusters,
|
||||
stream,
|
||||
)
|
||||
config_kernel_dict[kernel_cache_key] = compiled_gemm
|
||||
else:
|
||||
compiled_gemm = config_kernel_dict[kernel_cache_key]
|
||||
# end of cache the compiled GEMM kernel
|
||||
try:
|
||||
# define a benchmark function to measure the execution time of the compiled GEMM kernel
|
||||
cur_time = benchmark(
|
||||
partial(compiled_gemm, a, b, c, stream),
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Execution error: {e}")
|
||||
cur_time = float("inf")
|
||||
if cur_time < min_time:
|
||||
min_time = cur_time
|
||||
best_kernel = compiled_gemm
|
||||
if best_kernel is None:
|
||||
raise ValueError("No best kernel found")
|
||||
return best_kernel
|
||||
|
||||
This brute-force approach ensures we could find the optimal parameters, though at the cost of trying every possibilities.
|
||||
For more advanced use cases, users can explore sophisticated optimization
|
||||
techniques like search space pruning and genetic algorithms to reduce tuning overhead and discover better
|
||||
configurations more efficiently.
|
||||
|
||||
To further optimize tuning performance, we can utilize caching mechanisms to avoid redundant computations.
|
||||
We could cache the tuning results in a input-to-kernel dictionary (e.g., ``input_kernel_dict``).
|
||||
When processing inputs with matching ``config_key`` values, the cached kernel can be reused directly without re-tuning.
|
||||
The ``config_key`` is related with the input tensor's characteristics, such as the shape, data type, etc.
|
||||
The setup of ``config_key`` is very flexible, users can customize it based on their own application.
|
||||
For instance, if the data type is fixed in users' application, we could use the input tensor's shape as the key, i.e., ``(m, n, k)``.
|
||||
To further reduce tuning overhead, we could consider using a simplified key like ``config_key = (power_of_2(m), power_of_2(n), power_of_2(k))``,
|
||||
where ``m``, ``n``, and ``k`` are rounded up to the nearest power of 2. This simplification can significantly reduce the number
|
||||
of unique keys while still maintaining good performance in most cases. However, it's important to validate that this
|
||||
approximation doesn't negatively impact performance for your specific use case.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
config_key = (m, n, k)
|
||||
if config_key in input_kernel_dict:
|
||||
compiled_gemm = input_kernel_dict[config_key]
|
||||
else:
|
||||
compiled_gemm = autotune_gemm(...)
|
||||
input_kernel_dict[config_key] = compiled_gemm
|
||||
# launch gemm kernel
|
||||
compiled_gemm(a_tensor, b_tensor, c_tensor, stream)
|
||||
|
||||
By following the methods above, you can customize your own auto-tuner to find the optimal GEMM kernel configuration
|
||||
for specific matrix dimensions and data types, significantly improving computational performance for models.
|
||||
@@ -0,0 +1,133 @@
|
||||
.. _debugging:
|
||||
|
||||
Debugging
|
||||
=========
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
This page provides an overview of debugging techniques and tools for CuTe DSL programs.
|
||||
|
||||
|
||||
Getting Familiar with the Limitations
|
||||
-------------------------------------
|
||||
|
||||
Before diving into comprehensive debugging capabilities, it's important to understand the limitations of CuTe DSL.
|
||||
Understanding these limitations will help you avoid potential pitfalls from the start.
|
||||
|
||||
Please refer to :doc:`../limitations` for more details.
|
||||
|
||||
|
||||
DSL Debugging
|
||||
-------------
|
||||
|
||||
CuTe DSL provides built-in logging mechanisms to help you understand the code execution flow and
|
||||
some of the internal state.
|
||||
|
||||
Enabling Logging
|
||||
~~~~~~~~~~~~~~~~
|
||||
|
||||
CuTe DSL provides environment variables to control logging level:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# Enable console logging (default: False)
|
||||
export CUTE_DSL_LOG_TO_CONSOLE=1
|
||||
|
||||
# Log to file instead of console (default: False)
|
||||
export CUTE_DSL_LOG_TO_FILE=my_log.txt
|
||||
|
||||
# Control log verbosity (0, 10, 20, 30, 40, 50, default: 10)
|
||||
export CUTE_DSL_LOG_LEVEL=20
|
||||
|
||||
|
||||
Log Categories and Levels
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Similar to standard Python logging, different log levels provide varying degrees of detail:
|
||||
|
||||
+--------+-------------+
|
||||
| Level | Description |
|
||||
+========+=============+
|
||||
| 0 | Disabled |
|
||||
+--------+-------------+
|
||||
| 10 | Debug |
|
||||
+--------+-------------+
|
||||
| 20 | Info |
|
||||
+--------+-------------+
|
||||
| 30 | Warning |
|
||||
+--------+-------------+
|
||||
| 40 | Error |
|
||||
+--------+-------------+
|
||||
| 50 | Critical |
|
||||
+--------+-------------+
|
||||
|
||||
|
||||
Dump the generated IR
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
For users familiar with MLIR and compilers, CuTe DSL supports dumping the Intermediate Representation (IR).
|
||||
This helps you verify whether the IR is generated as expected.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# Dump Generated CuTe IR (default: False)
|
||||
export CUTE_DSL_PRINT_IR=1
|
||||
|
||||
# Keep Generated CuTe IR in a file (default: False)
|
||||
export CUTE_DSL_KEEP_IR=1
|
||||
|
||||
|
||||
|
||||
Kernel Functional Debugging
|
||||
----------------------------
|
||||
|
||||
Using Python's ``print`` and CuTe's ``cute.printf``
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
CuTe DSL programs can use both Python's native ``print()`` as well as our own ``cute.printf()`` to
|
||||
print debug information during kernel generation and execution. They differ in a few key ways:
|
||||
|
||||
- Python's ``print()`` executes during compile-time only (no effect on the generated kernel) and is
|
||||
typically used for printing static values (e.g. a fully static layouts).
|
||||
- ``cute.printf()`` executes at runtime on the GPU itself and changes the PTX being generated. This
|
||||
can be used for printing values of tensors at runtime for diagnostics, but comes at a performance
|
||||
overhead similar to that of `printf()` in CUDA C.
|
||||
|
||||
For detailed examples of using these functions for debugging, please refer to the associated
|
||||
notebook referenced in :doc:`notebooks`.
|
||||
|
||||
Handling Unresponsive/Hung Kernels
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
When a kernel becomes unresponsive and ``SIGINT`` (``CTRL+C``) fails to terminate it,
|
||||
you can follow these steps to forcefully terminate the process:
|
||||
|
||||
1. Use ``CTRL+Z`` to suspend the unresponsive kernel
|
||||
2. Execute the following command to terminate the suspended process:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# Terminate the most recently suspended process
|
||||
kill -9 $(jobs -p | tail -1)
|
||||
|
||||
|
||||
CuTe DSL can also be debugged using standard NVIDIA CUDA tools.
|
||||
|
||||
Using Compute-Sanitizer
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
For detecting memory errors and race conditions:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
compute-sanitizer --some_options python your_dsl_code.py
|
||||
|
||||
Please refer to the `compute-sanitizer documentation <https://developer.nvidia.com/compute-sanitizer>`_ for more details.
|
||||
|
||||
Conclusion
|
||||
----------
|
||||
|
||||
This page covered several key methods for debugging CuTe DSL programs. Effective debugging typically requires a combination of these approaches.
|
||||
If you encounter issues with DSL, you can enable logging and share the logs with the CUTLASS team as a GitHub issue to report a bug.
|
||||
@@ -0,0 +1,90 @@
|
||||
.. _dsl_code_generation:
|
||||
.. |DC| replace:: dynamic compilation
|
||||
.. |DSL| replace:: CuTe DSL
|
||||
.. |IR| replace:: intermediate representation (IR)
|
||||
|
||||
End-to-End Code Generation
|
||||
==========================
|
||||
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
|
||||
1. Techniques for Turning Python into |IR|
|
||||
------------------------------------------
|
||||
|
||||
1.1 AST rewrite
|
||||
^^^^^^^^^^^^^^^^
|
||||
The function’s abstract-syntax tree is analysed **before** execution.
|
||||
Python control-flow (``for``/``while``, ``if``/``else``) and built-ins are converted to structured |IR|
|
||||
constructs. Computation inside each region is left untouched at this stage.
|
||||
|
||||
*Advantages*
|
||||
|
||||
* Sees the entire program, so every branch and loop is preserved.
|
||||
* Keeps loop structure intact for optimization such as tiling, vectorisation
|
||||
or GPU thread mapping.
|
||||
|
||||
*Disadvantages*
|
||||
|
||||
* Requires a well-defined Python subset that the rewriter understands.
|
||||
|
||||
|
||||
1.2 Tracing
|
||||
^^^^^^^^^^^
|
||||
The decorated function is executed once with *proxy* arguments; overloaded
|
||||
operators record every tensor operation that actually runs and produce a flat
|
||||
trace that is lowered to |IR|.
|
||||
|
||||
*Advantages*
|
||||
|
||||
* Near-zero compile latency, ideal for straight-line arithmetic.
|
||||
* No need to parse Python source, so it supports many dynamic Python
|
||||
features, and Python has many features.
|
||||
|
||||
*Disadvantages*
|
||||
|
||||
* Untaken branches vanish, so the generated kernel may be wrong for other
|
||||
inputs.
|
||||
* Loops are flattened to the iteration count observed during tracing.
|
||||
* Data-dependent control-flow freezes to a single execution path.
|
||||
|
||||
|
||||
2. |DSL| Code-Generation Modes
|
||||
------------------------------
|
||||
|
||||
CuTe’s Python front-end combines the techniques above into **two mutually
|
||||
exclusive modes**, selectable with the ``preprocessor`` flag of the
|
||||
``@jit`` decorator:
|
||||
|
||||
1. Tracing mode ``@jit(preprocess=False)`` – tracing only.
|
||||
This results in the fastest compilation path and is recommended only for kernels that are guaranteed to be
|
||||
straight-line arithmetic. It suffers from all tracing limitations listed in the previous section.
|
||||
|
||||
2. Preprocessor mode (**default**) ``@jit(preprocess=True)`` – **AST rewrite + tracing**.
|
||||
The AST pass captures every loop and branch, eliminating the correctness and
|
||||
optimisation problems of pure tracing; tracing then fills in the arithmetic.
|
||||
This hybrid “preprocessor” pipeline is unique to |DSL| and was designed
|
||||
specifically to overcome the disadvantages identified above.
|
||||
|
||||
.. figure:: dsl_modes.png
|
||||
:width: 400
|
||||
:align: center
|
||||
|
||||
*Left*: tracing mode records only the path that executed.
|
||||
*Right*: preprocessor mode emits structured |IR| for every branch and loop
|
||||
before tracing the arithmetic.
|
||||
|
||||
|
||||
Why Tracing-Only Is Insufficient for Control-Flow
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
* **Branch loss** – The untaken side of an ``if``/``else`` is never lowered.
|
||||
* **Loop unrolling** – Loops are flattened to the iteration count observed,
|
||||
destroying structure needed for parallel mapping and tiling.
|
||||
* **Data-dependent paths** – Control-flow that depends on tensor values freezes
|
||||
to a single execution path at trace time.
|
||||
|
||||
The preprocessor mode fixes all of these by lowering control-flow first and delegating
|
||||
only the arithmetic to the tracer.
|
||||
@@ -0,0 +1,140 @@
|
||||
.. _dsl_control_flow:
|
||||
.. |DC| replace:: dynamic compilation
|
||||
.. |IR| replace:: intermediate representation (IR)
|
||||
.. |DSL| replace:: CuTe DSL
|
||||
.. |Constexpr| replace:: **Constexpr** (compile-time Python value)
|
||||
|
||||
|DSL| Control Flow
|
||||
==================
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
|
||||
Overview
|
||||
--------
|
||||
|DSL| walks Python’s AST and converts each control-flow construct it finds into
|
||||
structured |IR|. You can therefore write ordinary Python loops and branches
|
||||
while the compiler decides—statement by statement—whether to
|
||||
|
||||
* **evaluate at compile time** if the controlling value is a |Constexpr|, or
|
||||
* **emit intermediate representation (IR)** when the value is dynamic.
|
||||
|
||||
|
||||
For a high-level discussion of the overall pipeline, see
|
||||
:doc:`the code-generation overview <dsl_code_generation>`.
|
||||
|
||||
For Loops
|
||||
---------
|
||||
|DSL| recognises three kinds of ranges for ``for`` loops:
|
||||
|
||||
* ``range`` – the Python built-in
|
||||
* ``cutlass.range_dynamic`` – always lowers to |IR|
|
||||
* ``cutlass.range_constexpr`` – always unrolls at compile time
|
||||
|
||||
|
||||
range(...)
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
The AST rewriter inserts a small helper stub. At runtime the loop bounds are
|
||||
inspected:
|
||||
|
||||
* **Constant bounds** → the loop is unrolled at compile time.
|
||||
* **Dynamic bounds** → the loop is emitted as structured |IR|.
|
||||
|
||||
|
||||
cutlass.range_dynamic(...)
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
Use when you *always* want a loop in the generated |IR|, even if the bounds
|
||||
look constant.
|
||||
|
||||
|
||||
cutlass.range_constexpr(...)
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
Runs in the Python interpreter and is fully unrolled before code generation.
|
||||
All loop indices must be |Constexpr|.
|
||||
|
||||
|
||||
Limitations of Dynamic For Loops
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
* Early-exit ``break``, ``continue``, or raising exception are not yet supported.
|
||||
* Operations in the loop body are traced only when tracing is active in that
|
||||
region.
|
||||
|
||||
|
||||
**Example:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def loop_example():
|
||||
n = 10
|
||||
|
||||
# ❌ This loop is dynamic, early-exit isn't allowed.
|
||||
for i in cutlass.range_dynamic(n):
|
||||
if i == 5:
|
||||
break # Early-exit
|
||||
cute.printf("%d\\n", i)
|
||||
|
||||
# ✅ This loop is constexpr, early-exit is allowed.
|
||||
for i in cutlass.range_constexpr(n):
|
||||
if i == 5:
|
||||
break # Early-exit
|
||||
cute.printf("%d\\n", i)
|
||||
|
||||
If-Else Statements
|
||||
------------------
|
||||
|
||||
Standard Python ``if``/``else`` is supported.
|
||||
|
||||
* **Predicate is Constexpr (compile-time Python value)** → evaluated at compile time.
|
||||
* **Predicate is dynamic** → lowered to |IR|.
|
||||
|
||||
**Example:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def main(const_var: cutlass.Constexpr, dynamic_var: cutlass.Int32):
|
||||
if const_var: # compile-time branch
|
||||
cute.printf("Const branch\\n")
|
||||
else:
|
||||
cute.printf("Const else\\n")
|
||||
|
||||
if dynamic_var == 10: # dynamic branch
|
||||
cute.printf("Dynamic True\\n")
|
||||
else:
|
||||
cute.printf("Dynamic False\\n")
|
||||
|
||||
Similarly to for-loops, the ``if cutlass.const_expr`` and ``if cutlass.dynamic_expr`` constructs can
|
||||
be used to force the evaluation at compile-time or the generation of IR, respectively. Unstructured
|
||||
control flow is only supported when using ``if cutlass.const_expr``.
|
||||
|
||||
While Loops
|
||||
-----------
|
||||
|
||||
Python ``while`` loops are always treated as **dynamic** because the loop condition may become
|
||||
dynamic after the first iteration. Similarly to for-loops and ``if``/``else``, the
|
||||
``while cutlass.const_expr`` and ``while cutlass.dynamic_expr`` constructs are available.
|
||||
|
||||
Compile-Time Metaprogramming
|
||||
----------------------------
|
||||
|
||||
Mix compile-time constructs with normal |DSL| code to generate specialised
|
||||
kernels without runtime overhead. A compile-time flag can, for example, toggle
|
||||
an optional **ReLU** epilogue:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.kernel
|
||||
def gemm(..., do_relu: cutlass.Constexpr):
|
||||
# main GEMM work
|
||||
...
|
||||
if const_expr(do_relu): # compile-time guard
|
||||
# ReLU code is emitted only when do_relu is True
|
||||
...
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
gemm(..., False) # ReLU is omitted from the generated |IR|
|
||||
gemm(..., True) # ReLU is included
|
||||
@@ -0,0 +1,198 @@
|
||||
.. _dsl_dynamic_layout:
|
||||
.. |DSL| replace:: CuTe DSL
|
||||
.. |SLAY| replace:: static layout
|
||||
.. |DLAY| replace:: dynamic layout
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
Static vs Dynamic layouts
|
||||
=========================
|
||||
|
||||
Static Layout
|
||||
-------------
|
||||
|
||||
When integrating with popular deep learning frameworks, one question is how to deal with the layout of the converted ``cute.Tensor``.
|
||||
For example, when converting a ``torch.Tensor`` to a ``cute.Tensor``, the shape of the ``torch.Tensor`` is honored for the layout of
|
||||
``cute.Tensor``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import cutlass
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
@cute.jit
|
||||
def foo(tensor):
|
||||
print(f"tensor.layout: {tensor.layout}") # Prints tensor layout at compile time
|
||||
cute.printf("tensor: {}", tensor) # Prints tensor values at runtime
|
||||
|
||||
In this example, we define a JIT function ``foo`` that takes a ``cute.Tensor`` as input and prints its layout. Note
|
||||
that Python print is used to print the layout at compile time. This works fine for |SLAY| whose value is known at
|
||||
compile time.
|
||||
|
||||
Now let's try to run the JIT function ``foo`` with different shapes of the input ``torch.Tensor``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
a = torch.tensor([1, 2, 3], dtype=torch.uint16)
|
||||
a_pack = from_dlpack(a)
|
||||
compiled_func = cute.compile(foo, a_pack)
|
||||
compiled_func(a_pack)
|
||||
|
||||
Here we first convert a 1D ``torch.Tensor`` with 3 elements to a ``cute.Tensor`` using ``from_dlpack``. Then we compile
|
||||
the JIT function ``foo`` with the converted ``cute.Tensor`` and call the compiled function.
|
||||
|
||||
::
|
||||
|
||||
tensor.layout: (3):(1)
|
||||
tensor: raw_ptr(0x00000000079e5100: i16, generic, align<2>) o (3):(1) =
|
||||
( 1, 2, 3 )
|
||||
|
||||
It prints ``(3):(1)`` for the layout because the converted ``cute.Tensor`` has a |SLAY| with shape ``(3)`` which
|
||||
is the shape of the ``a``.
|
||||
|
||||
Now if we call the compiled function with a different shape of the input ``torch.Tensor``, it would result in an unexpected
|
||||
result at runtime due to the mismatch of the type since ``compiled_func`` expects a ``cute.Tensor`` with layout ``(3):(1)``
|
||||
while ``b`` has shape ``(5)``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
b = torch.tensor([11, 12, 13, 14, 15], dtype=torch.uint16)
|
||||
b_pack = from_dlpack(b)
|
||||
compiled_func(b_pack) # ❌ This results in an unexpected result at runtime due to type mismatch
|
||||
|
||||
Following is the output which is unexpected due to the type mismatch.
|
||||
|
||||
::
|
||||
|
||||
tensor: raw_ptr(0x00000000344804c0: i16, generic, align<2>) o (3):(1) =
|
||||
( 11, 12, 13 )
|
||||
|
||||
To fix that, we would have to trigger another code generation and compilation for the new shape for ``b``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
compiled_func_2 = cute.compile(foo, b_pack) # This would trigger another compilation
|
||||
compiled_func_2(b_pack) # ✅ Now this works fine
|
||||
|
||||
As shown in the example above, with the newly compiled ``compiled_func_2``, we can pass in ``b_pack`` to the compiled
|
||||
JIT function ``compiled_func_2``.
|
||||
|
||||
::
|
||||
|
||||
tensor.layout: (5):(1)
|
||||
tensor: raw_ptr(0x0000000034bb2840:: i16, generic, align<2>) o (5):(1) =
|
||||
( 11, 12, 13, 14, 15 )
|
||||
|
||||
Now it recompiles and prints the values of ``b`` correctly.
|
||||
|
||||
It's obvoius that we need distinct codes generated and compiled for different static layout. In this case, one for layout
|
||||
``(3):(1)`` and the other for layout ``(5):(1)``.
|
||||
|
||||
Dynamic Layout
|
||||
--------------
|
||||
|
||||
In order to avoid generating and compiling multiple times for different shapes of the input ``torch.Tensor``, |DSL| provides a way to
|
||||
generate and compile JIT function with |DLAY|.
|
||||
|
||||
To get dyanmic layout of the ``cute.Tensor``, a ``torch.Tensor`` object can be passed into the JIT function directly which instructs
|
||||
|DSL| to call ``cute.mark_layout_dynamic`` automatically on the converted ``cute.Tensor`` per the leading dimension of the layout.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import cutlass
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
@cute.jit
|
||||
def foo(tensor):
|
||||
print(tensor.layout) # Prints (?,?):(?,1) for dynamic layout
|
||||
|
||||
a = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint16)
|
||||
compiled_func = cute.compile(foo, a)
|
||||
compiled_func(a)
|
||||
|
||||
b = torch.tensor([[11, 12], [13, 14], [15, 16]], dtype=torch.uint16)
|
||||
compiled_func(b) # Reuse the same compiled function for different shape
|
||||
|
||||
In the example above, a single compilation of the JIT function ``foo`` is reused for different shapes of the input ``torch.Tensor``.
|
||||
This is possible because the converted ``cute.Tensor`` has a |DLAY| ``(?,?):(?,1)`` which is compatible with the shape of the
|
||||
input ``torch.Tensor`` of both calls.
|
||||
|
||||
Alternatively, for compact layout, ``cute.mark_compact_shape_dynamic`` can be called for a finer-grained control to specify the mode
|
||||
of the layout for dynamic and the divisibility constraint for the dynamic dimension.
|
||||
|
||||
Refer to :doc:`framework_integration` for more details on ``from_dlpack``, ``mark_layout_dynamic``,
|
||||
and ``mark_compact_shape_dynamic``.
|
||||
|
||||
Static Layout vs. Dynamic Layout
|
||||
--------------------------------
|
||||
|
||||
Per the previous sections, we have seen that |SLAY| leads to distinct JIT code generations while |DLAY| leads to a single
|
||||
compilation for different shapes.
|
||||
|
||||
That said, creating JIT function with |SLAY| is useful when the use cases targeting input data with fixed shapes.
|
||||
Since more information is available at compile time, the compiler would be able to kick in optimizations that otherwise would not
|
||||
be possible for the code generated for |DLAY|.
|
||||
|
||||
On the other hand, |DLAY| would be more flexible for the cases where the input data has varying shapes. This provides more
|
||||
scalability of the generated code to deal with varying input data of different shapes.
|
||||
|
||||
Programming with Static and Dynamic Layout
|
||||
------------------------------------------
|
||||
|
||||
|DSL| provides intuitive way to program with static and |DLAY| in the codes.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import cutlass
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
@cute.jit
|
||||
def foo(tensor, x: cutlass.Constexpr[int]):
|
||||
print(cute.size(tensor)) # Prints 3 for the 1st call
|
||||
# Prints ? for the 2nd call
|
||||
if cute.size(tensor) > x:
|
||||
cute.printf("tensor[2]: {}", tensor[2])
|
||||
else:
|
||||
cute.printf("tensor size <= {}", x)
|
||||
|
||||
a = torch.tensor([1, 2, 3], dtype=torch.uint16)
|
||||
foo(from_dlpack(a), 3) # First call with static layout
|
||||
|
||||
b = torch.tensor([1, 2, 3, 4, 5], dtype=torch.uint16)
|
||||
foo(b, 3) # Second call with dynamic layout
|
||||
|
||||
In this example, the JIT function ``foo`` is compiled with a |SLAY| ``(3):(1)`` for the first call, which means the
|
||||
size of the tensor is known at compile time. |DSL| makes good use of this and automatically handles the if condition at the
|
||||
compile time. Hence the generated codes are efficient without the if condition at all.
|
||||
|
||||
For the second call, the JIT function ``foo`` is compiled with a |DLAY| ``(?):(1)`` hence the tensor size is only
|
||||
evaluated at runtime. |DSL| automatically generates the code to handle the |DLAY| and the if condition at runtime.
|
||||
|
||||
The same applies to loop as well:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def foo(tensor, x: cutlass.Constexpr[int]):
|
||||
for i in range(cute.size(tensor)):
|
||||
cute.printf("tensor[{}]: {}", i, tensor[i])
|
||||
|
||||
a = torch.tensor([1, 2, 3], dtype=torch.uint16)
|
||||
foo(from_dlpack(a), 3) # First call with static layout
|
||||
|
||||
b = torch.tensor([1, 2, 3, 4, 5], dtype=torch.uint16)
|
||||
foo(b, 3) # Second call with dynamic layout
|
||||
|
||||
With the static layout in the first call, |DSL| is able to fully unroll the loop at compile time. While in the second call,
|
||||
the generated codes will have the loop executed at runtime based on the |DLAY|.
|
||||
|
||||
With the single JIT function implementation, |DSL| is able to handle control-flow constructs and automatically generate
|
||||
the optimized codes for different cases. This is all possible because |DSL| is able to walk the Python AST and convert
|
||||
each control-flow construct it finds accordingly.
|
||||
|
||||
Please refer to :doc:`dsl_control_flow` for more details.
|
||||
@@ -0,0 +1,128 @@
|
||||
.. _dsl_introduction:
|
||||
.. |DC| replace:: dynamic compilation
|
||||
.. |IR| replace:: IR
|
||||
.. |DSL| replace:: CuTe DSL
|
||||
|
||||
|
||||
|DSL|
|
||||
======================
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
Overview
|
||||
--------
|
||||
|
||||
|DSL| is a Python-based domain-specific language (DSL) designed for |DC| of numeric and GPU-oriented code. Its primary goals are:
|
||||
|
||||
- **Consistent with CuTe C++**, allowing users to express GPU kernels with full control of the hardware.
|
||||
- **JIT compilation** for both host and GPU execution.
|
||||
- `DLPack <https://github.com/dmlc/dlpack>`_ **integration**, enabling seamless interop with frameworks (e.g., PyTorch, JAX).
|
||||
- **JIT caching**, so that repeated calls to the same function benefit from cached |IR| modules.
|
||||
- **Native types and type inference** to reduce boilerplate and improve performance.
|
||||
- **Optional lower-level control**, offering direct access to GPU backends or specialized |IR| dialects.
|
||||
|
||||
Decorators
|
||||
----------
|
||||
|
||||
|DSL| provides two main Python decorators for generating optimized code via |DC|:
|
||||
|
||||
1. ``@jit`` — Host-side JIT-compiled functions
|
||||
2. ``@kernel`` — GPU kernel functions
|
||||
|
||||
Both decorators can optionally use a **preprocessor** that automatically expands Python control flow (loops, conditionals) into operations consumable by the underlying |IR|.
|
||||
|
||||
``@jit``
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
Declares JIT-compiled functions that can be invoked from Python or from other |DSL| functions.
|
||||
|
||||
**Decorator Parameters**:
|
||||
|
||||
* ``preprocessor``:
|
||||
|
||||
* ``True`` (default) — Automatically translate Python flow control (e.g., loops, if-statements) into |IR| operations.
|
||||
* ``False`` — No automatic expansion; Python flow control must be handled manually or avoided.
|
||||
|
||||
**Call-site Parameters**:
|
||||
|
||||
- ``no_cache``:
|
||||
|
||||
- ``True`` — Disables JIT caching, forcing a fresh compilation each call.
|
||||
- ``False`` (default) — Enables caching for faster subsequent calls.
|
||||
|
||||
``@kernel``
|
||||
~~~~~~~~~~~~~~~~
|
||||
|
||||
Defines GPU kernel functions, compiled as specialized GPU symbols through |DC|.
|
||||
|
||||
**Decorator Parameters**:
|
||||
|
||||
- ``preprocessor``:
|
||||
|
||||
- ``True`` (default) — Automatically expands Python loops/ifs into GPU-compatible |IR| operations.
|
||||
- ``False`` — Expects manual or simplified kernel implementations.
|
||||
|
||||
**Kernel Launch Parameters**:
|
||||
|
||||
- ``grid``
|
||||
Specifies the grid size as a list of integers.
|
||||
- ``block``
|
||||
Specifies the block size as a list of integers.
|
||||
- ``cluster``
|
||||
Specifies the cluster size as a list of integers.
|
||||
- ``smem``
|
||||
Specifies the size of shared memory in bytes (integer).
|
||||
|
||||
Calling Conventions
|
||||
-------------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 20 20 15 25
|
||||
|
||||
* - **Caller**
|
||||
- **Callee**
|
||||
- **Allowed**
|
||||
- **Compilation/Runtime**
|
||||
|
||||
* - Python function
|
||||
- ``@jit``
|
||||
- ✅
|
||||
- DSL runtime
|
||||
|
||||
* - Python function
|
||||
- ``@kernel``
|
||||
- ❌
|
||||
- N/A (error raised)
|
||||
|
||||
* - ``@jit``
|
||||
- ``@jit``
|
||||
- ✅
|
||||
- Compile-time call, inlined
|
||||
|
||||
* - ``@jit``
|
||||
- Python function
|
||||
- ✅
|
||||
- Compile-time call, inlined
|
||||
|
||||
* - ``@jit``
|
||||
- ``@kernel``
|
||||
- ✅
|
||||
- Dynamic call via GPU driver or runtime
|
||||
|
||||
* - ``@kernel``
|
||||
- ``@jit``
|
||||
- ✅
|
||||
- Compile-time call, inlined
|
||||
|
||||
* - ``@kernel``
|
||||
- Python function
|
||||
- ✅
|
||||
- Compile-time call, inlined
|
||||
|
||||
* - ``@kernel``
|
||||
- ``@kernel``
|
||||
- ❌
|
||||
- N/A (error raised)
|
||||
@@ -0,0 +1,196 @@
|
||||
.. _dsl_jit_arg_generation:
|
||||
.. |DSL| replace:: CuTe DSL
|
||||
.. |CUSTOM_TYPES| replace:: customized types
|
||||
|
||||
|DSL| JIT Function Argument Generation
|
||||
=======================================
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
In a nutshell
|
||||
--------------
|
||||
When using the ``@jit`` or ``@kernel`` decorators to define a JIT-compiled function, the arguments to the function are traced to determine the JIT function's signature.
|
||||
|DSL| provides a Pythonic way to write the arguments for JIT function as one normally would in Python, and the |DSL| will take care of the rest for you.
|
||||
|
||||
Specifically, |DSL| honors following when generating the JIT function's arguments:
|
||||
|
||||
- JIT function arguments are assumed to be **dynamic arguments** by default.
|
||||
- If an argument is explicitly type annotated with ``cutlass.Constexpr``, it is treated as a **compile-time constant**.
|
||||
- If type annotation is provided, |DSL| validates the argument type at compile time for **type safety**.
|
||||
- |DSL| provides **runtime checkable protocols** (``JitArgument`` and ``DynamicExpression``) for generating JIT function arguments for |CUSTOM_TYPES|.
|
||||
|
||||
More details below for each of the above.
|
||||
|
||||
Static argument vs. Dynamic argument
|
||||
------------------------------------
|
||||
|
||||
|DSL| supports both static and dynamic arguments for JIT functions.
|
||||
|
||||
1. **Static arguments** hold values that are known at compile time. It is not included in the generated JIT function signature.
|
||||
2. **Dynamic arguments** hold values that are only known at runtime.
|
||||
|
||||
By default, |DSL| assumes dynamic arguments and tries to infer the argument types from the call-site argument types. An explicit type annotation ``cutlass.Constexpr`` can be used to specify a static argument.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
|
||||
@cute.jit
|
||||
def foo(x: cutlass.Int32, y: cute.Constexpr):
|
||||
print("x = ", x) # Prints x = ?
|
||||
print("y = ", y) # Prints y = 2
|
||||
cute.printf("x: {}", x) # Prints x: 2
|
||||
cute.printf("y: {}", y) # Prints y: 2
|
||||
|
||||
foo(2, 2)
|
||||
|
||||
In the example above, ``x`` is a dynamic argument with type cutlass.Int32 and ``y`` is a static argument.
|
||||
|
||||
With the ``cutlass.Constexpr`` annotation, a more sophisticated uses case of static argument in the JIT functions can be something like:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
|
||||
@cute.kernel
|
||||
def kernel(
|
||||
self,
|
||||
tiled_mma: cute.TiledMma,
|
||||
tma_atom_a: cute.CopyAtom,
|
||||
mA_mkl: cute.Tensor,
|
||||
tma_atom_b: cute.CopyAtom,
|
||||
mB_nkl: cute.Tensor,
|
||||
tma_atom_c: Optional[cute.CopyAtom],
|
||||
mC_mnl: cute.Tensor,
|
||||
cluster_layout_vmnk: cute.Layout,
|
||||
a_smem_layout_staged: cute.ComposedLayout,
|
||||
b_smem_layout_staged: cute.ComposedLayout,
|
||||
c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None],
|
||||
epi_tile: cute.Tile,
|
||||
epilogue_op: cutlass.Constexpr,
|
||||
):
|
||||
...
|
||||
|
||||
# Perform epilogue op on accumulator and convert to C type
|
||||
acc_vec = tTR_rAcc.load()
|
||||
acc_vec = epilogue_op(acc_vec.to(self.c_dtype))
|
||||
tTR_rC.store(acc_vec)
|
||||
|
||||
In this example, ``epilogue_op`` is a static argument in the JIT kernel where the argument is used for the epilogue fusion. Upon calling the kernel,
|
||||
an elementwise lambda function can be passed in as the ``epilogue_op`` argument. For example, a ReLU can be applied for epilogue fusion by simply setting the
|
||||
``epilogue_op`` to ``lambda x: cute.where(x > 0, x, cute.full_like(x, 0))``
|
||||
|
||||
Refer to the `Blackwell dense GEMM example <https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py>`__ for a complete example.
|
||||
|
||||
Type safety
|
||||
-----------
|
||||
|
||||
|DSL| makes good use of type annotation in JIT function signature and validates the JIT function argument types at compile time for **type safety**.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import numpy as np
|
||||
|
||||
@cute.jit
|
||||
def foo(x: cute.Tensor, y: cutlass.Float16):
|
||||
...
|
||||
|
||||
a = np.random.randn(10, 10).astype(np.float16)
|
||||
b = 32
|
||||
|
||||
foo(a, b)
|
||||
foo(b, a) # This will fail at compile time due to type mismatch
|
||||
|
||||
The type safety check helps catch the type mismatch issue early at the compile time with clear error message to avoid tricky runtime errors which is usually more expensive to debug.
|
||||
In the example above, the second call to ``foo`` will fail at compile time due to the type mismatch with a clear error message:
|
||||
|
||||
::
|
||||
|
||||
cutlass.base_dsl.common.DSLRuntimeError: DSLRuntimeError: expects argument #1 (a) to be <class 'cutlass.cute.typing.Tensor'>, but got <class 'int'>
|
||||
|
||||
JIT function arguments with |CUSTOM_TYPES|
|
||||
--------------------------------------------
|
||||
|DSL| supports |CUSTOM_TYPES| for JIT function arguments by providing two runtime checkable protocols:
|
||||
|
||||
* ``JitArgument`` which is used for host JIT functions to be called from Python.
|
||||
- ``__c_pointers__``: Generate a list of ctypes pointers for the current object.
|
||||
- ``__get_mlir_types__``: Generate a list of MLIR types for the current object.
|
||||
- ``__new_from_mlir_values__``: Create a new object from MLIR values.
|
||||
|
||||
* ``DynamicExpression`` which is used for device JIT functions to be called from the host JIT functions.
|
||||
- ``__extract_mlir_values__``: Generate a dynamic expression for the current object.
|
||||
- ``__new_from_mlir_values__``: Create a new object from MLIR values.
|
||||
|
||||
Refer to `typing.py <https://github.com/NVIDIA/cutlass/tree/main/python/CuTeDSL/base_dsl/typing.py>`__ for more details on these protocol APIs.
|
||||
|
||||
Depending on different cases of the |CUSTOM_TYPES|, |DSL| provides easy ways to adopt |CUSTOM_TYPES| for JIT function arguments.
|
||||
|
||||
1. Direct protocol implementation in |CUSTOM_TYPES|
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
One way is to implement the protocol methods directly in the |CUSTOM_TYPES| to enable the protocol based JIT function argument generation.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
|
||||
# Customized type that implements the DynamicExpression protocol
|
||||
class MyDynamicExpression:
|
||||
def __init__(self, tensor, offset):
|
||||
self._tensor = tensor # Dynamic argument
|
||||
self._offset = offset # Dynamic argument
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
return [self._tensor.__extract_mlir_values__(), self._offset.__extract_mlir_values__()]
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
return MyDynamicExpression(values[0], values[1])
|
||||
|
||||
@cute.kernel
|
||||
def my_kernel(x: MyDynamicExpression):
|
||||
...
|
||||
|
||||
In the example above, the ``MyDynamicExpression`` implements the ``DynamicExpression`` protocol and |DSL| will generate the JIT function arguments for the JIT kernel ``my_kernel`` based on the protocol methods.
|
||||
|
||||
2. Adaptor based protocol implementation for |CUSTOM_TYPES|
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
For the case where directly changing the |CUSTOM_TYPES| to implement the protocol is not feasible, |DSL| provides adaptor based approach to adapt the |CUSTOM_TYPES| for JIT function argument generation.
|
||||
|
||||
The JIT function argument adaptor is a callable object that implements the desired protocol methods for the registered |CUSTOM_TYPES|. This way, |DSL| automatically queries the JIT argument adaptor registry
|
||||
to generate the JIT function arguments for the given |CUSTOM_TYPES|.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cutlass.register_jit_arg_adapter(MyFrameworkObject)
|
||||
class MyFrameworkObjectAdapter:
|
||||
"""
|
||||
Convert a 3rd party framework object to a JIT function argument with JitArgument protocol
|
||||
"""
|
||||
|
||||
def __init__(self, arg):
|
||||
self._arg = arg
|
||||
|
||||
def __c_pointers__(self):
|
||||
# Convert the framework object to a C-ABI compatible object
|
||||
# thru its C-ABI interface
|
||||
return [self._arg.get_cabi_pointer()]
|
||||
|
||||
def __get_mlir_types__(self):
|
||||
# Return the list of MLIR types the framework object represents
|
||||
return [self._arg.get_data().mlir_type]
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
# Convert the MLIR values back to the framework object
|
||||
return MyFrameworkObject(values[0])
|
||||
|
||||
In this example, the ``MyFrameworkObjectAdapter`` implements an adaptor class which bridges the |DSL| and the 3rd party framework type ``MyFrameworkObject``.
|
||||
The registration is done by just decorating the adaptor with ``cutlass.register_jit_arg_adapter`` for the customized type. With the registered adaptor,
|
||||
|DSL| will automatically use the adaptor to generate the JIT function arguments for ``MyFrameworkObject`` typed arguments.
|
||||
@@ -0,0 +1,152 @@
|
||||
.. _dsl_jit_caching:
|
||||
.. |DSL| replace:: CuTe DSL
|
||||
|
||||
.. _JIT_Caching:
|
||||
|
||||
|DSL| JIT Caching
|
||||
====================
|
||||
.. contents:: Table of Contents
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
Zero Compile and JIT Executor
|
||||
-----------------------------
|
||||
|
||||
Zero Compile is a feature that enables explicit kernel compilation on demand through ``cute.compile``.
|
||||
When ``cute.compile`` is called, it compiles the kernel and returns a JIT Executor instance.
|
||||
This JIT Executor instance can be cached and reused directly for subsequent executions without compiling the kernel again.
|
||||
|
||||
The JIT Executor is a component that independently executes compiled code.
|
||||
It can be created either through ``cute.compile`` or implicit compilation.
|
||||
The JIT Executor instance behaves like a callable object to execute the compiled code.
|
||||
Each JIT Executor instance maintains a single compiled host function.
|
||||
|
||||
It encompasses all necessary execution components:
|
||||
|
||||
* Host function pointer and its MLIR execution engine
|
||||
* CUDA modules (optional)
|
||||
* Argument specifications defining how Python arguments are converted to C ABI-compatible types. Note that arguments with the ``cutlass.Constexpr`` hint are excluded from argument specifications since they are evaluated at compile time rather than runtime.
|
||||
|
||||
For example, in the following code, ``print_result`` is a ``cutlass.Constexpr`` value that is **NOT** evaluated at runtime:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import cutlass.cute as cute
|
||||
|
||||
@cute.jit
|
||||
def add(a, b, print_result: cutlass.Constexpr):
|
||||
if print_result:
|
||||
cute.printf("Result: %d\n", a + b)
|
||||
return a + b
|
||||
|
||||
jit_executor = cute.compile(add, 1, 2, True)
|
||||
|
||||
jit_executor(1, 2) # output: ``Result: 3``
|
||||
|
||||
The JIT Executor ensures all components are properly initialized and loaded after compilation.
|
||||
|
||||
For example, all CUDA modules are loaded (via ``cuModuleLoad``) and kernel function pointers are extracted (via ``cuModuleGetFunction``).
|
||||
|
||||
When calling a JIT Executor instance, it:
|
||||
|
||||
* Parses Python runtime arguments and converts them to C ABI-compatible types according to argument specifications
|
||||
* Invokes the host function with the converted arguments
|
||||
|
||||
Custom Caching with ``cute.compile``
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
``cute.compile`` bypasses caching in |DSL| and always performs compilation, returning a fixed JIT Executor instance.
|
||||
This allows implementing custom caching strategies as shown below:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def add(b):
|
||||
return a + b
|
||||
|
||||
# Define a custom cache
|
||||
custom_cache = {}
|
||||
|
||||
a = 1
|
||||
compiled_add_1 = cute.compile(add, 2)
|
||||
custom_cache[1] = compiled_add_1
|
||||
compiled_add_1(2) # result = 3
|
||||
|
||||
a = 2
|
||||
compiled_add_2 = cute.compile(add, 2)
|
||||
custom_cache[2] = compiled_add_2
|
||||
compiled_add_2(2) # result = 4
|
||||
|
||||
# Use the custom cache
|
||||
custom_cache[1](2) # result = 3
|
||||
custom_cache[2](2) # result = 4
|
||||
|
||||
|
||||
Cache in |DSL|
|
||||
-----------------
|
||||
|
||||
By default, cache in |DSL| is implicitly enabled to avoid recompilation when kernels are called repeatedly without changes.
|
||||
|
||||
The cache is implemented as a map storing compiled JIT Executor instances within |DSL|.
|
||||
|
||||
The cache key combines hashes of:
|
||||
|
||||
* MLIR bytecode of the MLIR program generated by |DSL|
|
||||
* All |DSL| Python source files
|
||||
* All |DSL| shared libraries
|
||||
* All |DSL| environment variables
|
||||
|
||||
The cache value is a compiled JIT Executor instance.
|
||||
|
||||
On a cache hit, compilation is skipped and the cached JIT Executor instance is reused.
|
||||
|
||||
On a cache miss, the kernel is compiled and the new JIT Executor instance is stored in the cache.
|
||||
|
||||
Here is an example demonstrating automatic caching of the ``add`` kernel:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Global variable
|
||||
a = 1
|
||||
|
||||
@cute.jit
|
||||
def add(b):
|
||||
return a + b
|
||||
|
||||
# Cache is empty at beginning
|
||||
|
||||
# First call: cache miss triggers compilation
|
||||
result = add(2) # result = 3
|
||||
# Cache now has one instance
|
||||
|
||||
# Second call: cache hit reuses cached JIT Executor
|
||||
result = add(2) # result = 3
|
||||
|
||||
a = 2
|
||||
# Third call: cache miss due to changed IR code triggers recompilation
|
||||
result = add(2) # result = 4
|
||||
# Cache now has two instances
|
||||
|
||||
The cache can be serialized to files for subsequent runs.
|
||||
After serialization, compiled MLIR bytecode is stored in file.
|
||||
The cache directory is ``/tmp/{current_user}/cutlass_python_cache``.
|
||||
The cache loads from files into memory during |DSL| initialization and saves back to files when the process exits.
|
||||
|
||||
The following environment variables control file caching:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Disable file caching while keeping in-memory cache available, defaults to False.
|
||||
export CUTE_DSL_DISABLE_FILE_CACHING=True
|
||||
|
||||
# Maximum number of cache files allowed, defaults to 1000.
|
||||
export CUTE_DSL_FILE_CACHING_CAPACITY=1000
|
||||
|
||||
Limitations
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The intention of caching is to reduce the host launch overhead before each execution. As above example shows,
|
||||
the consistency between the original Python code and the MLIR program is hard to maintain because of the impact of dynamic factors such as global variables.
|
||||
Therefore, the MLIR program **MUST** always be generated to verify that the kernel content matches what was previously built.
|
||||
|
||||
For optimal host launch latency, we recommend using above custom caching method with ``cute.compile``.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
@@ -0,0 +1,412 @@
|
||||
.. _framework_integration:
|
||||
.. |DSL| replace:: CuTe DSL
|
||||
|
||||
Integration with Frameworks
|
||||
=============================
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
In order to facilitate the integration of CUTLASS Python with popular frameworks, we leverage the
|
||||
`DLPack protocol <https://github.com/dmlc/dlpack>`_ and transform tensors originating from these
|
||||
frameworks to CuTe tensors. The present page documents the conventions, the API available to the
|
||||
user, and provide example code snippets for common usage patterns.
|
||||
|
||||
Implicit Conversion
|
||||
-------------------
|
||||
|
||||
Tensors originating from frameworks supporting the DLPack protocol can be directly provided to a
|
||||
JIT function as a regular parameter. |DSL|'s runtime implicitly converts the original tensor to a
|
||||
CuTe tensor with a fully dynamic layout except for the stride element corresponding to the leading
|
||||
dimension. The example below demonstrates this use case.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import cutlass.cute as cute
|
||||
|
||||
@cute.jit
|
||||
def foo(src):
|
||||
"""
|
||||
The following lines print
|
||||
|
||||
ptr<f32, generic> o (?,?,?):(?,?,1)
|
||||
<class 'cutlass.cute.core._Tensor'>
|
||||
"""
|
||||
print(src)
|
||||
print(type(src))
|
||||
|
||||
a = torch.randn(30, 20, 32, device="cpu")
|
||||
foo(a)
|
||||
|
||||
|
||||
Explicit conversion using ``from_dlpack``
|
||||
------------------------------------------
|
||||
|
||||
|DSL|'s runtime provides an interface for converting DLPack-compatible tensors to CuTe tensors,
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
b = cute.runtime.from_dlpack(a)
|
||||
|
||||
where ``a`` is a tensor supporting the DLPack protocol with the ``__dlpack__``
|
||||
and ``__dlpack_device__`` methods. The resulting CuTe tensor ``b`` has a fully static layout. This
|
||||
conversion is performed without copying any tensor data, enabling seamless integration with major
|
||||
frameworks. Users can create tensors using NumPy, PyTorch, etc. and directly feed them into JIT
|
||||
functions writtnen using |DSL|.
|
||||
|
||||
The resulting CuTe tensor shares the same underlying memory buffer as the original tensor. This
|
||||
zero-copy approach maximizes performance by eliminating unnecessary data duplication. However, it is
|
||||
important to note that the CuTe tensor's validity is tied to the lifetime of the original tensor. If
|
||||
the source tensor is destroyed or goes out of scope, the corresponding CuTe tensor becomes invalid
|
||||
since it references the original memory location.
|
||||
|
||||
The full signature of from_dlpack is as follows:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def from_dlpack(tensor, assumed_align=None):
|
||||
|
||||
The ``assumed_align`` integer parameter specifies the alignment of the tensor in unit of bytes.
|
||||
The tensor's base address must be divisible by ``assumed_align``. When not provided explicitly,
|
||||
the alignment is set to the natural alignment of the tensor's element type. Note that the alignment
|
||||
information is part of the pointer type in the generated IR. Therefore, programs with different
|
||||
alignments have a different IR and identical IRs are required for hitting the kernel caching
|
||||
mechanism of |DSL|.
|
||||
|
||||
Code Example
|
||||
~~~~~~~~~~~~
|
||||
|
||||
The following code demonstrates how to convert a PyTorch tensor to a CuTe tensor using the
|
||||
``from_dlpack`` function with default parameters.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import cutlass
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
x = torch.randn(30, 20, device="cpu")
|
||||
y = from_dlpack(x)
|
||||
|
||||
Once converted, we can access the tensor's information through various
|
||||
attributes. The following list shows the attributes of the converted tensor:
|
||||
|
||||
- ``tensor.shape``: the tensor's shape
|
||||
- ``tensor.stride``: the tensor's stride
|
||||
- ``tensor.memspace``: the tensor's memory space
|
||||
- ``tensor.element_type``: the tensor's element data type
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import cutlass
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
x = torch.randn(30, 20, device="cpu")
|
||||
y = from_dlpack(x)
|
||||
|
||||
print(y.shape) # (30, 20)
|
||||
print(y.stride) # (20, 1)
|
||||
print(y.memspace) # generic (if torch tensor in on device memory, memspace will be gmem)
|
||||
print(y.element_type) # Float32
|
||||
print(y) # Tensor<0x000000000875f580@generic o (30, 20):(20, 1)>
|
||||
|
||||
The string format of the resulting CuTe tensor is
|
||||
|
||||
.. code-block::
|
||||
|
||||
Tensor<0x{tensor.data_ptr:016x}@{tensor.memspace} o {tensor.shape}:{tensor.stride}>
|
||||
|
||||
As can be seen in the example above, ``from_dlpack`` first results in a tensor with a static layout.
|
||||
To obtain dynamic or mixed static/dynamic layouts after calling ``from_dlpack``, the
|
||||
``mark_layout_dynamic`` and ``mark_compact_shape_dynamic`` functions are used and described in
|
||||
the following sections.
|
||||
|
||||
When to Use Explicit Conversion?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The DLPack protocol is a widely used protocol for interoperability between different frameworks.
|
||||
However, there is some associated overhead. Based on our benchmark, it usually takes between 2 to 3
|
||||
us per call to ``from_dlpack``.
|
||||
|
||||
Explicit conversion allows for caching the converted CuTe tensors in order to avoid the overhead of
|
||||
repeated calls to ``from_dlpack``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
x = torch.randn(30, 20, device="cpu")
|
||||
if key not in cached_tensors:
|
||||
# Do the conversion only for cache misses
|
||||
cached_tensors[key] = cute.runtime.from_dlpack(x)
|
||||
foo(cached_tensors[key])
|
||||
|
||||
Another use case for explicit conversion is to gain fine-grain control over which modes of a tensor
|
||||
are considered dynamic from the perspective of the generated program.
|
||||
|
||||
Mark the Tensor's Layout as Dynamic with ``mark_layout_dynamic``
|
||||
----------------------------------------------------------------
|
||||
|
||||
After calling this function, all shape modes become dynamic. The stride modes also become dynamic
|
||||
with the following two exceptions:
|
||||
|
||||
1. the leading dimension's stride remains fixed at 1;
|
||||
2. stride elements equal to 0 (which indicates broadcasting) are retained.
|
||||
|
||||
The full signature of ``mark_layout_dynamic`` is as follows:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def mark_layout_dynamic(self, leading_dim: int|None = None):
|
||||
|
||||
The ``leading_dim`` parameter specifies the leading dimension of the tensor. The leading dimension's
|
||||
stride is set to 1 unless inconsistent with the layout of the DLPack tensor. For example,
|
||||
|
||||
- For a tensor with layout ``(2,2,3,4):(2,1,4,12)``, if ``leading_dim`` is specified to be 1,
|
||||
the layout will be marked as ``(?,?,?,?):(?,1,?,?)``.
|
||||
- If ``leading_dim`` is specified to be 0, a deduction failure error is raised because the stride of
|
||||
dimension 0 is 2 (not 1).
|
||||
|
||||
The default value for ``leading_dim`` is ``None``. In such case, the system
|
||||
automatically deduces it from the tensor's layout using the following logic:
|
||||
|
||||
1. If a dimension's stride is 1, that dimension is marked as the leading dimension.
|
||||
2. If multiple dimensions satisfy condition 1, an error is thrown indicating deduction failure.
|
||||
Note that after converting a **PyTorch** tensor to the DLPack format, the stride for dimensions
|
||||
with size 1 are canonicalized to 1. This canonicalization can increase the likelihood of
|
||||
deduction failures. This behavior is specific to PyTorch and does not occur with NumPy for
|
||||
example.
|
||||
3. If no dimension satisfies condition 1, all strides are marked as dynamic.
|
||||
|
||||
For example:
|
||||
|
||||
- For a tensor with layout ``(2,2,3,4):(2,1,4,12)``, the leading dimension is 1.
|
||||
The layout will be marked as ``(?,?,?,?):(?,1,?,?)``.
|
||||
- For a tensor with layout ``(1,5,1):(1,1,1)``, if ``leading_dim`` is not specified,
|
||||
a deduction failure error is raised.
|
||||
- For a tensor with layout ``(2,2):(8,2)``, since no dimension has stride 1,
|
||||
all dimensions are marked as dynamic: ``(?,?):(?,?)``.
|
||||
|
||||
Code Example
|
||||
~~~~~~~~~~~~
|
||||
|
||||
The following example demonstrates how to use ``mark_layout_dynamic`` to specify dynamic tensor layouts.
|
||||
|
||||
* ``t0`` shows the usage of ``mark_layout_dynamic`` with unspecified ``leading_dim`` and the automatic deduction of leading dimension.
|
||||
* ``t1`` & ``t2`` shows the usage of ``mark_layout_dynamic`` with specified ``leading_dim``.
|
||||
* ``t3`` shows the usage of ``mark_layout_dynamic`` with no leading dimension.
|
||||
* ``t4`` shows the usage of ``mark_layout_dynamic`` with broadcasted dimensions.
|
||||
* ``t5`` demonstrates the deduction failure when the there're more than one dimensions with stride equals to 1.
|
||||
* ``t6`` & ``t7`` demonstrates incorrect settings for ``leading_dim`` and expected errors.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
# (8,4,16,2):(2,16,64,1)
|
||||
a = torch.empty(16, 4, 8, 2).permute(2, 1, 0, 3)
|
||||
# (1,4,1,32,1):(4,1,4,4,4) => torch tensor when dimension has shape 1, its stride is degenerated to 1,
|
||||
# resulting in (1,4,1,32,1):(1,1,1,4,1)
|
||||
b = torch.empty(32, 1, 1, 1, 4).permute(3, 4, 1, 0, 2)
|
||||
# (2,2):(8,2)
|
||||
c = torch.empty(3, 4)[::2, ::2]
|
||||
# (3,1,1,5):(5,0,0,1)
|
||||
d = torch.empty(3, 1, 1, 5).expand(3, 4, 2, 5)
|
||||
|
||||
# auto deduce the leading dimension to be 3
|
||||
t0 = from_dlpack(a).mark_layout_dynamic()
|
||||
print(t0)
|
||||
# (?,?,?,?):(?,?,?,1)
|
||||
|
||||
t1 = from_dlpack(b).mark_layout_dynamic(leading_dim=0)
|
||||
print(t2)
|
||||
# (?,?,?,?,?):(1,?,?,?,?)
|
||||
|
||||
t2 = from_dlpack(b).mark_layout_dynamic(leading_dim=2)
|
||||
print(t3)
|
||||
# (?,?,?,?,?):(?,?,1,?,?)
|
||||
|
||||
t3 = from_dlpack(c).mark_layout_dynamic()
|
||||
print(t3)
|
||||
# (?,?):(?,?)
|
||||
|
||||
t4 = from_dlpack(d).mark_layout_dynamic()
|
||||
print(t4)
|
||||
# (?,?,?,?):(?,0,0,1)
|
||||
|
||||
t5 = from_dlpack(b).mark_layout_dynamic()
|
||||
# Can't decude the leading dimension from layout, please specify the leading_dim explicitly.
|
||||
|
||||
t6 = from_dlpack(a).mark_layout_dynamic(leading_dim=1)
|
||||
# Expected strides[leading_dim] == 1, but got 16
|
||||
|
||||
t7 = from_dlpack(b).mark_layout_dynamic(leading_dim=3)
|
||||
# Expected strides[leading_dim] == 1, but got 4
|
||||
|
||||
Mark the Tensor's Layout as Dynamic with ``mark_compact_shape_dynamic``
|
||||
-----------------------------------------------------------------------
|
||||
|
||||
The ``mark_compact_shape_dynamic`` function provides fine-grain control over dynamic shapes for compact
|
||||
layouts. The full signature of ``mark_compact_shape_dynamic`` is as follows:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def mark_compact_shape_dynamic(self, mode: int, stride_order: tuple[int, ...]|None = None, divisibility: int = 1):
|
||||
|
||||
The ``mode`` parameter determines which shape dimension becomes dynamic. After calling this function,
|
||||
the specific shape dimension given by ``mode`` is marked as dynamic immediately. The stride will be
|
||||
updated accordingly but this process is delayed until the C ABI of the tensor is constructed.
|
||||
For modes that have a shape of size 1, their stride are canonicalized to 0.
|
||||
|
||||
The ``stride_order`` parameter specifies the ordering of strides in the tensor. It is consistent
|
||||
with ``torch.Tensor.dim_order()`` and defaults to ``None``. The parameter indicates the order of
|
||||
modes (dimensions) if the current layout were to be converted to row-major order. It starts from the
|
||||
outermost to the innermost dimension when reading it from left to right. This parameter must be
|
||||
explicitly set when the stride order cannot be automatically deduced from the tensor's layout, such
|
||||
as when multiple dimensions have a stride of 1.
|
||||
|
||||
For example:
|
||||
|
||||
- Layout ``(4,2):(1,4)`` has a ``stride_order`` of ``(1,0)`` indicates the innermost dimension is
|
||||
0 (``4:1``), the outermost dimension is 1 (``2:4``).
|
||||
- Layout ``(5,3,2,4):(3,1,15,30)`` has a ``stride_order`` of ``(3,2,0,1)`` indicates the innermost
|
||||
dimension is 1 (``3:1``), the outermost dimension is 3 (``4:30``).
|
||||
|
||||
If ``stride_order`` is not specified, the system automatically deduces it from the tensor's layout
|
||||
using the following logic:
|
||||
|
||||
1. Sort the strides in descending order.
|
||||
2. If multiple dimensions have a stride of 1, a deduction failure error is raised.
|
||||
|
||||
For example:
|
||||
|
||||
- For a tensor with layout ``(2,2,3,4):(2,1,4,12)``, the deduced ``stride_order`` is ``[3,2,0,1]``.
|
||||
- For a tensor with layout ``(1,5,1):(1,1,1)``, ``stride_order``'s deduction fails because
|
||||
all dimensions have an identical stride of 1, making it impossible to determine the correct ordering.
|
||||
|
||||
If ``stride_order`` is specified, the system validates that the order is consistent with the
|
||||
tensor's layout.
|
||||
|
||||
The ``divisibility`` parameter specifies the divisibility of the dynamic shape. It could be used to
|
||||
represent the assumption alignment of the input. Defaults to 1.
|
||||
|
||||
Note that this API is only available for compact tensors. For non-compact tensors, we can use
|
||||
``cute.assume`` to attach divisibility information to a specific shape mode in a host JIT function,
|
||||
as demonstrated in the following example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def foo(a: cute.Tensor):
|
||||
new_shape = a.shape
|
||||
# use cute.assume to set shape of mode=0 with divisibility=16
|
||||
new_shape[0] = cute.assume(new_shape[0], 16)
|
||||
new_layout = cute.make_layout(new_shape, stride=a.stride)
|
||||
new_a = cute.make_tensor(a.iterator, new_layout)
|
||||
|
||||
|
||||
Code Example
|
||||
~~~~~~~~~~~~
|
||||
|
||||
The following example demonstrates how to use ``mark_compact_shape_dynamic`` to specify dynamic tensor layouts.
|
||||
|
||||
* ``t0`` & ``t1`` show the usage of ``mark_compact_shape_dynamic`` with unspecified ``stride_order`` and different ``mode`` and ``divisibility``.
|
||||
* ``t2`` shows the usage of consecutive ``mark_compact_shape_dynamic`` with unspecified ``stride_order`` and different ``mode`` and ``divisibility``.
|
||||
* ``t3`` & ``t4`` show the usage of ``mark_compact_shape_dynamic`` with different specified ``stride_order``.
|
||||
* ``t5``, ``t6``, ``t7``, ``t8``, ``t9``, ``t10``, ``t11``, and ``t12`` demonstrate incorrect settings for parameters and expected errors.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
@cute.jit
|
||||
def kernel(t: cute.Tensor):
|
||||
pass
|
||||
|
||||
# (8,4,16,2):(2,16,64,1)
|
||||
a = torch.empty(16, 4, 8, 2).permute(2, 1, 0, 3)
|
||||
# (1,4,1,32,1):(4,1,4,4,4) => torch tensor when dimension has shape 1, its stride is degenerated to 1,
|
||||
# resulting in (1,4,1,32,1):(1,1,1,4,1)
|
||||
# b.dim_order() is (3,2,4,0,1)
|
||||
b = torch.empty(32, 1, 1, 1, 4).permute(3, 4, 1, 0, 2)
|
||||
|
||||
# auto deduce the stride order to be [2,1,0,3]
|
||||
t0 = from_dlpack(a).mark_compact_shape_dynamic(
|
||||
mode=0, divisibility=2
|
||||
)
|
||||
kernel(t0)
|
||||
# (?{div=2},4,16,2):(2,?{div=4},?{div=16},1)
|
||||
print(t0)
|
||||
|
||||
t1 = from_dlpack(a).mark_compact_shape_dynamic(
|
||||
mode=1, divisibility=2
|
||||
)
|
||||
kernel(t1)
|
||||
# (8,?{div=2},16,2):(2,16,?{div=32},1)
|
||||
print(t1)
|
||||
|
||||
t2 = from_dlpack(a).mark_compact_shape_dynamic(
|
||||
mode=1, divisibility=2
|
||||
).mark_compact_shape_dynamic(
|
||||
mode=3, divisibility=2
|
||||
)
|
||||
kernel(t2)
|
||||
# (8,?{div=2},16,?{div=2}):(?{div=2},?{div=16},?{div=32},1)
|
||||
print(t2)
|
||||
|
||||
t3 = from_dlpack(b).mark_compact_shape_dynamic(
|
||||
mode=2, divisibility=1, stride_order=(3, 0, 2, 4, 1)
|
||||
)
|
||||
kernel(t3)
|
||||
# (1,4,?,32,1):(0,1,4,?{div=4},0)
|
||||
print(t3)
|
||||
|
||||
t4 = from_dlpack(b).mark_compact_shape_dynamic(
|
||||
mode=2, divisibility=1, stride_order=(2, 3, 4, 0, 1)
|
||||
)
|
||||
kernel(t4)
|
||||
# (1,4,?,32,1):(0,1,128,4,0)
|
||||
print(t4)
|
||||
|
||||
t5 = t2.mark_compact_shape_dynamic(
|
||||
mode=3, divisibility=5, stride_order=(0, 1, 2, 3)
|
||||
)
|
||||
# The stride_order is not consistent with the last stride_order
|
||||
|
||||
t6 = from_dlpack(a).mark_compact_shape_dynamic(
|
||||
mode=3, divisibility=5, stride_order=(0, 1, 2, 3)
|
||||
)
|
||||
# The stride_order is not consistent with the deduced stride_order
|
||||
|
||||
t7 = from_dlpack(b).mark_compact_shape_dynamic(
|
||||
mode=0, divisibility=4
|
||||
)
|
||||
# The layout could not be deduced, please specify the stride_order explicitly
|
||||
|
||||
t8 = from_dlpack(b).mark_compact_shape_dynamic(
|
||||
mode=30, divisibility=5, stride_order=(3, 0, 2, 4, 1)
|
||||
)
|
||||
# Expected mode value to be in range [0, 5), but got 30
|
||||
|
||||
t9 = from_dlpack(b).mark_compact_shape_dynamic(
|
||||
mode=3, divisibility=5, stride_order=(2, 1, 2, 3, 4)
|
||||
)
|
||||
# Expected stride_order to contain all the dimensions of the tensor, but it doesn't contain 0.
|
||||
|
||||
t10 = from_dlpack(b).mark_compact_shape_dynamic(
|
||||
mode=3, divisibility=5, stride_order=(0, 1, 2, 3, 4, 5)
|
||||
)
|
||||
# Expected stride_order to have 5 elements, but got 6.
|
||||
|
||||
t11 = from_dlpack(b).mark_compact_shape_dynamic(
|
||||
mode=0, divisibility=4, stride_order=b.dim_order()
|
||||
)
|
||||
# The shape(1) of mode(0) is not divisible by the divisibility(4)
|
||||
|
||||
t12 = from_dlpack(b).mark_compact_shape_dynamic(
|
||||
mode=0, divisibility=1, stride_order=(2, 1, 3, 0, 4)
|
||||
)
|
||||
# The stride_order is not consistent with the layout
|
||||
@@ -0,0 +1,16 @@
|
||||
.. _notebooks:
|
||||
|
||||
Educational Notebooks
|
||||
=====================
|
||||
|
||||
A number of notebooks for educational purposes are provided in the `CUTLASS GitHub repository <https://github.com/NVIDIA/cutlass>`__.
|
||||
A list with handful links is given below:
|
||||
|
||||
- `"Hello world" <https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/hello_world.ipynb>`__
|
||||
- `Printing <https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/print.ipynb>`__
|
||||
- `Data Types Basics <https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/data_types.ipynb>`__
|
||||
- `Tensors <https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/tensor.ipynb>`__
|
||||
- `The TensorSSA Abstraction <https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/tensorssa.ipynb>`__
|
||||
- `Layout Algebra <https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/cute_layout_algebra.ipynb>`__
|
||||
- `Element-wise Add Tutorial <https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/elementwise_add.ipynb>`__
|
||||
- `Using CUDA Graphs <https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/cuda_graphs.ipynb>`__
|
||||
Reference in New Issue
Block a user