v4.3 tag release update. (#2789)
This commit is contained in:
Junkai-Wu
@@ -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()
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user