compiled_add = cute.compile(add, a_torch, b_torch, options="--enable-tvm-ffi")

Note that CuTe's convention is to write the stride order for dimensions from left to right,
where a lower order number means higher priority. In the context of the ``make_fake_compact_tensor`` API,
for shape ``(2, 3, 4)`` and stride order ``(0, 1, 2)``, the stride is ``(1, 2, 6)``.
This is commonly known as column-major order. If you want to create a fake tensor with compact row-major order,
you should explicitly pass in ``stride_order=tuple(reversed(range(len(shape))))``
to ``make_fake_compact_tensor``. Alternatively, you can always precisely control the
stride via the ``stride`` argument in the ``make_fake_tensor`` API.


``cute.Tensor`` adapter for TVM FFI
-----------------------------------

To adapt the ``cute.Tensor`` to the TVM FFI function, you can use the ``cute.runtime.from_dlpack`` function with the
``enable_tvm_ffi=True`` option or the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. For example:

.. code-block:: python

   def example_from_dlpack():
      a_cute = cute.runtime.from_dlpack(a_torch, enable_tvm_ffi=True).mark_layout_dynamic()
      b_cute = cute.runtime.from_dlpack(b_torch, enable_tvm_ffi=True).mark_layout_dynamic()

      compiled_add_one(a_cute, b_cute)

Note that because the ``cute.runtime.from_dlpack`` function performs an explicit DLPack conversion, it is less efficient than passing the ``torch.Tensor`` directly.
You can also use ``cute.Tensor`` as an argument hint for ``cute.compile``.

.. code-block:: python

   compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")


Working with torch Tensors
--------------------------

As you may have noticed in the examples above, TVM FFI-compiled functions can
directly accept ``torch.Tensor`` objects (and other DLPack-compatible tensors) as inputs.
The resulting functions add minimal overhead, enabling faster eager invocations
thanks to the optimized calling path.

Working with Streams
--------------------

In many cases, a CuTe kernel needs to run on a specific CUDA stream.
|DSL| provides two ways to work with streams through TVM FFI.
The first is to pass the stream explicitly as an argument.
The following example demonstrates this approach; the function accepts ``torch.cuda.Stream``,
``CUstream`` or any stream class that implements the CUDA stream protocol.

.. code-block:: python

   import cutlass.cute as cute
   import torch
   from cuda.bindings.driver import CUstream

   @cute.kernel
   def device_add_one(a: cute.Tensor, b: cute.Tensor):
      threads_per_block = 128
      cta_x_, _, _ = cute.arch.block_idx()
      tid_x, _, _ = cute.arch.thread_idx()
      tid = cta_x_ * threads_per_block + tid_x
      if tid < a.shape[0]:
         b[tid] = a[tid] + 1.0

   @cute.jit
   def add_one_with_stream(a: cute.Tensor, b: cute.Tensor, stream: CUstream):
      n = a.shape[0]
      threads_per_block = 128
      blocks = (n + threads_per_block - 1) // threads_per_block
      device_add_one(a, b).launch(
         grid=(blocks, 1, 1),
         block=(threads_per_block, 1, 1),
         stream=stream,
      )

   def example_add_one_with_stream():
      n = cute.sym_int()
      a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))
      b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))
      # Fake stream is a placeholder for stream argument
      stream = cute.runtime.make_fake_stream()
      compiled_add_one = cute.compile(
         add_one_with_stream, a_cute, b_cute, stream, options="--enable-tvm-ffi"
      )
      a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
      b_torch = torch.empty(10, dtype=torch.float32, device="cuda")
      torch_stream = torch.cuda.current_stream()
      compiled_add_one(a_torch, b_torch, torch_stream)
      torch_stream.synchronize()
      print("result of b_torch after compiled_add_one(a_torch, b_torch, torch_stream)")
      print(b_torch)

Using Environment Stream

Named tuples are also supported and help logically group related arguments together.
The example below shows how to use named tuples as arguments. Under the hood, named tuples
are passed as unnamed tuples at the ABI level. When errors occur, the function signature in
error messages will display unnamed tuple arguments.
Ensure that the compile-time CuTe named tuple type definition has the same fields
as the runtime PyTorch named tuple.
Currently, users need to explicitly unpack the named tuple outside of conditionals and then
use the unpacked variables inside the conditionals.

