mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Lei Wang
e07bd1d009
* [BugFix][Blackwell] Fix 256-bit vectorized codegen for fp4/int4/fp16
256-bit vectorization (sm_100+) broke nvcc compilation for several
packed dtypes because every (dtype x lanes) combination relied on
hand-written per-type glue:
- fp4 x 64 lanes: predicated copy materializes a typed vector temporary,
which needs `fp4_e2_64_t::operator=(ulonglong4)` (never defined; the
operator was copy-pasted onto the 128-bit fp4_e2_32_t with a broken
body) and `make_fp4_e2_64_t` (constructors stop at 32 lanes).
- int4/uint4 x 64 lanes: PrintType emitted the nonexistent CUDA type
`(u)int8`, and the packed-nibble broadcast only handled 16/32 lanes.
- fp16 x 16 lanes: broadcast calls tl::pack_float16x4 whose parameters
were `__half`, but codegen prints fp16 scalars as cutlass::half_t
(the bf16 variant already took bfloat16_t).
Instead of patching each cell, make the contracts generic:
- tl::load_global_256 / _conditional now return the pointee type
(static_assert sizeof == 32) and bit-cast internally, so assignment
type-checks for any 32-byte vector struct. The hand-written
operator=(ulonglong4) on fp8_e{4,5,8}_32_t become dead and are
removed, along with the misplaced fp4_e2_32_t one. The T.ldg256
builtin keeps its uint32x8 result via an explicit ulonglong4 cast.
- New tl::broadcast<V>(unit) replicates a packed unit across any vector
width; fp4 broadcast now emits one packed fp4x2 byte + tl::broadcast
for every lane count instead of a 64-arg make_ call.
- int4x64 maps to (u)longlong4 (like int8x32) and its broadcast packs
two 8-nibble fields per 64-bit component.
Adds a compile matrix test (12 dtypes x {copy, predicated copy, fill}
x {128b, 256b}) plus numeric checks for the previously broken cells, so
the next dtype/width gap fails in CI instead of on a user's machine.
On < sm_100 targets the 256b params compile at 128 bits and stay green.
* [Refactor][CUDA] Replace per-lane-count vector constructors with variadic packers
Follow-up to #3153: the make_<type>_<N>_t constructor ladders (5 for fp4,
15 for fp8) and several BroadcastNode special cases enumerated the same
(dtype x lanes) matrix by hand, so every new dtype or vector width needed
more glue. Replace the enumeration with two generic templates:
- tl::make_vec<V>(e0, e1, ...): builds a vector struct from one scalar
element per lane (byte-sized or larger elements).
- tl::make_fp4_vec<V>(e0, e1, ...): same for nibble-packed fp4.
and route the codegen through them:
- PrintVecConstructor emits tl::make_fp4_vec<V> / tl::make_vec<V> for
fp4/fp8 vectors, covering the Shuffle and element-wise assembly paths
(whose previous output for these types did not even compile).
- fp8 and int8 broadcasts of side-effect-free values emit
tl::broadcast<carrier>(byte) instead of per-lane make_ calls; the int4
16/32/64-lane packed broadcasts collapse into a single
tl::broadcast branch.
- Side-effecting broadcast values (e.g. casts of rng calls) keep their
evaluate-once-per-lane semantics: they fall through to the generic
per-lane path, guarded via tirx::SideEffect. The int32 rng broadcast
path and the fp16/bf16/f32x8 paths are intentionally untouched.
The make_fp4_e2_{2..32}_t and make_fp8_{e4,e5,e8}_{2..32}_t ladders are
deleted; the char-packing make_int* overloads in common.h stay because
the int8 Shuffle path still assembles vectors element-wise.
Verified on sm_103a: the 12-dtype x 3-op x 2-width compile matrix stays
green, new numeric tests cover non-constant fp4/fp8/int8 scalar fills,
and the ldg/ldst/vectorize/vectorized_cast/quantize suites pass
(228 tests) with TILELANG_DISABLE_CACHE=1.
* [Test] Update issue #2480 assertions for the tl::broadcast emission
The int8x32 broadcast is now emitted as tl::broadcast<(u)longlong4>(byte)
instead of make_(u)longlong4 with 32 scalar args; the exact-value checks
that guard the original issue (upper 32 bits of each field zeroed) are
unchanged.
248 lines
8.2 KiB
Python
248 lines
8.2 KiB
Python
"""Compile matrix over dtype x op x vector width for vectorized global
|
|
memory patterns.
|
|
|
|
Every cell used to be an independent hand-written case in the codegen
|
|
(make_<type> constructors, operator= from ulonglong4, pack helpers), so a
|
|
new dtype or a wider vector width silently created nvcc-only failures:
|
|
fp4 x 64 lanes, int4 x 64 lanes, and fp16 x 16 lanes all broke when 256-bit
|
|
vectorization was enabled on sm_100+. This matrix compiles every
|
|
combination so the next gap fails in CI instead of on a user's machine.
|
|
|
|
On targets without 256-bit vectorization (< sm_100) the "256" width params
|
|
compile at 128 bits, which keeps the test portable and still covers the
|
|
128-bit half of the matrix.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
MATRIX_DTYPES = [
|
|
"float4_e2m1fn",
|
|
"float8_e4m3fn",
|
|
"float8_e5m2",
|
|
"float8_e8m0fnu",
|
|
"float16",
|
|
"bfloat16",
|
|
"int8",
|
|
"uint8",
|
|
"int4",
|
|
"uint4",
|
|
"int32",
|
|
"float32",
|
|
]
|
|
|
|
# copy: pure vector load/store, no broadcast.
|
|
# pred_copy: loop-invariant predicate materializes a typed vector temporary
|
|
# (`condval = tl::load_global_...` assignment) and broadcasts the zero
|
|
# else-value -- the shape that exposed the fp4/fp16/int4 gaps.
|
|
# fill_one: broadcast of a non-zero constant (exercises packed-nibble
|
|
# constant folding for 4-bit dtypes).
|
|
MATRIX_OPS = ["copy", "pred_copy", "fill_one"]
|
|
|
|
|
|
def _matrix_kernel(dtype, op, disable_256):
|
|
@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_DISABLE_VECTORIZE_256: disable_256})
|
|
def factory():
|
|
@T.prim_func
|
|
def kern(
|
|
src: T.Tensor((128, 128), dtype),
|
|
predicate: T.Tensor((1,), T.int32),
|
|
dst: T.Tensor((128, 128), dtype),
|
|
):
|
|
with T.Kernel(1, threads=256):
|
|
for row, col in T.Parallel(128, 128):
|
|
if op == "copy":
|
|
dst[row, col] = src[row, col]
|
|
elif op == "pred_copy":
|
|
dst[row, col] = src[row, col] if predicate[0] != 0 else 0
|
|
else:
|
|
dst[row, col] = 1
|
|
|
|
return kern
|
|
|
|
return factory()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("dtype", MATRIX_DTYPES)
|
|
@pytest.mark.parametrize("op", MATRIX_OPS)
|
|
@pytest.mark.parametrize("width", ["128", "256"])
|
|
def test_vectorize_matrix_compiles(dtype, op, width):
|
|
_matrix_kernel(dtype, op, disable_256=(width == "128"))
|
|
|
|
|
|
def _pred_copy_kernel(dtype):
|
|
@tilelang.jit
|
|
def factory():
|
|
@T.prim_func
|
|
def kern(
|
|
src: T.Tensor((128, 128), dtype),
|
|
predicate: T.Tensor((1,), T.int32),
|
|
dst: T.Tensor((128, 128), dtype),
|
|
):
|
|
with T.Kernel(1, threads=256):
|
|
for row, col in T.Parallel(128, 128):
|
|
dst[row, col] = src[row, col] if predicate[0] != 0 else 0
|
|
|
|
return kern
|
|
|
|
return factory()
|
|
|
|
|
|
def _check_pred_copy_packed(kernel):
|
|
"""Run a predicated copy on byte-backed packed storage and check both
|
|
predicate branches. src/dst are uint8 views so equality is bitwise."""
|
|
src = torch.randint(0, 256, (128, 64), dtype=torch.uint8, device="cuda")
|
|
dst = torch.full((128, 64), 0xEE, dtype=torch.uint8, device="cuda")
|
|
packed = torch.float4_e2m1fn_x2
|
|
pred1 = torch.tensor([1], dtype=torch.int32, device="cuda")
|
|
pred0 = torch.tensor([0], dtype=torch.int32, device="cuda")
|
|
|
|
kernel(src.view(packed), pred1, dst.view(packed))
|
|
torch.cuda.synchronize()
|
|
assert torch.equal(dst, src)
|
|
|
|
kernel(src.view(packed), pred0, dst.view(packed))
|
|
torch.cuda.synchronize()
|
|
assert (dst == 0).all()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.skipif(
|
|
not hasattr(torch, "float4_e2m1fn_x2"),
|
|
reason="PyTorch float4_e2m1fn_x2 dtype is unavailable",
|
|
)
|
|
def test_pred_copy_fp4_numeric():
|
|
"""fp4 predicated copy: on sm_100+ this vectorizes to 64 lanes (256 bits)
|
|
and used to fail to compile (undefined make_fp4_e2_64_t, missing
|
|
ulonglong4 assignment)."""
|
|
_check_pred_copy_packed(_pred_copy_kernel("float4_e2m1fn"))
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_pred_copy_int4_numeric():
|
|
"""int4 predicated copy: on sm_100+ this vectorizes to 64 lanes and used
|
|
to emit the nonexistent CUDA type `int8` plus a 64-arg make_int8 call."""
|
|
kernel = _pred_copy_kernel("int4")
|
|
src = torch.randint(0, 256, (128, 64), dtype=torch.uint8, device="cuda")
|
|
dst = torch.full((128, 64), 0xEE, dtype=torch.uint8, device="cuda")
|
|
pred1 = torch.tensor([1], dtype=torch.int32, device="cuda")
|
|
pred0 = torch.tensor([0], dtype=torch.int32, device="cuda")
|
|
|
|
kernel(src, pred1, dst)
|
|
torch.cuda.synchronize()
|
|
assert torch.equal(dst, src)
|
|
|
|
kernel(src, pred0, dst)
|
|
torch.cuda.synchronize()
|
|
assert (dst == 0).all()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("fill_value", [1, 5])
|
|
def test_fill_int4_nonzero_numeric(fill_value):
|
|
"""int4 constant fill: checks the packed-nibble constant path (every
|
|
byte of the destination must hold the value in both nibbles)."""
|
|
|
|
@tilelang.jit
|
|
def factory():
|
|
@T.prim_func
|
|
def kern(dst: T.Tensor((128, 128), "int4")):
|
|
with T.Kernel(1, threads=256):
|
|
for row, col in T.Parallel(128, 128):
|
|
dst[row, col] = fill_value
|
|
|
|
return kern
|
|
|
|
kernel = factory()
|
|
dst = torch.full((128, 64), 0xEE, dtype=torch.uint8, device="cuda")
|
|
kernel(dst)
|
|
torch.cuda.synchronize()
|
|
expected = (fill_value & 0xF) | ((fill_value & 0xF) << 4)
|
|
assert (dst == expected).all()
|
|
|
|
|
|
def _fill_var_kernel(dtype):
|
|
@tilelang.jit
|
|
def factory():
|
|
@T.prim_func
|
|
def kern(scalar: T.Tensor((2,), dtype), dst: T.Tensor((128, 128), dtype)):
|
|
with T.Kernel(1, threads=256):
|
|
for row, col in T.Parallel(128, 128):
|
|
dst[row, col] = scalar[0]
|
|
|
|
return kern
|
|
|
|
return factory()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_fill_scalar_var_fp8_numeric():
|
|
"""fp8 fill from a runtime scalar: a non-constant broadcast replicated
|
|
across the 256-bit carrier (tl::broadcast byte replication)."""
|
|
kernel = _fill_var_kernel("float8_e4m3fn")
|
|
scalar = torch.full((2,), 1.5, device="cuda").to(torch.float8_e4m3fn)
|
|
dst = torch.zeros((128, 128), device="cuda").to(torch.float8_e4m3fn)
|
|
kernel(scalar, dst)
|
|
torch.cuda.synchronize()
|
|
assert (dst.view(torch.uint8) == scalar.view(torch.uint8)[0]).all()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_fill_scalar_var_int8_numeric():
|
|
"""int8 fill from a runtime scalar (non-constant byte broadcast)."""
|
|
kernel = _fill_var_kernel("int8")
|
|
scalar = torch.full((2,), -42, dtype=torch.int8, device="cuda")
|
|
dst = torch.zeros((128, 128), dtype=torch.int8, device="cuda")
|
|
kernel(scalar, dst)
|
|
torch.cuda.synchronize()
|
|
assert (dst == -42).all()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.skipif(
|
|
not hasattr(torch, "float4_e2m1fn_x2"),
|
|
reason="PyTorch float4_e2m1fn_x2 dtype is unavailable",
|
|
)
|
|
def test_fill_scalar_var_fp4_numeric():
|
|
"""fp4 fill from a runtime scalar: packed nibble read, then a
|
|
non-constant nibble-pair broadcast."""
|
|
kernel = _fill_var_kernel("float4_e2m1fn")
|
|
# Both nibbles hold 0x4 (fp4 e2m1 value 2.0), so every destination byte
|
|
# must equal 0x44.
|
|
scalar = torch.full((1,), 0x44, dtype=torch.uint8, device="cuda")
|
|
dst = torch.zeros((128, 64), dtype=torch.uint8, device="cuda")
|
|
packed = torch.float4_e2m1fn_x2
|
|
kernel(scalar.view(packed), dst.view(packed))
|
|
torch.cuda.synchronize()
|
|
assert (dst == 0x44).all()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_fill_fp16_numeric():
|
|
"""fp16 constant fill: at 16 lanes (256 bits) the broadcast goes through
|
|
tl::pack_float16x4, which used to reject cutlass::half_t arguments."""
|
|
|
|
@tilelang.jit
|
|
def factory():
|
|
@T.prim_func
|
|
def kern(dst: T.Tensor((128, 128), "float16")):
|
|
with T.Kernel(1, threads=256):
|
|
for row, col in T.Parallel(128, 128):
|
|
dst[row, col] = 1.5
|
|
|
|
return kern
|
|
|
|
kernel = factory()
|
|
dst = torch.zeros((128, 128), dtype=torch.float16, device="cuda")
|
|
kernel(dst)
|
|
torch.cuda.synchronize()
|
|
assert (dst == 1.5).all()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|