v4.4 update. (#2979)
This commit is contained in:
Junkai-Wu
@@ -0,0 +1,260 @@
|
||||
# Copyright (c) 2025 - 2026 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.
|
||||
|
||||
from functools import partial
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import cutlass.jax as cjax
|
||||
import cuda.bindings.driver as cuda
|
||||
|
||||
"""
|
||||
Examples of calling CuTe DSL from jax.jit function using cutlass_call.
|
||||
|
||||
cutlass_call is a Jax primitive the enables calling of CuTe DSL kernels within a
|
||||
a jit-compiled Jax function. During the lowering process cutlass_call will
|
||||
trigger compilation of the kernel and embed it into the HLO computation. It can
|
||||
then be efficiently launched by XLA without callback to Python.
|
||||
|
||||
This example assumes familiarity with CuTe DSL concepts such as layouts and
|
||||
dynamic shapes.
|
||||
|
||||
To run this example:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Run with addition operation
|
||||
python examples/jax/cutlass_call_basic.py
|
||||
"""
|
||||
|
||||
|
||||
# This is a typical CuTe DSL kernel function that accepts both tensor and scalar values.
|
||||
@cute.jit
|
||||
def launch(
|
||||
A: cute.Tensor,
|
||||
B: cute.Tensor,
|
||||
x: cute.Int32,
|
||||
y: cute.Int32,
|
||||
C: cute.Tensor,
|
||||
D: cute.Tensor,
|
||||
stream: cuda.CUstream,
|
||||
):
|
||||
# Print layouts
|
||||
print("A layout: ", A.layout)
|
||||
print("B layout: ", B.layout)
|
||||
print("C layout: ", C.layout)
|
||||
print("D layout: ", D.layout)
|
||||
cute.printf("A layout: {}", A.layout)
|
||||
cute.printf("B layout: {}", B.layout)
|
||||
cute.printf("C layout: {}", C.layout)
|
||||
cute.printf("D layout: {}", D.layout)
|
||||
cute.printf("")
|
||||
|
||||
# Print non-tensor values
|
||||
print("X is: ", x)
|
||||
print("Y is: ", y)
|
||||
cute.printf("X is: {}", x)
|
||||
cute.printf("Y is: {}", y)
|
||||
print()
|
||||
|
||||
|
||||
# cutlass_call uses a fixed function signature to pass arguments between Jax and CuTeDSL kernel.
|
||||
#
|
||||
# Function Signature Requirement:
|
||||
# stream, inputs, outputs, *, kwargs...
|
||||
#
|
||||
# The first argument must be the CUstream that the kernel is run. This stream is managed by the XLA runtime
|
||||
# and is necessary to schedule and synchronize launches with the rest of your computation.
|
||||
#
|
||||
# The second set of arguments are the Jax arrays for inputs and outputs. Inputs must be passed before
|
||||
# outputs.
|
||||
#
|
||||
# Lastly static arguments (i.e. static_argnums or static_argnames) values are passed as keyword only arguments
|
||||
# by name.
|
||||
#
|
||||
# The the kernel does not match this signature a wrapper functions like the one shown below can be written
|
||||
# or an inline lambda function can be used to rebind the arguments into the appropriate order.
|
||||
@cute.jit
|
||||
def launch_jax_wrapper(
|
||||
stream: cuda.CUstream,
|
||||
A: cute.Tensor,
|
||||
B: cute.Tensor,
|
||||
C: cute.Tensor,
|
||||
D: cute.Tensor,
|
||||
*,
|
||||
x: cute.Int32,
|
||||
y: cute.Int32,
|
||||
):
|
||||
launch(A, B, x, y, C, D, stream)
|
||||
|
||||
|
||||
@cute.jit
|
||||
def launch_aliased(
|
||||
A: cute.Tensor, B: cute.Tensor, x: cute.Int32, y: cute.Int32, stream: cuda.CUstream
|
||||
):
|
||||
# Print layouts
|
||||
print("A layout: ", A.layout)
|
||||
print("B layout: ", B.layout)
|
||||
cute.printf("A layout: {}", A.layout)
|
||||
cute.printf("B layout: {}", B.layout)
|
||||
cute.printf("")
|
||||
|
||||
# Print non-tensor values
|
||||
print("X is: ", x)
|
||||
print("Y is: ", y)
|
||||
cute.printf("X is: {}", x)
|
||||
cute.printf("Y is: {}", y)
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@partial(jax.jit, static_argnums=[2, 3])
|
||||
def run_cutlass_kernel(a, b, x, y):
|
||||
call = cjax.cutlass_call(
|
||||
launch_jax_wrapper,
|
||||
# Jax requires output shapes/dtype information for each output
|
||||
output_shape_dtype=(
|
||||
jax.ShapeDtypeStruct(a.shape, a.dtype),
|
||||
jax.ShapeDtypeStruct(b.shape, a.dtype),
|
||||
),
|
||||
# Static jit arguments are passed via additional keyword arguments
|
||||
x=x,
|
||||
y=y,
|
||||
)
|
||||
|
||||
# Returned value is a callable to invoke the kernel passing only jax arrays.
|
||||
return call(a, b)
|
||||
|
||||
print("\nExample: example_basic_call_from_jit")
|
||||
A = jnp.zeros((512, 32, 64))
|
||||
B = jnp.zeros((1, 256, 64, 128))
|
||||
C, D = run_cutlass_kernel(A, B, 0, 1)
|
||||
|
||||
@partial(jax.jit, static_argnums=[2, 3])
|
||||
def run_cutlass_kernel_lambda(a, b, x, y):
|
||||
call = cjax.cutlass_call(
|
||||
# A lambda function may be used to wrap and bind arguments passed by jax
|
||||
# to the kernel. Alternatively you can wrap using another separate cute.jit
|
||||
# function.
|
||||
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
|
||||
# Jax requires output shapes/dtype information for each output
|
||||
output_shape_dtype=(
|
||||
jax.ShapeDtypeStruct(a.shape, a.dtype),
|
||||
jax.ShapeDtypeStruct(b.shape, a.dtype),
|
||||
),
|
||||
# Static jit arguments are passed via additional keyword arguments
|
||||
x=x,
|
||||
y=y,
|
||||
)
|
||||
|
||||
# Returned value is a callable to invoke the kernel passing only jax arrays.
|
||||
return call(a, b)
|
||||
|
||||
print("\nExample: run_cutlass_kernel_lambda")
|
||||
A = jnp.zeros((512, 32, 64))
|
||||
B = jnp.zeros((1, 256, 64, 128))
|
||||
C, D = run_cutlass_kernel_lambda(A, B, 1, 2)
|
||||
|
||||
@partial(jax.jit, static_argnums=[2, 3])
|
||||
def run_cutlass_kernel_static_shapes(a, b, x, y):
|
||||
call = cjax.cutlass_call(
|
||||
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
|
||||
output_shape_dtype=(
|
||||
jax.ShapeDtypeStruct(a.shape, a.dtype),
|
||||
jax.ShapeDtypeStruct(b.shape, a.dtype),
|
||||
),
|
||||
# By default cutlass_call will treat all tensors as dynamic shape.
|
||||
# Dynamic shapes are often expected for kernels so this default ensures
|
||||
# the broadest support. If you know that a kernel can accept fully static
|
||||
# tensors then you can enable this flag to pass all tensors shapes and
|
||||
# layouts known at compile time.
|
||||
use_static_tensors=True,
|
||||
x=x,
|
||||
y=y,
|
||||
)
|
||||
return call(a, b)
|
||||
|
||||
print("\nExample: run_cutlass_kernel_static_shapes")
|
||||
A = jnp.zeros((512, 32, 64))
|
||||
B = jnp.zeros((1, 256, 64, 128))
|
||||
C, D = run_cutlass_kernel_static_shapes(A, B, 3, 4)
|
||||
|
||||
@partial(jax.jit, static_argnums=[2, 3])
|
||||
def run_cutlass_kernel_with_modes(a, b, x, y):
|
||||
call = cjax.cutlass_call(
|
||||
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
|
||||
output_shape_dtype=(
|
||||
jax.ShapeDtypeStruct(a.shape, a.dtype),
|
||||
jax.ShapeDtypeStruct(b.shape, a.dtype),
|
||||
),
|
||||
# The modes of the layout for each tensor can be specified using the
|
||||
# TensorSpec. By default modes will align with the physical layout
|
||||
# but can be mapped to specific index position. If None is passed
|
||||
# then the default mode is assumed for that tensor.
|
||||
#
|
||||
# Individual static/dynamic settings may also be applied. For example
|
||||
# a specific tensor can be marked to have static shape.
|
||||
input_spec=(
|
||||
cjax.TensorSpec(mode=(1, 0, 2), static=True),
|
||||
cjax.TensorSpec(mode=(3, 1, 2, 0)),
|
||||
),
|
||||
output_spec=(None, cjax.TensorSpec(mode=(0, 1, 3, 2))),
|
||||
x=x,
|
||||
y=y,
|
||||
)
|
||||
return call(a, b)
|
||||
|
||||
print("\nExample: run_cutlass_kernel_with_modes")
|
||||
A = jnp.zeros((512, 32, 64))
|
||||
B = jnp.zeros((1, 256, 64, 128))
|
||||
C, D = run_cutlass_kernel_with_modes(A, B, 5, 6)
|
||||
|
||||
@partial(jax.jit, static_argnums=[2, 3], donate_argnums=[0, 1])
|
||||
def run_cutlass_kernel_aliased_outputs(a, b, x, y):
|
||||
call = cjax.cutlass_call(
|
||||
lambda stream, a, b, *, x, y: launch_aliased(a, b, x, y, stream),
|
||||
output_shape_dtype=(
|
||||
jax.ShapeDtypeStruct(a.shape, a.dtype),
|
||||
jax.ShapeDtypeStruct(b.shape, b.dtype),
|
||||
),
|
||||
# Can specify the input tensors that are aliasing outputs of this call.
|
||||
# To avoid allocating separate output buffers. This is useful for kernels
|
||||
# that update a tensor.
|
||||
input_output_aliases={0: 0, 1: 1},
|
||||
x=x,
|
||||
y=y,
|
||||
)
|
||||
return call(a, b)
|
||||
|
||||
print("\nExample: run_cutlass_kernel_aliased_outputs")
|
||||
A = jnp.zeros((512, 32, 64))
|
||||
B = jnp.zeros((1, 256, 64, 128))
|
||||
A, B = run_cutlass_kernel_aliased_outputs(A, B, 7, 8)
|
||||
@@ -0,0 +1,168 @@
|
||||
# Copyright (c) 2025 - 2026 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 pytest
|
||||
from functools import partial
|
||||
import argparse
|
||||
|
||||
import cuda.bindings.driver as cuda
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
from jax import export
|
||||
|
||||
from cutlass.jax import cutlass_call, get_export_disabled_safety_checks
|
||||
from cutlass.jax.testing import create_tensor
|
||||
|
||||
"""
|
||||
Examples of using jax.export APIs with functions using cutlass_call.
|
||||
|
||||
This example demonstrates the use of jax.export with CuTe DSL kernel. It assumes
|
||||
familiarity with CuTe DSL concepts such as layouts and dynamic shapes as well as
|
||||
Jax's exporting and serialization features:
|
||||
https://docs.jax.dev/en/latest/export/index.html#export
|
||||
|
||||
To run this example:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Run with defaults
|
||||
python examples/jax/cutlass_call_export.py
|
||||
|
||||
# Run with shape (1024, 512)
|
||||
python examples/jax/cutlass_call_export.py --M 1024 --N 512
|
||||
|
||||
# Export with symbolic shapes.
|
||||
python examples/jax/cutlass_call_export.py --export_symbolic
|
||||
"""
|
||||
|
||||
|
||||
@cute.kernel
|
||||
def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
bidx, _, _ = cute.arch.block_idx()
|
||||
bdim, _, _ = cute.arch.block_dim()
|
||||
|
||||
thread_idx = bidx * bdim + tidx
|
||||
|
||||
m, n = gA.shape
|
||||
ni = thread_idx % n
|
||||
mi = thread_idx // n
|
||||
|
||||
a_val = gA[mi, ni]
|
||||
b_val = gB[mi, ni]
|
||||
gC[mi, ni] = a_val + b_val
|
||||
|
||||
|
||||
@cute.jit
|
||||
def launch(stream: cuda.CUstream, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
|
||||
print("mA: ", mA.layout)
|
||||
print("mB: ", mB.layout)
|
||||
print("mC: ", mC.layout)
|
||||
num_threads_per_block = 256
|
||||
m, n = mA.shape
|
||||
kernel(mA, mB, mC).launch(
|
||||
grid=((m * n) // num_threads_per_block, 1, 1),
|
||||
block=(num_threads_per_block, 1, 1),
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
|
||||
def run_example(M, N, export_symbolic_shapes):
|
||||
@jax.jit
|
||||
def f(a, b):
|
||||
call = cutlass_call(launch, output_shape_dtype=a)
|
||||
return jax.nn.sigmoid(call(a, b))
|
||||
|
||||
@jax.jit
|
||||
def ref_f(a, b):
|
||||
return jax.nn.sigmoid(a + b)
|
||||
|
||||
# Symbolic or partially shapes are supported by cutlass_call and cute.Tensor
|
||||
# This allows export of functions calling Cut eDSL kernels w/o having to re-compile
|
||||
# the kernel for each new shape.
|
||||
if export_symbolic_shapes:
|
||||
a, b = export.symbolic_shape("a, b")
|
||||
export_shape_dtype = jax.ShapeDtypeStruct((a, b), jnp.float32)
|
||||
else:
|
||||
export_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32)
|
||||
|
||||
print("Exporting with input signature: ")
|
||||
print(f"({export_shape_dtype}, {export_shape_dtype})")
|
||||
|
||||
# jax.export can be used to export a jit function containing cutlass_call.
|
||||
# The function get_export_disabled_safety_checks() returns a list of custom
|
||||
# call targets that are used by cutlass_call not part of Jax's built-in
|
||||
# list of stable custom calls.
|
||||
exported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())
|
||||
traced = exported(export_shape_dtype, export_shape_dtype)
|
||||
|
||||
# Serialize the computation to a byte blob.
|
||||
blob = traced.serialize()
|
||||
print(f"Serialized computation is {len(blob)} bytes.")
|
||||
|
||||
# Deserialize and run
|
||||
rehydrated = export.deserialize(blob)
|
||||
|
||||
key = jax.random.key(1123)
|
||||
a_key, b_key = jax.random.split(key, 2)
|
||||
|
||||
a = create_tensor((M, N), dtype=jnp.float32, key=a_key)
|
||||
b = create_tensor((M, N), dtype=jnp.float32, key=b_key)
|
||||
c = rehydrated.call(a, b)
|
||||
assert jnp.allclose(c, ref_f(a, b))
|
||||
|
||||
# If the computation was exported with dynamic shapes then we can also
|
||||
# call it with different shapes. The kernel will not be re-compiled
|
||||
# even though the shapes are changing.
|
||||
if export_symbolic_shapes:
|
||||
a = create_tensor((M * 2, N * 4), dtype=jnp.float32, key=a_key)
|
||||
b = create_tensor((M * 2, N * 4), dtype=jnp.float32, key=b_key)
|
||||
c = rehydrated.call(a, b)
|
||||
assert jnp.allclose(c, ref_f(a, b))
|
||||
|
||||
a = create_tensor((M * 4, N * 4), dtype=jnp.float32, key=a_key)
|
||||
b = create_tensor((M * 4, N * 4), dtype=jnp.float32, key=b_key)
|
||||
c = rehydrated.call(a, b)
|
||||
assert jnp.allclose(c, ref_f(a, b))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Demonstration of using jax.export with functions with cutlass_call"
|
||||
)
|
||||
parser.add_argument("--M", default=512, type=int)
|
||||
parser.add_argument("--N", default=256, type=int)
|
||||
parser.add_argument("--export_symbolic", action="store_true")
|
||||
|
||||
args = parser.parse_args()
|
||||
run_example(args.M, args.N, args.export_symbolic)
|
||||
print("PASS")
|
||||
@@ -0,0 +1,140 @@
|
||||
# Copyright (c) 2025 - 2026 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.
|
||||
|
||||
from functools import partial
|
||||
import argparse
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
|
||||
from jax.experimental.custom_partitioning import custom_partitioning
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import cutlass.jax as cjax
|
||||
from cutlass.jax.testing import create_tensor
|
||||
import cuda.bindings.driver as cuda
|
||||
|
||||
|
||||
"""
|
||||
Examples of combining jax.jit and jax.shard_map for sharding and executing kernels
|
||||
across multiple GPU devices.
|
||||
|
||||
To run this example:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Run with addition operation
|
||||
python examples/jax/cutlass_call_sharding.py
|
||||
"""
|
||||
|
||||
|
||||
@cute.kernel
|
||||
def kernel(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
bidx, _, _ = cute.arch.block_idx()
|
||||
|
||||
frgA = cute.make_rmem_tensor(cute.size(a, mode=[0]), a.element_type)
|
||||
frgB = cute.make_rmem_tensor(cute.size(b, mode=[0]), b.element_type)
|
||||
frgC = cute.make_rmem_tensor(cute.size(c, mode=[0]), c.element_type)
|
||||
|
||||
cute.autovec_copy(a[None, tidx, bidx], frgA)
|
||||
cute.autovec_copy(b[None, tidx, bidx], frgB)
|
||||
frgC.store(frgA.load() + frgB.load())
|
||||
cute.autovec_copy(frgC, c[None, tidx, bidx])
|
||||
|
||||
|
||||
@cute.jit
|
||||
def launch(
|
||||
stream: cuda.CUstream,
|
||||
a: cute.Tensor,
|
||||
b: cute.Tensor,
|
||||
c: cute.Tensor,
|
||||
):
|
||||
cute.printf("a: {}", a.layout)
|
||||
cute.printf("b: {}", b.layout)
|
||||
cute.printf("c: {}", c.layout)
|
||||
kernel(a, b, c).launch(
|
||||
grid=[a.shape[-1], 1, 1], block=[a.shape[-2], 1, 1], stream=stream
|
||||
)
|
||||
|
||||
|
||||
def run_example():
|
||||
# Create a device mesh with one axis b
|
||||
ngpu = jax.device_count()
|
||||
mesh = jax.make_mesh((ngpu,), "b")
|
||||
|
||||
if ngpu == 1:
|
||||
print("Note: only 1 GPU was detected.")
|
||||
|
||||
# We will shard our 3D tensors over b
|
||||
sharding = P("b", None, None)
|
||||
|
||||
@partial(jax.jit, static_argnums=[0, 1])
|
||||
def allocate_sharded_tensors(shape, dtype):
|
||||
key = jax.random.key(1123)
|
||||
a_key, b_keys = jax.random.split(key, 2)
|
||||
a = create_tensor(shape, dtype, a_key)
|
||||
b = create_tensor(shape, dtype, b_keys)
|
||||
a = jax.lax.with_sharding_constraint(a, NamedSharding(mesh, sharding))
|
||||
b = jax.lax.with_sharding_constraint(b, NamedSharding(mesh, sharding))
|
||||
return a, b
|
||||
|
||||
@jax.jit
|
||||
def compute(a, b):
|
||||
# This jax.shard_map partitions the cutlass_call over the mesh.
|
||||
@partial(
|
||||
jax.shard_map,
|
||||
mesh=mesh,
|
||||
in_specs=(sharding, sharding),
|
||||
out_specs=(sharding, sharding),
|
||||
)
|
||||
def sharded_call(a_block, b_block):
|
||||
call = cjax.cutlass_call(
|
||||
launch,
|
||||
use_static_tensors=True,
|
||||
output_shape_dtype=jax.ShapeDtypeStruct(a_block.shape, a_block.dtype),
|
||||
)
|
||||
ref_result = a_block + b_block
|
||||
return call(a_block, b_block), ref_result
|
||||
|
||||
return sharded_call(a, b)
|
||||
|
||||
# Allocate (32, 16, 64) on each GPU
|
||||
shape = (32 * ngpu, 16, 64)
|
||||
dtype = jnp.float32
|
||||
|
||||
a, b = allocate_sharded_tensors(shape, dtype)
|
||||
c, c_ref = compute(a, b)
|
||||
|
||||
assert jnp.allclose(c, c_ref)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_example()
|
||||
print("PASS")
|
||||
@@ -0,0 +1,329 @@
|
||||
# Copyright (c) 2025 - 2026 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 operator
|
||||
from functools import partial
|
||||
from typing import List, Type
|
||||
|
||||
import cuda.bindings.driver as cuda
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
|
||||
"""
|
||||
An Elementwise Apply Example using CuTe DSL with cutlass.jax.cutlass_call
|
||||
|
||||
This example is similar to examples/ampere/elementwise_apply.py but demonstrates
|
||||
how to run the code in a jax specific way using the cutlass_call primitive. It assumes
|
||||
familiarity with basic CuTe DSL concepts as well as the cutlass_call primitive.
|
||||
|
||||
To run this example:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Run with addition operation
|
||||
python examples/jax/elementwise_apply_example.py --M 1024 --N 512 --op add
|
||||
|
||||
# Run with multiplication operation
|
||||
python examples/ampere/elementwise_apply_example.py --M 1024 --N 512 --op mul
|
||||
|
||||
# Run with subtraction operation
|
||||
python examples/ampere/elementwise_apply_example.py --M 1024 --N 512 --op sub
|
||||
"""
|
||||
|
||||
|
||||
@cute.kernel
|
||||
def elementwise_apply_kernel(
|
||||
op: cutlass.Constexpr,
|
||||
mInputs: List[cute.Tensor],
|
||||
mC: cute.Tensor,
|
||||
cC: cute.Tensor, # coordinate tensor
|
||||
shape: cute.Shape,
|
||||
tv_layout: cute.Layout, # (tid, vid) -> logic coord
|
||||
):
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
bidx, bidy, _ = cute.arch.block_idx()
|
||||
|
||||
###############################################################################
|
||||
# Slice to local tile of thread block
|
||||
###############################################################################
|
||||
blk_crd = ((None, None), (bidx, bidy))
|
||||
|
||||
# Leverage the meta-programming capability of the DSL to slice the tensors for each input
|
||||
# All for loops below on input tensors would be fully unrolled automatically at compile time
|
||||
# logical coord -> memory address
|
||||
gInputs = [t[blk_crd] for t in mInputs] # (TileM, TileN)
|
||||
gC = mC[blk_crd] # (TileM, TileN)
|
||||
gCrd = cC[blk_crd] # (TileM, TileN)
|
||||
|
||||
print("[DSL INFO] Sliced Tensors per thread block:")
|
||||
for i in cutlass.range_constexpr(len(gInputs)):
|
||||
print(f"[DSL INFO] ctaInputs{i} = {gInputs[i].type}")
|
||||
print(f"[DSL INFO] gC = {gC.type}")
|
||||
print(f"[DSL INFO] gCrd = {gCrd.type}")
|
||||
|
||||
###############################################################################
|
||||
# Compose with thread block TV layout to map thread & value indices to memory address
|
||||
###############################################################################
|
||||
# (tid, vid) -> memory address
|
||||
tidfrgInputs = [cute.composition(t, tv_layout) for t in gInputs]
|
||||
tidfrgC = cute.composition(gC, tv_layout)
|
||||
tidfrgCrd = cute.composition(gCrd, tv_layout)
|
||||
|
||||
# repeat None like vid to remove hierarchy of layout
|
||||
thr_crd = (tidx, cute.repeat_like(None, tidfrgInputs[0][1]))
|
||||
|
||||
###############################################################################
|
||||
# Slice to local tile of thread
|
||||
###############################################################################
|
||||
# vid -> address
|
||||
thrInputs = [t[thr_crd] for t in tidfrgInputs] # (V)
|
||||
thrC = tidfrgC[thr_crd] # (V)
|
||||
thrCrd = tidfrgCrd[thr_crd]
|
||||
|
||||
print("[DSL INFO] Sliced Tensors per thread:")
|
||||
for i in cutlass.range_constexpr(len(thrInputs)):
|
||||
print(f"[DSL INFO] thrInputs{i} = {thrInputs[i].type}")
|
||||
print(f"[DSL INFO] thrC = {thrC.type}")
|
||||
print(f"[DSL INFO] thrCrd = {thrCrd.type}")
|
||||
|
||||
###############################################################################
|
||||
# Compute predicate for out of boundary checks
|
||||
###############################################################################
|
||||
frgPred = cute.make_fragment(thrCrd.shape, cutlass.Boolean)
|
||||
print(f"[DSL INFO] frgPred = {frgPred.type}")
|
||||
|
||||
for i in cutlass.range_constexpr(cute.size(frgPred)):
|
||||
frgPred[i] = cute.elem_less(thrCrd[i], shape)
|
||||
|
||||
# if tidx == 0 and bidx == 0:
|
||||
# cute.print_tensor(frgPred)
|
||||
|
||||
##########################################################
|
||||
# Load data and compute result
|
||||
##########################################################
|
||||
|
||||
# Load data before use. The compiler will optimize the copy and load
|
||||
# operations to convert some memory ld/st into register uses.
|
||||
result = op(*[thrInput.load() for thrInput in thrInputs])
|
||||
thrC.store(result)
|
||||
|
||||
|
||||
@cute.jit
|
||||
def elementwise_apply(
|
||||
op: cutlass.Constexpr, inputs, result: cute.Tensor, stream: cuda.CUstream
|
||||
):
|
||||
"""CUDA kernel applying binary operator on each element of two n-D input tensors in
|
||||
CuTe Python and store to result tensor.
|
||||
|
||||
:param op: Binary operator or lambda function to apply element-wise
|
||||
:type op: cutlass.Constexpr
|
||||
:param a: First input tensor
|
||||
:type a: cute.Tensor
|
||||
:param b: Second input tensor
|
||||
:type b: cute.Tensor
|
||||
:param result: Output tensor to store the results of op(a, b)
|
||||
:type result: cute.Tensor
|
||||
:return: None
|
||||
:rtype: None
|
||||
"""
|
||||
|
||||
# Baseline: naive TV layout
|
||||
# * mA layout: (4096, 4096):(4096, 1)
|
||||
# * TV layout map to (512, 4) tile
|
||||
# * tidx maps to mode-0 but input layout is contiguous on mode-1, performance will be bad
|
||||
# tv_layout = cute.make_layout((128, (4, 4)), stride=(4, (512, 1)))
|
||||
# cta_tiler = (512, 4)
|
||||
|
||||
# Opt-1: better TV layout with better 1D thread layout (SOL with 1D thread layout)
|
||||
# * mA layout: (4096, 4096):(4096, 1)
|
||||
# * TV layout map to (4, 512) tile
|
||||
# * tidx maps to mode-1 which is leading mode of input tensor for coalesced load
|
||||
# tv_layout = cute.make_layout((128, (4, 4)), stride=(16, (4, 1)))
|
||||
# cta_tiler = (4, 512)
|
||||
|
||||
# Opt-2: 2D tile but worse
|
||||
# * mA layout: (4096, 4096):(4096, 1)
|
||||
# * TV layout map to (128, 16) logical tile
|
||||
# * V layout is bad as contiguous mode is not on right-most
|
||||
# * `cute.copy` only supports vectorize when stride-1 of v-layout on right-most )
|
||||
# tv_layout = cute.make_layout(((32, 4), (4, 4)), stride=((4, 512), (1, 128)))
|
||||
# cta_tiler = (128, 16)
|
||||
|
||||
# Opt-3: SOL with 2D thread tile
|
||||
# * mA layout: (4096, 4096):(4096, 1)
|
||||
# * TV layout map to (64, 256) logical tile
|
||||
# * tidx maps to mode-1 and input layout is contiguous on mode-1 for coalesced load-store
|
||||
|
||||
# Use 128bit(16B) load as canonicalized form of val_layout then recast to target element-type
|
||||
coalesced_ldst_bytes = 16
|
||||
|
||||
# Compile time validation: expect same element type for all input tensors
|
||||
assert all(t.element_type == inputs[0].element_type for t in inputs)
|
||||
dtype = inputs[0].element_type
|
||||
|
||||
thr_layout = cute.make_ordered_layout((4, 64), order=(1, 0))
|
||||
val_layout = cute.make_ordered_layout((16, coalesced_ldst_bytes), order=(1, 0))
|
||||
val_layout = cute.recast_layout(dtype.width, 8, val_layout)
|
||||
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
|
||||
|
||||
print("[DSL INFO] Input Tensors:")
|
||||
for i, t in enumerate(inputs):
|
||||
print(f"[DSL INFO] inputs{i} = {t}")
|
||||
print(f"[DSL INFO] result = {result}")
|
||||
|
||||
print("[DSL INFO] Tiling Parameters:")
|
||||
print(f"[DSL INFO] tiler_mn = {tiler_mn} per thread block")
|
||||
print(f"[DSL INFO] tv_layout = {tv_layout}")
|
||||
|
||||
print("[DSL INFO] Tiled Tensors:")
|
||||
mInputs = [cute.zipped_divide(input, tiler_mn) for input in inputs]
|
||||
# ((TileM, TileN), (RestM, RestN))
|
||||
mC = cute.zipped_divide(result, tiler_mn)
|
||||
|
||||
# (RestM, RestN) -> (RestN, RestM)
|
||||
remap_block = cute.make_ordered_layout(
|
||||
cute.select(mInputs[0].shape[1], mode=[1, 0]), order=(1, 0)
|
||||
)
|
||||
for i, t in enumerate(mInputs):
|
||||
print(f"[DSL INFO] gInputs{i} = {mInputs[i]}")
|
||||
mInputs[i] = cute.composition(t, (None, remap_block))
|
||||
print(f"[DSL INFO] gInputs{i} (remapped) = {mInputs[i]}")
|
||||
|
||||
mC = cute.composition(mC, (None, remap_block))
|
||||
print(f"[DSL INFO] gC = {mC}")
|
||||
|
||||
idC = cute.make_identity_tensor(result.shape)
|
||||
cC = cute.zipped_divide(idC, tiler=tiler_mn)
|
||||
print(f"[DSL INFO] coord tensor = {cC}")
|
||||
|
||||
# Launch the kernel asynchronously
|
||||
# Group input tensors into a list as a single argument
|
||||
elementwise_apply_kernel(op, mInputs, mC, cC, result.shape, tv_layout).launch(
|
||||
# Compute production at each mode of mC.shape[1] to get multi-dimensional grid size
|
||||
grid=cute.product_each(mC.shape[1]),
|
||||
block=[cute.size(tv_layout, mode=[0]), 1, 1],
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
|
||||
@cutlass.dsl_user_op
|
||||
def leaky_relu(x, alpha, *, loc=None, ip=None):
|
||||
return cute.where(x > 0, x, alpha * x, loc=loc, ip=ip)
|
||||
|
||||
|
||||
def leaky_relu_ref(x, alpha):
|
||||
import jax.numpy as jnp
|
||||
|
||||
return jnp.where(x > 0, x, alpha * x)
|
||||
|
||||
|
||||
def run_and_verify(op, M, N, dtype, skip_ref_check=False):
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import cutlass.jax as cjax
|
||||
import cutlass.jax.testing as testing
|
||||
|
||||
if op == "leaky_relu":
|
||||
op = partial(leaky_relu, alpha=0.01)
|
||||
ref_op = partial(leaky_relu_ref, alpha=0.01)
|
||||
num_inputs = 1
|
||||
else:
|
||||
op = getattr(operator, op)
|
||||
ref_op = op
|
||||
num_inputs = 2
|
||||
|
||||
# This jax function is transformed using jax.jit to compile its contents
|
||||
# into an efficient HLO executable.
|
||||
@partial(jax.jit, static_argnums=[1])
|
||||
def jax_function(inputs, op):
|
||||
call = cjax.cutlass_call(
|
||||
# Bind jax arguments to kernel signature
|
||||
lambda stream, inputs, output, *, op: elementwise_apply(
|
||||
op, inputs, output, stream
|
||||
),
|
||||
# Specify output shape/dtype of result
|
||||
output_shape_dtype=jax.ShapeDtypeStruct(inputs[0].shape, inputs[0].dtype),
|
||||
# Pass static/constexpr values as kwargs
|
||||
op=op,
|
||||
)
|
||||
|
||||
# Call the kernel!
|
||||
return call(inputs)
|
||||
|
||||
@partial(jax.jit, static_argnums=[1])
|
||||
def jax_ref_function(inputs, op):
|
||||
return op(*inputs)
|
||||
|
||||
print("\nRunning Elementwise Apply test with:")
|
||||
print(f"Tensor dimensions: [{M}, {N}]")
|
||||
print(f"Input and Output Data type: {dtype}")
|
||||
|
||||
jax_dtype = cjax.cutlass_to_jax_dtype(dtype)
|
||||
keys = jax.random.split(jax.random.key(1435), num_inputs)
|
||||
inputs = [testing.create_tensor((M, N), jax_dtype, key) for key in keys]
|
||||
|
||||
print("Input tensor shapes:")
|
||||
for i in range(num_inputs):
|
||||
print(f"inputs[{i}]: {inputs[i].shape}, dtype: {inputs[i].dtype}")
|
||||
|
||||
epsilon = 1.2
|
||||
if op in (operator.truediv, operator.floordiv):
|
||||
inputs[1] = jnp.where(inputs[1] == 0, epsilon, inputs[1])
|
||||
|
||||
# Call the jax.jit function which will compile the kernel
|
||||
c = jax_function(inputs, op)
|
||||
|
||||
if not skip_ref_check:
|
||||
print("Executing elementwise apply kernel...")
|
||||
c = jax_function(inputs, op)
|
||||
print("Verifying results...")
|
||||
assert jnp.allclose(ref_op(*inputs), c)
|
||||
print("Results verified successfully!")
|
||||
print(f"First few elements of result: \n{c[:3, :3]}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Demonstration of calling a kernel with cutlass_call"
|
||||
)
|
||||
parser.add_argument("--M", default=4096, type=int)
|
||||
parser.add_argument("--N", default=4096, type=int)
|
||||
parser.add_argument("--op", default="add", type=str)
|
||||
parser.add_argument("--skip_ref_check", action="store_true")
|
||||
|
||||
args = parser.parse_args()
|
||||
run_and_verify(
|
||||
args.op,
|
||||
args.M,
|
||||
args.N,
|
||||
dtype=cutlass.Float32,
|
||||
skip_ref_check=args.skip_ref_check,
|
||||
)
|
||||
print("\nPASS")
|
||||
Reference in New Issue
Block a user