.. code-block:: python

   from typing import NamedTuple
   from cutlass import cute
   import torch

   class CuteNamedTuple(NamedTuple):
      a: cute.Tensor
      b: cute.Tensor
      c: cute.Float32 = cute.Float32(1)

      def __new_from_mlir_values__(self, values):
         return CuteNamedTuple(*values)

   class TorchNamedTuple(NamedTuple):
      a: torch.Tensor
      b: torch.Tensor
      c: float = 1

   @cute.kernel
   def device_add_one_named_tuple(value: CuteNamedTuple):
      tid = cute.arch.block_idx()[0] * 128 + cute.arch.thread_idx()[0]
      # need to unpack namedtuple outside conditionals
      a = value.a
      b = value.b
      c = value.c
      if tid < a.shape[0]:
         b[tid] = a[tid] + c

   @cute.jit
   def add_one_with_named_tuple(value: CuteNamedTuple):
      n = value.a.shape[0]
      threads_per_block = 128
      blocks = (n + threads_per_block - 1) // threads_per_block
      device_add_one_named_tuple(value).launch(grid=(blocks, 1, 1), block=(threads_per_block, 1, 1))

   def example_add_one_with_named_tuple():
      n = cute.sym_int()
      a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))
      b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))

      compiled_add_one = cute.compile(
         add_one_with_named_tuple, CuteNamedTuple(a=a_cute, b=b_cute),
         options="--enable-tvm-ffi"
      )
      a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
      b_torch = torch.empty(10, dtype=torch.float32, device="cuda")
      compiled_add_one(TorchNamedTuple(a=a_torch, b=b_torch))
      print("result of b_torch")
      print(b_torch)

   example_add_one_with_named_tuple()


Supported types
---------------

The TVM FFI function supports the following |DSL|-specific types as arguments:

- ``cute.Tensor``
- ``cutlass.Boolean``, ``cutlass.Int8``, ``cutlass.Int16``, ``cutlass.Int32``, ``cutlass.Int64``, ``cutlass.Uint8``, ``cutlass.Uint16``, ``cutlass.Uint32``, ``cutlass.Uint64``, ``cutlass.Float32``,  ``cutlass.Float64``
- ``cute.Shape``, ``cute.Stride``, ``cute.Coord``, ``cute.Tile``, ``cute.IntTuple``

.. list-table::
   :header-rows: 1
   :widths: 30 70

   * - Compile-time type
     - Call-time type
   * - ``cute.Pointer``
     - ``ctypes.c_void_p`` or a class that implements ``__tvm_ffi_opaque_ptr__`` protocol.
   * - ``cute.runtime.FakeTensor``
     - ``torch.Tensor`` and other DLPack-compatible tensors.
   * - Scalar types (e.g. ``cutlass.Boolean``, ``cutlass.Int32``)
     - Python scalars (e.g. True, 123).
   * - CuTe algebra types (e.g. ``cute.Shape``, ``cute.Stride``)
     - ``tvm_ffi.Shape`` or python tuple of ints.
   * - CUDA stream ``cuda.CUstream``
     - A stream class that implements the CUDA stream protocol (e.g. ``torch.cuda.Stream``, ``cuda.CUstream``).
   * - Tuple of types (e.g. ``Tuple[cute.Tensor, cute.Tensor, cutlass.Int32]``)
     - Python tuple of corresponding call-time types.

Error handling
--------------

TVM FFI functions will enable validation of arguments to make sure they match the expected type
and value constraints declared by the user. These checks are compiled into the function, run very fast,
and have no observable overhead during function invocation. Each of those errors will translate
into a proper Python exception that can be caught and handled. The example below shows some
example error cases that can be checked:

.. code-block:: python

   def example_constraint_checks():
      n = cute.sym_int(divisibility=16)
      # assume align to 16 bytes (4 int32), both should share same shape variable n
      a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,), assumed_align=16)
      b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,), assumed_align=16)
      compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")
      a = torch.zeros(128, dtype=torch.float32, device="cuda")
      b = torch.zeros(128, dtype=torch.float32, device="cuda")

      try:
         # raises type mismatch error because we expect a and b to be float32
         compiled_add_one(a, 1)
      except TypeError as e:
         # Mismatched type on argument #1 when calling:
         # `add_one(a: Tensor([n0], float32), b: Tensor([n0], float32))`,
         # expected Tensor
         print(f"TypeError: {e}")

      try:
         # raises shape mismatch error because we expect both a and b have shap [n]
         compiled_add_one(a, b[:126])
      except ValueError as e:
         # Mismatched b.shape[0] on argument #1 when calling:
         # `add_one(a: Tensor([n0], float32), b: Tensor([n0], float32))`,
         # expected to match a.shape[0]
         print(f"ValueError: {e}")

      try:
         # triggers divisibility mismatch error because 126 is not divisible by 16
         compiled_add_one(a[:126], b[:126])
      except ValueError as e:
         # Invalid a.shape[0] on argument #0 when calling:
         # `add_one(a: Tensor([n0], float32), b: Tensor([n0], float32)`,
         # expected to be divisible by 16
         print(f"ValueError: {e}")

      try:
         a = torch.zeros(129, dtype=torch.float32, device="cuda")
         b = torch.zeros(129, dtype=torch.float32, device="cuda")
         # triggers data alignment mismatch error because x and y are not aligned to 16 bytes
         compiled_add_one(a[1:], b[1:])
      except ValueError as e:
         # raises: Misaligned Tensor data on argument #0 when calling:
         # `add_one(a: Tensor([n0], float32), b: Tensor([n0], float32)`,
         # expected data alignment=16 bytes
         print(f"ValueError: {e}")

Any CUDA errors encountered will also be automatically converted into Python exceptions by the TVM FFI function.

.. code-block:: python

   @cute.jit
   def add_one_invalid_launch(a: cute.Tensor, b: cute.Tensor):
      # Intentionally exceed the maximum block dimension (1024 threads) so the
      # CUDA runtime reports an invalid configuration error.
      device_add_one(a, b).launch(grid=(1, 1, 1), block=(4096, 1, 1))

   def example_error_cuda_error():
      a_torch = torch.zeros((10,), dtype=torch.float32, device="cuda")
      b_torch = torch.zeros((10,), dtype=torch.float32, device="cuda")

      a_cute = cute.runtime.from_dlpack(a_torch, enable_tvm_ffi=True)
      b_cute = cute.runtime.from_dlpack(b_torch, enable_tvm_ffi=True)
      compiled_add_one_invalid_launch = cute.compile(
         add_one_invalid_launch, a_cute, b_cute, options="--enable-tvm-ffi"
      )

      try:
         compiled_add_one_invalid_launch(a_torch, b_torch)
      except RuntimeError as e:
         # raises RuntimeError: CUDA Error: cudaErrorInvalidValue
         print(f"RuntimeError: {e}")


Working with Devices
--------------------
TVM FFI-compiled functions naturally work across GPU devices.
The device index of the first input GPU tensor determines the kernel's device context.
The TVM FFI function calls ``cudaSetDevice`` to set the correct device
before launching the kernel based on that tensor's device index.
For advanced scenarios that pass raw pointers instead of tensors, you should call
``cudaSetDevice`` explicitly through the CUDA Python API.

Call a compiled function from C++ via the TVM FFI registry
----------------------------------------------------------

The object returned by ``cute.compile(..., options="--enable-tvm-ffi")`` is itself a
``tvm_ffi.Function``: a native callable that follows the TVM FFI calling convention.
You can publish it in TVM FFI's process-global function registry under a string name with
``tvm_ffi.register_global_func``. Once registered, the same compiled kernel can be
looked up by name and invoked from any TVM FFI-supported language (e.g. C++) running
in the same process. This allows you to bypass the Python interpreter entirely and avoid
the CPU overhead by staying in the C++ environment only.
In such a case, Python is only used for expressiveness as a DSL to describe the kernel and register it in the shared registry,
and C++ is used for the actual execution with efficiency.

The following is a minimal example of how to call a compiled CuTeDSL function in C++.

This is the C++ code we will compile into a PyTorch extension. We name it ``extension.cpp`` here:

.. code-block:: cpp

   #include <ATen/DLConvertor.h>     // at::toDLPackNonOwning
   #include <torch/extension.h>      // pybind11 + at::Tensor
   #include <tvm/ffi/container/tensor.h>  // tvm::ffi::TensorView
   #include <tvm/ffi/function.h>     // tvm::ffi::Function

   #include <string>

   void apply_tvm_function(const std::string& name, at::Tensor &x, at::Tensor &y, at::Tensor &z) {
      tvm::ffi::Function fn = tvm::ffi::Function::GetGlobalRequired(name);
      DLTensor dl_x = {};
      DLTensor dl_y = {};
      DLTensor dl_z = {};
      at::toDLPackNonOwning(x, &dl_x);
      at::toDLPackNonOwning(y, &dl_y);
      at::toDLPackNonOwning(z, &dl_z);
      fn(&dl_x, &dl_y, &dl_z);
   }

   PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("apply_tvm_function", &apply_tvm_function,
         "Look up a tvm-ffi global function by name and call it with three tensors.");
   }


