leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Merge pull request #2813 from fengxie/ftse/fix/example

Browse Source 
Refactor TVM FFI examples and update doc
This commit is contained in:
drazi
2025-11-28 09:07:15 +08:00
committed by GitHub
parent 1de3a576cc 8588d099e4
commit 5016493cc0
7 changed files with 249 additions and 219 deletions
Download Patch File Download Diff File Expand all files Collapse all files

128
examples/python/CuTeDSL/ampere/call_with_tvm_ffi.py
View File

@@ -1,128 +0,0 @@
# 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 sys
import os
import torch
import time
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 ampere.tensorop_gemm import TensorOpGemm
def compile_op(use_tvm_ffi: bool = True):
from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_tensor
a_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
b_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
c_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
a = make_fake_compact_tensor(
cutlass.Float16, a_shape, stride_order=(1, 0, 2), assumed_align=16
)
b = make_fake_compact_tensor(
cutlass.Float16, b_shape, stride_order=(1, 0, 2), assumed_align=16
)
c = make_fake_compact_tensor(
cutlass.Float16, c_shape, stride_order=(1, 0, 2), assumed_align=16
)
tensor_op_gemm = TensorOpGemm(
cutlass.Float16, cutlass.Float16, cutlass.Float32, (2, 2, 1)
)
compiled_fn = cute.compile(
tensor_op_gemm, a, b, c, options="--enable-tvm-ffi" if use_tvm_ffi else ""
)
return compiled_fn
def run_op(compiled_fn, mnkl, *, use_tvm_ffi: bool = True):
print("\nRunning TensorOpGemm test with:")
print(f"Tensor dimensions: {mnkl}")
torch.manual_seed(1112)
# (M,K,L)
a_torch = torch.randn(
mnkl[3], mnkl[0], mnkl[2], dtype=torch.float16, device="cuda"
).permute(1, 2, 0)
# (N,K,L)
b_torch = torch.randn(
mnkl[3], mnkl[1], mnkl[2], dtype=torch.float16, device="cuda"
).permute(1, 2, 0)
# (N,M,L)
c_torch = 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_torch.shape}, dtype: {a_torch.dtype}")
print(f"b: {b_torch.shape}, dtype: {b_torch.dtype}")
print(f"c: {c_torch.shape}, dtype: {c_torch.dtype}\n")
if not use_tvm_ffi:
a = from_dlpack(a_torch).mark_layout_dynamic(leading_dim=1)
b = from_dlpack(b_torch).mark_layout_dynamic(leading_dim=1)
c = from_dlpack(c_torch).mark_layout_dynamic(leading_dim=1)
else:
a = a_torch
b = b_torch
c = c_torch
# pass in torch tensor as input
compiled_fn(a, b, c)
torch.cuda.synchronize()
# measure the launch overhead of tvm ffi function
repeat = 100
start_time = time.time()
for i in range(repeat):
compiled_fn(a, b, c)
end_time = time.time()
print(
f"Launch overhead of tvm ffi function: {(end_time - start_time) / repeat} seconds"
)
ref = torch.einsum("mkl,nkl->mnl", a_torch, b_torch)
torch.testing.assert_close(c_torch, ref, atol=1e-05, rtol=1e-05)
print("\n[DSL INFO] Results verified successfully!")
print(f"First few elements of result: \n{c_torch[:3, :3, :3]}")
if __name__ == "__main__":
compiled_fn = compile_op(use_tvm_ffi=False)
run_op(compiled_fn, [512, 512, 256, 1], use_tvm_ffi=False)
compiled_fn = compile_op(use_tvm_ffi=True)
run_op(compiled_fn, [512, 512, 256, 1], use_tvm_ffi=True)

109
examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
View File

