v4.5 tag update (#3202)

* Python DSL examples reorganization.

* v4.5 tag update.

examples/python/CuTeDSL/dsl_tutorials/call_bypass_dlpack.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2025 - 2026 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.
+
+import sys
+import os
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import make_ptr
+
+
+"""
+An Example demonstrating how to call off-the-shelf kernel by-passing dlpack protocol
+
+The example shows how to directly pass pointers from PyTorch tensors to off-the-shelf kernels
+written by CuTe DSL with a thin customized wrapper jit function. The jit function will be
+compiled with inline without introducing overhead.
+
+To run this example:
+
+.. code-block:: bash
+
+    python examples/ampere/call_bypass_dlpack.py
+
+
+It's worth to mention that by-passing dlpack protocol can resolve the issue that dlpack doesn't handle shape-1
+mode correctly. For example, the following code will fail, because dlpack will convert the shape-1 mode
+with stride-1 which propagate alignment incorrectly.
+
+.. code-block:: python
+
+    @cute.kernel
+    def fails_kernel(gX: cute.Tensor):
+        bidx, _, _ = cute.arch.block_idx()
+        mX = gX[None, bidx, None]  # We wish to retain alignment
+        # assert mX.iterator.alignment == 16
+
+
+    @cute.jit
+    def fails(gX_: cute.Tensor):
+        gX = gX_
+        fails_kernel(gX).launch(grid=(1, 1, 1), block=(128, 1, 1))
+
+
+    gX_torch = torch.rand((128, 1, 128), device="cuda", dtype=torch.bfloat16)
+    fails(from_dlpack(gX_torch, assumed_align=16))
+
+"""
+
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, ".."))
+
+from cute.ampere.kernel.dense_gemm.tensorop_gemm import TensorOpGemm
+
+
+@cute.jit
+def tensor_op_gemm_wrapper(
+    a_ptr: cute.Pointer,
+    b_ptr: cute.Pointer,
+    c_ptr: cute.Pointer,
+    m: cutlass.Int32,
+    n: cutlass.Int32,
+    k: cutlass.Int32,
+    l: cutlass.Int32,
+):
+    print("\n[DSL INFO] Input Parameters:")
+    print(f"[DSL INFO]   mnkl: {(m, n, k, l)}")
+
+    # Assume alignment of shape to call tensorop_gemm example
+    m = cute.assume(m, divby=8)
+    n = cute.assume(n, divby=8)
+
+    # Torch is row major
+    a_layout = cute.make_ordered_layout((m, k, l), order=(0, 1, 2))
+    b_layout = cute.make_ordered_layout((n, k, l), order=(0, 1, 2))
+    c_layout = cute.make_ordered_layout((m, n, l), order=(1, 0, 2))
+    mA = cute.make_tensor(a_ptr, layout=a_layout)
+    mB = cute.make_tensor(b_ptr, layout=b_layout)
+    mC = cute.make_tensor(c_ptr, layout=c_layout)
+
+    print(f"[DSL INFO]   mA: {mA}")
+    print(f"[DSL INFO]   mB: {mB}")
+    print(f"[DSL INFO]   mC: {mC}")
+
+    tensor_op_gemm = TensorOpGemm(
+        a_ptr.value_type, c_ptr.value_type, cutlass.Float32, (2, 2, 1)
+    )
+    print("\n[DSL INFO] Created TensorOpGemm instance")
+    print(f"[DSL INFO]   Input dtype: {a_ptr.value_type}")
+    print(f"[DSL INFO]   Output dtype: {c_ptr.value_type}")
+    print(f"[DSL INFO]   Accumulation dtype: {cutlass.Float32}")
+    print(f"[DSL INFO]   Atom layout: {(2, 2, 1)}")
+
+    # No need to compile inside jit function
+    tensor_op_gemm(mA, mB, mC)
+    print("\n[DSL INFO] Executed TensorOpGemm")
+
+
+def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
+    import torch
+
+    print("\nRunning TensorOpGemm test with:")
+    print(f"Tensor dimensions: {mnkl}")
+
+    # (M,K,L)
+    a = torch.randn(
+        mnkl[3], mnkl[2], mnkl[0], dtype=torch.float16, device="cuda"
+    ).permute(2, 1, 0)
+    # (N,K,L)
+    b = torch.randn(
+        mnkl[3], mnkl[2], mnkl[1], dtype=torch.float16, device="cuda"
+    ).permute(2, 1, 0)
+    # (N,M,L)
+    c = torch.randn(
+        mnkl[3], mnkl[0], mnkl[1], dtype=torch.float16, device="cuda"
+    ).permute(1, 2, 0)
+
+    print("Input tensor shapes:")
+    print(f"a: {a.shape}, dtype: {a.dtype}")
+    print(f"b: {b.shape}, dtype: {b.dtype}")
+    print(f"c: {c.shape}, dtype: {c.dtype}\n")
+
+    a_ptr = make_ptr(
+        cutlass.Float16, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
+    )
+    b_ptr = make_ptr(
+        cutlass.Float16, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
+    )
+    c_ptr = make_ptr(
+        cutlass.Float16, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
+    )
+    tensor_op_gemm_wrapper(a_ptr, b_ptr, c_ptr, *mnkl)
+    torch.cuda.synchronize()
+
+    ref = torch.einsum("mkl,nkl->mnl", a, b)
+    torch.testing.assert_close(c, ref, atol=1e-05, rtol=1e-05)
+    print("\n[DSL INFO] Results verified successfully!")
+    print(f"First few elements of result: \n{c[:3, :3, :3]}")
+
+
+if __name__ == "__main__":
+    run_tensor_op_gemm_wrapper((512, 256, 128, 16))

examples/python/CuTeDSL/dsl_tutorials/call_from_jit.py

@@ -0,0 +1,260 @@
+# Copyright (c) 2025 - 2026 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.
+
+"""
+Demonstrating JIT GEMM Implementation with Static Shape Wrapper
+
+This example illustrates how to invoke a JIT-compiled GEMM implementation through a wrapper function
+with static shapes. It showcases the integration between PyTorch and CuTe tensors in a JIT context.
+
+Key features demonstrated:
+1. Seamless conversion between PyTorch and CuTe tensors using the JitArgument protocol
+2. Integration of static shape GEMM operations within a JIT-compiled wrapper function
+
+Core components:
+- BufferWithLayout: Handles memory buffer management with configurable stride ordering
+- tensor_op_gemm_wrapper: JIT-compiled entry point that orchestrates the GEMM operation
+
+Usage:
+
+.. code-block:: bash
+
+    python examples/ampere/call_from_jit.py
+
+Default configuration:
+- Batch dimension (L): 16
+- Matrix dimensions: M=512, N=256, K=128
+- Precision: Float16 inputs with Float32 accumulation
+
+Requirements:
+- CUDA-capable GPU
+- PyTorch with CUDA support
+"""
+
+import os
+import sys
+from typing import Type, Tuple
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import make_ptr
+
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, ".."))
+
+from cute.ampere.kernel.dense_gemm.tensorop_gemm import TensorOpGemm
+
+
+class BufferWithLayout:
+    def __init__(self, ptr: cute.Pointer, stride_order: tuple[int, int, int]):
+        self.ptr = ptr
+
+        # static properties
+        self.stride_order = stride_order
+
+    def to_tensor(
+        self, shape: tuple[int, int, int], *, loc=None, ip=None
+    ) -> cute.Tensor:
+        assert len(shape) == len(self.stride_order), (
+            f"Shape {shape} and stride_order {self.stride_order} must have the "
+            "same rank."
+        )
+        layout = cute.make_ordered_layout(shape, self.stride_order)
+        # permute (l, mn, k) -> (mn, k, l)
+        res = cute.make_tensor(self.ptr, cute.select(layout, mode=[1, 2, 0]))
+        return res
+
+    # Implement JitArgument Protocol and DynamicExpression Protocol
+
+    def __c_pointers__(self):
+        """Get the C pointers for the underlying pointer.
+
+        This method is part of the JitArgument Protocol and returns the C pointers
+        from the underlying pointer object.
+
+        This is required for user to define a custom data type which can pass to JIT function.
+        When JIT compiled function is called, JIT executor will call this method to get raw pointers
+        to underlying data object.
+
+        Following condition must be satisfied:
+
+        len(__c_pointers__()) == len(__get_mlir_types__()) == len(__extract_mlir_values__())
+
+        :return: The C pointers from the underlying pointer object
+        :rtype: Any
+        """
+        return self.ptr.__c_pointers__()
+
+    def __get_mlir_types__(self):
+        """Get the MLIR types for the underlying pointer.
+
+        This method is part of the JitArgument Protocol and returns the MLIR types
+        used for compiler to generate code. It must match the type of the underlying pointers
+        returned by __c_pointers__().
+
+        :return: The MLIR types from the underlying pointer object
+        :rtype: Any
+        """
+        return self.ptr.__get_mlir_types__()
+
+    def __extract_mlir_values__(self):
+        """Extract MLIR values from the underlying pointer.
+
+        This method is part of the DynamicExpression Protocol and extracts MLIR values
+        from the underlying pointer object.
+
+        It is used by compiler to generate function call in MLIR to another JIT function.
+        It must match the types returned by __get_mlir_types__().
+
+        :return: The MLIR values extracted from the underlying pointer object
+        :rtype: Any
+        """
+        return self.ptr.__extract_mlir_values__()
+
+    def __new_from_mlir_values__(self, values):
+        """Create a new BufferWithLayout instance from MLIR values.
+
+        This method is part of the JitArgument & DynamicExpression Protocol and creates a new
+        BufferWithLayout instance with pointer initialized from the given MLIR values.
+
+        It is used by compiler to generate function body in MLIR called by JIT function.
+        It must match the types returned by __c_pointers__() and __get_mlir_types__().
+        code generator takes function arguments and reconstructs python object which is legal
+        inside function body.
+
+        :param values: MLIR values to initialize the underlying pointer
+        :type values: Any
+        :return: A new BufferWithLayout instance with pointer initialized from values
+        :rtype: BufferWithLayout
+        """
+        return BufferWithLayout(
+            self.ptr.__new_from_mlir_values__(values), self.stride_order
+        )
+
+
+@cute.jit
+def tensor_op_gemm_wrapper(
+    buffer_a: BufferWithLayout,
+    buffer_b: BufferWithLayout,
+    buffer_c: BufferWithLayout,
+    mnkl: cutlass.Constexpr[tuple[int, int, int, int]],
+    acc_dtype: Type[cutlass.Numeric],
+    atom_layout_mnk: cutlass.Constexpr[tuple[int, int, int]],
+):
+    print("\n[DSL INFO] Input Parameters:")
+    print(f"[DSL INFO]   mnkl: {mnkl}")
+    print(f"[DSL INFO]   buffer_a: {buffer_a}")
+    print(f"[DSL INFO]   buffer_b: {buffer_b}")
+    print(f"[DSL INFO]   buffer_c: {buffer_c}")
+    print(f"[DSL INFO]   acc_dtype: {acc_dtype}")
+    print(f"[DSL INFO]   atom_layout_mnk: {atom_layout_mnk}")
+
+    mA = buffer_a.to_tensor(cute.select(mnkl, mode=[3, 0, 2]))
+    mB = buffer_b.to_tensor(cute.select(mnkl, mode=[3, 1, 2]))
+    mC = buffer_c.to_tensor(cute.select(mnkl, mode=[3, 0, 1]))
+
+    print("\n[DSL INFO] Created Tensors:")
+    print(f"[DSL INFO]   mA = {mA}")
+    print(f"[DSL INFO]   mB = {mB}")
+    print(f"[DSL INFO]   mC = {mC}")
+
+    tensor_op_gemm = TensorOpGemm(
+        buffer_a.ptr.value_type,
+        buffer_c.ptr.value_type,
+        acc_dtype,
+        atom_layout_mnk,
+    )
+    print("\n[DSL INFO] Created TensorOpGemm instance")
+    print(f"[DSL INFO]   Input dtype: {buffer_a.ptr.value_type}")
+    print(f"[DSL INFO]   Output dtype: {buffer_c.ptr.value_type}")
+    print(f"[DSL INFO]   Accumulation dtype: {acc_dtype}")
+    print(f"[DSL INFO]   Atom layout: {atom_layout_mnk}")
+
+    # No need to compile inside jit function
+    tensor_op_gemm(mA, mB, mC)
+    print("\n[DSL INFO] Executed TensorOpGemm")
+
+
+def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
+    import torch
+    from cutlass.torch import dtype as torch_dtype
+
+    print("\nRunning TensorOpGemm test with:")
+    print(f"Tensor dimensions: {mnkl}")
+
+    ab_dtype = cutlass.Float16
+    c_dtype = cutlass.Float16
+
+    a = torch.randn(
+        mnkl[3], mnkl[0], mnkl[2], dtype=torch_dtype(ab_dtype), device="cuda"
+    )
+    b = torch.randn(
+        mnkl[3], mnkl[1], mnkl[2], dtype=torch_dtype(ab_dtype), device="cuda"
+    )
+    c = torch.randn(
+        mnkl[3], mnkl[0], mnkl[1], dtype=torch_dtype(c_dtype), device="cuda"
+    )
+
+    print("Input tensor shapes:")
+    print(f"a: {a.shape}, dtype: {a.dtype}")
+    print(f"b: {b.shape}, dtype: {b.dtype}")
+    print(f"c: {c.shape}, dtype: {c.dtype}\n")
+
+    buffer_a = BufferWithLayout(
+        make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
+        (2, 1, 0),
+    )
+    buffer_b = BufferWithLayout(
+        make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
+        (2, 1, 0),
+    )
+    buffer_c = BufferWithLayout(
+        make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
+        (2, 1, 0),
+    )
+
+    tensor_op_gemm_wrapper(
+        buffer_a,
+        buffer_b,
+        buffer_c,
+        mnkl,  # pass shape as static value
+        # no stride passing
+        cutlass.Float32,
+        (2, 2, 1),
+    )
+    torch.cuda.synchronize()
+
+    ref = torch.einsum("lmk,lnk->lmn", a, b)
+    torch.testing.assert_close(c, ref, atol=1e-05, rtol=1e-05)
+    print("\n[DSL INFO] Results verified successfully!")
+    print(f"First few elements of result: \n{c[:3, :3, :3]}")
+
+
+if __name__ == "__main__":
+    run_tensor_op_gemm_wrapper((512, 256, 128, 16))

examples/python/CuTeDSL/dsl_tutorials/cooperative_launch.py

