mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +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.