@@ -34,6 +34,7 @@ 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
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
@@ -1704,55 +1705,6 @@ def bmm(
gemm_op(a, b, c, max_active_clusters, stream, epilogue_op)
def compile_bmm(
gemm_op: PersistentDenseGemmKernel,
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
c_major: str,
max_active_clusters: cutlass.Constexpr,
stream: cuda.CUstream,
epilogue_op: cutlass.Constexpr = lambda x: x,
options: str = "",
):
from cutlass.cute.runtime import make_fake_compact_tensor
a_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
b_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
c_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
if a_major == "k":
a_order = (2, 1, 0) # k is leading dimension
elif a_major == "m":
a_order = (2, 0, 1) # m is leading dimension
if b_major == "n":
b_order = (2, 1, 0) # n is leading dimension
elif b_major == "k":
b_order = (2, 0, 1) # k is leading dimension
if c_major == "n":
c_order = (2, 1, 0) # n is leading dimension
elif c_major == "m":
c_order = (2, 0, 1) # m is leading dimension
a = make_fake_compact_tensor(
a_dtype, a_shape, stride_order=a_order, assumed_align=16
)
b = make_fake_compact_tensor(
b_dtype, b_shape, stride_order=b_order, assumed_align=16
)
c = make_fake_compact_tensor(
c_dtype, c_shape, stride_order=c_order, assumed_align=16
)
return cute.compile(
bmm, gemm_op, a, b, c, max_active_clusters, stream, epilogue_op, options=options
)
def prepare_tensors(
mnkl: Tuple[int, int, int, int],
ab_dtype: Type[cutlass.Numeric],
@@ -1811,7 +1763,6 @@ def run(
iterations: int = 1,
skip_ref_check: bool = False,
use_cold_l2: bool = False,
use_tvm_ffi: bool = False,
benchmark: bool = False,
**kwargs,
):
@@ -1853,8 +1804,6 @@ def run(
:type skip_ref_check: bool, optional
:param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False.
:type use_cold_l2: bool, optional
:param use_tvm_ffi: Whether to use TVM FFI for the kernel, defaults to False.
:type use_tvm_ffi: bool, optional
:param benchmark: Whether to only benchmark the kernel, defaults to False.
:type benchmark: bool, optional
:raises RuntimeError: If CUDA GPU is not available.
@@ -1874,16 +1823,15 @@ def run(
print(f"Iterations: {iterations}")
print(f"Skip reference checking: {skip_ref_check}")
print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}")
print(f"Use TVM FFI: {'True' if use_tvm_ffi else 'False'}")
import torch
from cutlass.torch import dtype as torch_dtype
# Build GEMM object
gemm = PersistentDenseGemmKernel(
gemm_op = PersistentDenseGemmKernel(
acc_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store
)
can_implement = gemm.can_implement(
can_implement = gemm_op.can_implement(
mnkl, ab_dtype, c_dtype, a_major, b_major, c_major
)
if not can_implement:
@@ -1906,29 +1854,32 @@ def run(
cluster_shape_mn[0] * cluster_shape_mn[1]
)
options = []
if use_tvm_ffi:
options.append("--enable-tvm-ffi")
compiled_fn = compile_bmm(
gemm,
ab_dtype,
ab_dtype,
c_dtype,
a_major,
b_major,
c_major,
max_active_clusters,
current_stream,
options=",".join(options),
)
# Run and verify BMM with torch
a, b, c = prepare_tensors(mnkl, ab_dtype, c_dtype, a_major, b_major, c_major)
# Leading dim is 2
leading_dim_a = 2 if a_major == "k" else 1
leading_dim_b = 1 if b_major == "k" else 2
leading_dim_c = 2 if c_major == "n" else 1
a_ = from_dlpack(a).mark_layout_dynamic(leading_dim=leading_dim_a)
b_ = from_dlpack(b).mark_layout_dynamic(leading_dim=leading_dim_b)
c_ = from_dlpack(c).mark_layout_dynamic(leading_dim=leading_dim_c)
compiled_fn = cute.compile(
bmm,
gemm_op,
a_,
b_,
c_,
max_active_clusters,
current_stream,
epilogue_op=lambda x: x,
)
if not skip_ref_check:
# Use small random number for deterministic result for reference check
compiled_fn(a, b, c, torch_stream)
compiled_fn(a_, b_, c_, current_stream)
# Manually quantize to be comparable
ref = (
@@ -1954,7 +1905,10 @@ def run(
c_major,
init_random=not init_normal,
)
return testing.JitArguments(a, b, c, torch_stream)
a_ = from_dlpack(a).mark_layout_dynamic(leading_dim=leading_dim_a)
b_ = from_dlpack(b).mark_layout_dynamic(leading_dim=leading_dim_b)
c_ = from_dlpack(c).mark_layout_dynamic(leading_dim=leading_dim_c)
return testing.JitArguments(a_, b_, c_, current_stream)
workspace_count = 1
if use_cold_l2:
@@ -2043,12 +1997,6 @@ def prepare_parser():
default=False,
help="Use circular buffer tensor sets to ensure L2 cold cache",
)
parser.add_argument(
"--use_tvm_ffi",
action="store_true",
default=False,
help="Enable TVM FFI for the kernel, defaults to False using CuTe DSL's native runtime",
)
return parser
@@ -2090,7 +2038,6 @@ if __name__ == "__main__":
args.iterations,
args.skip_ref_check,
args.use_cold_l2,
args.use_tvm_ffi,
args.benchmark,
)
print("PASS")

170
examples/python/CuTeDSL/cute/tvm_ffi/ampere_gemm_with_fake_tensor.py Normal file
View File

@@ -0,0 +1,170 @@
# 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 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 ampere.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}")

10
examples/python/CuTeDSL/cute/fake_tensor.py → examples/python/CuTeDSL/cute/tvm_ffi/compile_with_fake_tensor.py
View File

@@ -22,26 +22,26 @@ def run():
shape = (3, 4)
a = make_fake_compact_tensor(cutlass.Float16, (3, 4), stride_order=(1, 0))
cute.compile(print_tensor_type, a)
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)
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)
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)
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)
cute.compile(print_tensor_type, a, options="--enable-tvm-ffi")
if __name__ == "__main__":