@@ -0,0 +1,627 @@
+# Copyright (c) 2025 - 2026 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.
+
+"""
+Cooperative Launch Example:
+
+This module demonstrates CUDA Cooperative Launch functionality. It implements a
+global barrier that synchronizes ALL threads across the entire GPU grid.
+
+In traditional CUDA kernel launches, there is no guarantee that all thread blocks
+will be resident on the GPU simultaneously. This means that thread blocks may
+execute in waves (some finish before others start) and attempting to synchronize
+across blocks can cause deadlock.
+
+**Cooperative Launch** solves this by guaranteeing that all thread blocks launch
+atomically and simultaneously.
+
+For more details, see the CUDA Programming Guide official documentation:
+https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/cooperative-groups.html#when-to-use-cudalaunchcooperativekernel
+
+Cooperative Launch Limitations:
+
+Cooperative launch has strict grid size constraints.
+If you exceed this limit, cudaLaunchCooperativeKernel returns
+cudaErrorCooperativeLaunchTooLarge.
+
+This example demonstrates both a successful cooperative launch with a small grid
+and an expected failure when exceeding the grid size limit.
+
+Usage:
+
+Run directly:
+    $ python cooperative_launch.py
+
+This will:
+    1. Demonstrate expected failure with too many thread blocks
+    2. Successfully run a cooperative kernel with grid-wide barrier
+    3. Print confirmation that all threads synchronized successfully
+
+"""
+
+from typing import List, Optional
+import cutlass
+import cutlass.cute as cute
+from cutlass._mlir import ir
+from cutlass.cutlass_dsl import (
+    dsl_user_op,  # Decorator for user-defined device operations
+    DSLCudaRuntimeError,  # Exception type for CUDA runtime errors
+    extract_mlir_values,  # Extract MLIR values from the object
+    new_from_mlir_values,  # Create a new instance from MLIR values
+)
+
+# Function to check cuda errors
+from cutlass.base_dsl.runtime.cuda import checkCudaErrors
+
+# LLVM dialect for inline PTX assembly generation
+from cutlass._mlir.dialects import llvm
+
+# CUDA Python bindings for runtime API (memory allocation, synchronization, etc.)
+import cuda.bindings.runtime as cuda_runtime
+
+
+class GlobalBarrier:
+    """
+    A grid-wide barrier for synchronizing ALL thread blocks on the GPU.
+
+    This class implements a cooperative barrier that enables grid-wide
+    synchronization. It requires cooperative launch to function correctly.
+
+    Design Overview:
+
+    The barrier uses a single 32-bit integer in global memory with the
+    following bit layout:
+
+        ┌──────────────────────────────────────────────────────────────────┐
+        │ Bit 31     │ Bits 30-0                                          │
+        │ ────────── │ ───────────────────────────────────────────────────│
+        │ Phase Bit  │ Arrival Counter (supports up to 2^31 - 1 blocks)   │
+        └──────────────────────────────────────────────────────────────────┘
+
+    Capacity:
+
+    - Maximum thread blocks: 2^31 - 1 = 2,147,483,647 blocks
+
+    Memory Ordering:
+
+    The barrier uses specific memory ordering semantics:
+
+    - Release semantics on arrival (atom.add.release.gpu)
+
+    - Acquire semantics on wait (ld.global.acquire.gpu)
+
+    Usage Example:
+
+    Host-side setup:
+
+        >>> barrier_ptr = GlobalBarrier.allocate()  # Allocate barrier memory
+
+    Device-side usage (inside a kernel):
+
+        >>> barrier = GlobalBarrier(barrier_ptr)
+        >>>
+        >>> # Do some work...
+        >>>
+        >>> barrier.arrive_and_wait()  # Synchronize all blocks
+        >>>
+        >>> # All blocks proceed together after this point
+
+    Warning:
+
+    This barrier requires cooperative launch! Using it with a regular launch
+    can result in a deadlock because not all thread blocks may be resident
+    simultaneously.
+    """
+
+    @staticmethod
+    def allocate() -> cute.runtime.Pointer:
+        """
+        Allocate and initialize barrier memory on the GPU.
+
+        This function allocates device memory for the barrier.
+        It must be called before launching any kernel that uses the barrier.
+        """
+        ptr = checkCudaErrors(cuda_runtime.cudaMalloc(4))
+
+        # This sets all 32 bits to 0:
+        #   - Phase bit (bit 31) = 0
+        #   - Counter (bits 30-0) = 0
+        checkCudaErrors(cuda_runtime.cudaMemset(ptr, 0, 4))
+
+        # Create a pointer with the following properties:
+        #   - Type: Uint32 (32-bit unsigned integer)
+        #   - Address: the allocated device pointer
+        #   - Address Space: gmem (global memory)
+        barrier_ptr = cute.runtime.make_ptr(
+            cutlass.Uint32,  # Element type
+            ptr,  # Raw CUDA pointer
+            cute.AddressSpace.gmem,  # Memory address space
+        )
+
+        return barrier_ptr
+
+    @staticmethod
+    def free(barrier_ptr: cute.Pointer):
+        """
+        Free the barrier memory on the GPU.
+
+        This function frees the device memory for the barrier.
+        It must be called after the barrier is no longer needed.
+        """
+        checkCudaErrors(cuda_runtime.cudaFree(barrier_ptr._pointer))
+
+    @dsl_user_op
+    def __init__(
+        self,
+        barrier_ptr: cute.Pointer,
+        *,
+        phase: Optional[cutlass.Uint32] = None,
+        is_leader: Optional[cutlass.Boolean] = None,
+        number_of_thread_blocks: Optional[cutlass.Uint32] = None,
+        loc=None,
+        ip=None,
+    ):
+        """
+        Initialize a GlobalBarrier instance on the device.
+
+        This constructor is called by each thread when the kernel
+        starts. It sets up the barrier state for this thread's participation
+        in grid-wide synchronization.
+
+        Each thread stores the following:
+
+        - A reference to the shared barrier memory
+        - Whether it's the leader thread of its block
+        - The current phase for barrier tracking
+        - The total number of thread blocks in the grid
+        """
+        # The barrier is shared across ALL thread blocks, so it must be in
+        # global memory. Shared memory (smem) is block-local and wouldn't work.
+        if barrier_ptr.memspace != cute.AddressSpace.gmem:
+            raise ValueError(
+                "GlobalBarrier requires barrier_ptr to be in global memory (gmem)"
+            )
+
+        # Store barrier pointer reference
+        self.barrier_ptr = barrier_ptr
+
+        # Initialize phase tracking
+        # Phase starts at 0 for the first barrier, then alternates:
+        #   First barrier:  wait for phase 1
+        #   Second barrier: wait for phase 0
+        #   Third barrier:  wait for phase 1
+        #   ... and so on
+        if phase is not None:
+            self.phase = phase
+        else:
+            self.phase = cutlass.Uint32(0)
+
+        if is_leader is not None:
+            self.is_leader = is_leader
+        else:
+            # Determine if this thread is the block leader
+            # Get this thread's position within its block
+            tidx, tidy, tidz = cute.arch.thread_idx()
+
+            # Leader is the thread at position (0, 0, 0) in the block
+            # We use bitwise AND to combine the three conditions efficiently
+            self.is_leader = (
+                cutlass.Boolean(tidx == 0)  # First in X dimension
+                & cutlass.Boolean(tidy == 0)  # First in Y dimension
+                & cutlass.Boolean(tidz == 0)  # First in Z dimension
+            )
+
+        if number_of_thread_blocks is not None:
+            self.number_of_thread_blocks = number_of_thread_blocks
+        else:
+            # Calculate total number of thread blocks in the grid
+            # Get grid dimensions (how many blocks in each dimension)
+            gidx, gidy, gidz = cute.arch.grid_dim()
+
+            # Total blocks = gridDim.x × gridDim.y × gridDim.z
+            # This is needed to know when ALL blocks have arrived
+            self.number_of_thread_blocks = cutlass.Uint32(gidx * gidy * gidz)
+
+    @dsl_user_op
+    @cute.jit
+    def arrive(self, *, loc=None, ip=None):
+        """
+        Arrive at the barrier without waiting.
+
+        This signals that the calling thread block has reached the barrier
+        point, but does not wait for other blocks. Use this when you want
+        to overlap computation with barrier synchronization.
+
+        This method must be called by ALL threads in the block,
+        not just the leader. The internal block-level sync ensures all
+        threads in the block agree before the leader signals arrival.
+        """
+        # Ensure ALL threads in this block have reached this point before
+        # the leader signals arrival. This is critical for correctness!
+        cute.arch.sync_threads(loc=loc, ip=ip)
+
+        # Only the leader thread performs atomic operations to minimize
+        # contention on the barrier memory location
+        if self.is_leader:
+            # Atomically increment the arrival counter by 1
+            # The atomic add returns the value before the add, so we add 1
+            # to get the current value after our arrival
+            barrier_value = (
+                self._increment_barrier(cutlass.Uint32(1), loc=loc, ip=ip) + 1
+            )
+
+            # Check if we're the last block to arrive
+            # Mask out the phase bit (bit 31) to get just the counter value
+            # Compare against total number of thread blocks
+            if (barrier_value & ~(1 << 31)) == self.number_of_thread_blocks:
+                # Flip phase and reset counter
+                # We add a value that simultaneously:
+                #   1. Flips bit 31 (adds 2^31)
+                #   2. Resets counter to 0 (subtracts N, where N was the count)
+                #
+                # Example with 8 blocks:
+                #   Current: 0x00000008 (phase=0, counter=8)
+                #   Add:     0x80000000 - 8 = 0x7FFFFFF8
+                #   Result:  0x80000000 (phase=1, counter=0) ✓
+                #
+                # This works because we're doing modular arithmetic and the
+                # counter wraps correctly
+                self._increment_barrier(
+                    cutlass.Uint32((1 << 31) - self.number_of_thread_blocks),
+                    loc=loc,
+                    ip=ip,
+                )
+
+    def _read_barrier(self, *, loc=None, ip=None) -> cutlass.Uint32:
+        """
+        Read the barrier value with acquire memory semantics.
+
+        This is an internal method that reads the 32-bit barrier value from
+        global memory using GPU-scope acquire semantics.
+
+        Notes
+        -----
+        PTX Instruction:
+            Uses ld.global.acquire.gpu.b32 which is a:
+
+            - Global memory load (ld.global)
+            - With acquire semantics (.acquire)
+            - At GPU scope (.gpu) - visible across all thread blocks
+            - For 32-bit data (.b32)
+
+        Inline Assembly:
+            We use LLVM inline assembly because CuTe DSL may not have a direct
+            high-level API for acquire loads. The assembly string format:
+
+            - $0: Output operand (the loaded value)
+            - $1: Input operand (the address to load from)
+
+        """
+        # Use inline PTX assembly for the acquire-semantics load
+        return cutlass.Uint32(
+            llvm.inline_asm(
+                # Return type: 32-bit unsigned integer
+                cutlass.Uint32.mlir_type,
+                # Input arguments: barrier pointer address
+                # We convert the pointer to an integer (64-bit address)
+                [self.barrier_ptr.toint().ir_value(loc=loc, ip=ip)],
+                # PTX instruction
+                "ld.global.acquire.gpu.b32 $0, [$1];",
+                # Constraint string
+                # "=r" : Output is a 32-bit register (write-only)
+                # "l"  : Input is a 64-bit register (pointer address)
+                "=r,l",
+                # Assembly attributes
+                # Mark as having side effects
+                has_side_effects=True,
+                # No special stack alignment needed
+                is_align_stack=False,
+                # Use AT&T syntax (required for LLVM inline asm)
+                asm_dialect=llvm.AsmDialect.AD_ATT,
+                # MLIR location and insertion point
+                loc=loc,
+                ip=ip,
+            )
+        )
+
+    def _increment_barrier(
+        self, value: cutlass.Uint32, *, loc=None, ip=None
+    ) -> cutlass.Uint32:
+        """
+        Atomically increment the barrier with release memory semantics.
+
+        This is an internal method that performs an atomic add on the barrier
+        value using GPU-scope release semantics.
+
+        Notes
+        -----
+        PTX Instruction:
+            Uses atom.add.release.gpu.u32 which is a:
+
+            - Atomic operation (atom)
+            - Addition (.add)
+            - With release semantics (.release)
+            - At GPU scope (.gpu)
+            - For unsigned 32-bit integers (.u32)
+
+        Atomicity:
+            The atomic add is guaranteed to be indivisible - no other thread
+            can see a partial update or interleave with this operation.
+
+        Return Value:
+            Atomic operations return the OLD value, not the new value.
+            This is why the caller adds 1 to get the current count.
+        """
+        # Atomic add using inline PTX assembly with release semantics
+        return cutlass.Uint32(
+            llvm.inline_asm(
+                # Return type: 32-bit unsigned integer (the old value)
+                cutlass.Uint32.mlir_type,
+                # Input arguments: (barrier address, value to add)
+                [
+                    self.barrier_ptr.toint().ir_value(
+                        loc=loc, ip=ip
+                    ),  # Barrier address
+                    value.ir_value(loc=loc, ip=ip),  # Value to add
+                ],
+                # PTX  instruction
+                "atom.add.release.gpu.u32 $0, [$1], $2;",
+                # Constraint string
+                # "=r" : Output is a 32-bit register
+                # "l"  : First input is 64-bit (pointer)
+                # "r"  : Second input is 32-bit (value)
+                "=r,l,r",
+                # Assembly attributes
+                has_side_effects=True,
+                is_align_stack=False,
+                asm_dialect=llvm.AsmDialect.AD_ATT,
+                # MLIR metadata
+                loc=loc,
+                ip=ip,
+            )
+        )
+
+    @dsl_user_op
+    @cute.jit
+    def wait(self, *, loc=None, ip=None):
+        """
+        Wait for all thread blocks to arrive at the barrier.
+
+        This method blocks (spins) until all thread blocks have called
+        arrive() on the barrier. It does NOT signal arrival itself - use
+        arrive_and_wait() if you need to both arrive and wait.
+
+        IMPORTANT: This method MUST be called by ALL threads in the block.
+        The internal sync_threads ensures all threads proceed together.
+
+        Algorithm:
+        - Leader thread spins, reading barrier with acquire semantics
+        - Waits until phase bit matches expected value
+        - Block-level sync ensures all threads proceed together
+        - Update local phase tracking for next barrier
+
+        """
+        # Leader thread: spin-wait for phase flip
+        if self.is_leader:
+            # Calculate expected phase (opposite of current phase)
+            # XOR with 1 flips: 0→1, 1→0
+            expected = self.phase ^ 1
+
+            # Initial read of barrier value
+            barrier_value = self._read_barrier(loc=loc, ip=ip)
+
+            # Spin loop: wait until phase matches expected
+            # Extract phase bit (bit 31)
+            # Compare against expected phase value
+            while (barrier_value >> 31) != expected:
+                # Keep reading barrier until phase flips
+                # The acquire semantics ensure memory ordering
+                barrier_value = self._read_barrier(loc=loc, ip=ip)
+
+        # Block-level synchronization
+        # Ensure all threads in the block wait for the leader to see the
+        # phase flip before any thread proceeds
+        cute.arch.sync_threads(loc=loc, ip=ip)
+
+        # Update phase for next barrier
+        # Flip local phase: 0→1 or 1→0
+        # This prepares for the next barrier synchronization
+        self.phase = self.phase ^ 1
+
+    @dsl_user_op
+    def arrive_and_wait(self, *, loc=None, ip=None):
+        """
+        Arrive at the barrier AND wait for all other thread blocks.
+
+        This is the most common barrier operation - it combines arrive()
+        and wait() into a single call. All thread blocks will be synchronized
+        after this call returns.
+
+        IMPORTANT: This method MUST be called by ALL threads in the block.
+
+        Semantics
+        ---------
+
+        Logically equivalent to:
+
+            >>> barrier.arrive()   # Signal we've reached this point
+            >>> barrier.wait()     # Wait for everyone else
+        """
+        # Execute both phases: arrive then wait
+        self.arrive(loc=loc, ip=ip)
+        self.wait(loc=loc, ip=ip)
+
+    def __extract_mlir_values__(self) -> List[ir.Value]:
+        """
+        Extract MLIR values from the GlobalBarrier instance.
+        """
+
+        assert len(extract_mlir_values(self.barrier_ptr)) == 1
+        assert len(extract_mlir_values(self.is_leader)) == 1
+        assert len(extract_mlir_values(self.phase)) == 1
+        assert len(extract_mlir_values(self.number_of_thread_blocks)) == 1
+
+        return (
+            extract_mlir_values(self.barrier_ptr)
+            + extract_mlir_values(self.is_leader)
+            + extract_mlir_values(self.phase)
+            + extract_mlir_values(self.number_of_thread_blocks)
+        )
+
+    def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GlobalBarrier":
+        """
+        Create a new GlobalBarrier instance from MLIR values.
+        """
+        assert len(values) == 4, f"Expected 4 IR values, but got {len(values)}"
+        return GlobalBarrier(
+            barrier_ptr=new_from_mlir_values(self.barrier_ptr, [values[0]]),
+            is_leader=new_from_mlir_values(self.is_leader, [values[1]]),
+            phase=new_from_mlir_values(self.phase, [values[2]]),
+            number_of_thread_blocks=new_from_mlir_values(
+                self.number_of_thread_blocks, [values[3]]
+            ),
+        )
+
+
+@cute.kernel
+def cooperative_kernel(barrier_ptr: cute.Pointer):
+    """
+    Example kernel demonstrating cooperative launch with grid-wide barrier.
+
+    This kernel shows how to use the GlobalBarrier class to synchronize all
+    thread blocks in a grid. It performs 10 iterations, with a barrier
+    synchronization after each iteration.
+
+    Launch Requirements: This kernel MUST be launched with cooperative=True.
+
+    """
+    # Initialize the barrier for this thread
+    # Each thread creates its own GlobalBarrier instance, all sharing the
+    # same underlying barrier_ptr in global memory
+    barrier = GlobalBarrier(barrier_ptr=barrier_ptr)
+
+    for i in range(10):
+        # Synchronize all thread blocks across the entire grid
+        # After this call, ALL blocks have completed iterations 0..i
+        barrier.arrive_and_wait()
+
+        # Get block and thread indices
+        bidx, bidy, bidz = cute.arch.block_idx()
+        tidx, tidy, tidz = cute.arch.thread_idx()
+
+        # Check if this is the leader block (first block in the grid)
+        leader_cluster = bidx == 0 and bidy == 0 and bidz == 0
+
+        # Check if this is the leader thread (first thread in the block)
+        leader_thread = tidx == 0 and tidy == 0 and tidz == 0
+
+        # Only the single leader thread of the leader block prints
+        if leader_cluster and leader_thread:
+            cute.printf("All threads arrived at barrier for the %dth iteration\n", i)
+
+
+# =============================================================================
+#                           KERNEL LAUNCH WRAPPERS
+# =============================================================================
+
+
+@cute.jit
+def run_cooperative_kernel(barrier_ptr: cute.runtime.Pointer):
+    """
+    Launch the cooperative kernel with a reasonable grid size.
+
+    This wrapper launches the cooperative_kernel with a grid of 8 thread blocks
+    (2×2×2), where each block contains 128 threads (32×2×2).
+
+    Notes: The cooperative=True flag is ESSENTIAL. It tells CUDA to:
+    - Verify the grid fits within hardware limits
+    - Launch all blocks atomically
+    """
+    cooperative_kernel(barrier_ptr).launch(
+        grid=(2, 2, 2),  # 8 thread blocks
+        block=(32, 2, 2),  # 128 threads per block
+        cooperative=True,  # Enable cooperative launch semantics
+    )
+
+
+@cute.jit
+def xfail_run_cooperative_kernel(barrier_ptr: cute.runtime.Pointer):
+    """
+    Demonstrate cooperative launch failure with an oversized grid.
+
+    This wrapper intentionally launches with a grid that exceeds
+    the limits, demonstrating how cooperative launch fails.
+
+    This launch is expected to fail with cudaErrorCooperativeLaunchTooLarge.
+
+    This demonstrates proper error handling for cooperative launch.
+
+    See Also
+    --------
+    The main() function shows how to properly catch and handle this error.
+    """
+    # Attempt to launch with way too many blocks
+    cooperative_kernel(barrier_ptr).launch(
+        grid=(10000, 1, 1),  # 10,000 blocks
+        block=(1024, 1, 1),  # 1,024 threads per block
+        cooperative=True,  # Cooperative launch will reject this
+    )
+
+
+if __name__ == "__main__":
+    # Initialize CUDA context
+    cutlass.cuda.initialize_cuda_context()
+
+    # Allocate barrier memory
+    # Allocate 4 bytes in device global memory for the barrier state
+    barrier_ptr = GlobalBarrier.allocate()
+
+    # Demonstrate expected failure (grid too large)
+    expectedly_failed = False
+    try:
+        # Attempt to launch with 10,000 blocks - this WILL fail
+        xfail_run_cooperative_kernel(barrier_ptr)
+    except DSLCudaRuntimeError as e:
+        # Verify we got the expected error code
+        assert (
+            e.error_code == cuda_runtime.cudaError_t.cudaErrorCooperativeLaunchTooLarge
+        )
+        expectedly_failed = True
+    finally:
+        # Ensure the failure actually happened (test validation)
+        assert expectedly_failed
+
+    # Run successful cooperative kernel
+    # Launch with a reasonable grid size that fits hardware constraints
+    run_cooperative_kernel(barrier_ptr)
+
+    # Synchronize and clean up
+    checkCudaErrors(cuda_runtime.cudaDeviceSynchronize())
+
+    # Free the barrier memory
+    GlobalBarrier.free(barrier_ptr)

examples/python/CuTeDSL/dsl_tutorials/dataclass_immutable.py

