v4.5 tag update (#3202)

* Python DSL examples reorganization.

* v4.5 tag update.

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