v4.5 tag update (#3202)
* Python DSL examples reorganization. * v4.5 tag update.
This commit is contained in:
Junkai-Wu
@@ -39,10 +39,11 @@ CuTe DSL provides environment variables to control logging level:
|
||||
# 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
|
||||
# Log to file instead of console (default: False).
|
||||
# Set to 1/True to enable; the log file path is chosen automatically by the DSL.
|
||||
export CUTE_DSL_LOG_TO_FILE=1
|
||||
|
||||
# Control log verbosity (0, 10, 20, 30, 40, 50, default: 10)
|
||||
# Control log verbosity (0=disabled, 1=all messages (debug and above), 10=debug, 20=info, 30=warning, 40=error, 50=critical; default: 1)
|
||||
export CUTE_DSL_LOG_LEVEL=20
|
||||
|
||||
|
||||
@@ -68,40 +69,53 @@ Similar to standard Python logging, different log levels provide varying degrees
|
||||
+--------+-------------+
|
||||
|
||||
|
||||
Dump the generated IR
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
Save generated artifacts to files
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
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.
|
||||
CuTe DSL can save generated artifacts (IR, PTX, CUBIN, …) to files for offline inspection.
|
||||
Use ``CUTE_DSL_KEEP`` with a comma-separated list of artifact tokens:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# Dump Generated CuTe IR (default: False)
|
||||
# Save clean IR (after canonicalize+cse, human-readable) to a .mlir file
|
||||
export CUTE_DSL_KEEP=ir
|
||||
|
||||
# Save raw IR (before any passes) to a .mlir file
|
||||
export CUTE_DSL_KEEP=ir-debug
|
||||
|
||||
# Save PTX assembly to a .ptx file
|
||||
export CUTE_DSL_KEEP=ptx
|
||||
|
||||
# Save CUBIN binary to a .cubin file
|
||||
export CUTE_DSL_KEEP=cubin
|
||||
|
||||
# Save LLVM IR to a file
|
||||
export CUTE_DSL_KEEP=llvm
|
||||
|
||||
# Save multiple artifacts at once
|
||||
export CUTE_DSL_KEEP=ir,ptx,cubin
|
||||
|
||||
# Save all supported artifacts
|
||||
export CUTE_DSL_KEEP=all
|
||||
|
||||
Files are written to the current working directory by default. Use ``CUTE_DSL_DUMP_DIR``
|
||||
to redirect them (see `Change the dump directory`_ below).
|
||||
|
||||
.. note::
|
||||
|
||||
The ``sass`` token requires ``nvdisasm`` (or ``nvdisasm_internal``) to be available
|
||||
in your ``PATH``. It is usually installed with the CUDA toolkit.
|
||||
|
||||
Print the generated IR to the console
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
To print the IR directly to the console (without writing a file):
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# Print generated IR to stdout (default: False)
|
||||
export CUTE_DSL_PRINT_IR=1
|
||||
|
||||
# Keep Generated CuTe IR in a file (default: False)
|
||||
export CUTE_DSL_KEEP_IR=1
|
||||
|
||||
|
||||
Dump the generated PTX & CUBIN
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
For users familiar with PTX and SASS, CuTe DSL supports dumping the generated PTX and CUBIN.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# Dump generated PTX in a .ptx file (default: False)
|
||||
export CUTE_DSL_KEEP_PTX=1
|
||||
|
||||
# Dump generated cubin in a .cubin file (default: False)
|
||||
export CUTE_DSL_KEEP_CUBIN=1
|
||||
|
||||
To further get SASS from cubin, users can use ``nvdisasm`` (usually installed with CUDA toolkit) to disassemble the cubin.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
nvdisasm your_dsl_code.cubin > your_dsl_code.sass
|
||||
|
||||
|
||||
Access the dumped contents programmatically
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
@@ -101,18 +101,11 @@ The result:
|
||||
|
||||
|DSL| bridges Python and GPU hardware through a three-stage pipeline.
|
||||
|
||||
.. _fig-dsl-modes:
|
||||
|
||||
.. figure:: dsl_modes.png
|
||||
:width: 400
|
||||
.. figure:: dsl_compilation.png
|
||||
:width: 600
|
||||
: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.
|
||||
|
||||
|
||||
The default |DSL| compilation pipeline (mode 2): Python source flows through AST preprocessing
|
||||
The |DSL| compilation pipeline: Python source flows through AST preprocessing
|
||||
and interpreter-driven tracing to produce |IR|, which is then lowered and
|
||||
compiled to device code.
|
||||
|
||||
@@ -258,8 +251,8 @@ Practical Implications
|
||||
4. |DSL| Code-Generation Modes
|
||||
------------------------------
|
||||
|
||||
CuTe's Python front-end combines the techniques above into **two mutually
|
||||
exclusive modes** (see :ref:`fig-dsl-modes`), selectable with the ``preprocessor`` flag of the
|
||||
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.
|
||||
@@ -272,3 +265,10 @@ 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.
|
||||
|
||||
@@ -117,6 +117,12 @@ Defines GPU kernel functions, compiled as specialized GPU symbols through |DC|.
|
||||
- ``False`` (default) — Standard kernel launch.
|
||||
- ``True`` — Cooperative kernel launch.
|
||||
|
||||
- ``smem_merge_branch_allocs``
|
||||
Enables mutually exclusive control flow branches (sequentially executed if-else) to reuse the same shared memory.
|
||||
|
||||
- ``False`` (default) — Shared memory is allocated additively across all branches (default CUDA C++ behavior).
|
||||
- ``True`` — Merge shared-memory allocations across branches (experimental feature, recommended for mega-kernels).
|
||||
|
||||
Calling Conventions
|
||||
-------------------
|
||||
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
.. _dsl_struct_types:
|
||||
|
||||
Struct-like JIT Arguments
|
||||
=========================
|
||||
|
||||
|DSL| supports several struct-like Python types as JIT function arguments.
|
||||
Each provides a different trade-off between mutability, syntax convenience,
|
||||
and low-level control.
|
||||
|
||||
.. |DSL| replace:: CuTe DSL
|
||||
|
||||
.. contents:: On this page
|
||||
:local:
|
||||
:depth: 2
|
||||
|
||||
|
||||
Overview
|
||||
--------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 25 15 60
|
||||
|
||||
* - Type
|
||||
- Mutable fields?
|
||||
- Notes
|
||||
* - ``typing.NamedTuple``
|
||||
- **No**
|
||||
- Tuple subclass — fields fixed at construction.
|
||||
Flattened field-by-field through the pytree system.
|
||||
* - ``@dataclass(frozen=True)``
|
||||
- **No**
|
||||
- Frozen dataclass — treated as a read-only pytree container,
|
||||
similar to ``NamedTuple``.
|
||||
|
||||
|
||||
NamedTuple
|
||||
----------
|
||||
|
||||
A ``typing.NamedTuple`` whose fields are DSL scalar types (``Int32``,
|
||||
``Float32``, etc.) can be passed directly to ``@cute.jit`` /
|
||||
``cute.compile`` without any boilerplate or protocol implementation.
|
||||
|
||||
**How it works.** NamedTuples are registered as pytree containers in the DSL
|
||||
tree system. Each field is flattened individually through the existing DSL
|
||||
type paths and reconstructed by calling the NamedTuple constructor on the way
|
||||
into the kernel body. Field attribute access (``tup.a``, ``tup.b``, …)
|
||||
works exactly as in native Python.
|
||||
|
||||
Basic usage
|
||||
^^^^^^^^^^^
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from typing import NamedTuple
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
|
||||
class Vec3(NamedTuple):
|
||||
x: cutlass.Int32
|
||||
y: cutlass.Int32
|
||||
z: cutlass.Int32
|
||||
|
||||
@cute.jit
|
||||
def print_vec(v: Vec3):
|
||||
cute.printf("x=%d y=%d z=%d\n", v.x, v.y, v.z)
|
||||
|
||||
v = Vec3(x=cutlass.Int32(1), y=cutlass.Int32(2), z=cutlass.Int32(3))
|
||||
cute.compile(print_vec, v)(v)
|
||||
|
||||
Control flow on fields
|
||||
^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Fields are DSL values inside the kernel, so they work in ``if``/``else``
|
||||
branches and ``for`` loops:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def clamp_positive(v: Vec3, out: cute.Tensor):
|
||||
"""Write max(field, 0) for each component."""
|
||||
out[0] = cutlass.Int32(0) if v.x < cutlass.Int32(0) else v.x
|
||||
out[1] = cutlass.Int32(0) if v.y < cutlass.Int32(0) else v.y
|
||||
out[2] = cutlass.Int32(0) if v.z < cutlass.Int32(0) else v.z
|
||||
|
||||
@cute.jit
|
||||
def triangular_sum(v: Vec3, out: cute.Tensor):
|
||||
"""Sum 0..v.x-1 into out[0], and so on."""
|
||||
s = cutlass.Int32(0)
|
||||
for i in range(v.x):
|
||||
s = s + i
|
||||
out[0] = s
|
||||
|
||||
Creating a new NamedTuple value inside the kernel
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
NamedTuple fields are **immutable** — the same constraint as native Python
|
||||
tuples. Assigning ``tup.x = ...`` inside a kernel raises ``AttributeError``.
|
||||
To "update" a field, construct a replacement NamedTuple:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def scale(v: Vec3, factor: cutlass.Int32, out: cute.Tensor):
|
||||
# Construct a new Vec3 with all fields scaled
|
||||
scaled = Vec3(x=v.x * factor, y=v.y * factor, z=v.z * factor)
|
||||
out[0] = scaled.x
|
||||
out[1] = scaled.y
|
||||
out[2] = scaled.z
|
||||
|
||||
Choosing the right type
|
||||
-----------------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 65
|
||||
|
||||
* - Use case
|
||||
- Recommended type
|
||||
* - Read-only config / parameters passed into a kernel
|
||||
- ``NamedTuple`` or ``@dataclass(frozen=True)``
|
||||
* - Accumulator or running state updated inside a kernel
|
||||
- ``@native_struct``
|
||||
* - Want Python-native immutable semantics (hashable, unpackable)
|
||||
- ``NamedTuple``
|
||||
* - Need fine-grained LLVM struct control (packing, zero-init)
|
||||
- ``@native_struct``
|
||||
|
||||
|
||||
See also
|
||||
--------
|
||||
|
||||
* :doc:`dsl_jit_arg_generation` — overview of JIT function argument protocols
|
||||
* :doc:`dsl_dynamic_layout` — passing ``Layout`` objects as JIT arguments
|
||||
@@ -79,6 +79,12 @@ by reducing register usage and the number of address calculation instructions. W
|
||||
to True, a runtime check is performed to ensure that the layout does not overflow. Please note that this parameter
|
||||
only has an effect when the tensor's layout is marked as dynamic.
|
||||
|
||||
For packed subbyte torch dtypes such as ``torch.float4_e2m1fn_x2``, ``from_dlpack`` exposes the
|
||||
logical element layout expected by CuTe instead of the packed storage layout. For example, a torch
|
||||
tensor with shape ``(128, 128)`` and dtype ``torch.float4_e2m1fn_x2`` is exposed as a logical FP4
|
||||
tensor with shape ``(128, 256)``. The same logical reinterpretation also applies when the leading
|
||||
dimension is not the last mode.
|
||||
|
||||
Code Example
|
||||
~~~~~~~~~~~~
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ Conference Talks
|
||||
An introduction to the |DSL| architecture, covering the hybrid AST-rewrite and
|
||||
tracing approach, MLIR code generation, and integration with CUTLASS.
|
||||
|
||||
* `LLVM Video <https://www.youtube.com/watch?v=5NXd6MbKYNQ>`_
|
||||
* `Video <https://www.youtube.com/watch?v=5NXd6MbKYNQ>`__
|
||||
* `Slides (PDF) <https://llvm.org/devmtg/2025-10/slides/technical_talks/ozen.pdf>`_
|
||||
|
||||
----
|
||||
@@ -25,4 +25,4 @@ tracing approach, MLIR code generation, and integration with CUTLASS.
|
||||
Learn how to leverage Tensor Cores directly from Python using CUTLASS 4.0's
|
||||
new DSL front-end, enabling rapid kernel development without writing CUDA C++.
|
||||
|
||||
* `GTC Video <https://www.nvidia.com/en-us/on-demand/session/gtc25-s74639/>`_
|
||||
* `Video <https://www.nvidia.com/en-us/on-demand/session/gtc25-s74639/>`__
|
||||
|
||||
Reference in New Issue
Block a user