367 lines
14 KiB
Python
367 lines
14 KiB
Python
# 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 cutlass
|
|
import cutlass.cute as cute
|
|
import cutlass.jax as cjax
|
|
import cuda.bindings.driver as cuda
|
|
|
|
"""
|
|
CuTe DSL kernels used by the ``cute_dsl_jax.ipynb`` notebook.
|
|
|
|
This module defines GPU kernels written in CuTe DSL (CUTLASS 4.x Python DSL)
|
|
that are called from JAX via ``cutlass.jax.cutlass_call``. ``cutlass_call`` is a
|
|
JAX primitive that triggers compilation of the kernel during lowering and embeds
|
|
it into the HLO computation, so XLA can launch it efficiently without callback
|
|
to Python.
|
|
|
|
Kernels provided:
|
|
|
|
- ``vector_add`` — element-wise c = a + b (3-D CuTe layout)
|
|
- ``saxpy`` — y = alpha * x + y
|
|
- ``relu`` — element-wise ReLU with flat indexing
|
|
- ``fused_bias_relu`` — fused bias addition + ReLU
|
|
- ``gemm`` — tiled matrix multiplication
|
|
- ``elementwise_add`` — 2-D element-wise add (flat indexing, ``jax.export``-compatible)
|
|
|
|
The notebook imports these kernels and wraps each one with ``cutlass_call``
|
|
inside ``@jax.jit`` functions. See ``cute_dsl_jax.ipynb`` for usage, validation,
|
|
and step-by-step explanations.
|
|
|
|
This module is imported by the notebook and by ``cute_dsl_jax.py``. It can also
|
|
be run directly to validate every kernel:
|
|
|
|
.. code-block:: bash
|
|
|
|
# Interactive notebook (recommended for learning)
|
|
jupyter lab cute_dsl_jax.ipynb
|
|
|
|
# Full demo as a standalone script
|
|
python cute_dsl_jax_kernels.py
|
|
"""
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Vector Add: c = a + b #
|
|
# ------------------------------------------------------------------ #
|
|
@cute.kernel
|
|
def vector_add_kernel(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
|
|
"""Per-thread kernel: each thread adds one element."""
|
|
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_vector_add(
|
|
stream: cuda.CUstream,
|
|
a: cute.Tensor, b: cute.Tensor, c: cute.Tensor,
|
|
):
|
|
vector_add_kernel(a, b, c).launch(
|
|
grid=[a.shape[-1], 1, 1],
|
|
block=[a.shape[-2], 1, 1],
|
|
stream=stream,
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# SAXPY: y = alpha * x + y #
|
|
# ------------------------------------------------------------------ #
|
|
@cute.kernel
|
|
def saxpy_kernel(x: cute.Tensor, y: cute.Tensor, out: cute.Tensor, alpha: float):
|
|
"""SAXPY: out[i] = alpha * x[i] + y[i]."""
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
bidx, _, _ = cute.arch.block_idx()
|
|
|
|
frgX = cute.make_rmem_tensor(cute.size(x, mode=[0]), x.element_type)
|
|
frgY = cute.make_rmem_tensor(cute.size(y, mode=[0]), y.element_type)
|
|
frgO = cute.make_rmem_tensor(cute.size(out, mode=[0]), out.element_type)
|
|
|
|
cute.autovec_copy(x[None, tidx, bidx], frgX)
|
|
cute.autovec_copy(y[None, tidx, bidx], frgY)
|
|
frgO.store(alpha * frgX.load() + frgY.load())
|
|
cute.autovec_copy(frgO, out[None, tidx, bidx])
|
|
|
|
|
|
@cute.jit
|
|
def launch_saxpy(
|
|
stream: cuda.CUstream,
|
|
x: cute.Tensor, y: cute.Tensor, out: cute.Tensor,
|
|
*, alpha: float,
|
|
):
|
|
saxpy_kernel(x, y, out, alpha).launch(
|
|
grid=[x.shape[-1], 1, 1],
|
|
block=[x.shape[-2], 1, 1],
|
|
stream=stream,
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# ReLU: out = max(0, x) #
|
|
# ------------------------------------------------------------------ #
|
|
@cute.kernel
|
|
def relu_kernel(x: cute.Tensor, out: cute.Tensor, N: int):
|
|
"""Per-thread kernel: each thread computes ReLU of one element."""
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
bidx, _, _ = cute.arch.block_idx()
|
|
bdx, _, _ = cute.arch.block_dim()
|
|
|
|
idx = bidx * bdx + tidx
|
|
if idx < N:
|
|
val = x[idx]
|
|
out[idx] = cutlass.max(val, cutlass.Float32(0.0))
|
|
|
|
|
|
@cute.jit
|
|
def launch_relu(
|
|
stream: cuda.CUstream,
|
|
x: cute.Tensor, out: cute.Tensor,
|
|
*, N: int,
|
|
):
|
|
BLOCK_SIZE = 256
|
|
grid_size = (N + BLOCK_SIZE - 1) // BLOCK_SIZE
|
|
relu_kernel(x, out, N).launch(
|
|
grid=[grid_size, 1, 1],
|
|
block=[BLOCK_SIZE, 1, 1],
|
|
stream=stream,
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Fused Bias + ReLU: out = max(0, x + bias[col]) #
|
|
# ------------------------------------------------------------------ #
|
|
@cute.kernel
|
|
def fused_bias_relu_kernel(
|
|
x: cute.Tensor, bias: cute.Tensor, out: cute.Tensor, N: int, width: int,
|
|
):
|
|
"""Per-thread: out[i] = max(0, x[i] + bias[i % width])."""
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
bidx, _, _ = cute.arch.block_idx()
|
|
bdx, _, _ = cute.arch.block_dim()
|
|
|
|
idx = bidx * bdx + tidx
|
|
if idx < N:
|
|
col = idx % width
|
|
val = x[idx] + bias[col]
|
|
out[idx] = cutlass.max(val, cutlass.Float32(0.0))
|
|
|
|
|
|
@cute.jit
|
|
def launch_fused_bias_relu(
|
|
stream: cuda.CUstream,
|
|
x: cute.Tensor, bias: cute.Tensor, out: cute.Tensor,
|
|
*, N: int, width: int,
|
|
):
|
|
BLOCK_SIZE = 256
|
|
grid_size = (N + BLOCK_SIZE - 1) // BLOCK_SIZE
|
|
fused_bias_relu_kernel(x, bias, out, N, width).launch(
|
|
grid=[grid_size, 1, 1],
|
|
block=[BLOCK_SIZE, 1, 1],
|
|
stream=stream,
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# GEMM: D = A @ B #
|
|
# ------------------------------------------------------------------ #
|
|
@cute.kernel
|
|
def gemm_kernel(
|
|
A: cute.Tensor, B: cute.Tensor, D: cute.Tensor,
|
|
M: int, N: int, K: int, BLOCK_M: int, BLOCK_N: int,
|
|
):
|
|
"""Tiled GEMM: each thread accumulates output elements."""
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
bm, bn, _ = cute.arch.block_idx()
|
|
bdx, _, _ = cute.arch.block_dim()
|
|
|
|
for i in cutlass.range(tidx, BLOCK_M * BLOCK_N, bdx):
|
|
row = i // BLOCK_N
|
|
col = i % BLOCK_N
|
|
m_idx = bm * BLOCK_M + row
|
|
n_idx = bn * BLOCK_N + col
|
|
if m_idx < M and n_idx < N:
|
|
acc = cutlass.Float32(0.0)
|
|
for k in cutlass.range(K):
|
|
acc += A[m_idx * K + k] * B[k * N + n_idx]
|
|
D[m_idx * N + n_idx] = acc
|
|
|
|
|
|
@cute.jit
|
|
def launch_gemm(
|
|
stream: cuda.CUstream,
|
|
A: cute.Tensor, B: cute.Tensor, D: cute.Tensor,
|
|
*, M: int, N: int, K: int,
|
|
):
|
|
BLOCK_M, BLOCK_N = 64, 64
|
|
grid_m = (M + BLOCK_M - 1) // BLOCK_M
|
|
grid_n = (N + BLOCK_N - 1) // BLOCK_N
|
|
gemm_kernel(A, B, D, M, N, K, BLOCK_M, BLOCK_N).launch(
|
|
grid=[grid_m, grid_n, 1],
|
|
block=[256, 1, 1],
|
|
stream=stream,
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Element-wise Add (2-D, flat indexing) #
|
|
# ------------------------------------------------------------------ #
|
|
@cute.kernel
|
|
def elementwise_add_kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
|
|
"""Per-thread kernel: 2-D element-wise add using flat indexing."""
|
|
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_elementwise_add(
|
|
stream: cuda.CUstream,
|
|
mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor,
|
|
):
|
|
num_threads_per_block = 256
|
|
m, n = mA.shape
|
|
elementwise_add_kernel(mA, mB, mC).launch(
|
|
grid=((m * n) // num_threads_per_block, 1, 1),
|
|
block=(num_threads_per_block, 1, 1),
|
|
stream=stream,
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Self-tests #
|
|
# ------------------------------------------------------------------ #
|
|
if __name__ == '__main__':
|
|
import os
|
|
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
|
|
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import numpy as np
|
|
|
|
BLOCK = 256
|
|
N_BLOCKS = 4
|
|
|
|
# ── Vector Add ────────────────────────────────────────────────────
|
|
# 3-D CuTe layout: (elems_per_thread, threads_per_block, num_blocks)
|
|
a = jax.random.normal(jax.random.PRNGKey(0), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
|
|
b = jax.random.normal(jax.random.PRNGKey(1), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
|
|
call = cjax.cutlass_call(
|
|
launch_vector_add,
|
|
output_shape_dtype=jax.ShapeDtypeStruct(a.shape, a.dtype),
|
|
use_static_tensors=True,
|
|
)
|
|
c = jax.jit(call)(a, b)
|
|
np.testing.assert_allclose(np.array(c), np.array(a + b), rtol=1e-5, atol=1e-5)
|
|
print('vector_add: PASSED')
|
|
|
|
# ── SAXPY ─────────────────────────────────────────────────────────
|
|
x = jax.random.normal(jax.random.PRNGKey(2), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
|
|
y = jax.random.normal(jax.random.PRNGKey(3), (1, BLOCK, N_BLOCKS), dtype=jnp.float32)
|
|
alpha = 2.5
|
|
call = cjax.cutlass_call(
|
|
launch_saxpy,
|
|
output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype),
|
|
use_static_tensors=True,
|
|
alpha=alpha,
|
|
)
|
|
out = jax.jit(call)(x, y)
|
|
np.testing.assert_allclose(np.array(out), np.array(alpha * x + y), rtol=1e-5, atol=1e-5)
|
|
print('saxpy: PASSED')
|
|
|
|
# ── ReLU ──────────────────────────────────────────────────────────
|
|
N_ELEM = BLOCK * N_BLOCKS
|
|
x = jax.random.normal(jax.random.PRNGKey(4), (N_ELEM,), dtype=jnp.float32)
|
|
call = cjax.cutlass_call(
|
|
launch_relu,
|
|
output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype),
|
|
N=N_ELEM,
|
|
)
|
|
out = jax.jit(call)(x)
|
|
np.testing.assert_allclose(np.array(out), np.array(jnp.maximum(x, 0)), rtol=1e-5, atol=1e-5)
|
|
print('relu: PASSED')
|
|
|
|
# ── Fused Bias + ReLU ─────────────────────────────────────────────
|
|
ROWS, COLS = 16, 64
|
|
x = jax.random.normal(jax.random.PRNGKey(5), (ROWS * COLS,), dtype=jnp.float32)
|
|
bias = jax.random.normal(jax.random.PRNGKey(6), (COLS,), dtype=jnp.float32)
|
|
call = cjax.cutlass_call(
|
|
launch_fused_bias_relu,
|
|
output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype),
|
|
N=ROWS * COLS, width=COLS,
|
|
)
|
|
out = jax.jit(call)(x, bias)
|
|
ref = jnp.maximum(x.reshape(ROWS, COLS) + bias, 0).reshape(-1)
|
|
np.testing.assert_allclose(np.array(out), np.array(ref), rtol=1e-5, atol=1e-5)
|
|
print('fused_bias_relu: PASSED')
|
|
|
|
# ── GEMM ──────────────────────────────────────────────────────────
|
|
M, N, K = 128, 128, 64
|
|
A = jax.random.normal(jax.random.PRNGKey(7), (M * K,), dtype=jnp.float32)
|
|
B = jax.random.normal(jax.random.PRNGKey(8), (K * N,), dtype=jnp.float32)
|
|
call = cjax.cutlass_call(
|
|
launch_gemm,
|
|
output_shape_dtype=jax.ShapeDtypeStruct((M * N,), A.dtype),
|
|
M=M, N=N, K=K,
|
|
)
|
|
D = jax.jit(call)(A, B)
|
|
ref = A.reshape(M, K) @ B.reshape(K, N)
|
|
np.testing.assert_allclose(np.array(D.reshape(M, N)), np.array(ref), rtol=1e-2, atol=1e-2)
|
|
print('gemm: PASSED')
|
|
|
|
# ── Elementwise Add (2-D) ─────────────────────────────────────────
|
|
M, N = 16, 256
|
|
a = jax.random.normal(jax.random.PRNGKey(9), (M, N), dtype=jnp.float32)
|
|
b = jax.random.normal(jax.random.PRNGKey(10), (M, N), dtype=jnp.float32)
|
|
call = cjax.cutlass_call(
|
|
launch_elementwise_add,
|
|
output_shape_dtype=jax.ShapeDtypeStruct(a.shape, a.dtype),
|
|
)
|
|
c = jax.jit(call)(a, b)
|
|
np.testing.assert_allclose(np.array(c), np.array(a + b), rtol=1e-5, atol=1e-5)
|
|
print('elementwise_add: PASSED')
|
|
|
|
print('\nAll kernels passed.')
|