Then we need to compile and load this extension into PyTorch:

.. code-block:: python

   import subprocess
   import sys

   from torch.utils.cpp_extension import load

   def _tvm_ffi_config(flag: str) -> str:
      """Ask the installed apache-tvm-ffi where its headers and lib live."""
      out = subprocess.check_output([sys.executable, "-m", "tvm_ffi.config", flag])
      return out.decode().strip()


   def build_extension():
      include_dir = _tvm_ffi_config("--includedir")
      dlpack_include_dir = _tvm_ffi_config("--dlpack-includedir")
      lib_dir = _tvm_ffi_config("--libdir")
      return load(
         name="tvm_ffi_demo_ext",
         sources=["extension.cpp"],
         extra_include_paths=[include_dir, dlpack_include_dir],
         extra_cflags=["-std=c++17"],
         # -ltvm_ffi to link, and -rpath so the .so is found at runtime. This is
         # the same libtvm_ffi.so that `import tvm_ffi` loads -> shared registry.
         extra_ldflags=[f"-L{lib_dir}", "-ltvm_ffi", f"-Wl,-rpath,{lib_dir}"],
         verbose=True,
      )

With all the boilerplate code ready, now let's write a CuTeDSL kernel and use the C++ extension to call it.
In practice you might want to call the CuTeDSL function from C++ directly without going back to Python. We call from Python here just for demonstration purposes.

.. code-block:: python

   import cutlass
   import torch
   import tvm_ffi
   import cutlass.cute as cute

   @cute.jit
   def add(x: cute.Tensor, y: cute.Tensor, z: cute.Tensor):
      add_kernel(x, y, z).launch(grid=[1, 1, 1], block=[16, 1, 1])

   @cute.kernel
   def add_kernel(x: cute.Tensor, y: cute.Tensor, z: cute.Tensor):
      tidx, _, _ = cute.arch.thread_idx()
      if tidx < 16:
         z[tidx] = x[tidx] + y[tidx]

   def main() -> None:
      ext = build_extension()

      fake_x = cute.runtime.make_fake_compact_tensor(cutlass.BFloat16, (4, 4), stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4)
      fake_y = cute.runtime.make_fake_compact_tensor(cutlass.BFloat16, (4, 4), stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4)
      fake_z = cute.runtime.make_fake_compact_tensor(cutlass.BFloat16, (4, 4), stride_order=(1, 0), memspace=cute.AddressSpace.gmem, assumed_align=4)
      compiled = cute.compile(add, fake_x, fake_y, fake_z, options="--enable-tvm-ffi",)
      tvm_ffi.register_global_func("CuTeDSL_add", compiled, override=True)

      x = torch.randn((4, 4), dtype=torch.bfloat16, device="cuda")
      y = torch.randn((4, 4), dtype=torch.bfloat16, device="cuda")
      z = torch.randn((4, 4), dtype=torch.bfloat16, device="cuda")

      # Launch the C++ function. This is only for demonstration because it's the easiest way to run our C++ function.
      # In practice you might be working with some C++ heavy framework and you should call the C++ function from C++ directly without going through Python.
      ext.apply_tvm_function("CuTeDSL_add", x, y, z)
      assert torch.allclose(x + y, z, atol=1e-8, rtol=1e-8)
      print("Successfully called CuTeDSL function from C++!")

   if __name__ == "__main__":
      main()

Calling convention of TVM-FFI in C++

The function returned by ``cute.compile`` supports keyword arguments and defaults.
The example below shows how to use keyword arguments and defaults:

.. code-block:: python

   import torch
   from cutlass import cute

   @cute.kernel
   def device_add_scalar(a: cute.Tensor, b: cute.Tensor, offset: cutlass.Float32):
      threads_per_block = 128
      cta_x_, _, _ = cute.arch.block_idx()
      tid_x, _, _ = cute.arch.thread_idx()
      tid = cta_x_ * threads_per_block + tid_x
      if tid < a.shape[0]:
         b[tid] = a[tid] + offset

   @cute.jit
   def add_constant(a: cute.Tensor, b: cute.Tensor, offset: cutlass.Float32=cutlass.Float32(1)):
      n = a.shape[0]
      threads_per_block = 128
      blocks = (n + threads_per_block - 1) // threads_per_block
      device_add_scalar(a, b, offset).launch(grid=(blocks, 1, 1), block=(threads_per_block, 1, 1))

   def example_kwargs_and_defaults():
      n = cute.sym_int()
      a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))
      b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))
      compiled_add_constant = cute.compile(add_constant, a_cute, b_cute, options="--enable-tvm-ffi")
      a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
      b_torch = torch.empty(10, dtype=torch.float32, device="cuda")
      compiled_add_constant(a_torch, b_torch)
      print("result of b_torch after compiled_add_constant(a_torch, b_torch)")
      print(b_torch)
      compiled_add_constant(a_torch, b_torch, offset=4)
      print("result of b_torch after compiled_add_constant(a_torch, b_torch, offset=4)")
      print(b_torch)

