v4.5 dev update. (#3153)
This commit is contained in:
Junkai-Wu
@@ -26,45 +26,45 @@
|
||||
# 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:
|
||||
This example demonstrates three export modes:
|
||||
|
||||
1. Concrete shapes -- shapes are fixed constants baked into the export.
|
||||
2. Unconstrained symbolic shapes ("a, b")
|
||||
3. Constrained symbolic shapes ("32*M, 16*N")
|
||||
|
||||
The JAX function being exported is the same in all three cases; only the
|
||||
shape specification passed to jax.export differs.
|
||||
|
||||
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
|
||||
python examples/jax/cutlass_call_export.py --M 512 --N 256
|
||||
|
||||
# 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
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
||||
import cuda.bindings.driver as cuda
|
||||
|
||||
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, TensorSpec
|
||||
from cutlass.jax.testing import create_tensor
|
||||
|
||||
|
||||
# Simple element-wise addition kernel: gC[i,j] = gA[i,j] + gB[i,j]
|
||||
@cute.kernel
|
||||
def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
@@ -84,9 +84,6 @@ def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
|
||||
|
||||
@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(
|
||||
@@ -96,63 +93,100 @@ def launch(stream: cuda.CUstream, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Ten
|
||||
)
|
||||
|
||||
|
||||
def run_example(M, N, export_symbolic_shapes):
|
||||
def _export_and_run(f, ref_f, input_shape_dtype, run_shapes):
|
||||
"""Export f, serialize/deserialize, then run on each shape in run_shapes.
|
||||
|
||||
Both inputs (a, b) are assumed to share the same input_shape_dtype.
|
||||
"""
|
||||
print(f"Exporting with input signature: ({input_shape_dtype}, {input_shape_dtype})")
|
||||
|
||||
# jax.export can be used to export a jit function containing cutlass_call.
|
||||
# CUTLASS custom call targets are not on JAX's built-in stable custom-call
|
||||
# allowlist, so we pass them via disabled_checks to suppress that safety check.
|
||||
exported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())
|
||||
traced = exported(input_shape_dtype, input_shape_dtype)
|
||||
|
||||
blob = traced.serialize()
|
||||
print(f"Serialized computation is {len(blob)} bytes.")
|
||||
|
||||
rehydrated = export.deserialize(blob)
|
||||
|
||||
key = jax.random.key(1123)
|
||||
a_key, b_key = jax.random.split(key, 2)
|
||||
for shape in run_shapes:
|
||||
a = create_tensor(shape, dtype=jnp.float32, key=a_key)
|
||||
b = create_tensor(shape, dtype=jnp.float32, key=b_key)
|
||||
c = rehydrated.call(a, b)
|
||||
assert jnp.allclose(c, ref_f(a, b)), f"Mismatch at shape {shape}"
|
||||
print(f" shape {shape}: OK")
|
||||
|
||||
|
||||
def run_example(M, N):
|
||||
@jax.jit
|
||||
def ref_f(a, b):
|
||||
return jax.nn.sigmoid(a + b)
|
||||
|
||||
# The same JAX function is used in all three examples below. The export
|
||||
# mode is determined entirely by the shape spec passed to jax.export.
|
||||
@jax.jit
|
||||
def f(a, b):
|
||||
call = cutlass_call(launch, output_shape_dtype=a)
|
||||
return jax.nn.sigmoid(call(a, b))
|
||||
|
||||
# ── 1. Concrete shapes ────────────────────────────────────────────────────
|
||||
# Shapes are fixed constants baked into the export. The deserialized
|
||||
# computation only accepts exactly these dimensions at runtime.
|
||||
print("\nConcrete shapes:")
|
||||
|
||||
input_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32)
|
||||
_export_and_run(
|
||||
f,
|
||||
ref_f,
|
||||
input_shape_dtype,
|
||||
run_shapes=[(M, N)], # concrete exports reject any other shape
|
||||
)
|
||||
|
||||
# ── 2. Unconstrained symbolic shapes ─────────────────────────────────────
|
||||
# Both dimensions are fully dynamic. The exported computation accepts any
|
||||
# (M, N) at runtime without recompilation.
|
||||
print("\nUnconstrained symbolic shapes:")
|
||||
|
||||
a_sym, b_sym = export.symbolic_shape("a, b")
|
||||
input_shape_dtype = jax.ShapeDtypeStruct((a_sym, b_sym), jnp.float32)
|
||||
_export_and_run(
|
||||
f,
|
||||
ref_f,
|
||||
input_shape_dtype,
|
||||
run_shapes=[(M, N), (M * 2, N * 4), (M * 4, N * 4)],
|
||||
)
|
||||
|
||||
# ── 3. Constrained symbolic shapes (divisibility) ─────────────────────────
|
||||
# Shapes are declared as multiples of a tile size via TensorSpec.divisibility.
|
||||
# The symbolic expression "32*M, 16*N" tells jax.export that dim 0 is always
|
||||
# a multiple of 32 and dim 1 is always a multiple of 16. This lets the
|
||||
# compiler generate more efficient code (e.g. no remainder handling).
|
||||
# Runtime shapes must satisfy these divisibility constraints.
|
||||
print("\nConstrained symbolic shapes:")
|
||||
|
||||
@jax.jit
|
||||
def ref_f(a, b):
|
||||
return jax.nn.sigmoid(a + b)
|
||||
def f_divisible(a, b):
|
||||
spec = TensorSpec(divisibility=(32, 16))
|
||||
call = cutlass_call(
|
||||
launch,
|
||||
output_shape_dtype=a,
|
||||
input_spec=(spec, spec),
|
||||
output_spec=spec,
|
||||
)
|
||||
return jax.nn.sigmoid(call(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))
|
||||
m_sym, n_sym = export.symbolic_shape("32*M, 16*N")
|
||||
input_shape_dtype = jax.ShapeDtypeStruct((m_sym, n_sym), jnp.float32)
|
||||
_export_and_run(
|
||||
f_divisible,
|
||||
ref_f,
|
||||
input_shape_dtype,
|
||||
run_shapes=[(M, N), (M * 2, N * 2), (M * 4, N * 4)],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
@@ -161,8 +195,7 @@ if __name__ == "__main__":
|
||||
)
|
||||
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)
|
||||
run_example(args.M, args.N)
|
||||
print("PASS")
|
||||
|
||||
Reference in New Issue
Block a user