v4.3 tag release update. (#2789)

2025-11-21 09:49:44 +08:00
parent 406e078b29
commit 8cd5bef43a
225 changed files with 23229 additions and 2813 deletions
--- a/examples/python/CuTeDSL/cute/fake_tensor.py
+++ b/examples/python/CuTeDSL/cute/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)
+
+    # 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)
+
+    # with static stride
+    a = make_fake_tensor(cutlass.Float16, shape, stride=(4, 1))
+    cute.compile(print_tensor_type, a)
+
+    # 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)
+
+    # 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)
+
+
+if __name__ == "__main__":
+    run()
--- a/examples/python/CuTeDSL/cute/ffi/CMakeLists.txt
+++ b/examples/python/CuTeDSL/cute/ffi/CMakeLists.txt
@@ -35,20 +35,26 @@ find_package(Python3 COMPONENTS Interpreter Development REQUIRED)

 # Get Python site-packages directory using Python
 execute_process(
-    COMMAND ${Python3_EXECUTABLE} -c "import site; print(site.getsitepackages()[0])"
-    OUTPUT_VARIABLE Python_SITE_PACKAGES
+    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 directory: ${Python_SITE_PACKAGES}")
+message(STATUS "Python site-packages directories: ${Python_SITE_PACKAGES_PATHS}")

 # Add nanobind path to CMAKE_PREFIX_PATH
-list(APPEND CMAKE_PREFIX_PATH ${Python_SITE_PACKAGES}/nanobind/cmake)
+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 
+    message(FATAL_ERROR
        "nanobind not found!\n"
        "Please install nanobind with: pip install nanobind\n"
    )
--- a/examples/python/CuTeDSL/cute/ffi/jit_argument.py
+++ b/examples/python/CuTeDSL/cute/ffi/jit_argument.py
@@ -243,12 +243,7 @@ import tempfile
 import torch


-def run_test(tmpdir=None, cmake_args=""):
-    # Skip cleanup if user provides tmpdir
-    cleanup = tmpdir is None
-    # Initialize temporary build directory
-    tmpdir = tmpdir or tempfile.mkdtemp()
-
+def run_test(tmpdir=None, cmake_args="", cleanup=True):
    try:
        current_dir = os.path.dirname(os.path.abspath(__file__))

@@ -256,8 +251,6 @@ def run_test(tmpdir=None, cmake_args=""):
        subprocess.run(["cmake", "-B", tmpdir, current_dir] + cmake_args, check=True)
        subprocess.run(["cmake", "--build", tmpdir], check=True)

-        sys.path.append(tmpdir)
-
        from tensor import make_tensor, pycapsule_get_pointer

        # Mock test tensor and corresponding C structure for this example
@@ -314,4 +307,13 @@ if __name__ == "__main__":
    )
    args = parser.parse_args()

-    run_test(tmpdir=args.tmp_dir, cmake_args=args.cmake_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)
--- a/examples/python/CuTeDSL/cute/print_latex.py
+++ b/examples/python/CuTeDSL/cute/print_latex.py
@@ -0,0 +1,72 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+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)
--- a/examples/python/CuTeDSL/cute/torch_fake_tensor.py
+++ b/examples/python/CuTeDSL/cute/torch_fake_tensor.py
@@ -65,7 +65,7 @@ def print_tensor(t: cute.Tensor):
    cute.print_tensor(t)


-if __name__ == "__main__":
+def run():
    from torch._subclasses.fake_tensor import FakeTensorMode

    shape = (3, 4)
@@ -75,3 +75,7 @@ if __name__ == "__main__":

    real_tensor = torch.randn(shape, dtype=torch.float32)
    compiled_fn(from_dlpack(real_tensor))
+
+
+if __name__ == "__main__":
+    run()
--- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py
+++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py
@@ -0,0 +1,92 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""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 examples/cute/tvm_ffi/aot_export.py
+    # run example to use in torch
+    python examples/cute/tvm_ffi/aot_use_in_torch.py
+    # run example to use in jax
+    python examples/cute/tvm_ffi/aot_use_in_jax.py
+    # run example to use in c++ bundle
+    bash examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
+"""
+
+from pathlib import Path
+import torch
+import os
+import subprocess
+import tvm_ffi
+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).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")
+
+    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()
--- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp
+++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp
@@ -0,0 +1,97 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023 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
+// examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
+
+#include <cuda_runtime.h>
+#include <iostream>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/error.h>
+#include <tvm/ffi/extra/module.h>
+#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;
+}
--- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
+++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
@@ -0,0 +1,48 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#!/bin/bash
+CUDA_DIALECT_PATH="build/lib/"
+export LD_LIBRARY_PATH=${CUDA_DIALECT_PATH}:`tvm-ffi-config --libdir`
+
+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 \
+  -L${CUDA_DIALECT_PATH} \
+  -L${CUDA_HOME}/lib64 \
+  -lcuda_dialect_runtime -lcuda -lcudart \
+  `tvm-ffi-config --ldflags` \
+  `tvm-ffi-config --libs`
+
+echo "Running the executable..."
+./build/aot_use_in_cpp_bundle
--- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py
+++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py
@@ -0,0 +1,52 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+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)
+    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()
--- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py
+++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py
@@ -0,0 +1,44 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import cutlass.cute as cute
+import torch
+
+# now load it back
+def main():
+    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)
+    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()
--- a/examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py
+++ b/examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py
@@ -0,0 +1,71 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""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 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()
--- a/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py
+++ b/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py
@@ -0,0 +1,93 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""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 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()
--- a/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py
+++ b/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py
@@ -0,0 +1,83 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""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 examples/cute/tvm_ffi/jit_and_use_in_torch.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).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()