For efficiency and portability reasons, TVM FFI ABI supports functions with positional-only arguments.
If you export the compiled module to an object file and then load it back, the function
will only accept positional arguments in the order of the arguments in the function signature.
You can rewrap the function or use the TVM FFI wrapper generator to generate a kwargs wrapper.
The code block below shows how to do this:

.. code-block:: python

   def example_kwargs_and_defaults():
      n = cute.sym_int()
      a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))
      b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))
      compiled_add_constant = cute.compile(add_constant, a_cute, b_cute, options="--enable-tvm-ffi")
      # export the compiled module to object file
      compiled_add_constant.export_to_c("./add_constant.o", function_name="add_constant")
      # obtain necessary runtime libs for loading the shared library
      runtime_libs = cute.runtime.find_runtime_libraries(enable_tvm_ffi=True)
      # compile the object file to a shared library
      cmd = ["gcc", "-shared", "-o", "./add_constant.so", "./add_constant.o", *runtime_libs]
      subprocess.run(cmd, check=True)

      a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
      b_torch = torch.empty(10, dtype=torch.float32, device="cuda")

      mod = cute.runtime.load_module("./add_constant.so")
      try:
         mod.add_constant(a_torch, b_torch)
      except Exception as e:
         # Raises a missing arguments error because kwargs and default information are lost
         print(e)
      # We rewrap the function to regain argument and kwargs support.
      # Alternatively, use the TVM FFI wrapper generator to generate a kwargs wrapper function.
      from tvm_ffi.utils import kwargs_wrapper
      # arg_defaults are aligned to the end of the argument list
      wrapped_func = kwargs_wrapper.make_kwargs_wrapper(
         mod.add_constant, arg_names=["a", "b", "offset"], arg_defaults=(1,)
      )
      wrapped_func(a_torch, b_torch)
      print("result of b_torch after wrapped_func(a_torch, b_torch)")
      print(b_torch)
      # You can also use the signature of the original function
      # to generate a kwargs wrapper function. Make sure to exclude
      # arguments that are not included in the runtime,
      # such as 'self', constexpr, and env stream arguments.
      wrapped_func = kwargs_wrapper.make_kwargs_wrapper_from_signature(
         mod.add_constant, signature=inspect.signature(add_constant),
         exclude_arg_names=["self"]
      )
      wrapped_func(a_torch, b_torch, offset=4)
      print("result of b_torch after wrapped_func(a_torch, b_torch, offset=4)")
      print(b_torch)


Limitations
-----------

The Fake Tensor flow is ONLY compatible with TVM FFI because TVM FFI supports more flexible constraints on Tensor arguments.
For instance, fake tensor can specify per-mode static shape or constraints on shape and strides which are not supported by
existing ``from_dlpack`` flow. It's expected that JIT function compiled with fake tensor will have different ABI compared to
tensor converted by ``from_dlpack``.

.. code-block:: python

   import cutlass.cute as cute
   import torch

   n = cute.sym_int()
   # Dynamic Shape
   fake_a = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))

   # Compile without tvm-ffi
   compiled_fn = cute.compile(foo, fake_a)

   # Wrong, in compatible ABI
   compiled_fn(from_dlpack(a))


In order to avoid such issue, it's recommended to use fake tensor only with TVM FFI backend. Practically speaking,
as we only want to call ``from_dlpack`` once and reuse for both compilation and runtime, the benefit of
using fake tensor is limited in this case.

Host cross-compilation is not supported on the TVM FFI path: the ``--host-target`` option cannot be combined with
``--enable-tvm-ffi`` (doing so raises an error). To cross-compile the host object for another CPU architecture (for example
AArch64), use the low-level cute ABI AOT path described in :ref:`dsl_aot_host_cross_compilation`.