@@ -0,0 +1,98 @@
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Dataclass Example: Passing Tensors via Frozen Dataclass
+=======================================================
+
+Demonstrates passing tensors from `@cute.jit` to `@cute.kernel`
+using `@dataclass(frozen=True)`.
+
+Uses fake tensors for compilation and TVM-FFI for efficient runtime dispatch.
+
+To run:
+
+.. code-block:: bash
+
+    python examples/python/CuTeDSL/cute/dataclass_immutable.py
+
+"""
+
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32
+
+
+@dataclass(frozen=True)
+class CopyParams:
+    """
+    Dataclass with one static attribute and two tensors.
+
+    Attributes:
+        num_threads: Threads per block (static - used for loop unrolling)
+        src: Source tensor (dynamic)
+        dst: Destination tensor (dynamic)
+    """
+
+    num_threads: int  # static
+    src: cute.Tensor  # dynamic
+    dst: cute.Tensor  # dynamic
+
+
+@cute.jit
+def memcpy_jit(mSrc: cute.Tensor, mDst: cute.Tensor, num_threads: cutlass.Constexpr[int]):
+    """Host function - passes tensors via dataclass."""
+    params = CopyParams(
+        num_threads=num_threads,
+        src=mSrc,
+        dst=mDst,
+    )
+
+    N = mSrc.shape[0]
+    num_blocks = cute.ceil_div(N, num_threads)
+
+    memcpy_kernel(params, N).launch(
+        grid=[num_blocks, 1, 1],
+        block=[num_threads, 1, 1],
+    )
+
+
+@cute.kernel
+def memcpy_kernel(params: CopyParams, N: Int32):
+    """Device kernel - receives tensors via CopyParams dataclass."""
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    idx = bidx * params.num_threads + tidx
+
+    if idx < N:
+        params.dst[idx] = params.src[idx]
+
+
+if __name__ == "__main__":
+    import torch
+    from cutlass.cute.runtime import make_fake_compact_tensor
+
+    n = cute.sym_int64()
+    # 1. Create fake tensors for compilation (no GPU memory allocated)
+    fake_src = make_fake_compact_tensor(cutlass.Float32, (n,), stride_order=(0,))
+    fake_dst = make_fake_compact_tensor(cutlass.Float32, (n,), stride_order=(0,))
+
+    # 2. Compile with --enable-tvm-ffi
+    compiled_fn = cute.compile(
+        memcpy_jit, fake_src, fake_dst, num_threads=128, options="--enable-tvm-ffi"
+    )
+
+    # 3. Create real torch tensors
+    src = torch.randn(1024, device="cuda", dtype=torch.float32)
+    dst = torch.zeros_like(src)
+
+    # 4. Run compiled kernel
+    compiled_fn(src, dst)
+
+    # 5. Verify dst == src
+    torch.cuda.synchronize()
+    torch.testing.assert_close(dst, src)
+    print("PASS")

examples/python/CuTeDSL/dsl_tutorials/dynamic_smem_size.py

@@ -0,0 +1,145 @@
+# Copyright (c) 2025 - 2026 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.
+
+
+import cutlass.cute as cute
+import cutlass
+
+
+"""
+Example of automatic shared memory size computation for configuring kernel launch
+
+This example demonstrates how to let the DSL automatically set shared memory
+  size for a kernel launch rather explicitly configuring it at launch time,
+  provided that developers are using `SmemAllocator` for all allocations.
+
+Usage:
+    python dynamic_smem_size.py         # Show auto inference
+"""
+
+
+@cute.struct
+class SharedData:
+    """A struct to demonstrate shared memory allocation."""
+
+    values: cute.struct.MemRange[cutlass.Float32, 64]  # 256 bytes
+    counter: cutlass.Int32  # 4 bytes
+    flag: cutlass.Int8  # 1 byte
+
+
+@cute.kernel
+def kernel_static():
+    """
+    Example kernel that allocates shared memory.
+    The total allocation will be automatically calculated when smem=None.
+    """
+    tidx, _, _ = cute.arch.block_idx()
+    if tidx == 0:
+        cute.printf("Running kernel_static")
+
+    allocator = cutlass.utils.SmemAllocator()
+
+    # Allocate various types of shared memory
+    shared_data = allocator.allocate(SharedData)
+    raw_buffer = allocator.allocate(512, byte_alignment=64)
+    int_array = allocator.allocate_array(element_type=cutlass.Int32, num_elems=128)
+    tensor_smem = allocator.allocate_tensor(
+        element_type=cutlass.Float16,
+        layout=cute.make_layout((32, 16)),
+        byte_alignment=16,
+        swizzle=None,
+    )
+
+    cute.printf("Kernel launch smem size: {}", cute.arch.dynamic_smem_size())
+    return
+
+
+@cute.kernel
+def kernel_no_smem():
+    """
+    Example kernel that does not allocates shared memory.
+    The total allocation will be automatically calculated as 0 when smem=None.
+    """
+    tidx, _, _ = cute.arch.block_idx()
+    if tidx == 0:
+        cute.printf("Running kernel_no_smem")
+    return
+
+
+if __name__ == "__main__":
+    # Initialize CUDA context
+    cutlass.cuda.initialize_cuda_context()
+
+    print("Launching kernel with auto smem size. (launch config `smem=None`)")
+
+    # Compile the static example
+    @cute.jit
+    def launch_kernelno_smem():
+        kernel_no_smem().launch(
+            grid=(1, 1, 1),
+            block=(1, 1, 1),
+        )
+
+    # --------
+    print(f" > Run {kernel_no_smem.__name__}")
+    func = cute.compile(launch_kernelno_smem)
+    func()
+    cutlass.cuda.stream_sync(cutlass.cuda.default_stream())
+
+    @cute.jit
+    def launch_kernel_static():
+        kernel_static().launch(
+            grid=(1, 1, 1),
+            block=(1, 1, 1),
+            # smem=None
+            # auto infer launch kernel static smem usage
+        )
+
+    # --------
+    print(f" > Run {kernel_static.__name__} with sufficient smem")
+    func = cute.compile(launch_kernel_static)
+    func()
+    cutlass.cuda.stream_sync(cutlass.cuda.default_stream())
+
+    @cute.jit
+    def launch_kernel_static_insufficient():
+        kernel_static().launch(
+            grid=(1, 1, 1),
+            block=(1, 1, 1),
+            # launch kernel with static smem usage exceeds cfg
+            # show warning
+            smem=16,
+        )
+
+    # --------
+    print(f" > Run {kernel_static.__name__} with insufficient smem, show warning:")
+    func = cute.compile(launch_kernel_static_insufficient)
+    func()
+    cutlass.cuda.stream_sync(cutlass.cuda.default_stream())
+
+    print("PASS")

examples/python/CuTeDSL/dsl_tutorials/export/export_to_c.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2025 - 2026 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.
+
+import cutlass
+import cutlass.cute as cute
+
+import os
+import subprocess
+import cuda.bindings.driver as cuda
+
+"""Example demonstrating how to export a CuTe function.
+
+This example shows how to:
+1. Compile a CuTe function
+2. Export the compiled function to a object file and a C header file
+3. Compile the object file to a shared library
+
+To run this example:
+
+.. code-block:: bash
+
+    # export the compiled function to a object file and a C header file
+    python cutlass_ir/compiler/python/examples/cute/export/export_to_c.py
+"""
+
+
+@cute.kernel
+def print_tensor_kernel(a: cute.Tensor):
+    cute.printf("a: {}", a)
+
+
+@cute.jit
+def print_tensor(a: cute.Tensor, stream: cuda.CUstream):
+    print_tensor_kernel(a).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream)
+
+
+@cute.kernel
+def add_one_kernel(a: cute.Tensor, b: cute.Tensor):
+    a[0] = b[0] + 1
+
+
+@cute.jit
+def add_one(a: cute.Tensor, b: cute.Tensor, stream: cuda.CUstream):
+    add_one_kernel(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream)
+
+
+def run():
+    from cutlass.cute.runtime import make_fake_compact_tensor
+
+    shape = (cute.SymInt(divisibility=16), cute.SymInt())
+    a = make_fake_compact_tensor(cutlass.Float32, shape, stride_order=(1, 0))
+    b = make_fake_compact_tensor(cutlass.Float32, shape, stride_order=(1, 0))
+    stream = cuda.CUstream(0)
+    compiled_print_tensor = cute.compile(print_tensor, a, stream=stream)
+    compiled_add_one = cute.compile(add_one, a, b, stream=stream)
+
+    os.makedirs("./build", exist_ok=True)
+    compiled_print_tensor.export_to_c(
+        file_path="./build",
+        file_name="print_tensor_example",
+        function_prefix="print_tensor",
+    )
+    compiled_add_one.export_to_c(
+        file_path="./build", file_name="add_one_example", function_prefix="add_one"
+    )
+
+    cc = os.environ.get("CC", "gcc")
+
+    # compile the object file to a shared library
+    cmd = [
+        cc,
+        "-shared",
+        "-o",
+        "./build/libexport_example.so",
+        "./build/print_tensor_example.o",
+        "./build/add_one_example.o",
+    ]
+    subprocess.run(cmd, check=True)
+
+
+if __name__ == "__main__":
+    run()

examples/python/CuTeDSL/dsl_tutorials/export/load_in_python.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2025 - 2026 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.
+
+import cutlass
+import cutlass.cute as cute
+import cuda.bindings.driver as cuda
+
+
+"""Example demonstrating how to load a CuTe module/function in Python.
+
+This example shows how to:
+1. Load a CuTe module from a object file or a shared library
+2. Extract the function from the module and call it in Python
+
+To run this example:
+
+.. code-block:: bash
+
+    # prerequesites: export the compiled functions to object files and compile them into a shared library
+    python examples/cute/export/export_to_c.py
+    # load the module from a object file or a shared library
+    python examples/cute/export/load_in_python.py
+"""
+
+
+def run():
+    import torch
+    from cutlass.cute.runtime import from_dlpack
+
+    a = (
+        torch.arange(16 * 10, dtype=torch.float32, device="cuda")
+        .reshape(16, 10)
+        .permute(1, 0)
+    )
+    b = torch.ones(16, 10, dtype=torch.float32, device="cuda").permute(1, 0)
+    a_cute = from_dlpack(a).mark_layout_dynamic()
+    b_cute = from_dlpack(b).mark_layout_dynamic()
+    stream = cuda.CUstream(0)
+
+    print_tensor_mod = cute.runtime.load_module("./build/print_tensor_example.o")
+    print_tensor_mod.print_tensor(a_cute, stream=stream)
+
+    add_one_mod = cute.runtime.load_module("./build/add_one_example.o")
+    add_one_mod.add_one(a_cute, b_cute, stream=stream)
+    assert a[0, 0] == b[0, 0] + 1
+
+    shared_mod = cute.runtime.load_module("./build/libexport_example.so")
+    shared_mod.print_tensor(a_cute, stream=stream)
+    shared_mod.add_one(a_cute, b_cute, stream=stream)
+    assert a[0, 0] == b[0, 0] + 1
+
+
+if __name__ == "__main__":
+    run()

examples/python/CuTeDSL/dsl_tutorials/export/run_with_dynamic_loading.cpp

@@ -0,0 +1,254 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved. SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+ *
+ * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+ * property and proprietary rights in and to this material, related
+ * documentation and any modifications thereto. Any use, reproduction,
+ * disclosure or distribution of this material and related documentation
+ * without an express license agreement from NVIDIA CORPORATION or
+ * its affiliates is strictly prohibited.
+ */
+/**
+ * This example demonstrates how to load a CuTe module/function from a object
+ * file or a shared library and run it in a CUDA kernel.
+ *
+ * To run this example:
+ *
+ * .. code-block:: bash
+ *
+ *    # prerequesites: export the compiled functions to object files and compile
+ * them into a shared library python examples/cute/export/export_to_c.py # run
+ * the example bash ./examples/cute/export/run_with_dynamic_loading.sh
+ */
+
+#include "CuteDSLRuntime.h"
+#include <cuda_runtime.h>
+#include <fstream>
+#include <vector>
+
+std::vector<unsigned char> read_file(const std::string &filename) {
+  std::ifstream file(filename, std::ios::binary);
+  std::vector<unsigned char> content((std::istreambuf_iterator<char>(file)),
+                                     std::istreambuf_iterator<char>());
+  return content;
+}
+
+void initialize_cuda_context() {
+  // Initialize cuda context
+  cudaSetDevice(0);
+}
+
+void check_error(CuteDSLRT_Error_t error) {
+  if (error != CuteDSLRT_Error_Success) {
+    printf("Got runtime error: %d\n", error);
+    exit(1);
+  }
+}
+
+// Copy the definition of the tensor from `print_tensor_example.h` to here
+typedef struct {
+  void *data;
+  int32_t dynamic_shapes[2];
+  int64_t dynamic_strides[1];
+} print_tensor_Tensor_a_t;
+
+// Copy the definition of the tensor from `add_one_example.h` to here
+typedef struct {
+  void *data;
+  int32_t dynamic_shapes[2];
+  int64_t dynamic_strides[1];
+} add_one_Tensor_a_t;
+typedef struct {
+  void *data;
+  int32_t dynamic_shapes[2];
+  int64_t dynamic_strides[1];
+} add_one_Tensor_b_t;
+
+print_tensor_Tensor_a_t prepare_print_tensor_tensor() {
+  print_tensor_Tensor_a_t tensor;
+  tensor.data = NULL;
+  tensor.dynamic_shapes[0] = 32;
+  tensor.dynamic_shapes[1] = 16;
+  tensor.dynamic_strides[0] = 16;
+  return tensor;
+}
+
+add_one_Tensor_a_t prepare_add_one_tensor_a() {
+  float a_host_ptr[32 * 16];
+  for (int i = 0; i < 32 * 16; i++) {
+    a_host_ptr[i] = i;
+  }
+  void *a_device_ptr;
+  cudaMalloc(&a_device_ptr, sizeof(float) * 32 * 16);
+  cudaMemcpy(a_device_ptr, a_host_ptr, sizeof(float) * 32 * 16,
+             cudaMemcpyHostToDevice);
+  add_one_Tensor_a_t tensor;
+  tensor.data = (void *)a_device_ptr;
+  tensor.dynamic_shapes[0] = 32;
+  tensor.dynamic_shapes[1] = 16;
+  tensor.dynamic_strides[0] = 16;
+  return tensor;
+}
+
+add_one_Tensor_b_t prepare_add_one_tensor_b() {
+  float b_host_ptr[32 * 16];
+  for (int i = 0; i < 32 * 16; i++) {
+    b_host_ptr[i] = 32 * 16 - i;
+  }
+  void *b_device_ptr;
+  cudaMalloc(&b_device_ptr, sizeof(float) * 32 * 16);
+  cudaMemcpy(b_device_ptr, b_host_ptr, sizeof(float) * 32 * 16,
+             cudaMemcpyHostToDevice);
+  add_one_Tensor_b_t tensor;
+  tensor.data = (void *)b_device_ptr;
+  tensor.dynamic_shapes[0] = 32;
+  tensor.dynamic_shapes[1] = 16;
+  tensor.dynamic_strides[0] = 16;
+  return tensor;
+}
+
+void run_print_tensor_with_object_file() {
+
+  // prepare tensor
+  print_tensor_Tensor_a_t tensor = prepare_print_tensor_tensor();
+
+  // prepare stream
+  cudaStream_t stream;
+  cudaStreamCreate(&stream);
+
+  // load module and function
+  CuteDSLRT_Module_t *module = nullptr;
+  std::vector<unsigned char> print_tensor_example_o_bytes =
+      read_file("build/print_tensor_example.o");
+  size_t print_tensor_example_o_size = print_tensor_example_o_bytes.size();
+  CuteDSLRT_Error_t error = CuteDSLRT_Module_Create_From_Bytes(
+      &module, print_tensor_example_o_bytes.data(), print_tensor_example_o_size,
+      nullptr, 0);
+  check_error(error);
+
+  CuteDSLRT_Function_t *function = nullptr;
+  error = CuteDSLRT_Module_Get_Function(&function, module, "print_tensor");
+  check_error(error);
+
+  // run kernel, refer to the wrapper function in `print_tensor_example.h`
+  int32_t cuda_result;
+  void *args[3] = {&tensor, &stream, &cuda_result};
+  error = CuteDSLRT_Function_Run(function, args, 3);
+  check_error(error);
+
+  // synchronize stream
+  cudaStreamSynchronize(stream);
+
+  // unload module
+  error = CuteDSLRT_Module_Destroy(module);
+  check_error(error);
+  cudaStreamDestroy(stream);
+}
+
+void run_add_one_with_object_file() {
+  // prepare a tensor
+  add_one_Tensor_a_t tensor_a = prepare_add_one_tensor_a();
+
+  // prepare b tensor
+  add_one_Tensor_b_t tensor_b = prepare_add_one_tensor_b();
+
+  // prepare stream
+  cudaStream_t stream;
+  cudaStreamCreate(&stream);
+
+  // load module
+  CuteDSLRT_Module_t *module = nullptr;
+  std::vector<unsigned char> add_one_example_o_bytes =
+      read_file("build/add_one_example.o");
+  size_t add_one_example_o_size = add_one_example_o_bytes.size();
+  CuteDSLRT_Error_t error = CuteDSLRT_Module_Create_From_Bytes(
+      &module, add_one_example_o_bytes.data(), add_one_example_o_size, nullptr,
+      0);
+  check_error(error);
+
+  CuteDSLRT_Function_t *function = nullptr;
+  error = CuteDSLRT_Module_Get_Function(&function, module, "add_one");
+  check_error(error);
+
+  // run kernel, refer to the wrapper function in `add_one_example.h`
+  int32_t cuda_result;
+  void *args[4] = {&tensor_a, &tensor_b, &stream, &cuda_result};
+  error = CuteDSLRT_Function_Run(function, args, 4);
+  check_error(error);
+  cudaStreamSynchronize(stream);
+
+  // unload module
+  error = CuteDSLRT_Module_Destroy(module);
+  check_error(error);
+  cudaStreamDestroy(stream);
+
+  // check result
+  float a_host_ptr[32 * 16];
+  cudaMemcpy(a_host_ptr, tensor_a.data, sizeof(float) * 32 * 16,
+             cudaMemcpyDeviceToHost);
+
+  if (a_host_ptr[0] != 32 * 16 + 1) {
+    printf("Error: a_host_ptr[0] = %f, expected %d\n", a_host_ptr[0], 32 * 16);
+    exit(1);
+  }
+}
+
+void run_example_with_shared_library() {
+  // prepare tensor
+  print_tensor_Tensor_a_t tensor = prepare_print_tensor_tensor();
+
+  // prepare a tensor
+  add_one_Tensor_a_t tensor_a = prepare_add_one_tensor_a();
+
+  // prepare b tensor
+  add_one_Tensor_b_t tensor_b = prepare_add_one_tensor_b();
+
+  // prepare stream
+  cudaStream_t stream;
+  cudaStreamCreate(&stream);
+
+  // load module
+  CuteDSLRT_Module_t *module = nullptr;
+  const char *shared_libs[] = {"build/libexport_example.so"};
+  CuteDSLRT_Error_t error =
+      CuteDSLRT_Module_Create_From_Bytes(&module, nullptr, 0, shared_libs, 1);
+  check_error(error);
+
+  // get print tensor function
+  CuteDSLRT_Function_t *print_tensor_function = nullptr;
+  error = CuteDSLRT_Module_Get_Function(&print_tensor_function, module,
+                                        "print_tensor");
+  check_error(error);
+
+  // get add one function
+  CuteDSLRT_Function_t *add_one_function = nullptr;
+  error = CuteDSLRT_Module_Get_Function(&add_one_function, module, "add_one");
+  check_error(error);
+
+  // run print tensor kernel
+  int32_t cuda_result;
+  void *print_tensor_args[3] = {&tensor, &stream, &cuda_result};
+  error = CuteDSLRT_Function_Run(print_tensor_function, print_tensor_args, 3);
+  check_error(error);
+  cudaStreamSynchronize(stream);
+
+  // run add one kernel
+  void *add_one_args[4] = {&tensor_a, &tensor_b, &stream, &cuda_result};
+  error = CuteDSLRT_Function_Run(add_one_function, add_one_args, 4);
+  check_error(error);
+  cudaStreamSynchronize(stream);
+
+  // unload module
+  error = CuteDSLRT_Module_Destroy(module);
+  check_error(error);
+  cudaStreamDestroy(stream);
+}
+
+int main() {
+  initialize_cuda_context();
+  run_print_tensor_with_object_file();
+  run_add_one_with_object_file();
+  run_example_with_shared_library();
+  return 0;
+}

examples/python/CuTeDSL/dsl_tutorials/export/run_with_dynamic_loading.sh

@@ -0,0 +1,80 @@
+# Copyright (c) 2025 - 2026 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.
+
+#!/bin/bash
+set -eu
+
+# Try to find the wheel path of nvidia-cutlass-dsl
+WHEEL_PATH=$(python3 -c "import cutlass, os; print(os.path.dirname(cutlass.__file__))" 2>/dev/null)/../..
+
+if [[ -z "$WHEEL_PATH" ]]; then
+    echo "nvidia-cutlass-dsl wheel not found in the current Python environment."
+    exit 1
+else
+    echo "nvidia-cutlass-dsl wheel path found at: $WHEEL_PATH"
+fi
+
+CUTE_DSL_LIB_PATH="${WHEEL_PATH}/lib/"
+export LD_LIBRARY_PATH=${CUTE_DSL_LIB_PATH}:${CUDA_HOME}/lib64:./build
+
+if [ -z "$CUDA_HOME" ]; then
+    CUDA_HOME=/usr/local/cuda
+    echo "CUDA_HOME not set, using default: $CUDA_HOME"
+else
+    echo "CUDA_HOME found: $CUDA_HOME"
+fi
+SOURCE_FILE="$(dirname "$0")/run_with_dynamic_loading.cpp"
+
+echo "Compiling the executable..."
+# Search for a common C++ compiler: g++, clang++, or c++
+if [ -n "${CXX-}" ] && command -v "$CXX" &> /dev/null; then
+    CXX="$CXX"
+elif command -v g++ &> /dev/null; then
+    CXX="g++"
+elif command -v clang++ &> /dev/null; then
+    CXX="clang++"
+elif command -v c++ &> /dev/null; then
+    CXX="c++"
+else
+    echo "Error: No common C++ compiler found (g++, clang++, or c++). Please install a C++ compiler to continue."
+    exit 1
+fi
+
+$CXX -o build/run_with_dynamic_loading \
+  -I${CUDA_HOME}/include \
+  -I${WHEEL_PATH}/include \
+  ${SOURCE_FILE} \
+  -L${CUTE_DSL_LIB_PATH} \
+  -L${CUDA_HOME}/lib64 \
+  -L./build \
+  -lcudart \
+  -lcute_dsl_runtime \
+  -lexport_example
+
+echo "Running the executable..."
+./build/run_with_dynamic_loading

examples/python/CuTeDSL/dsl_tutorials/export/run_with_static_linking.cpp

@@ -0,0 +1,124 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved. SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+ *
+ * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+ * property and proprietary rights in and to this material, related
+ * documentation and any modifications thereto. Any use, reproduction,
+ * disclosure or distribution of this material and related documentation
+ * without an express license agreement from NVIDIA CORPORATION or
+ * its affiliates is strictly prohibited.
+ */
+/**
+ * This example demonstrates how to load a CuTe module/function from compilation
+ * and static linking and run it in a CUDA kernel.
+ *
+ * To run this example:
+ *
+ * .. code-block:: bash
+ *
+ *    # prerequesites: export the compiled functions to object files and compile
+ * them into a shared library python examples/cute/export/export_to_c.py # run
+ * the example bash ./examples/cute/export/run_with_static_linking.sh
+ */
+
+#include "add_one_example.h"
+#include "print_tensor_example.h"
+#include <cuda_runtime.h>
+
+void initialize_cuda_context() {
+  // Initialize cuda context
+  int device_id = 0;
+  cudaSetDevice(device_id);
+}
+
+void run_print_tensor() {
+  // prepare tensor
+  print_tensor_Tensor_a_t tensor;
+  tensor.data = NULL;
+  tensor.dynamic_shapes[0] = 32;
+  tensor.dynamic_shapes[1] = 16;
+  tensor.dynamic_strides[0] = 16;
+
+  // prepare stream
+  cudaStream_t stream;
+  cudaStreamCreate(&stream);
+
+  // load module
+  print_tensor_Kernel_Module_t module;
+  print_tensor_Kernel_Module_Load(&module);
+
+  // run kernel
+  cute_dsl_print_tensor_wrapper(&module, &tensor, stream);
+
+  // synchronize stream
+  cudaStreamSynchronize(stream);
+
+  // unload module
+  print_tensor_Kernel_Module_Unload(&module);
+  cudaStreamDestroy(stream);
+}
+
+void run_add_one() {
+  // prepare a tensor
+  add_one_Tensor_a_t a_tensor;
+  float a_host_ptr[32 * 16];
+  for (int i = 0; i < 32 * 16; i++) {
+    a_host_ptr[i] = i;
+  }
+  void *a_device_ptr;
+  cudaMalloc(&a_device_ptr, sizeof(float) * 32 * 16);
+  cudaMemcpy(a_device_ptr, a_host_ptr, sizeof(float) * 32 * 16,
+             cudaMemcpyHostToDevice);
+  a_tensor.data = (void *)a_device_ptr;
+  a_tensor.dynamic_shapes[0] = 32;
+  a_tensor.dynamic_shapes[1] = 16;
+  a_tensor.dynamic_strides[0] = 16;
+
+  // prepare b tensor
+  add_one_Tensor_b_t b_tensor;
+  float b_host_ptr[32 * 16];
+  for (int i = 0; i < 32 * 16; i++) {
+    b_host_ptr[i] = 32 * 16 - i;
+  }
+  void *b_device_ptr;
+  cudaMalloc(&b_device_ptr, sizeof(float) * 32 * 16);
+  cudaMemcpy(b_device_ptr, b_host_ptr, sizeof(float) * 32 * 16,
+             cudaMemcpyHostToDevice);
+
+  b_tensor.data = (void *)b_device_ptr;
+  b_tensor.dynamic_shapes[0] = 32;
+  b_tensor.dynamic_shapes[1] = 16;
+  b_tensor.dynamic_strides[0] = 16;
+
+  // prepare stream
+  cudaStream_t stream;
+  cudaStreamCreate(&stream);
+
+  // load module
+  add_one_Kernel_Module_t module;
+  add_one_Kernel_Module_Load(&module);
+
+  // run kernel
+  cute_dsl_add_one_wrapper(&module, &a_tensor, &b_tensor, stream);
+  cudaStreamSynchronize(stream);
+
+  // unload module
+  add_one_Kernel_Module_Unload(&module);
+  cudaStreamDestroy(stream);
+
+  // check result
+  cudaMemcpy(a_host_ptr, a_device_ptr, sizeof(float) * 32 * 16,
+             cudaMemcpyDeviceToHost);
+
+  if (a_host_ptr[0] != 32 * 16 + 1) {
+    printf("Error: a_host_ptr[0] = %f, expected %d\n", a_host_ptr[0], 32 * 16);
+  }
+}
+
+int main() {
+  initialize_cuda_context();
+  run_print_tensor();
+  run_add_one();
+  return 0;
+}

examples/python/CuTeDSL/dsl_tutorials/export/run_with_static_linking.sh

@@ -0,0 +1,76 @@
+# Copyright (c) 2025 - 2026 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.
+
+#!/bin/bash
+set -eu
+
+# Try to find the wheel path of nvidia-cutlass-dsl
+WHEEL_PATH=$(python3 -c "import cutlass, os; print(os.path.dirname(cutlass.__file__))" 2>/dev/null)/../..
+
+if [[ -z "$WHEEL_PATH" ]]; then
+    echo "nvidia-cutlass-dsl wheel not found in the current Python environment."
+    exit 1
+else
+    echo "nvidia-cutlass-dsl wheel path found at: $WHEEL_PATH"
+fi
+CUTE_DSL_LIB_PATH="${WHEEL_PATH}/lib/"
+export LD_LIBRARY_PATH=${CUTE_DSL_LIB_PATH}
+
+if [ -z "$CUDA_HOME" ]; then
+    CUDA_HOME=/usr/local/cuda
+    echo "CUDA_HOME not set, using default: $CUDA_HOME"
+else
+    echo "CUDA_HOME found: $CUDA_HOME"
+fi
+SOURCE_FILE="$(dirname "$0")/run_with_static_linking.cpp"
+
+echo "Compiling the executable..."
+# Search for a common C++ compiler: g++, clang++, or c++
+if [ -n "${CXX-}" ] && command -v "$CXX" &> /dev/null; then
+    CXX="$CXX"
+elif command -v g++ &> /dev/null; then
+    CXX="g++"
+elif command -v clang++ &> /dev/null; then
+    CXX="clang++"
+elif command -v c++ &> /dev/null; then
+    CXX="c++"
+else
+    echo "Error: No common C++ compiler found (g++, clang++, or c++). Please install a C++ compiler to continue."
+    exit 1
+fi
+
+# Note: The options -ldl, -lrt, and -lpthread are optional to satisfy symbol dependencies required by `libcudart_static.a`.
+$CXX -o build/run_with_static_linking \
+  -I${CUDA_HOME}/include \
+  -I./build \
+  ${SOURCE_FILE} build/add_one_example.o build/print_tensor_example.o \
+  ${CUTE_DSL_LIB_PATH}/libcuda_dialect_runtime_static.a \
+  ${CUDA_HOME}/lib64/libcudart_static.a -ldl -lrt -lpthread
+
+echo "Running the executable..."
+./build/run_with_static_linking

examples/python/CuTeDSL/dsl_tutorials/ffi/CMakeLists.txt

@@ -0,0 +1,64 @@
+# Copyright (c) 2025 - 2026 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.
+
+cmake_minimum_required(VERSION 3.15)
+project(tensor)
+
+# Find Python
+find_package(Python COMPONENTS Interpreter Development REQUIRED)
+find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
+
+# Get Python site-packages directory using Python
+execute_process(
+    COMMAND ${Python3_EXECUTABLE} -c "import sys, sysconfig; print(';'.join([sysconfig.get_paths()['purelib'],sysconfig.get_paths(vars={'base': sys.base_prefix, 'platbase': sys.base_prefix})['purelib']]))"
+    OUTPUT_VARIABLE Python_SITE_PACKAGES_PATHS
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+
+message(STATUS "Python site-packages directories: ${Python_SITE_PACKAGES_PATHS}")
+
+# Add nanobind path to CMAKE_PREFIX_PATH
+foreach(path IN LISTS Python_SITE_PACKAGES_PATHS)
+    if(EXISTS "${path}/nanobind/cmake")
+        message(STATUS "Adding nanobind cmake path: ${path}/nanobind/cmake")
+        list(APPEND CMAKE_PREFIX_PATH "${path}/nanobind/cmake")
+        break()
+    endif()
+endforeach()
+
+# Find nanobind
+find_package(nanobind)
+if(NOT nanobind_FOUND)
+    message(FATAL_ERROR
+        "nanobind not found!\n"
+        "Please install nanobind with: pip install nanobind\n"
+    )
+endif()
+
+# Add the module
+nanobind_add_module(tensor tensor.cpp)

examples/python/CuTeDSL/dsl_tutorials/ffi/jit_argument.py

@@ -0,0 +1,320 @@
+# Copyright (c) 2025 - 2026 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.
+
+"""Example of accessing POD (Plain Old Data) from C or other languages via LLVM operations.
+
+This example demonstrates a basic approach to building customized interfaces as C-structures between user code
+and JIT compiled functions. It provides a minimal-cost solution for calling JIT functions
+and can be used to build AOT (Ahead-of-Time) launchers for JIT compiled functions.
+
+The C-structure is defined as:
+
+.. code-block:: c
+
+    struct Tensor {
+        void *ptr;          // Pointer to tensor data
+        int32_t shape[3];   // Tensor dimensions
+        int32_t strides[3]; // Memory strides for each dimension
+    };
+
+The example defines Tensor and TensorValue classes that wrap C structs for view of a tensor with its data pointer,
+shape, and strides, enabling efficient data passing between different language boundaries.
+
+.. note::
+   Future development may include automated code generation flows.
+"""
+
+import cutlass
+import cutlass.cute as cute
+
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm
+
+
+class ExampleTensorValue(ir.Value):
+    """A wrapper class for tensor values in MLIR.
+
+    This class extends ir.Value to provide convenient access to tensor data pointer,
+    shape, and strides through MLIR operations.
+
+    :type: ir.Value
+    """
+
+    def __init__(self, v):
+        """Initialize a new TensorValue.
+
+        :param v: The underlying MLIR value to wrap
+        :type v: ir.Value
+        """
+        super().__init__(v)
+
+    @property
+    def data_ptr(self, *, loc=None, ip=None):
+        """Get the data pointer from the tensor value.
+
+        Extracts the data pointer (first field) from the LLVM struct value.
+
+        :param loc: Optional location information for MLIR operations
+        :type loc: Optional[ir.Location]
+        :param ip: Optional insertion point for MLIR operations
+        :type ip: Optional[ir.InsertionPoint]
+        :return: An integer value representing the data pointer
+        :rtype: ir.Value
+        """
+        # Extract the data pointer from the LLVM struct value
+        # The data pointer is the first field (index 0) in the struct
+
+        # Use llvm.extractvalue to get the pointer field from the struct
+        ptr_val = llvm.extractvalue(
+            llvm.PointerType.get(),
+            self,
+            [0],  # Extract the first field (index 0)
+            loc=loc,
+            ip=ip,
+        )
+
+        return cute.make_ptr(cutlass.Float32, ptr_val)
+
+    @property
+    def shape(self):
+        """Get the shape of the tensor.
+
+        Extracts the shape (second field) from the LLVM struct value.
+
+        :return: A tuple of integers representing the tensor dimensions
+        :rtype: tuple[ir.Value, ...]
+        """
+        i32_type = ir.IntegerType.get_signless(32)
+        # Extract the shape field from the LLVM struct value
+        # The shape is the second field (index 1) in the struct
+        shape_val = llvm.extractvalue(
+            llvm.StructType.get_literal([i32_type] * 3),
+            self,
+            [1],  # Extract the second field (index 1)
+        )
+
+        # Extract each dimension from the shape struct
+        return tuple(llvm.extractvalue(i32_type, shape_val, [i]) for i in range(3))
+
+    @property
+    def stride(self):
+        """Get the strides of the tensor.
+
+        Extracts the strides (third field) from the LLVM struct value.
+
+        :return: A tuple of integers representing the tensor strides
+        :rtype: tuple[ir.Value, ...]
+        """
+        i32_type = ir.IntegerType.get_signless(32)
+        # Extract the strides field from the LLVM struct value
+        # The strides are the third field (index 2) in the struct
+        strides_val = llvm.extractvalue(
+            llvm.StructType.get_literal([i32_type] * 3),
+            self,
+            [2],  # Extract the third field (index 2)
+        )
+
+        # Extract each dimension from the strides struct
+        return tuple(llvm.extractvalue(i32_type, strides_val, [i]) for i in range(3))
+
+
+class ExampleTensor:
+    """A class representing a tensor with its data pointer, shape, and strides.
+
+    This class provides a Python interface to create and manipulate tensor structures
+    that can be passed to CUTE JIT compiled functions.
+
+    :ivar _c_struct_p: The C struct pointer for the tensor
+    :ivar _rank: The number of dimensions in the tensor
+    """
+
+    def __init__(self, c_struct_p, rank):
+        """Initialize a new Tensor.
+
+        :param c_struct_p: The C struct pointer for the tensor
+        :type c_struct_p: int
+        :param rank: The number of dimensions in the tensor
+        :type rank: int
+        """
+        self._c_struct_p = c_struct_p
+        self._rank = rank
+
+    def __get_mlir_types__(self):
+        """Get the MLIR types for this tensor.
+
+        Creates an LLVM structure type representing a C-structure with:
+
+        .. code-block:: c
+
+            struct Tensor {
+                void *ptr;
+                int32_t shape[3];
+                int32_t strides[3];
+            };
+
+        :return: A list containing the MLIR struct type
+        :rtype: list[llvm.StructType]
+
+        Create an LLVM structure type that represents a C-structure like:
+        """
+
+        # Get the number of dimensions from the shape
+        ndim = self._rank
+
+        # Create the pointer type (void*)
+        ptr_type = llvm.PointerType.get()
+
+        # Create array types for shape and strides (int32_t[ndim])
+        int32_type = ir.IntegerType.get_signless(32)
+        shape_type = llvm.StructType.get_literal([int32_type] * ndim)
+        strides_type = llvm.StructType.get_literal([int32_type] * ndim)
+
+        # Create the structure type
+        struct_type = llvm.StructType.get_literal([ptr_type, shape_type, strides_type])
+
+        return [struct_type]
+
+    def __new_from_mlir_values__(self, values):
+        """Create a new TensorValue from MLIR values.
+
+        :param values: A list of MLIR values
+        :type values: list[ir.Value]
+        :return: A new TensorValue instance
+        :rtype: TensorValue
+        """
+        return ExampleTensorValue(values[0])
+
+    def __c_pointers__(self):
+        """Get the C pointers for this tensor.
+
+        :return: A list containing the C struct pointer
+        :rtype: list[int]
+        """
+        return [self._c_struct_p]
+
+
+@cute.jit
+def foo(tensor):
+    """Example JIT function that prints tensor information.
+
+    :param tensor: A Tensor instance to print information about
+    :type tensor: Tensor
+    """
+    cute.printf("data_ptr: {}", tensor.data_ptr)
+    cute.printf("shape: {}", tensor.shape)
+    cute.printf("stride: {}", tensor.stride)
+
+    mA = cute.make_tensor(
+        tensor.data_ptr, cute.make_layout(tensor.shape, stride=tensor.stride)
+    )
+    cute.print_tensor(mA)
+
+
+import sys
+import os
+import subprocess
+import shutil
+import tempfile
+
+
+def run_test(tmpdir=None, cmake_args="", cleanup=True):
+    import torch
+
+    try:
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+
+        cmake_args = cmake_args.split()
+        subprocess.run(["cmake", "-B", tmpdir, current_dir] + cmake_args, check=True)
+        subprocess.run(["cmake", "--build", tmpdir], check=True)
+
+        from tensor import make_tensor, pycapsule_get_pointer
+
+        # Mock test tensor and corresponding C structure for this example
+        # In production, this may come from external library
+        x = torch.arange(2 * 8 * 4).to(torch.float32).reshape(2, 8, 4)
+        c_struct = make_tensor(x.data_ptr(), x.shape, x.stride())
+        c_struct_p = pycapsule_get_pointer(c_struct)
+
+        # Initialize tensor wrapper and compile test function
+        tensor = ExampleTensor(c_struct_p, len(x.shape))
+        compiled_func = cute.compile(foo, tensor)
+
+        # Benchmark pointer access performance
+        from time import time
+
+        start = time()
+        # Measure performance of critical path pointer access
+        # get C pointers is on critical path to call JIT compiled function
+        for _ in range(1000):
+            tensor.__c_pointers__()
+        end = time()
+        print(f"__c_pointers__: {(end - start) * 1000} us")
+
+        # Execute compiled function
+        compiled_func(tensor)
+    except Exception as e:
+        import traceback
+
+        traceback.print_exception(type(e), e, e.__traceback__)
+        raise e
+    finally:
+        if cleanup:
+            # Clean up the temporary directory
+            shutil.rmtree(tmpdir)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Set temporary directory for building C modules"
+    )
+    parser.add_argument(
+        "--tmp-dir",
+        type=str,
+        default=None,
+        help="Temporary directory path for building C modules",
+    )
+    parser.add_argument(
+        "--cmake-args",
+        type=str,
+        default="",
+        help="Extra CMake arguments for building C modules",
+    )
+    args = parser.parse_args()
+
+    if args.tmp_dir:
+        tmp_dir = args.tmp_dir
+        cleanup = False
+    else:
+        tmp_dir = tempfile.mkdtemp()
+        cleanup = True
+
+    sys.path.append(tmp_dir)
+
+    run_test(tmp_dir, args.cmake_args, cleanup)

examples/python/CuTeDSL/dsl_tutorials/ffi/tensor.cpp

@@ -0,0 +1,82 @@
+// Copyright (c) 2025 - 2026 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.
+
+#include <cstdint>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/string.h>
+#include <nanobind/stl/vector.h>
+
+namespace nb = nanobind;
+
+// Forward declaration of the MockTensor struct for testing only
+struct MockTensor {
+  void *ptr;
+  struct {
+    int32_t shape[3];
+  } shape;
+
+  struct {
+    int32_t strides[3];
+  } strides;
+};
+
+NB_MODULE(tensor, m) {
+  // create a tensor for testing
+  m.def("make_tensor", [](int64_t ptr, std::vector<int32_t> shape,
+                          std::vector<int32_t> strides) {
+    auto *tensor = new MockTensor();
+    tensor->ptr = reinterpret_cast<void *>(ptr);
+
+    assert(shape.size() == 3 && "shape must have 3 elements");
+    assert(strides.size() == 3 && "strides must have 3 elements");
+
+    for (size_t i = 0; i < shape.size(); i++) {
+      tensor->shape.shape[i] = shape[i];
+      tensor->strides.strides[i] = strides[i];
+    }
+
+    return nb::steal(PyCapsule_New(tensor, "tensor", [](PyObject *capsule) {
+      auto n = PyCapsule_GetName(capsule);
+      if (void *p = PyCapsule_GetPointer(capsule, n)) {
+        delete reinterpret_cast<MockTensor *>(p);
+      }
+    }));
+  });
+
+  m.def(
+      "pycapsule_get_pointer",
+      [](nb::object &capsule) {
+        void *ptr = PyCapsule_GetPointer(capsule.ptr(), "tensor");
+        if (!ptr) {
+          throw std::runtime_error("Invalid tensor capsule");
+        }
+        return reinterpret_cast<uintptr_t>(ptr);
+      },
+      "Get pointer from PyCapsule");
+}

examples/python/CuTeDSL/dsl_tutorials/inline_ptx.py

@@ -0,0 +1,243 @@
+# Copyright (c) 2025 - 2026 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.
+
+from functools import partial
+from typing import Union
+
+import cutlass.cute as cute
+from cutlass import Constexpr
+from cutlass.cute.runtime import from_dlpack
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.typing import Boolean, Int32, Int
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass.cute.arch.nvvm_wrappers import FULL_MASK, WARP_SIZE
+
+"""
+A simple example to show how to wrap PTX instructions by using inline_asm op in llvm dialect.
+
+Situations like:
+
+1. Instructions that are not already exposed by CuTe DSL via `nvvm` module
+2. Sequences of instructions that the compiler otherwise does not generate optimally
+
+motivate developers to inline PTX themselves.
+
+In this example, we inline the vote.sync.ballot.b32, vote.sync.any.pred, vote.sync.all.pred,
+vote.sync.uni.pred, and use the corresponding ops in nvvm dialect for the test.
+
+You can refer to the documentation of `inline_asm op in llvm dialect <https://mlir.llvm.org/docs/Dialects/LLVM/#llvminline_asm-llvminlineasmop>`_
+and `vote.sync <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync>`_
+for more details.
+
+To run this example:
+
+.. code-block:: bash
+
+    python examples/dsl/inline_ptx.py
+
+The example will run the vote kernel with inline ptx and nvvm dialect separately.
+The results from inline ptx and nvvm dialect will be verified correspondingly.
+
+"""
+
+
+@dsl_user_op
+def ptx_vote_sync_op(
+    pred: Boolean, kind: str, mask: Int = FULL_MASK, *, loc=None, ip=None
+) -> Union[Int32, Boolean]:
+    return_type = Boolean
+    return_type_str = "pred"
+    return return_type(
+        llvm.inline_asm(
+            T.bool(),
+            [
+                Boolean(pred).ir_value(loc=loc, ip=ip),
+                Int32(mask).ir_value(loc=loc, ip=ip),
+            ],
+            f"""{{\n\t
+            .reg .pred ps;\n\t
+            .reg .pred pd;\n\t
+            setp.ne.b32 ps, $1, 0;\n\t
+            vote.sync.{kind}.{return_type_str} pd, ps, $2;\n\t
+            selp.b32 $0, 1, 0, pd;\n\t
+            }}""",
+            "=r,r,i",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+ptx_vote_any_sync = partial(ptx_vote_sync_op, kind="any")
+ptx_vote_all_sync = partial(ptx_vote_sync_op, kind="all")
+ptx_vote_uni_sync = partial(ptx_vote_sync_op, kind="uni")
+
+
+@dsl_user_op
+def ptx_vote_ballot_sync(
+    pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None
+) -> Union[Int32, Boolean]:
+    return_type = Int32
+    return_type_str = "b32"
+    return return_type(
+        llvm.inline_asm(
+            T.i32(),
+            [
+                Boolean(pred).ir_value(loc=loc, ip=ip),
+                Int32(mask).ir_value(loc=loc, ip=ip),
+            ],
+            f"""{{\n\t
+            .reg .pred p;\n\t
+            setp.ne.b32 p, $1, 0;\n\t
+            vote.sync.ballot.{return_type_str} $0, p, $2;\n\t
+            }}""",
+            "=r,r,i",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@cute.kernel
+def vote_kernel(
+    mBallot: cute.Tensor,
+    mAny: cute.Tensor,
+    mAll: cute.Tensor,
+    mUni: cute.Tensor,
+    use_inline_ptx: Constexpr[bool],
+):
+    tidx, _, _ = cute.arch.thread_idx()
+
+    vote_ballot = (
+        ptx_vote_ballot_sync(tidx < 10)
+        if use_inline_ptx
+        else cute.arch.vote_ballot_sync(tidx < 10)
+    )
+    vote_any = (
+        ptx_vote_any_sync(tidx < 10)
+        if use_inline_ptx
+        else cute.arch.vote_any_sync(tidx < 10)
+    )
+    vote_all = (
+        ptx_vote_all_sync(tidx < 10)
+        if use_inline_ptx
+        else cute.arch.vote_all_sync(tidx < 10)
+    )
+    vote_uni = (
+        ptx_vote_uni_sync(tidx < 10)
+        if use_inline_ptx
+        else cute.arch.vote_uni_sync(tidx < 10)
+    )
+
+    mBallot[tidx] = vote_ballot
+    mAny[tidx] = vote_any
+    mAll[tidx] = vote_all
+    mUni[tidx] = vote_uni
+
+
+@cute.jit
+def vote(
+    mBallot: cute.Tensor,
+    mAny: cute.Tensor,
+    mAll: cute.Tensor,
+    mUni: cute.Tensor,
+    use_inline_ptx: Constexpr[bool],
+):
+    vote_kernel(
+        mBallot,
+        mAny,
+        mAll,
+        mUni,
+        use_inline_ptx,
+    ).launch(
+        grid=[1, 1, 1],
+        block=[cute.size(WARP_SIZE, mode=[0]), 1, 1],
+    )
+
+
+def run():
+    import torch
+
+    ballot_ptx = torch.randint(
+        0, 100, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.int32
+    )
+    any_ptx = torch.randint(
+        0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
+    )
+    all_ptx = torch.randint(
+        0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
+    )
+    uni_ptx = torch.randint(
+        0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
+    )
+
+    mBallotPTX = from_dlpack(ballot_ptx).mark_layout_dynamic()
+    mAnyPTX = from_dlpack(any_ptx).mark_layout_dynamic()
+    mAllPTX = from_dlpack(all_ptx).mark_layout_dynamic()
+    mUniPTX = from_dlpack(uni_ptx).mark_layout_dynamic()
+
+    # get the results from ptx
+    vote(mBallotPTX, mAnyPTX, mAllPTX, mUniPTX, use_inline_ptx=True)
+
+    ballot_nvvm = torch.randint(
+        0, 100, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.int32
+    )
+    any_nvvm = torch.randint(
+        0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
+    )
+    all_nvvm = torch.randint(
+        0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
+    )
+    uni_nvvm = torch.randint(
+        0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
+    )
+
+    mBallotNVVM = from_dlpack(ballot_nvvm).mark_layout_dynamic()
+    mAnyNVVM = from_dlpack(any_nvvm).mark_layout_dynamic()
+    mAllNVVM = from_dlpack(all_nvvm).mark_layout_dynamic()
+    mUniNVVM = from_dlpack(uni_nvvm).mark_layout_dynamic()
+
+    # get the results from nvvm
+    vote(mBallotNVVM, mAnyNVVM, mAllNVVM, mUniNVVM, use_inline_ptx=False)
+
+    print("Verifying ballot results...")
+    torch.testing.assert_close(ballot_ptx, ballot_nvvm)
+    print("Verifying any results...")
+    torch.testing.assert_close(any_ptx, any_nvvm)
+    print("Verifying all results...")
+    torch.testing.assert_close(all_ptx, all_nvvm)
+    print("Verifying uni results...")
+    torch.testing.assert_close(uni_ptx, uni_nvvm)
+
+    print("Results verified successfully!")
+
+
+if __name__ == "__main__":
+    run()

examples/python/CuTeDSL/dsl_tutorials/jax/cute_dsl_jax_kernels.py

@@ -0,0 +1,366 @@
+# Copyright (c) 2025 - 2026 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.
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.jax as cjax
+import cuda.bindings.driver as cuda
+
+"""
+CuTe DSL kernels used by the ``cute_dsl_jax.ipynb`` notebook.
+
+This module defines GPU kernels written in CuTe DSL (CUTLASS 4.x Python DSL)
+that are called from JAX via ``cutlass.jax.cutlass_call``. ``cutlass_call`` is a
+JAX primitive that triggers compilation of the kernel during lowering and embeds
+it into the HLO computation, so XLA can launch it efficiently without callback
+to Python.
+
+Kernels provided:
+
+- ``vector_add``        — element-wise c = a + b (3-D CuTe layout)
+- ``saxpy``             — y = alpha * x + y
+- ``relu``              — element-wise ReLU with flat indexing
+- ``fused_bias_relu``   — fused bias addition + ReLU
+- ``gemm``              — tiled matrix multiplication
+- ``elementwise_add``   — 2-D element-wise add (flat indexing, ``jax.export``-compatible)
+
+The notebook imports these kernels and wraps each one with ``cutlass_call``
+inside ``@jax.jit`` functions. See ``cute_dsl_jax.ipynb`` for usage, validation,
+and step-by-step explanations.
+
+This module is imported by the notebook and by ``cute_dsl_jax.py``. It can also
+be run directly to validate every kernel:
+
+.. code-block:: bash
+
+    # Interactive notebook (recommended for learning)
+    jupyter lab cute_dsl_jax.ipynb
+
+    # Full demo as a standalone script
+    python cute_dsl_jax_kernels.py
+"""
+
+
+# ------------------------------------------------------------------ #
+#  Vector Add: c = a + b                                             #
+# ------------------------------------------------------------------ #
+@cute.kernel
+def vector_add_kernel(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
+    """Per-thread kernel: each thread adds one element."""
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    frgA = cute.make_rmem_tensor(cute.size(a, mode=[0]), a.element_type)
+    frgB = cute.make_rmem_tensor(cute.size(b, mode=[0]), b.element_type)
+    frgC = cute.make_rmem_tensor(cute.size(c, mode=[0]), c.element_type)
+
+    cute.autovec_copy(a[None, tidx, bidx], frgA)
+    cute.autovec_copy(b[None, tidx, bidx], frgB)
+    frgC.store(frgA.load() + frgB.load())
+    cute.autovec_copy(frgC, c[None, tidx, bidx])
+
+
+@cute.jit
+def launch_vector_add(
+    stream: cuda.CUstream,
+    a: cute.Tensor, b: cute.Tensor, c: cute.Tensor,
+):
+    vector_add_kernel(a, b, c).launch(
+        grid=[a.shape[-1], 1, 1],
+        block=[a.shape[-2], 1, 1],
+        stream=stream,
+    )
+
+
+# ------------------------------------------------------------------ #
+#  SAXPY: y = alpha * x + y                                          #
+# ------------------------------------------------------------------ #
+@cute.kernel
+def saxpy_kernel(x: cute.Tensor, y: cute.Tensor, out: cute.Tensor, alpha: float):
+    """SAXPY: out[i] = alpha * x[i] + y[i]."""
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    frgX = cute.make_rmem_tensor(cute.size(x, mode=[0]), x.element_type)
+    frgY = cute.make_rmem_tensor(cute.size(y, mode=[0]), y.element_type)
+    frgO = cute.make_rmem_tensor(cute.size(out, mode=[0]), out.element_type)
+
+    cute.autovec_copy(x[None, tidx, bidx], frgX)
+    cute.autovec_copy(y[None, tidx, bidx], frgY)
+    frgO.store(alpha * frgX.load() + frgY.load())
+    cute.autovec_copy(frgO, out[None, tidx, bidx])
+
+
+@cute.jit
+def launch_saxpy(
+    stream: cuda.CUstream,
+    x: cute.Tensor, y: cute.Tensor, out: cute.Tensor,
+    *, alpha: float,
+):
+    saxpy_kernel(x, y, out, alpha).launch(
+        grid=[x.shape[-1], 1, 1],
+        block=[x.shape[-2], 1, 1],
+        stream=stream,
+    )
+
+
+# ------------------------------------------------------------------ #
+#  ReLU: out = max(0, x)                                             #
+# ------------------------------------------------------------------ #
+@cute.kernel
+def relu_kernel(x: cute.Tensor, out: cute.Tensor, N: int):
+    """Per-thread kernel: each thread computes ReLU of one element."""
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+    bdx, _, _ = cute.arch.block_dim()
+
+    idx = bidx * bdx + tidx
+    if idx < N:
+        val = x[idx]
+        out[idx] = cutlass.max(val, cutlass.Float32(0.0))
+
+
+@cute.jit
+def launch_relu(
+    stream: cuda.CUstream,
+    x: cute.Tensor, out: cute.Tensor,
+    *, N: int,
+):
+    BLOCK_SIZE = 256
+    grid_size = (N + BLOCK_SIZE - 1) // BLOCK_SIZE
+    relu_kernel(x, out, N).launch(
+        grid=[grid_size, 1, 1],
+        block=[BLOCK_SIZE, 1, 1],
+        stream=stream,
+    )
+
+
+# ------------------------------------------------------------------ #
+#  Fused Bias + ReLU: out = max(0, x + bias[col])                    #
+# ------------------------------------------------------------------ #
+@cute.kernel
+def fused_bias_relu_kernel(
+    x: cute.Tensor, bias: cute.Tensor, out: cute.Tensor, N: int, width: int,
+):
+    """Per-thread: out[i] = max(0, x[i] + bias[i % width])."""
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+    bdx, _, _ = cute.arch.block_dim()
+
+    idx = bidx * bdx + tidx
+    if idx < N:
+        col = idx % width
+        val = x[idx] + bias[col]
+        out[idx] = cutlass.max(val, cutlass.Float32(0.0))
+
+
+@cute.jit
+def launch_fused_bias_relu(
+    stream: cuda.CUstream,
+    x: cute.Tensor, bias: cute.Tensor, out: cute.Tensor,
+    *, N: int, width: int,
+):
+    BLOCK_SIZE = 256
+    grid_size = (N + BLOCK_SIZE - 1) // BLOCK_SIZE
+    fused_bias_relu_kernel(x, bias, out, N, width).launch(
+        grid=[grid_size, 1, 1],
+        block=[BLOCK_SIZE, 1, 1],
+        stream=stream,
+    )
+
+
+# ------------------------------------------------------------------ #
+#  GEMM: D = A @ B                                                   #
+# ------------------------------------------------------------------ #
+@cute.kernel
+def gemm_kernel(
+    A: cute.Tensor, B: cute.Tensor, D: cute.Tensor,
+    M: int, N: int, K: int, BLOCK_M: int, BLOCK_N: int,
+):
+    """Tiled GEMM: each thread accumulates output elements."""
+    tidx, _, _ = cute.arch.thread_idx()
+    bm, bn, _ = cute.arch.block_idx()
+    bdx, _, _ = cute.arch.block_dim()
+
+    for i in cutlass.range(tidx, BLOCK_M * BLOCK_N, bdx):
+        row = i // BLOCK_N
+        col = i % BLOCK_N
+        m_idx = bm * BLOCK_M + row
+        n_idx = bn * BLOCK_N + col
+        if m_idx < M and n_idx < N:
+            acc = cutlass.Float32(0.0)
+            for k in cutlass.range(K):
+                acc += A[m_idx * K + k] * B[k * N + n_idx]
+            D[m_idx * N + n_idx] = acc
+
+
+@cute.jit
+def launch_gemm(
+    stream: cuda.CUstream,
+    A: cute.Tensor, B: cute.Tensor, D: cute.Tensor,
+    *, M: int, N: int, K: int,
+):
+    BLOCK_M, BLOCK_N = 64, 64
+    grid_m = (M + BLOCK_M - 1) // BLOCK_M
+    grid_n = (N + BLOCK_N - 1) // BLOCK_N
+    gemm_kernel(A, B, D, M, N, K, BLOCK_M, BLOCK_N).launch(
+        grid=[grid_m, grid_n, 1],
+        block=[256, 1, 1],
+        stream=stream,
+    )
+
+    
+# ------------------------------------------------------------------ #
+#  Element-wise Add (2-D, flat indexing)                             #
+# ------------------------------------------------------------------ #
+@cute.kernel
+def elementwise_add_kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
+    """Per-thread kernel: 2-D element-wise add using flat indexing."""
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+    bdim, _, _ = cute.arch.block_dim()
+
+    thread_idx = bidx * bdim + tidx
+
+    m, n = gA.shape
+    ni = thread_idx % n
+    mi = thread_idx // n
+
+    a_val = gA[mi, ni]
+    b_val = gB[mi, ni]
+    gC[mi, ni] = a_val + b_val
+
+
+@cute.jit
+def launch_elementwise_add(
+    stream: cuda.CUstream,
+    mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor,
+):
+    num_threads_per_block = 256
+    m, n = mA.shape
+    elementwise_add_kernel(mA, mB, mC).launch(
+        grid=((m * n) // num_threads_per_block, 1, 1),
+        block=(num_threads_per_block, 1, 1),
+        stream=stream,
+    )
+
+
+# ------------------------------------------------------------------ #
+#  Self-tests                                                         #
+# ------------------------------------------------------------------ #
+if __name__ == '__main__':
+    import os
+    os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
+
+    import jax
+    import jax.numpy as jnp
+    import numpy as np
+
+    BLOCK = 256
+    N_BLOCKS = 4
+
+    # ── Vector Add ────────────────────────────────────────────────────
+    # 3-D CuTe layout: (elems_per_thread, threads_per_block, num_blocks)
+    a = jax.random.normal(jax.random.PRNGKey(0), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
+    b = jax.random.normal(jax.random.PRNGKey(1), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
+    call = cjax.cutlass_call(
+        launch_vector_add,
+        output_shape_dtype=jax.ShapeDtypeStruct(a.shape, a.dtype),
+        use_static_tensors=True,
+    )
+    c = jax.jit(call)(a, b)
+    np.testing.assert_allclose(np.array(c), np.array(a + b), rtol=1e-5, atol=1e-5)
+    print('vector_add:       PASSED')
+
+    # ── SAXPY ─────────────────────────────────────────────────────────
+    x = jax.random.normal(jax.random.PRNGKey(2), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
+    y = jax.random.normal(jax.random.PRNGKey(3), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
+    alpha = 2.5
+    call = cjax.cutlass_call(
+        launch_saxpy,
+        output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype),
+        use_static_tensors=True,
+        alpha=alpha,
+    )
+    out = jax.jit(call)(x, y)
+    np.testing.assert_allclose(np.array(out), np.array(alpha * x + y), rtol=1e-5, atol=1e-5)
+    print('saxpy:            PASSED')
+
+    # ── ReLU ──────────────────────────────────────────────────────────
+    N_ELEM = BLOCK * N_BLOCKS
+    x = jax.random.normal(jax.random.PRNGKey(4), (N_ELEM,), dtype=jnp.float32)
+    call = cjax.cutlass_call(
+        launch_relu,
+        output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype),
+        N=N_ELEM,
+    )
+    out = jax.jit(call)(x)
+    np.testing.assert_allclose(np.array(out), np.array(jnp.maximum(x, 0)), rtol=1e-5, atol=1e-5)
+    print('relu:             PASSED')
+
+    # ── Fused Bias + ReLU ─────────────────────────────────────────────
+    ROWS, COLS = 16, 64
+    x = jax.random.normal(jax.random.PRNGKey(5), (ROWS * COLS,), dtype=jnp.float32)
+    bias = jax.random.normal(jax.random.PRNGKey(6), (COLS,), dtype=jnp.float32)
+    call = cjax.cutlass_call(
+        launch_fused_bias_relu,
+        output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype),
+        N=ROWS * COLS, width=COLS,
+    )
+    out = jax.jit(call)(x, bias)
+    ref = jnp.maximum(x.reshape(ROWS, COLS) + bias, 0).reshape(-1)
+    np.testing.assert_allclose(np.array(out), np.array(ref), rtol=1e-5, atol=1e-5)
+    print('fused_bias_relu:  PASSED')
+
+    # ── GEMM ──────────────────────────────────────────────────────────
+    M, N, K = 128, 128, 64
+    A = jax.random.normal(jax.random.PRNGKey(7), (M * K,), dtype=jnp.float32)
+    B = jax.random.normal(jax.random.PRNGKey(8), (K * N,), dtype=jnp.float32)
+    call = cjax.cutlass_call(
+        launch_gemm,
+        output_shape_dtype=jax.ShapeDtypeStruct((M * N,), A.dtype),
+        M=M, N=N, K=K,
+    )
+    D = jax.jit(call)(A, B)
+    ref = A.reshape(M, K) @ B.reshape(K, N)
+    np.testing.assert_allclose(np.array(D.reshape(M, N)), np.array(ref), rtol=1e-2, atol=1e-2)
+    print('gemm:             PASSED')
+
+    # ── Elementwise Add (2-D) ─────────────────────────────────────────
+    M, N = 16, 256
+    a = jax.random.normal(jax.random.PRNGKey(9), (M, N), dtype=jnp.float32)
+    b = jax.random.normal(jax.random.PRNGKey(10), (M, N), dtype=jnp.float32)
+    call = cjax.cutlass_call(
+        launch_elementwise_add,
+        output_shape_dtype=jax.ShapeDtypeStruct(a.shape, a.dtype),
+    )
+    c = jax.jit(call)(a, b)
+    np.testing.assert_allclose(np.array(c), np.array(a + b), rtol=1e-5, atol=1e-5)
+    print('elementwise_add:  PASSED')
+
+    print('\nAll kernels passed.')

examples/python/CuTeDSL/dsl_tutorials/jax/cutlass_call_basic.py

@@ -0,0 +1,254 @@
+# Copyright (c) 2025 - 2026 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.
+
+from functools import partial
+import jax
+import jax.numpy as jnp
+
+import cutlass.cute as cute
+import cutlass.jax as cjax
+import cuda.bindings.driver as cuda
+
+"""
+Examples of calling CuTe DSL from jax.jit function using cutlass_call.
+
+cutlass_call is a Jax primitive the enables calling of CuTe DSL kernels within a
+a jit-compiled Jax function. During the lowering process cutlass_call will
+trigger compilation of the kernel and embed it into the HLO computation. It can
+then be efficiently launched by XLA without callback to Python.
+
+This example assumes familiarity with CuTe DSL concepts such as layouts and
+dynamic shapes.
+
+To run this example:
+
+.. code-block:: bash
+
+    # Run with addition operation
+    python examples/jax/cutlass_call_basic.py
+"""
+
+
+# This is a typical CuTe DSL kernel function that accepts both tensor and scalar values.
+@cute.jit
+def launch(
+    A: cute.Tensor,
+    B: cute.Tensor,
+    x: cute.Int32,
+    y: cute.Int32,
+    C: cute.Tensor,
+    D: cute.Tensor,
+    stream: cuda.CUstream,
+):
+    # Print layouts
+    print("A layout: ", A.layout)
+    print("B layout: ", B.layout)
+    print("C layout: ", C.layout)
+    print("D layout: ", D.layout)
+    cute.printf("A layout: {}", A.layout)
+    cute.printf("B layout: {}", B.layout)
+    cute.printf("C layout: {}", C.layout)
+    cute.printf("D layout: {}", D.layout)
+    cute.printf("")
+
+    # Print non-tensor values
+    print("X is: ", x)
+    print("Y is: ", y)
+    cute.printf("X is: {}", x)
+    cute.printf("Y is: {}", y)
+    print()
+
+
+# cutlass_call uses a fixed function signature to pass arguments between Jax and CuTeDSL kernel.
+#
+# Function Signature Requirement:
+#   stream, inputs, outputs, *, kwargs...
+#
+# The first argument must be the CUstream that the kernel is run. This stream is managed by the XLA runtime
+# and is necessary to schedule and synchronize launches with the rest of your computation.
+#
+# The second set of arguments are the Jax arrays for inputs and outputs. Inputs must be passed before
+# outputs.
+#
+# Lastly static arguments (i.e. static_argnums or static_argnames) values are passed as keyword only arguments
+# by name.
+#
+# The the kernel does not match this signature a wrapper functions like the one shown below can be written
+# or an inline lambda function can be used to rebind the arguments into the appropriate order.
+@cute.jit
+def launch_jax_wrapper(
+    stream: cuda.CUstream,
+    A: cute.Tensor,
+    B: cute.Tensor,
+    C: cute.Tensor,
+    D: cute.Tensor,
+    *,
+    x: cute.Int32,
+    y: cute.Int32,
+):
+    launch(A, B, x, y, C, D, stream)
+
+
+@cute.jit
+def launch_aliased(
+    A: cute.Tensor, B: cute.Tensor, x: cute.Int32, y: cute.Int32, stream: cuda.CUstream
+):
+    # Print layouts
+    print("A layout: ", A.layout)
+    print("B layout: ", B.layout)
+    cute.printf("A layout: {}", A.layout)
+    cute.printf("B layout: {}", B.layout)
+    cute.printf("")
+
+    # Print non-tensor values
+    print("X is: ", x)
+    print("Y is: ", y)
+    cute.printf("X is: {}", x)
+    cute.printf("Y is: {}", y)
+    print()
+
+
+if __name__ == "__main__":
+
+    @partial(jax.jit, static_argnums=[2, 3])
+    def run_cutlass_kernel(a, b, x, y):
+        call = cjax.cutlass_call(
+            launch_jax_wrapper,
+            # Describe the shape and dtype of each output buffer.
+            output_shape_dtype=(
+                jax.ShapeDtypeStruct(a.shape, a.dtype),
+                jax.ShapeDtypeStruct(b.shape, a.dtype),
+            ),
+            # Static jit arguments are passed via additional keyword arguments.
+            x=x,
+            y=y,
+        )
+
+        # Returned value is a callable to invoke the kernel passing only jax arrays.
+        return call(a, b)
+
+    print("\nExample: example_basic_call_from_jit")
+    A = jnp.zeros((512, 32, 64))
+    B = jnp.zeros((1, 256, 64, 128))
+    C, D = run_cutlass_kernel(A, B, 0, 1)
+
+    @partial(jax.jit, static_argnums=[2, 3])
+    def run_cutlass_kernel_lambda(a, b, x, y):
+        call = cjax.cutlass_call(
+            # A lambda function may be used to wrap and bind arguments passed by jax
+            # to the kernel. Alternatively you can wrap using another separate cute.jit
+            # function.
+            lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
+            output_shape_dtype=(
+                jax.ShapeDtypeStruct(a.shape, a.dtype),
+                jax.ShapeDtypeStruct(b.shape, a.dtype),
+            ),
+            # Static jit arguments are passed via additional keyword arguments.
+            x=x,
+            y=y,
+        )
+
+        # Returned value is a callable to invoke the kernel passing only jax arrays.
+        return call(a, b)
+
+    print("\nExample: run_cutlass_kernel_lambda")
+    A = jnp.zeros((512, 32, 64))
+    B = jnp.zeros((1, 256, 64, 128))
+    C, D = run_cutlass_kernel_lambda(A, B, 1, 2)
+
+    @partial(jax.jit, static_argnums=[2, 3])
+    def run_cutlass_kernel_static_shapes(a, b, x, y):
+        call = cjax.cutlass_call(
+            lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
+            output_shape_dtype=(
+                jax.ShapeDtypeStruct(a.shape, a.dtype),
+                jax.ShapeDtypeStruct(b.shape, a.dtype),
+            ),
+            # By default cutlass_call treats all tensors as dynamic shape.
+            # Dynamic shapes are often expected for kernels so this default ensures
+            # the broadest support. If you know that a kernel can accept fully static
+            # tensors then you can enable this flag to compile all tensor shapes and
+            # layouts as constexpr values known at compile time.
+            # Individual tensors may opt out via .mark_layout_dynamic().
+            use_static_tensors=True,
+            x=x,
+            y=y,
+        )
+        return call(a, b)
+
+    print("\nExample: run_cutlass_kernel_static_shapes")
+    A = jnp.zeros((512, 32, 64))
+    B = jnp.zeros((1, 256, 64, 128))
+    C, D = run_cutlass_kernel_static_shapes(A, B, 3, 4)
+
+    @partial(jax.jit, static_argnums=[2, 3])
+    def run_cutlass_kernel_with_modes(a, b, x, y):
+        # input_spec and output_spec accept TensorSpec values to attach layout
+        # metadata to tensors.  mode remaps the logical dimension order seen by
+        # the kernel.  static=True compiles that tensor's layout as constexpr.
+        call = cjax.cutlass_call(
+            lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
+            output_shape_dtype=(
+                jax.ShapeDtypeStruct(a.shape, a.dtype),
+                jax.ShapeDtypeStruct(b.shape, a.dtype),
+            ),
+            input_spec=(
+                cjax.TensorSpec(mode=(1, 0, 2), static=True),
+                cjax.TensorSpec(mode=(3, 1, 2, 0)),
+            ),
+            output_spec=(None, cjax.TensorSpec(mode=(0, 1, 3, 2))),
+            x=x,
+            y=y,
+        )
+        return call(a, b)
+
+    print("\nExample: run_cutlass_kernel_with_modes")
+    A = jnp.zeros((512, 32, 64))
+    B = jnp.zeros((1, 256, 64, 128))
+    C, D = run_cutlass_kernel_with_modes(A, B, 5, 6)
+
+    @partial(jax.jit, static_argnums=[2, 3], donate_argnums=[0, 1])
+    def run_cutlass_kernel_aliased_outputs(a, b, x, y):
+        call = cjax.cutlass_call(
+            lambda stream, a, b, *, x, y: launch_aliased(a, b, x, y, stream),
+            output_shape_dtype=(
+                jax.ShapeDtypeStruct(a.shape, a.dtype),
+                jax.ShapeDtypeStruct(b.shape, b.dtype),
+            ),
+            # Map input indices to output indices so XLA can reuse the input
+            # buffers for the outputs, avoiding extra allocations.
+            input_output_aliases={0: 0, 1: 1},
+            x=x,
+            y=y,
+        )
+        return call(a, b)
+
+    print("\nExample: run_cutlass_kernel_aliased_outputs")
+    A = jnp.zeros((512, 32, 64))
+    B = jnp.zeros((1, 256, 64, 128))
+    A, B = run_cutlass_kernel_aliased_outputs(A, B, 7, 8)

examples/python/CuTeDSL/dsl_tutorials/jax/cutlass_call_export.py

@@ -0,0 +1,201 @@
+# Copyright (c) 2025 - 2026 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.
+
+"""
+Examples of using jax.export APIs with functions using cutlass_call.
+
+This example demonstrates three export modes:
+
+1. Concrete shapes  -- shapes are fixed constants baked into the export.
+2. Unconstrained symbolic shapes ("a, b")
+3. Constrained symbolic shapes ("32*M, 16*N")
+
+The JAX function being exported is the same in all three cases; only the
+shape specification passed to jax.export differs.
+
+It assumes familiarity with CuTe DSL concepts such as layouts and dynamic shapes
+as well as JAX's exporting and serialization features:
+https://docs.jax.dev/en/latest/export/index.html#export
+
+To run this example:
+
+.. code-block:: bash
+
+    python examples/jax/cutlass_call_export.py --M 512 --N 256
+
+"""
+
+import argparse
+
+import cuda.bindings.driver as cuda
+
+import cutlass.cute as cute
+
+import jax
+import jax.numpy as jnp
+from jax import export
+
+from cutlass.jax import cutlass_call, get_export_disabled_safety_checks, TensorSpec
+from cutlass.jax.testing import create_tensor
+
+
+# Simple element-wise addition kernel: gC[i,j] = gA[i,j] + gB[i,j]
+@cute.kernel
+def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+    bdim, _, _ = cute.arch.block_dim()
+
+    thread_idx = bidx * bdim + tidx
+
+    m, n = gA.shape
+    ni = thread_idx % n
+    mi = thread_idx // n
+
+    a_val = gA[mi, ni]
+    b_val = gB[mi, ni]
+    gC[mi, ni] = a_val + b_val
+
+
+@cute.jit
+def launch(stream: cuda.CUstream, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
+    num_threads_per_block = 256
+    m, n = mA.shape
+    kernel(mA, mB, mC).launch(
+        grid=((m * n) // num_threads_per_block, 1, 1),
+        block=(num_threads_per_block, 1, 1),
+        stream=stream,
+    )
+
+
+def _export_and_run(f, ref_f, input_shape_dtype, run_shapes):
+    """Export f, serialize/deserialize, then run on each shape in run_shapes.
+
+    Both inputs (a, b) are assumed to share the same input_shape_dtype.
+    """
+    print(f"Exporting with input signature: ({input_shape_dtype}, {input_shape_dtype})")
+
+    # jax.export can be used to export a jit function containing cutlass_call.
+    # CUTLASS custom call targets are not on JAX's built-in stable custom-call
+    # allowlist, so we pass them via disabled_checks to suppress that safety check.
+    exported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())
+    traced = exported(input_shape_dtype, input_shape_dtype)
+
+    blob = traced.serialize()
+    print(f"Serialized computation is {len(blob)} bytes.")
+
+    rehydrated = export.deserialize(blob)
+
+    key = jax.random.key(1123)
+    a_key, b_key = jax.random.split(key, 2)
+    for shape in run_shapes:
+        a = create_tensor(shape, dtype=jnp.float32, key=a_key)
+        b = create_tensor(shape, dtype=jnp.float32, key=b_key)
+        c = rehydrated.call(a, b)
+        assert jnp.allclose(c, ref_f(a, b)), f"Mismatch at shape {shape}"
+        print(f"  shape {shape}: OK")
+
+
+def run_example(M, N):
+    @jax.jit
+    def ref_f(a, b):
+        return jax.nn.sigmoid(a + b)
+
+    # The same JAX function is used in all three examples below. The export
+    # mode is determined entirely by the shape spec passed to jax.export.
+    @jax.jit
+    def f(a, b):
+        call = cutlass_call(launch, output_shape_dtype=a)
+        return jax.nn.sigmoid(call(a, b))
+
+    # ── 1. Concrete shapes ────────────────────────────────────────────────────
+    # Shapes are fixed constants baked into the export. The deserialized
+    # computation only accepts exactly these dimensions at runtime.
+    print("\nConcrete shapes:")
+
+    input_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32)
+    _export_and_run(
+        f,
+        ref_f,
+        input_shape_dtype,
+        run_shapes=[(M, N)],  # concrete exports reject any other shape
+    )
+
+    # ── 2. Unconstrained symbolic shapes ─────────────────────────────────────
+    # Both dimensions are fully dynamic. The exported computation accepts any
+    # (M, N) at runtime without recompilation.
+    print("\nUnconstrained symbolic shapes:")
+
+    a_sym, b_sym = export.symbolic_shape("a, b")
+    input_shape_dtype = jax.ShapeDtypeStruct((a_sym, b_sym), jnp.float32)
+    _export_and_run(
+        f,
+        ref_f,
+        input_shape_dtype,
+        run_shapes=[(M, N), (M * 2, N * 4), (M * 4, N * 4)],
+    )
+
+    # ── 3. Constrained symbolic shapes (divisibility) ─────────────────────────
+    # Shapes are declared as multiples of a tile size via TensorSpec.divisibility.
+    # The symbolic expression "32*M, 16*N" tells jax.export that dim 0 is always
+    # a multiple of 32 and dim 1 is always a multiple of 16. This lets the
+    # compiler generate more efficient code (e.g. no remainder handling).
+    # Runtime shapes must satisfy these divisibility constraints.
+    print("\nConstrained symbolic shapes:")
+
+    @jax.jit
+    def f_divisible(a, b):
+        spec = TensorSpec(divisibility=(32, 16))
+        call = cutlass_call(
+            launch,
+            output_shape_dtype=a,
+            input_spec=(spec, spec),
+            output_spec=spec,
+        )
+        return jax.nn.sigmoid(call(a, b))
+
+    m_sym, n_sym = export.symbolic_shape("32*M, 16*N")
+    input_shape_dtype = jax.ShapeDtypeStruct((m_sym, n_sym), jnp.float32)
+    _export_and_run(
+        f_divisible,
+        ref_f,
+        input_shape_dtype,
+        run_shapes=[(M, N), (M * 2, N * 2), (M * 4, N * 4)],
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Demonstration of using jax.export with functions with cutlass_call"
+    )
+    parser.add_argument("--M", default=512, type=int)
+    parser.add_argument("--N", default=256, type=int)
+
+    args = parser.parse_args()
+    run_example(args.M, args.N)
+    print("PASS")

examples/python/CuTeDSL/dsl_tutorials/jax/cutlass_call_sharding.py

@@ -0,0 +1,175 @@
+# Copyright (c) 2025 - 2026 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.
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax.sharding import NamedSharding, PartitionSpec as P, AxisType
+from jax.experimental.custom_partitioning import custom_partitioning
+
+import cutlass.cute as cute
+import cutlass.jax as cjax
+from cutlass.jax.testing import create_tensor
+import cuda.bindings.driver as cuda
+
+
+"""
+Examples of combining jax.jit, jax.shard_map and custom_partitioning for sharding
+and executing kernels across multiple GPU devices.
+
+To run this example:
+
+.. code-block:: bash
+
+    # Run with addition operation
+    python examples/jax/cutlass_call_sharding.py
+"""
+
+
+@cute.kernel
+def kernel(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    frgA = cute.make_rmem_tensor(cute.size(a, mode=[0]), a.element_type)
+    frgB = cute.make_rmem_tensor(cute.size(b, mode=[0]), b.element_type)
+    frgC = cute.make_rmem_tensor(cute.size(c, mode=[0]), c.element_type)
+
+    cute.autovec_copy(a[None, tidx, bidx], frgA)
+    cute.autovec_copy(b[None, tidx, bidx], frgB)
+    frgC.store(frgA.load() + frgB.load())
+    cute.autovec_copy(frgC, c[None, tidx, bidx])
+
+
+@cute.jit
+def launch(
+    stream: cuda.CUstream,
+    a: cute.Tensor,
+    b: cute.Tensor,
+    c: cute.Tensor,
+):
+    cute.printf("a: {}", a.layout)
+    cute.printf("b: {}", b.layout)
+    cute.printf("c: {}", c.layout)
+    kernel(a, b, c).launch(
+        grid=[a.shape[-1], 1, 1], block=[a.shape[-2], 1, 1], stream=stream
+    )
+
+
+def sharded_cutlass_call_impl(a_block, b_block):
+    """The sharded implementation that operates on a single device."""
+    call = cjax.cutlass_call(
+        launch,
+        use_static_tensors=True,
+        output_shape_dtype=jax.ShapeDtypeStruct(a_block.shape, a_block.dtype),
+    )
+    ref_result = a_block + b_block
+    return call(a_block, b_block), ref_result
+
+
+@custom_partitioning
+def custom_shared_call(a, b):
+    return sharded_cutlass_call_impl(a, b)
+
+
+def custom_shared_call_partitioner(mesh, arg_shapes, result_shape):
+    arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)
+    result_shardings = tuple([arg_shardings[0]] * len(result_shape))
+
+    def lower_fn(*args):
+        return sharded_cutlass_call_impl(*args)
+
+    return mesh, lower_fn, result_shardings, arg_shardings
+
+
+custom_shared_call.def_partition(custom_shared_call_partitioner)
+
+
+def run_example():
+    # Create a device mesh with one axis b
+    ngpu = jax.device_count()
+    mesh = jax.make_mesh((ngpu,), "b", axis_types=(AxisType.Explicit,))
+
+    if ngpu == 1:
+        print("Note: only 1 GPU was detected.")
+
+    # We will shard our 3D tensors over b
+    sharding = P("b", None, None)
+    named_sharding = NamedSharding(mesh, sharding)
+
+    print("Testing shard_map...")
+
+    @partial(
+        jax.jit, static_argnums=[0, 1], out_shardings=(named_sharding, named_sharding)
+    )
+    def allocate_sharded_tensors(shape, dtype):
+        key = jax.random.key(1123)
+        a_key, b_key = jax.random.split(key, 2)
+        a = create_tensor(shape, dtype, a_key)
+        b = create_tensor(shape, dtype, b_key)
+        return a, b
+
+    @jax.jit
+    def compute(a, b):
+        # This jax.shard_map partitions the cutlass_call over the mesh.
+        @partial(
+            jax.shard_map,
+            mesh=mesh,
+            in_specs=(sharding, sharding),
+            out_specs=(sharding, sharding),
+        )
+        def sharded_call(a_block, b_block):
+            return sharded_cutlass_call_impl(a_block, b_block)
+
+        return sharded_call(a, b)
+
+    # Allocate (32, 16, 64) on each GPU
+    shape = (32 * ngpu, 16, 64)
+    dtype = jnp.float32
+
+    a, b = allocate_sharded_tensors(shape, dtype)
+    c, c_ref = compute(a, b)
+
+    assert jnp.allclose(c, c_ref)
+
+    print("Testing custom_partitioning...")
+
+    # Test custom_partitioning implementation which should produce identical results
+    @jax.jit
+    def compute_cp(a, b):
+        return custom_shared_call(a, b)
+
+    c, c_ref = compute_cp(a, b)
+
+    assert jnp.allclose(c, c_ref)
+
+
+if __name__ == "__main__":
+    run_example()
+    print("PASS")

examples/python/CuTeDSL/dsl_tutorials/jax/elementwise_apply_example.py

@@ -0,0 +1,329 @@
+# Copyright (c) 2025 - 2026 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.
+
+
+import argparse
+import operator
+from functools import partial
+from typing import List
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+
+"""
+An Elementwise Apply Example using CuTe DSL with cutlass.jax.cutlass_call
+
+This example is similar to examples/ampere/elementwise_apply.py but demonstrates
+how to run the code in a jax specific way using the cutlass_call primitive. It assumes
+familiarity with basic CuTe DSL concepts as well as the cutlass_call primitive.
+
+To run this example:
+
+.. code-block:: bash
+
+    # Run with addition operation
+    python examples/jax/elementwise_apply_example.py --M 1024 --N 512 --op add
+
+    # Run with multiplication operation
+    python examples/ampere/elementwise_apply_example.py --M 1024 --N 512 --op mul
+
+    # Run with subtraction operation
+    python examples/ampere/elementwise_apply_example.py --M 1024 --N 512 --op sub
+"""
+
+
+@cute.kernel
+def elementwise_apply_kernel(
+    op: cutlass.Constexpr,
+    mInputs: List[cute.Tensor],
+    mC: cute.Tensor,
+    cC: cute.Tensor,  # coordinate tensor
+    shape: cute.Shape,
+    tv_layout: cute.Layout,  # (tid, vid) -> logic coord
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, bidy, _ = cute.arch.block_idx()
+
+    ###############################################################################
+    # Slice to local tile of thread block
+    ###############################################################################
+    blk_crd = ((None, None), (bidx, bidy))
+
+    # Leverage the meta-programming capability of the DSL to slice the tensors for each input
+    # All for loops below on input tensors would be fully unrolled automatically at compile time
+    # logical coord -> memory address
+    gInputs = [t[blk_crd] for t in mInputs]  # (TileM, TileN)
+    gC = mC[blk_crd]  # (TileM, TileN)
+    gCrd = cC[blk_crd]  # (TileM, TileN)
+
+    print("[DSL INFO] Sliced Tensors per thread block:")
+    for i in cutlass.range_constexpr(len(gInputs)):
+        print(f"[DSL INFO]   ctaInputs{i} = {gInputs[i].type}")
+    print(f"[DSL INFO]   gC = {gC.type}")
+    print(f"[DSL INFO]   gCrd = {gCrd.type}")
+
+    ###############################################################################
+    # Compose with thread block TV layout to map thread & value indices to memory address
+    ###############################################################################
+    # (tid, vid) -> memory address
+    tidfrgInputs = [cute.composition(t, tv_layout) for t in gInputs]
+    tidfrgC = cute.composition(gC, tv_layout)
+    tidfrgCrd = cute.composition(gCrd, tv_layout)
+
+    # repeat None like vid to remove hierarchy of layout
+    thr_crd = (tidx, cute.repeat_like(None, tidfrgInputs[0][1]))
+
+    ###############################################################################
+    # Slice to local tile of thread
+    ###############################################################################
+    # vid -> address
+    thrInputs = [t[thr_crd] for t in tidfrgInputs]  # (V)
+    thrC = tidfrgC[thr_crd]  # (V)
+    thrCrd = tidfrgCrd[thr_crd]
+
+    print("[DSL INFO] Sliced Tensors per thread:")
+    for i in cutlass.range_constexpr(len(thrInputs)):
+        print(f"[DSL INFO]   thrInputs{i} = {thrInputs[i].type}")
+    print(f"[DSL INFO]   thrC = {thrC.type}")
+    print(f"[DSL INFO]   thrCrd = {thrCrd.type}")
+
+    ###############################################################################
+    # Compute predicate for out of boundary checks
+    ###############################################################################
+    frgPred = cute.make_fragment(thrCrd.shape, cutlass.Boolean)
+    print(f"[DSL INFO]   frgPred = {frgPred.type}")
+
+    for i in cutlass.range_constexpr(cute.size(frgPred)):
+        frgPred[i] = cute.elem_less(thrCrd[i], shape)
+
+    # if tidx == 0 and bidx == 0:
+    #     cute.print_tensor(frgPred)
+
+    ##########################################################
+    # Load data and compute result
+    ##########################################################
+
+    # Load data before use. The compiler will optimize the copy and load
+    # operations to convert some memory ld/st into register uses.
+    result = op(*[thrInput.load() for thrInput in thrInputs])
+    thrC.store(result)
+
+
+@cute.jit
+def elementwise_apply(
+    op: cutlass.Constexpr, inputs, result: cute.Tensor, stream: cuda.CUstream
+):
+    """CUDA kernel applying binary operator on each element of two n-D input tensors in
+    CuTe Python and store to result tensor.
+
+    :param op: Binary operator or lambda function to apply element-wise
+    :type op: cutlass.Constexpr
+    :param a: First input tensor
+    :type a: cute.Tensor
+    :param b: Second input tensor
+    :type b: cute.Tensor
+    :param result: Output tensor to store the results of op(a, b)
+    :type result: cute.Tensor
+    :return: None
+    :rtype: None
+    """
+
+    # Baseline: naive TV layout
+    #   * mA layout: (4096, 4096):(4096, 1)
+    #   * TV layout map to (512, 4) tile
+    #   * tidx maps to mode-0 but input layout is contiguous on mode-1, performance will be bad
+    # tv_layout = cute.make_layout((128, (4, 4)), stride=(4, (512, 1)))
+    # cta_tiler = (512, 4)
+
+    # Opt-1: better TV layout with better 1D thread layout (SOL with 1D thread layout)
+    #   * mA layout: (4096, 4096):(4096, 1)
+    #   * TV layout map to (4, 512) tile
+    #   * tidx maps to mode-1 which is leading mode of input tensor for coalesced load
+    # tv_layout = cute.make_layout((128, (4, 4)), stride=(16, (4, 1)))
+    # cta_tiler = (4, 512)
+
+    # Opt-2: 2D tile but worse
+    #   * mA layout: (4096, 4096):(4096, 1)
+    #   * TV layout map to (128, 16) logical tile
+    #   * V layout is bad as contiguous mode is not on right-most
+    #     * `cute.copy` only supports vectorize when stride-1 of v-layout on right-most )
+    # tv_layout = cute.make_layout(((32, 4), (4, 4)), stride=((4, 512), (1, 128)))
+    # cta_tiler = (128, 16)
+
+    # Opt-3: SOL with 2D thread tile
+    #   * mA layout: (4096, 4096):(4096, 1)
+    #   * TV layout map to (64, 256) logical tile
+    #   * tidx maps to mode-1 and input layout is contiguous on mode-1 for coalesced load-store
+
+    # Use 128bit(16B) load as canonicalized form of val_layout then recast to target element-type
+    coalesced_ldst_bytes = 16
+
+    # Compile time validation: expect same element type for all input tensors
+    assert all(t.element_type == inputs[0].element_type for t in inputs)
+    dtype = inputs[0].element_type
+
+    thr_layout = cute.make_ordered_layout((4, 64), order=(1, 0))
+    val_layout = cute.make_ordered_layout((16, coalesced_ldst_bytes), order=(1, 0))
+    val_layout = cute.recast_layout(dtype.width, 8, val_layout)
+    tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+
+    print("[DSL INFO] Input Tensors:")
+    for i, t in enumerate(inputs):
+        print(f"[DSL INFO]   inputs{i} = {t}")
+    print(f"[DSL INFO]   result = {result}")
+
+    print("[DSL INFO] Tiling Parameters:")
+    print(f"[DSL INFO]   tiler_mn = {tiler_mn} per thread block")
+    print(f"[DSL INFO]   tv_layout = {tv_layout}")
+
+    print("[DSL INFO] Tiled Tensors:")
+    mInputs = [cute.zipped_divide(input, tiler_mn) for input in inputs]
+    # ((TileM, TileN), (RestM, RestN))
+    mC = cute.zipped_divide(result, tiler_mn)
+
+    # (RestM, RestN) -> (RestN, RestM)
+    remap_block = cute.make_ordered_layout(
+        cute.select(mInputs[0].shape[1], mode=[1, 0]), order=(1, 0)
+    )
+    for i, t in enumerate(mInputs):
+        print(f"[DSL INFO]   gInputs{i}            = {mInputs[i]}")
+        mInputs[i] = cute.composition(t, (None, remap_block))
+        print(f"[DSL INFO]   gInputs{i} (remapped) = {mInputs[i]}")
+
+    mC = cute.composition(mC, (None, remap_block))
+    print(f"[DSL INFO]   gC = {mC}")
+
+    idC = cute.make_identity_tensor(result.shape)
+    cC = cute.zipped_divide(idC, tiler=tiler_mn)
+    print(f"[DSL INFO]   coord tensor = {cC}")
+
+    # Launch the kernel asynchronously
+    # Group input tensors into a list as a single argument
+    elementwise_apply_kernel(op, mInputs, mC, cC, result.shape, tv_layout).launch(
+        # Compute production at each mode of mC.shape[1] to get multi-dimensional grid size
+        grid=cute.product_each(mC.shape[1]),
+        block=[cute.size(tv_layout, mode=[0]), 1, 1],
+        stream=stream,
+    )
+
+
+@cutlass.dsl_user_op
+def leaky_relu(x, alpha, *, loc=None, ip=None):
+    return cute.where(x > 0, x, alpha * x, loc=loc, ip=ip)
+
+
+def leaky_relu_ref(x, alpha):
+    import jax.numpy as jnp
+
+    return jnp.where(x > 0, x, alpha * x)
+
+
+def run_and_verify(op, M, N, dtype, skip_ref_check=False):
+    import jax
+    import jax.numpy as jnp
+    import cutlass.jax as cjax
+    import cutlass.jax.testing as testing
+
+    if op == "leaky_relu":
+        op = partial(leaky_relu, alpha=0.01)
+        ref_op = partial(leaky_relu_ref, alpha=0.01)
+        num_inputs = 1
+    else:
+        op = getattr(operator, op)
+        ref_op = op
+        num_inputs = 2
+
+    # This jax function is transformed using jax.jit to compile its contents
+    # into an efficient HLO executable.
+    @partial(jax.jit, static_argnums=[1])
+    def jax_function(inputs, op):
+        call = cjax.cutlass_call(
+            # Bind jax arguments to kernel signature
+            lambda stream, inputs, output, *, op: elementwise_apply(
+                op, inputs, output, stream
+            ),
+            # Specify output shape/dtype of result
+            output_shape_dtype=jax.ShapeDtypeStruct(inputs[0].shape, inputs[0].dtype),
+            # Pass static/constexpr values as kwargs
+            op=op,
+        )
+
+        # Call the kernel!
+        return call(inputs)
+
+    @partial(jax.jit, static_argnums=[1])
+    def jax_ref_function(inputs, op):
+        return op(*inputs)
+
+    print("\nRunning Elementwise Apply test with:")
+    print(f"Tensor dimensions: [{M}, {N}]")
+    print(f"Input and Output Data type: {dtype}")
+
+    jax_dtype = cjax.cutlass_to_jax_dtype(dtype)
+    keys = jax.random.split(jax.random.key(1435), num_inputs)
+    inputs = [testing.create_tensor((M, N), jax_dtype, key) for key in keys]
+
+    print("Input tensor shapes:")
+    for i in range(num_inputs):
+        print(f"inputs[{i}]: {inputs[i].shape}, dtype: {inputs[i].dtype}")
+
+    epsilon = 1.2
+    if op in (operator.truediv, operator.floordiv):
+        inputs[1] = jnp.where(inputs[1] == 0, epsilon, inputs[1])
+
+    # Call the jax.jit function which will compile the kernel
+    c = jax_function(inputs, op)
+
+    if not skip_ref_check:
+        print("Executing elementwise apply kernel...")
+        c = jax_function(inputs, op)
+        print("Verifying results...")
+        assert jnp.allclose(ref_op(*inputs), c)
+        print("Results verified successfully!")
+        print(f"First few elements of result: \n{c[:3, :3]}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Demonstration of calling a kernel with cutlass_call"
+    )
+    parser.add_argument("--M", default=4096, type=int)
+    parser.add_argument("--N", default=4096, type=int)
+    parser.add_argument("--op", default="add", type=str)
+    parser.add_argument("--skip_ref_check", action="store_true")
+
+    args = parser.parse_args()
+    run_and_verify(
+        args.op,
+        args.M,
+        args.N,
+        dtype=cutlass.Float32,
+        skip_ref_check=args.skip_ref_check,
+    )
+    print("\nPASS")

examples/python/CuTeDSL/dsl_tutorials/print_latex.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2025 - 2026 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.
+
+import argparse
+import cutlass
+import cutlass.cute as cute
+from cutlass.utils import print_latex, print_latex_tv
+from cutlass import for_generate, yield_out
+
+"""
+A Latex Printing Example using CuTe DSL.
+This example prints latex for a given layout or thread value layout.
+The primary goal for this example is to demonstrate how to dump latex, which can then be
+turned into an image in your favorite latex compiler.
+To run this example:
+.. code-block:: bash
+    python examples/python/CuteDSL/cute/print_latex.py
+    python examples/python/CuteDSL/cute/print_latex.py --tv_layout
+To compile, pipe the output to a file and use a tool like pdflatex:
+.. code-block:: bash
+    python examples/python/CuTeDSL/cute/print_latex.py > latex.tex
+    pdflatex latex.tex
+"""
+
+
+@cute.jit
+def main(print_tv_layout: cutlass.Constexpr[bool]):
+    # Note: only support compile time printing layouts
+    if cutlass.const_expr(print_tv_layout):
+        thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
+        val_layout = cute.make_ordered_layout((4, 1), order=(1, 0))
+        tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+        print_latex_tv(tv_layout, tiler_mn)
+    else:
+        layout = cute.make_layout((10, 10))
+        print_latex(layout)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="example of print latex and print latex tv"
+    )
+    parser.add_argument("--tv_layout", action="store_true")
+
+    args = parser.parse_args()
+
+    main(args.tv_layout)

examples/python/CuTeDSL/dsl_tutorials/programmatic_dependent_launch.py

@@ -0,0 +1,368 @@
+# Copyright (c) 2025 - 2026 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.
+
+
+import argparse
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.testing as testing
+from cutlass.cute.runtime import from_dlpack
+
+
+def supports_pdl():
+    import torch
+
+    return torch.cuda.get_device_capability()[0] >= 9
+
+
+"""
+This example demonstrates the use of Programmatic Dependent Launch (PDL) using
+CuTe DSL.
+
+PDL is a mechanism which allows for overlapping execution of back-to-back kernels
+within the same stream.
+For example, consider the following two elementwise add operations, where the second
+operation's first operand is the result of the first operation. While performing
+``w = u + v`` we will load u and v, add them, and then store the result. Once we
+have finished loading data, we are no longer utilizing the read bandwidth.
+To effectively utilize the read bandwidth, we can start loading ``x``
+immediately upon finishing reading. This is what PDL enables us to do.
+
+.. code-block:: bash
+
+w = u + v
+y = w + x
+
+To enable PDL, we need to do two things:
+
+1. Insert the ``griddepcontrol.launch_dependents`` and ``griddepcontrol.wait`` instructions in the kernel.
+2. Set the PDL launch attribute when launching the kernel.
+
+The ``griddepcontrol.launch_dependents`` and ``griddepcontrol.wait``
+instructions enable fine-grained control over kernel execution in PDL.
+Once all thread blocks execute the ``griddepcontrol.launch_dependents``
+instruction, the dependent kernels can opportunistically be early-launched.
+``griddepcontrol.wait`` functions as a synchronization barrier - any warp
+executing this instruction will block until the previous kernel finishes
+execution. This allows precise control over data dependencies between kernels.
+
+The following diagram shows the overlapping execution of two dependent kernels.
+We call the instructions before ``griddepcontrol.wait`` as prologue (``P0``),
+which may include barrier initialization and loading of independent data, etc.
+We call the instructions after ``griddepcontrol.launch_dependents`` as epilogue
+(``P2``), which may include math operations, data stores, etc. PDL enables
+these prologue and epilogue phases to execute concurrently across dependent
+kernels, improving GPU resource utilization. This is particularly beneficial
+when prologue and epilogue are bound by different resources (e.g., memory
+bandwidth vs compute throughput).
+
+      # P0: Prologue, P1: Main compute, P2: Epilogue
+
+         P0    P1    P2
+  K1: |=====|+++++|-----|
+
+                  <-----> K2 can start early
+                         (K1's P2 overlaps with K2's P0)
+
+                     P0        P1    P2
+  K2:             |=====|   |+++++|-----|
+                          ^
+                          |
+                          wait for K1 to complete
+Time ------------------------------------------------------>
+
+We could run this example with and without PDL:
+
+.. code-block:: bash
+
+    python examples/blackwell/programmatic_dependent_launch.py --benchmark
+    python examples/blackwell/programmatic_dependent_launch.py --benchmark --use_pdl
+
+From the benchmark results, you can see some speedups for the PDL version in most cases, benefiting from
+the overlapping execution of consecutive kernels. Moreover, you can use nsys to observe the overlapping execution.
+
+.. code-block:: bash
+
+    nsys profile python examples/blackwell/programmatic_dependent_launch.py --benchmark --use_pdl
+
+Note, PDL feature is supported on Hopper and later GPUs.
+
+See [the programming guide](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization)
+and the [PTX documentation](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-griddepcontrol)
+for more details.
+"""
+
+
+@cute.kernel
+def elementwise_add_kernel(
+    gA: cute.Tensor,
+    gB: cute.Tensor,
+    gC: cute.Tensor,
+    cC: cute.Tensor,  # coordinate tensor
+    shape: cute.Shape,
+    thr_layout: cute.Layout,
+    val_layout: cute.Layout,
+    is_first_kernel: cutlass.Constexpr = True,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    blk_coord = ((None, None), bidx)
+    blkA = gA[blk_coord]  # (TileM,TileN)
+    blkB = gB[blk_coord]  # (TileM,TileN)
+    blkC = gC[blk_coord]  # (TileM,TileN)
+    blkCrd = cC[blk_coord]  # (TileM, TileN)
+
+    copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gA.element_type)
+    copy_atom_store = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gC.element_type)
+
+    tiled_copy_A = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
+    tiled_copy_B = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
+    tiled_copy_C = cute.make_tiled_copy_tv(copy_atom_store, thr_layout, val_layout)
+
+    thr_copy_A = tiled_copy_A.get_slice(tidx)
+    thr_copy_B = tiled_copy_B.get_slice(tidx)
+    thr_copy_C = tiled_copy_C.get_slice(tidx)
+
+    thrA = thr_copy_A.partition_S(blkA)
+    thrB = thr_copy_B.partition_S(blkB)
+    thrC = thr_copy_C.partition_S(blkC)
+
+    frgA = cute.make_fragment_like(thrA)
+    frgB = cute.make_fragment_like(thrB)
+    frgC = cute.make_fragment_like(thrC)
+
+    thrCrd = thr_copy_C.partition_S(blkCrd)
+    frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)
+
+    for i in range(cute.size(frgPred)):
+        val = cute.elem_less(thrCrd[i], shape)
+        frgPred[i] = val
+
+    # Note: when not using cuda-graph, the kernel execution may be blocked by the host overhead.
+    # In this case we won't see overlapping even when pdl is enabled.
+    # In this example, we add a loop (10 times) for all the copy and compute operations in the following code
+    # to make kernel running longer and make pdl benefits observable for both cuda-graph enabled and disabled cases.
+    if is_first_kernel:
+        for _ in range(10):
+            cute.copy(copy_atom_load, thrA, frgA, pred=frgPred)
+            cute.copy(copy_atom_load, thrB, frgB, pred=frgPred)
+        # Here we add the launch dependents instruction for the first kernel as a hint to the runtime to early-launch
+        # the next kernel. If the next kernel becomes concurrent, we will have overlap where the second kernel
+        # can start reading x to ensure an E2E speedup. Note the placement of launch dependents has no implication
+        # on correctness, only performance.
+        cute.arch.griddepcontrol_launch_dependents()
+    else:
+        # In this example, the second kernel's second operand ``gB`` has no dependencies, its loading can overlap
+        # with the computation of ``gC`` from the first kernel.
+        for _ in range(10):
+            cute.copy(copy_atom_load, thrB, frgB, pred=frgPred)
+
+        # For the second kernel, its first operand ``gA`` is dependent on the previous kernel, we must call
+        # griddepcontrol.wait to assure correctness. This instruction will block until the prior kernels finishes
+        # and its memory operations are visible. Since gA is written by the prior kernel, this will block until gA
+        # is visible to our kernel. Without it, we would have undefined behavior due to a race condition.
+        cute.arch.griddepcontrol_wait()
+
+        for _ in range(10):
+            cute.copy(copy_atom_load, thrA, frgA, pred=frgPred)
+
+    for _ in range(10):
+        result = frgA.load() + frgB.load()
+        frgC.store(result)
+    cute.copy(copy_atom_store, frgC, thrC, pred=frgPred)
+
+
+@cute.jit
+def elementwise_add(
+    mA,
+    mB,
+    mC,
+    stream: cuda.CUstream,
+    use_pdl: cutlass.Constexpr = True,
+    is_first_kernel: cutlass.Constexpr = True,
+):
+    dtype = mA.element_type
+    # copy_bits for a thread is 128 bits, and we use 128 // dtype.width to get the vector size
+    vector_size = 128 // dtype.width
+
+    thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
+    val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0))
+    tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+
+    gA = cute.zipped_divide(mA, tiler_mn)  # ((TileM,TileN),(RestM,RestN))
+    gB = cute.zipped_divide(mB, tiler_mn)  # ((TileM,TileN),(RestM,RestN))
+    gC = cute.zipped_divide(mC, tiler_mn)  # ((TileM,TileN),(RestM,RestN))
+
+    idC = cute.make_identity_tensor(mC.shape)
+    cC = cute.zipped_divide(idC, tiler=tiler_mn)
+
+    elementwise_add_kernel(
+        gA, gB, gC, cC, mC.shape, thr_layout, val_layout, is_first_kernel
+    ).launch(
+        grid=[cute.size(gC, mode=[1]), 1, 1],
+        block=[cute.size(tv_layout, mode=[0]), 1, 1],
+        stream=stream,
+        use_pdl=use_pdl,
+    )
+
+
+def run_pdl_example(
+    M,
+    N,
+    skip_ref_check=False,
+    benchmark=False,
+    warmup_iterations=5,
+    iterations=100,
+    use_pdl=True,
+):
+    import torch
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("Blackwell/Hopper GPU is required to run this example!")
+
+    print("\nRunning Elementwise Add test with:")
+    print(f"Tensor dimensions: [{M}, {N}]")
+    print(f"Use PDL: {use_pdl}")
+
+    u = torch.randn(M, N, dtype=torch.float32, device="cuda")
+    v = torch.randn(M, N, dtype=torch.float32, device="cuda")
+    w = torch.randn(M, N, dtype=torch.float32, device="cuda")
+    x = torch.randn(M, N, dtype=torch.float32, device="cuda")
+    y = torch.empty(M, N, dtype=torch.float32, device="cuda")
+
+    u_tensor = from_dlpack(u).mark_layout_dynamic()
+    v_tensor = from_dlpack(v).mark_layout_dynamic()
+    w_tensor = from_dlpack(w).mark_layout_dynamic()
+    x_tensor = from_dlpack(x).mark_layout_dynamic()
+    y_tensor = from_dlpack(y).mark_layout_dynamic()
+
+    stream = torch.cuda.Stream()
+    current_stream = cuda.CUstream(stream.cuda_stream)
+    # Since is_first_kernel is cutlass.Constexpr, we need to compile for
+    # the first and second kernel separately.
+    compiled_func_first_kernel = cute.compile(
+        elementwise_add,
+        u_tensor,
+        v_tensor,
+        w_tensor,
+        current_stream,
+        use_pdl,
+        is_first_kernel=True,
+        options="--enable-tvm-ffi",
+    )
+    compiled_func_second_kernel = cute.compile(
+        elementwise_add,
+        w_tensor,
+        x_tensor,
+        y_tensor,
+        current_stream,
+        use_pdl,
+        is_first_kernel=False,
+        options="--enable-tvm-ffi",
+    )
+
+    # launch and run the two consecutive kernels in a same stream.
+    def run_func(current_stream, u, v, w, x, y):
+        # Run first operation: w_tensor = u_tensor + v_tensor
+        compiled_func_first_kernel(
+            u,
+            v,
+            w,
+            current_stream,
+        )
+        # Run second operation: y_tensor = w_tensor + x_tensor
+        # its first operand ``w_tensor`` is the result of the first operation,
+        # they use the same memory space.
+        compiled_func_second_kernel(
+            w,
+            x,
+            y,
+            current_stream,
+        )
+
+    if not skip_ref_check:
+        run_func(current_stream, u, v, w, x, y)
+        print("Verifying results...")
+        torch.testing.assert_close(u.cpu() + v.cpu() + x.cpu(), y.cpu())
+        print("Results verified successfully!")
+
+    if not benchmark:
+        return
+
+    def generate_kernel_arguments():
+        u = torch.randn(M, N, dtype=torch.float32, device="cuda")
+        v = torch.randn(M, N, dtype=torch.float32, device="cuda")
+        w = torch.randn(M, N, dtype=torch.float32, device="cuda")
+        x = torch.randn(M, N, dtype=torch.float32, device="cuda")
+        y = torch.empty(M, N, dtype=torch.float32, device="cuda")
+
+        return testing.JitArguments(current_stream, u, v, w, x, y)
+
+    avg_time_us = testing.benchmark(
+        run_func,
+        workspace_generator=generate_kernel_arguments,
+        workspace_count=10,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+        stream=current_stream,
+        use_cuda_graphs=True,
+    )
+    print(f"Execution time: {avg_time_us:.4f} us")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="example of Programmatic Dependent Launch (PDL) using CuTe DSL"
+    )
+    parser.add_argument("--M", default=256, type=int)
+    parser.add_argument("--N", default=256, type=int)
+    parser.add_argument("--warmup_iterations", default=5, type=int)
+    parser.add_argument("--iterations", default=10, type=int)
+    parser.add_argument("--skip_ref_check", action="store_true")
+    parser.add_argument("--benchmark", action="store_true")
+    parser.add_argument("--use_pdl", action="store_true")
+
+    args = parser.parse_args()
+    if supports_pdl():
+        run_pdl_example(
+            args.M,
+            args.N,
+            skip_ref_check=args.skip_ref_check,
+            benchmark=args.benchmark,
+            warmup_iterations=args.warmup_iterations,
+            iterations=args.iterations,
+            use_pdl=args.use_pdl,
+        )
+        print("\nPASS")
+    else:
+        print(
+            "PDL is not supported on this device, it requires Hopper or newer generations"
+        )

examples/python/CuTeDSL/dsl_tutorials/smem_allocator.py

@@ -0,0 +1,223 @@
+# Copyright (c) 2025 - 2026 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.
+
+import cutlass.cute as cute
+import cutlass
+import numpy as np
+from cutlass.cute.runtime import from_dlpack
+
+"""
+A Shared Memory Allocator Example on NVIDIA Ampere architecture using CuTe DSL.
+
+This example demonstrates how to allocate and manage shared memory in JIT kernels by using the SmemAllocator in CuTe DSL.
+It shows various ways to allocate different data structures in shared memory:
+
+1. Struct allocation with natural and strict alignment
+2. Raw memory block allocation with custom alignment
+3. Array allocation with automatic alignment
+4. Tensor allocation with layout specification
+
+The example includes:
+- Shared storage struct with mixed alignment requirements
+- Memory allocation patterns for different data types
+- Tensor operations on allocated memory
+
+To run this example:
+
+.. code-block:: bash
+
+    python examples/ampere/smem_allocator.py
+
+The example will allocate shared memory, perform tensor operations, and verify the results.
+"""
+
+
+@cute.struct
+class complex:
+    real: cutlass.Float32
+    imag: cutlass.Float32
+
+
+# SharedStorage size is 512, alignment is 128
+@cute.struct
+class SharedStorage:
+    # struct elements with natural alignment
+    a: cute.struct.MemRange[cutlass.Float32, 32]  # array
+    b: cutlass.Int64  # scalar
+    c: complex  # nested struct
+    # struct elements with strict alignment
+    x: cute.struct.Align[
+        cute.struct.MemRange[cutlass.Float32, 32],
+        128,
+    ]
+    y: cute.struct.Align[cutlass.Int32, 8]
+    z: cute.struct.Align[complex, 16]
+
+
+@cute.kernel
+def kernel(
+    const_a: cutlass.Constexpr,
+    dst_a: cute.Tensor,
+    const_b: cutlass.Constexpr,
+    dst_b: cute.Tensor,
+    const_c: cutlass.Constexpr,
+    dst_c: cute.Tensor,
+):
+    # Note: SMEM_SIZE bytes (specified in kernel().launch(smem=...)) can be reserved for developer to utilize
+    # Note: alignment of initial allocator base ptr is 1024
+    allocator = cutlass.utils.SmemAllocator()
+    # base ptr of allocator points at: SMEM_ADDR_START (the starting address of available shared memory)
+
+    # -- Allocate a scalar
+    int_ptr = allocator.allocate(cutlass.Int32)
+    # base ptr of allocator now points at: SMEM_ADDR_AFTER_INT = SMEM_ADDR_START + aligned_size(int)
+    assert int_ptr.dtype == cutlass.Int32, "Expected Int32, but got {}".format(
+        int_ptr.dtype
+    )
+
+    # -- Allocate a struct --
+    # Note: when specified alignment, max(alignment, alignof(struct)) will be applied
+    # reserves the section of struct in smem, elements in the struct can be accessed by ptr
+    struct_in_smem = allocator.allocate(SharedStorage)
+    # base ptr of allocator now points at: SMEM_ADDR_AFTER_STRUCT = SMEM_ADDR_START + aligned_size(struct)
+
+    # -- Allocate a block of memory --
+    # reserves a section of 64 bytes in smem, align to 128 bytes, returns the section base ptr
+    section_in_smem = allocator.allocate(64, byte_alignment=128)
+    # base ptr of allocator now points at: SMEM_ADDR_AFTER_SECTION = SMEM_ADDR_AFTER_STRUCT + aligned_size(section)
+
+    # -- Allocate an array --
+    # reserves an int64 array of size 14 in smem, returns the array base ptr
+    array_in_smem = allocator.allocate_array(element_type=cutlass.Int64, num_elems=14)
+    # base ptr of allocator now points at: SMEM_ADDR_AFTER_ARRAY = SMEM_ADDR_AFTER_SECTION + aligned_size(array)
+
+    # -- Allocate a tensor --
+    # Note: use cute.ComposedLayout or cute.Layout to specify layout of tensor
+    # Note: iterator swizzle with swizzle layout is currently not supported
+    layout = cute.make_layout((16, 2))
+    tensor_in_smem = allocator.allocate_tensor(
+        element_type=cutlass.Float32, layout=layout, byte_alignment=32, swizzle=None
+    )
+    # base ptr of allocator now points at: SMEM_ADDR_AFTER_TENSOR = SMEM_ADDR_AFTER_ARRAY + aligned_size(tensor)
+
+    # ptr<f16, smem, align<1024>>
+    # ptr<i64, smem, align<128>>
+    # ptr<f32, smem, align<8>>
+    print(struct_in_smem.a.data_ptr())
+    print(struct_in_smem.b.ptr)
+    print(struct_in_smem.c.real.ptr)
+    # ptr<i8, smem, align<512>>
+    print(section_in_smem)
+    # ptr<i64, smem, align<64>>
+    print(array_in_smem)
+    # tensor<ptr<f16, smem, align<32>> o (16,4):(1,16)>
+    print(tensor_in_smem)
+
+    # assign struct member array element
+    cute.printf("struct_in_smem.a[0] = {}", struct_in_smem.a[0])
+    struct_in_smem.a[0] = 2
+    cute.printf("struct_in_smem.a[0] = {}", struct_in_smem.a[0])
+
+    # assign struct member scalar
+    cute.printf("struct_in_smem.b.ptr = {}", struct_in_smem.b.ptr)
+    cute.printf("struct_in_smem.b: value = {}", struct_in_smem.b.ptr.load())
+    struct_in_smem.b = 16
+    cute.printf("struct_in_smem.b: value = {}", struct_in_smem.b.ptr.load())
+
+    # fill MemRange tensor in struct and copy to dst
+    a_tensor = struct_in_smem.a.get_tensor(cute.make_layout((8, 4)))
+    a_tensor.fill(const_a)
+    cute.printf("cute.struct.MemRange: {}", a_tensor)
+    dst_a.store(a_tensor.load())
+
+    # convert block of smem to fill tensor and copy to dst
+    layout = cute.make_layout((8, 2))
+    sec_ptr = cute.recast_ptr(section_in_smem, dtype=cutlass.Float32)
+    sec_tensor = cute.make_tensor(sec_ptr, layout)
+    sec_tensor.fill(const_b)
+    cute.printf("block of memory: {}", sec_tensor)
+    dst_b.store(sec_tensor.load())
+
+    # fill allocated tensor in smem and copy to dst
+    tensor_in_smem.fill(const_c)
+    cute.printf("tensor in smem: {}", tensor_in_smem)
+    dst_c.store(tensor_in_smem.load())
+
+
+@cute.jit
+def host(
+    const_a: cutlass.Constexpr,
+    dst_a: cute.Tensor,
+    const_b: cutlass.Constexpr,
+    dst_b: cute.Tensor,
+    const_c: cutlass.Constexpr,
+    dst_c: cute.Tensor,
+):
+    # Note: Shared Memory size is automatically calculated now
+    kernel(const_a, dst_a, const_b, dst_b, const_c, dst_c).launch(
+        grid=(1, 1, 1),
+        block=(1, 1, 1),
+        # Automatically calculate the launch kernel shared memory usage when `smem=None`
+    )
+
+
+def run_and_verify(const_a, const_b, const_c):
+    import torch
+
+    dst_a = torch.zeros((8, 4), dtype=torch.float32, device="cuda")
+    dst_b = torch.zeros((8, 2), dtype=torch.float32, device="cuda")
+    dst_c = torch.zeros((16, 2), dtype=torch.float32, device="cuda")
+
+    host(
+        const_a,
+        from_dlpack(dst_a),
+        const_b,
+        from_dlpack(dst_b),
+        const_c,
+        from_dlpack(dst_c),
+    )
+
+    assert const_a == dst_a.cpu()[0, 0], (
+        f"Expected {const_a}, but got {dst_a.cpu()[0, 0]}"
+    )
+    assert const_b == dst_b.cpu()[0, 0], (
+        f"Expected {const_b}, but got {dst_b.cpu()[0, 0]}"
+    )
+    assert const_c == dst_c.cpu()[0, 0], (
+        f"Expected {const_c}, but got {dst_c.cpu()[0, 0]}"
+    )
+
+
+if __name__ == "__main__":
+    # prepare cuda context
+    cutlass.cuda.initialize_cuda_context()
+    # An example for shared memory allocation
+    const_a = 0.5
+    const_b = 1.0
+    const_c = 2.0
+    run_and_verify(const_a, const_b, const_c)

examples/python/CuTeDSL/dsl_tutorials/torch_fake_tensor.py

@@ -0,0 +1,80 @@
+# Copyright (c) 2025 - 2026 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.
+
+
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+
+"""Example demonstrating how to use CuTe with PyTorch's FakeTensor mode.
+
+This example shows how to:
+1. Use PyTorch's FakeTensor mode to compile a CuTe function without real data
+2. Execute the compiled function on real data later
+
+FakeTensor mode allows compiling code without allocating real memory, which is useful
+for ahead-of-time compilation scenarios. The compiled function can then be executed
+on real tensors that match the expected shapes and dtypes.
+
+Primary goals of this example are to demonstrate: How to use PyTorch's FakeTensor mode with CuTe
+to enable ahead-of-time compilation without real data allocation.
+
+The example:
+1. Creates a fake tensor in PyTorch using FakeTensor mode
+2. Compiles a CuTe function using the fake tensor without allocating real memory
+3. Creates a real tensor with matching shape and dtype
+4. Executes the compiled function on the real tensor
+
+To run this example:
+
+.. code-block:: bash
+
+    python examples/cute/torch_fake_tensor.py
+"""
+
+
+@cute.jit
+def print_tensor(t: cute.Tensor):
+    cute.print_tensor(t)
+
+
+def run():
+    import torch
+    from torch._subclasses.fake_tensor import FakeTensorMode
+
+    shape = (3, 4)
+    with FakeTensorMode():
+        fake_tensor = torch.zeros(shape, dtype=torch.float32)
+        compiled_fn = cute.compile(print_tensor, from_dlpack(fake_tensor))
+
+    real_tensor = torch.randn(shape, dtype=torch.float32)
+    compiled_fn(from_dlpack(real_tensor))
+
+
+if __name__ == "__main__":
+    run()

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/ampere_gemm_with_fake_tensor.py

@@ -0,0 +1,170 @@
+# Copyright (c) 2025 - 2026 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.
+
+import sys
+import os
+import torch
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+"""Demonstrates calling off-the-shelf kernels with TVM FFI without DLPack.
+
+This example shows how to compile CuTe JIT function with fake tensors then run it with TVM FFI.
+"""
+
+if __name__ == "__main__":
+    # Add the current directory to sys.path
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, "..", ".."))
+from cute.ampere.kernel.dense_gemm.tensorop_gemm import TensorOpGemm
+
+
+@cute.jit
+def bmm(
+    a: cute.Tensor,  # (l, m, k)
+    b: cute.Tensor,  # (l, k, n)
+    c: cute.Tensor,  # (l, m, n)
+):
+    gemm_op = TensorOpGemm(cutlass.Float16, cutlass.Float16, cutlass.Float32, (2, 2, 1))
+
+    # Permute to follow convention of CuTe
+
+    # (l, m, k) -> (m, k, l)
+    a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0]))
+    # (l, k, n) -> (n, k, l)
+    b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0]))
+    # (l, m, n) -> (m, n, l)
+    c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0]))
+
+    gemm_op(a, b, c)
+
+
+def compile_bmm_dynamic_layout():
+    from cutlass.cute.runtime import make_fake_compact_tensor
+
+    m = cute.sym_int()
+    n = cute.sym_int(divisibility=16)
+    k = cute.sym_int(divisibility=16)
+    l = cute.sym_int()
+
+    # Contiguous on K
+    fake_a = make_fake_compact_tensor(
+        cutlass.Float16, (l, m, k), stride_order=(2, 1, 0), assumed_align=16
+    )
+    # Contiguous on N
+    fake_b = make_fake_compact_tensor(
+        cutlass.Float16, (l, k, n), stride_order=(2, 1, 0), assumed_align=16
+    )
+    # Contiguous on N
+    fake_c = make_fake_compact_tensor(
+        cutlass.Float16, (l, m, n), stride_order=(2, 1, 0), assumed_align=16
+    )
+
+    compiled_fn = cute.compile(bmm, fake_a, fake_b, fake_c, options="--enable-tvm-ffi")
+    return compiled_fn
+
+
+def compile_bmm_static_layout(m, n, k, l):
+    from cutlass.cute.runtime import make_fake_compact_tensor
+
+    fake_a = make_fake_compact_tensor(
+        cutlass.Float16, (l, m, k), stride_order=(2, 1, 0), assumed_align=16
+    )
+    fake_b = make_fake_compact_tensor(
+        cutlass.Float16, (l, k, n), stride_order=(2, 1, 0), assumed_align=16
+    )
+    fake_c = make_fake_compact_tensor(
+        cutlass.Float16, (l, m, n), stride_order=(2, 1, 0), assumed_align=16
+    )
+
+    compiled_fn = cute.compile(bmm, fake_a, fake_b, fake_c, options="--enable-tvm-ffi")
+    return compiled_fn
+
+
+def run_bmm_and_verify(compiled_fn, m, n, k, l):
+    torch.manual_seed(1112)
+
+    a = torch.randn(l, m, k, dtype=torch.float16, device="cuda")
+    b = torch.randn(l, k, n, dtype=torch.float16, device="cuda")
+    c = torch.randn(l, m, n, dtype=torch.float16, device="cuda")
+
+    print("[Runtime INFO] Input tensor shapes:")
+    print(f"a: {a.shape=}, {a.stride()=}, {a.dtype=}")
+    print(f"b: {b.shape=}, {b.stride()=}, {b.dtype=}")
+    print(f"c: {c.shape=}, {c.stride()=}, {c.dtype=}\n")
+
+    # pass in torch tensor as input
+    compiled_fn(a, b, c)
+    torch.cuda.synchronize()
+
+    ref = torch.bmm(a, b)
+    torch.testing.assert_close(c, ref, atol=1e-05, rtol=1e-05)
+    print("[Runtime INFO] Verification successful!")
+    print(f"  First few elements of result: \n{c[:3, :3, :3]}")
+
+
+if __name__ == "__main__":
+    m, n, k, l = (512, 512, 256, 2)
+
+    compiled_fn_dynamic = compile_bmm_dynamic_layout()
+    run_bmm_and_verify(compiled_fn_dynamic, m, n, k, l)
+
+    compiled_fn_static = compile_bmm_static_layout(m, n, k, l)
+    run_bmm_and_verify(compiled_fn_static, m, n, k, l)
+
+    # Error Check:
+    #   1. mis-matched tensor dim raise error
+    a = torch.randn(l, m, k, dtype=torch.float16, device="cuda")
+    b = torch.randn(l, 2 * k, n, dtype=torch.float16, device="cuda")
+    c = torch.randn(l, m, n, dtype=torch.float16, device="cuda")
+    try:
+        compiled_fn_dynamic(a, b, c)
+    except Exception as e:
+        print(f"\n[Runtime Error]: {e}")
+
+    #   2. mis-matched divisibility
+    a = torch.randn(l, m, k + 1, dtype=torch.float16, device="cuda")
+    b = torch.randn(l, k + 1, n, dtype=torch.float16, device="cuda")
+    c = torch.randn(l, m, n, dtype=torch.float16, device="cuda")
+
+    try:
+        compiled_fn_dynamic(a, b, c)
+    except Exception as e:
+        print(f"\n[Runtime Error]: {e}")
+
+    # 3. mis-matched static shape constraint
+    a = torch.randn(l * 2, m, k, dtype=torch.float16, device="cuda")
+    b = torch.randn(l * 2, k, n, dtype=torch.float16, device="cuda")
+    c = torch.randn(l * 2, m, n, dtype=torch.float16, device="cuda")
+
+    try:
+        compiled_fn_static(a, b, c)
+    except Exception as e:
+        print(f"\n[Runtime Error]: {e}")

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/aot_export.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2025 - 2026 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.
+
+"""Example demonstrating how to use TVM-FFI ABI with CuTe.
+
+This example shows how to:
+1. Compile a CuTe function with "--enable-tvm-ffi" option
+2. Export the compiled function to a shared library
+3. Load the shared library and use the compiled function to work with torch.Tensor
+
+To run this example:
+
+.. code-block:: bash
+
+    python cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_export.py
+    # run example to use in torch
+    python cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_torch.py
+    # run example to use in jax
+    python cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_jax.py
+    # run example to use in c++ bundle
+    bash cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
+"""
+
+from pathlib import Path
+import os
+import subprocess
+import tvm_ffi
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+
+@cute.kernel
+def device_add_one(a: cute.Tensor, b: cute.Tensor):
+    for i in range(a.shape[0]):
+        b[i] = a[i] + 1
+
+
+@cute.jit
+def add_one(a: cute.Tensor, b: cute.Tensor):
+    """b = a + 1"""
+    device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1))
+
+
+def main():
+    import torch
+
+    # compile the kernel with "--enable-tvm-ffi" option
+    a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
+    b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
+    a_cute = from_dlpack(a_torch, enable_tvm_ffi=True).mark_layout_dynamic()
+    b_cute = from_dlpack(b_torch, enable_tvm_ffi=True).mark_layout_dynamic()
+    # compile the kernel with "--enable-tvm-ffi" option
+    compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")
+
+    os.makedirs("./build", exist_ok=True)
+    object_file_path = "./build/add_one.o"
+    lib_path = "./build/add_one.so"
+    compiled_add_one.export_to_c(object_file_path, function_name="add_one")
+    shared_libs = cute.runtime.find_runtime_libraries(enable_tvm_ffi=True)
+    # compile the object file to a shared library
+    cmd = ["gcc", "-shared", "-o", lib_path, object_file_path, *shared_libs]
+    print(cmd)
+    subprocess.run(cmd, check=True)
+    print(f"Successfully created shared library: {lib_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/aot_use_in_cpp_bundle.cpp

@@ -0,0 +1,94 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+ *
+ * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+ * property and proprietary rights in and to this material, related
+ * documentation and any modifications thereto. Any use, reproduction,
+ * disclosure or distribution of this material and related documentation
+ * without an express license agreement from NVIDIA CORPORATION or
+ * its affiliates is strictly prohibited.
+ */
+
+// clang-format on
+// This example shows how to interface with an AOT compiled function in a C++
+// bundle. to build and run the example, run the following command in project
+// root bash
+// cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
+
+#include <cuda_runtime.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/error.h>
+#include <tvm/ffi/extra/module.h>
+#include <iostream>
+#include <vector>
+
+
+namespace ffi = tvm::ffi;
+
+struct CUDANDAlloc {
+  void AllocData(DLTensor* tensor) {
+    size_t data_size = ffi::GetDataSize(*tensor);
+    void* ptr = nullptr;
+    cudaError_t err = cudaMalloc(&ptr, data_size);
+    TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaMalloc failed: " << cudaGetErrorString(err);
+    tensor->data = ptr;
+  }
+
+  void FreeData(DLTensor* tensor) {
+    if (tensor->data != nullptr) {
+      cudaError_t err = cudaFree(tensor->data);
+      TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaFree failed: " << cudaGetErrorString(err);
+      tensor->data = nullptr;
+    }
+  }
+};
+
+inline ffi::Tensor Empty(ffi::Shape shape, DLDataType dtype, DLDevice device) {
+  return ffi::Tensor::FromNDAlloc(CUDANDAlloc(), shape, dtype, device);
+}
+
+// symbol from the shared library
+extern "C" int __tvm_ffi_add_one(void*, const TVMFFIAny*, int32_t, TVMFFIAny*);
+
+// Redirects into the exported function in object
+void CallAddOne(ffi::TensorView x, ffi::TensorView y) {
+    tvm::ffi::Function::InvokeExternC(nullptr, __tvm_ffi_add_one, x, y);
+}
+
+int main() {
+    DLDataType f32_dtype{kDLFloat, 32, 1};
+    DLDevice cuda_device{kDLCUDA, 0};
+
+    constexpr int ARRAY_SIZE = 10;
+
+    ffi::Tensor x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
+    ffi::Tensor y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
+
+    std::vector<float> host_x(ARRAY_SIZE);
+    for (int i = 0; i < ARRAY_SIZE; ++i) {
+      host_x[i] = static_cast<float>(i);
+    }
+
+    size_t nbytes = host_x.size() * sizeof(float);
+    cudaError_t err = cudaMemcpy(x.data_ptr(), host_x.data(), nbytes, cudaMemcpyHostToDevice);
+    TVM_FFI_ICHECK_EQ(err, cudaSuccess)
+        << "cudaMemcpy host to device failed: " << cudaGetErrorString(err);
+
+    // Call into the FFI function; tensors remain on device because they carry a
+    // kDLCUDA device tag.
+    CallAddOne(x, y);
+
+    std::vector<float> host_y(host_x.size());
+    err = cudaMemcpy(host_y.data(), y.data_ptr(), nbytes, cudaMemcpyDeviceToHost);
+    TVM_FFI_ICHECK_EQ(err, cudaSuccess)
+        << "cudaMemcpy device to host failed: " << cudaGetErrorString(err);
+
+    std::cout << "y after add_one_cuda(x, y)" << std::endl;
+    for (float value : host_y) {
+        std::cout << value << " ";
+    }
+    std::cout << std::endl;
+    return 0;
+}

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/aot_use_in_cpp_bundle.sh

@@ -0,0 +1,48 @@
+# Copyright (c) 2025 - 2026 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.
+
+#!/bin/bash
+# Set up library paths for runtime
+export LD_LIBRARY_PATH=$(python3 -m cutlass.cute.export.aot_config --libdir):$(tvm-ffi-config --libdir):$LD_LIBRARY_PATH
+
+CUDA_HOME=/usr/local/cuda
+SOURCE_FILE="$(dirname "$0")/aot_use_in_cpp_bundle.cpp"
+
+echo "Compiling the executable..."
+g++ -o build/aot_use_in_cpp_bundle \
+  -I${CUDA_HOME}/include \
+  `tvm-ffi-config --cxxflags`  \
+  ${SOURCE_FILE} build/add_one.o \
+  $(python3 -m cutlass.cute.export.aot_config --ldflags) \
+  -L${CUDA_HOME}/lib64 \
+  $(python3 -m cutlass.cute.export.aot_config --libs) -lcuda -lcudart \
+  $(tvm-ffi-config --ldflags) \
+  $(tvm-ffi-config --libs)
+
+echo "Running the executable..."
+./build/aot_use_in_cpp_bundle

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/aot_use_in_jax.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2025 - 2026 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.
+
+import jax
+import jax.numpy as jnp
+import jax_tvm_ffi
+import cutlass.cute as cute
+
+
+# now load it back
+def main():
+    a_jax = jnp.arange(10, dtype=jnp.float32)
+    b_jax = jnp.zeros(10, dtype=jnp.float32)
+    lib_path = "./build/add_one.so"
+    aot_mod = cute.runtime.load_module(lib_path, enable_tvm_ffi=True)
+    jax_tvm_ffi.register_ffi_target("add_one_cute", aot_mod.add_one, platform="gpu")
+    b_jax = jax.ffi.ffi_call(
+        "add_one_cute",
+        jax.ShapeDtypeStruct(a_jax.shape, a_jax.dtype),
+        vmap_method="broadcast_all",
+    )(a_jax)
+    print("result of b after aot_mod.add_one(a, b)")
+    print(b_jax)
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/aot_use_in_torch.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2025 - 2026 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.
+
+import cutlass.cute as cute
+
+
+# now load it back
+def main():
+    import torch
+
+    a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
+    b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
+    lib_path = "./build/add_one.so"
+    aot_mod = cute.runtime.load_module(lib_path, enable_tvm_ffi=True)
+    aot_mod.add_one(a_torch, b_torch)
+    print("result of b after aot_mod.add_one(a, b)")
+    print(b_torch)
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/compile_with_fake_tensor.py

@@ -0,0 +1,48 @@
+import cutlass
+import cutlass.cute as cute
+
+
+"""
+Example of using fake tensors in CuTe.
+
+This script demonstrates how to use fake tensors in CuTe to drive compilation without creating actual tensors
+from frameworks like PyTorch or TensorFlow.
+
+Run this file directly to see the output type information.
+"""
+
+
+@cute.jit
+def print_tensor_type(t: cute.Tensor):
+    print(t)
+
+
+def run():
+    from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_tensor
+
+    shape = (3, 4)
+    a = make_fake_compact_tensor(cutlass.Float16, (3, 4), stride_order=(1, 0))
+    cute.compile(print_tensor_type, a, options="--enable-tvm-ffi")
+
+    # 32-bit symbolic integer with divisibility 8
+    shape = (3, cute.sym_int32(divisibility=8))
+    a = make_fake_compact_tensor(cutlass.Float16, shape, stride_order=(1, 0))
+    cute.compile(print_tensor_type, a, options="--enable-tvm-ffi")
+
+    # with static stride
+    a = make_fake_tensor(cutlass.Float16, shape, stride=(4, 1))
+    cute.compile(print_tensor_type, a, options="--enable-tvm-ffi")
+
+    # with dynamic stride using 32bit integer
+    stride = (cute.sym_int32(divisibility=8), 1)
+    a = make_fake_tensor(cutlass.Float16, shape, stride=stride)
+    cute.compile(print_tensor_type, a, options="--enable-tvm-ffi")
+
+    # with dynamic stride using 64bit integer
+    stride = (cute.sym_int64(divisibility=8), 1)
+    a = make_fake_tensor(cutlass.Float16, shape, stride=stride)
+    cute.compile(print_tensor_type, a, options="--enable-tvm-ffi")
+
+
+if __name__ == "__main__":
+    run()

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/error_reporting.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2025 - 2026 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.
+
+"""Example demonstrating how to use TVM-FFI ABI with CuTe.
+
+This example shows how to:
+1. Compile a CuTe function with "--enable-tvm-ffi" option
+2. Directly use the compiled function to work with torch.Tensor
+
+To run this example:
+
+.. code-block:: bash
+
+    python cutlass_ir/compiler/python/examples/cute/tvm_ffi/error_reporting.py
+"""
+
+import torch
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+
+@cute.kernel
+def device_add_one(a: cute.Tensor, b: cute.Tensor):
+    for i in range(a.shape[0]):
+        b[i] = a[i] + 1
+
+
+@cute.jit
+def add_one(a: cute.Tensor, b: cute.Tensor):
+    """b = a + 1"""
+    device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1))
+
+
+def main():
+    # compile the kernel with "--enable-tvm-ffi" option
+    a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
+    b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
+    a_cute = from_dlpack(a_torch, enable_tvm_ffi=True)
+    b_cute = from_dlpack(b_torch, enable_tvm_ffi=True)
+    # compile the kernel with "--enable-tvm-ffi" option
+    compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")
+    # should raise an error because of shape mismatch
+    compiled_add_one(torch.arange(5, dtype=torch.float32, device="cuda"), b_cute)
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/jit_and_use_in_jax.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2025 - 2026 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.
+
+"""Example demonstrating how to use TVM-FFI ABI with CuTe.
+
+This example shows how to:
+1. Compile a CuTe function with "--enable-tvm-ffi" option
+2. Directly use the compiled function to work with JAX
+
+To run this example:
+
+.. code-block:: bash
+
+    pip install jax-tvm-ffi
+    pip install jax[cuda13]
+
+    python cutlass_ir/compiler/python/examples/cute/tvm_ffi/jit_and_use_in_jax.py
+"""
+
+import jax
+from jax import numpy as jnp
+import jax_tvm_ffi
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+
+@cute.kernel
+def device_add_one(a: cute.Tensor, b: cute.Tensor):
+    for i in range(a.shape[0]):
+        b[i] = a[i] + 1
+
+
+@cute.jit
+def add_one(a: cute.Tensor, b: cute.Tensor):
+    """b = a + 1"""
+    device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1))
+
+
+def main():
+    # compile the kernel with "--enable-tvm-ffi" option
+    a_jax = jnp.arange(
+        10,
+        dtype=jnp.float32,
+    )
+    b_jax = jnp.zeros(
+        10,
+        dtype=jnp.float32,
+    )
+    a_cute = from_dlpack(a_jax, enable_tvm_ffi=True).mark_layout_dynamic()
+    b_cute = from_dlpack(b_jax, enable_tvm_ffi=True).mark_layout_dynamic()
+    # compile the kernel with "--enable-tvm-ffi" option
+    compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")
+
+    # register the compiled function to JAX as a FFI target
+    jax_tvm_ffi.register_ffi_target("add_one_cute", compiled_add_one, platform="gpu")
+    a_jax = jnp.arange(10, dtype=jnp.float32)
+    # call the compiled function using JAX FFI
+    b_jax = jax.ffi.ffi_call(
+        "add_one_cute",
+        jax.ShapeDtypeStruct(a_jax.shape, a_jax.dtype),
+        vmap_method="broadcast_all",
+    )(a_jax)
+    print("result of b_jax after add_one_cute")
+    print(b_jax)
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/jit_and_use_in_torch.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2025 - 2026 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.
+
+"""Example demonstrating how to use TVM-FFI ABI with CuTe.
+
+This example shows how to:
+1. Compile a CuTe function with "--enable-tvm-ffi" option
+2. Directly use the compiled function to work with torch.Tensor
+
+To run this example:
+
+.. code-block:: bash
+
+    python cutlass_ir/compiler/python/examples/cute/tvm_ffi/jit_and_use_in_torch.py
+"""
+
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+
+@cute.kernel
+def device_add_one(a: cute.Tensor, b: cute.Tensor):
+    for i in range(a.shape[0]):
+        b[i] = a[i] + 1
+
+
+@cute.jit
+def add_one(a: cute.Tensor, b: cute.Tensor):
+    """b = a + 1"""
+    device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1))
+
+
+def main():
+    import torch
+
+    # compile the kernel with "--enable-tvm-ffi" option
+    a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
+    b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
+    a_cute = from_dlpack(a_torch, enable_tvm_ffi=True).mark_layout_dynamic()
+    b_cute = from_dlpack(b_torch, enable_tvm_ffi=True).mark_layout_dynamic()
+    # compile the kernel with "--enable-tvm-ffi" option
+    compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")
+    # run the compiled function by passing in cute.Tensor as input
+    # you need to set enable_tvm_ffi=True for now
+    compiled_add_one(a_cute, b_cute)
+    # print the result
+    print("result of b after compiled_add_one(a, b)")
+    print(b_torch)
+    a_torch = a_torch + 1
+    # We can directly pass in torch.Tensor as input
+    # the call overhead is optimized so it is very fast to pass in torch.Tensor as input
+    # takes about less than 0.5us per call likely in terms of API overhead
+    compiled_add_one(a_torch, b_torch)
+    # print the result
+    print("result of b after compiled_add_one(a, b)")
+    print(b_torch)
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/dsl_tutorials/tvm_ffi/requirements.txt

@@ -0,0 +1,2 @@
+apache-tvm-ffi
+torch-c-dlpack-ext