2
examples/python/CuTeDSL/cute/tvm_ffi/requirements.txt Normal file
View File

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

46
media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst
View File

@@ -4,7 +4,8 @@
Compile with TVM FFI
====================
Apache TVM FFI is an open ABI and FFI for machine learning systems. More information can be found in the `official documentation <https://tvm.apache.org/ffi/>`_.
Apache TVM FFI is an open ABI and FFI for machine learning systems. More information can be found in
the `official documentation <https://tvm.apache.org/ffi/>`_.
To install TVM FFI, you can run the following command:
@@ -14,7 +15,9 @@ To install TVM FFI, you can run the following command:
# optional package for improved torch tensor calling performance
pip install torch-c-dlpack-ext
In |DSL|, TVM FFI can be enabled as an option for JIT-compiled functions. Using TVM FFI can lead to faster JIT function invocation and provides better interoperability with machine learning frameworks (e.g., directly take ``torch.Tensor`` as arguments).
In |DSL|, TVM FFI can be enabled as an option for JIT-compiled functions. Using TVM FFI can lead to faster
JIT function invocation and provides better interoperability with machine learning frameworks
(e.g., directly take ``torch.Tensor`` as arguments).
Enable Apache TVM FFI in |DSL|
@@ -40,7 +43,8 @@ There are two ways to enable TVM FFI in |DSL|:
Note that the object returned by ``cute.compile`` is a Python function specific to TVM FFI.
2. Alternatively, you can enable TVM FFI globally by setting the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. Please note that this setting will apply to all JIT compilations within the environment.
2. Alternatively, you can enable TVM FFI globally by setting the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``.
Please note that this setting will apply to all JIT compilations within the environment.
Minimizing Host Overhead
@@ -114,6 +118,37 @@ The fake tensor is a placeholder that mimics the interface of a real tensor but
It is used in compilation or testing scenarios where only shape/type/layout information is needed.
All attempts to access or mutate data will raise errors.
Interoperability with `from_dlpack`
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The Fake Tensor flow supports more flexible constraints on Tensor arguments than the `from_dlpack` flow.
When fake tensor is used, it's recommended to use TVM FFI backend as it supports more flexible constraints on
Tensor arguments than the `from_dlpack` flow.
For instance, fake tensor can specify per-mode static shape or constraints on shape and strides which is not supported by
`from_dlpack`. It's expected that JIT function compiled with fake tensor may have different ABI with tensor converted
with `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 mismatched ABI, it's recommended to use TVM FFI when fake tensor is used for compilation.
Note on Stride Order
~~~~~~~~~~~~~~~~~~~~
@@ -129,7 +164,8 @@ 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:
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
@@ -288,6 +324,7 @@ composed of the types that are supported by TVM FFI. The example below shows how
example_add_one_with_tuple()
Supported types
---------------
@@ -316,6 +353,7 @@ The TVM FFI function supports the following |DSL|-specific types as arguments:
* - Tuple of types (e.g. ``Tuple[cute.Tensor, cute.Tensor, cutlass.Int32]``)
- Python tuple of corresponding call-time types.
Error handling
--------------

3
python/CuTeDSL/cutlass/cute/runtime.py
View File

@@ -617,7 +617,8 @@ def make_fake_compact_tensor(
:param shape: Shape of the tensor.
:type shape: tuple[int, ...]
:param stride_order: Order in which strides (memory layout) are assigned to the tensor dimensions.
If None, the default layout is col-major. Otherwise, it should be a permutation of the dimension indices.
If None, the default layout is left-to-right order (known as column-major order for flatten layout).
Otherwise, it should be a permutation order of the dimension indices.
:type stride_order: tuple[int, ...], optional
:param memspace: Memory space where the fake tensor resides. Optional.
:type memspace: str, optional