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

v4.4 update. (#2979)

Browse Source
This commit is contained in:
Junkai-Wu
2026-01-24 11:46:17 -05:00
committed by GitHub
parent 2fafefb7b9
commit 9fba3195f9
293 changed files with 46343 additions and 2995 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py
+6 -6
View File
@@ -37,21 +37,19 @@ To run this example:
.. code-block:: bash
python examples/cute/tvm_ffi/aot_export.py
python cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_export.py
# run example to use in torch
python examples/cute/tvm_ffi/aot_use_in_torch.py
python cutlass_ir/compiler/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
python cutlass_ir/compiler/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
bash cutlass_ir/compiler/python/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
@@ -69,6 +67,8 @@ def add_one(a: cute.Tensor, b: cute.Tensor):
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")