From 5cc9ca9b6629df72cfcc402222f22a3919f669d9 Mon Sep 17 00:00:00 2001 From: Saiyam Bharara Date: Thu, 10 Sep 2026 01:32:24 +0530 Subject: [PATCH 1/7] Fix nvcc inff crash on fp8_e4m3/float4 reduce init reduce_max/min build their -INFINITY/+INFINITY init via make_const, bypassing T.infinity()'s guard. e4m3/float4 have no inf encoding, so PrintConst printed it as the invalid literal `inff`. Clamp to the format's max-finite value instead, since (unlike e5m2 in #2671) there's no bit pattern to bitcast to. Fixes #2998 --- src/cuda/codegen/codegen_cuda.cc | 27 ++++++++++++++-- .../language/test_tilelang_language_reduce.py | 32 +++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/src/cuda/codegen/codegen_cuda.cc b/src/cuda/codegen/codegen_cuda.cc index e17df44331..4f8cf4cffa 100644 --- a/src/cuda/codegen/codegen_cuda.cc +++ b/src/cuda/codegen/codegen_cuda.cc @@ -6053,6 +6053,21 @@ void CodeGenTileLangCUDA::VisitExpr_(const BroadcastNode *op, os << ')'; } +namespace { +// Max finite magnitude for fp8/fp4 dtypes with no (or no non-finite) +// infinity encoding, e.g. the e4m3 family and float4_e2m1fn. +double MaxFiniteFp8Fp4(const DataType &dtype) { + if (dtype.is_float8_e4m3fnuz()) { + return 240.0; + } + if (dtype.is_float8_e4m3() || dtype.is_float8_e4m3fn()) { + return 448.0; + } + // float4_e2m1fn / float4_e2m1_unpacked + return 6.0; +} +} // namespace + inline void PrintConst(const FloatImmNode *op, std::ostream &os, CodeGenTileLangCUDA *p) { // NOLINT(*) // Type code is kBFloat/kFloat16 @@ -6115,9 +6130,17 @@ inline void PrintConst(const FloatImmNode *op, std::ostream &os, os << "::bitcast(0x7e)"; return; } + // No literal spelling for non-finite values here (e.g. -inff/nanf are + // invalid identifiers); clamp to the format's max-finite magnitude. + double value = op->value; + if (std::isnan(value)) { + value = MaxFiniteFp8Fp4(op->dtype); + } else if (std::isinf(value)) { + value = std::copysign(MaxFiniteFp8Fp4(op->dtype), value); + } p->PrintType(op->dtype, os); - os << '(' << FlexibleHexFormat(op->value) << 'f'; - os << "/*" << std::scientific << op->value << "*/"; + os << '(' << FlexibleHexFormat(value) << 'f'; + os << "/*" << std::scientific << value << "*/"; os << ')'; return; } diff --git a/testing/python/language/test_tilelang_language_reduce.py b/testing/python/language/test_tilelang_language_reduce.py index 3ef7ddef9b..b512f2b0f1 100644 --- a/testing/python/language/test_tilelang_language_reduce.py +++ b/testing/python/language/test_tilelang_language_reduce.py @@ -1134,6 +1134,38 @@ def kernel(A: T.Tensor((4, 32), T.float8_e4m3fn), B: T.Tensor((4,), T.float16)): torch.testing.assert_close(B, ref, atol=0, rtol=0) +def _make_fp8_e4m3_reduce_kernel(reduce_fn, M, N): + @T.prim_func + def kernel(A: T.Tensor((M, N), T.float8_e4m3), B: T.Tensor((M,), T.float8_e4m3)): + with T.Kernel(1, threads=128): + src = T.alloc_fragment((M, N), T.float8_e4m3) + dst = T.alloc_fragment((M,), T.float8_e4m3) + T.copy(A, src) + reduce_fn(src, dst, dim=1) + T.copy(dst, B) + + return kernel + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(8, 9) +@pytest.mark.parametrize(("reduce_fn", "torch_op"), [(T.reduce_max, "max"), (T.reduce_min, "min")], ids=["max", "min"]) +def test_reduce_fp8_e4m3_dst_runtime(reduce_fn, torch_op): + """reduce_max/reduce_min into a float8_e4m3 dst used to crash nvcc with + an undefined `inff` identifier (GH-2998).""" + if not hasattr(torch, "float8_e4m3fn"): + pytest.skip("torch.float8_e4m3fn is not available") + + M, N = 32, 64 + k = _compile(_make_fp8_e4m3_reduce_kernel(reduce_fn, M, N)) + + a = torch.rand(M, N, device="cuda").to(torch.float8_e4m3fn) + out = k(a) + # torch has no min/max reduction kernel for float8_e4m3fn; reduce in fp32. + ref = getattr(a.float(), torch_op)(dim=1).values + assert torch.allclose(out.float(), ref, atol=1e-1) + + @tilelang.testing.requires_cuda def test_reduce_packed_max_nan_propagate_uses_nan_intrinsics(): k = _compile(_make_nan_reduce_kernel(T.reduce_max, 128, 128, T.float16, threads=256, nan_propagate=True)) From 52a5b54ab01ffe36f8e6bceb72319a7fa0c26136 Mon Sep 17 00:00:00 2001 From: Saiyam Bharara Date: Sat, 3 Oct 2026 13:30:48 +0530 Subject: [PATCH 2/7] Fix fp4 reduce and e8m0 non-finite constants --- src/cuda/codegen/codegen_cuda.cc | 11 +++- src/tl_templates/cuda/common.h | 11 ++-- src/tl_templates/cuda/cuda_fp4.h | 25 +++++++++ .../language/test_tilelang_language_reduce.py | 53 +++++++++++++++++++ 4 files changed, 93 insertions(+), 7 deletions(-) diff --git a/src/cuda/codegen/codegen_cuda.cc b/src/cuda/codegen/codegen_cuda.cc index 4f8cf4cffa..d5e5eecb9e 100644 --- a/src/cuda/codegen/codegen_cuda.cc +++ b/src/cuda/codegen/codegen_cuda.cc @@ -6055,7 +6055,7 @@ void CodeGenTileLangCUDA::VisitExpr_(const BroadcastNode *op, namespace { // Max finite magnitude for fp8/fp4 dtypes with no (or no non-finite) -// infinity encoding, e.g. the e4m3 family and float4_e2m1fn. +// infinity encoding, e.g. the e4m3 family, e8m0fnu and float4_e2m1fn. double MaxFiniteFp8Fp4(const DataType &dtype) { if (dtype.is_float8_e4m3fnuz()) { return 240.0; @@ -6063,7 +6063,10 @@ double MaxFiniteFp8Fp4(const DataType &dtype) { if (dtype.is_float8_e4m3() || dtype.is_float8_e4m3fn()) { return 448.0; } - // float4_e2m1fn / float4_e2m1_unpacked + if (dtype.is_float8_e8m0fnu()) { + return std::ldexp(1.0, 127); + } + ICHECK(dtype.is_float4()) << "No max-finite value for " << dtype; return 6.0; } } // namespace @@ -6135,6 +6138,10 @@ inline void PrintConst(const FloatImmNode *op, std::ostream &os, double value = op->value; if (std::isnan(value)) { value = MaxFiniteFp8Fp4(op->dtype); + } else if (std::isinf(value) && value < 0 && + op->dtype.is_float8_e8m0fnu()) { + // e8m0fnu is unsigned: its smallest value, 2^-127, is the lower bound. + value = std::ldexp(1.0, -127); } else if (std::isinf(value)) { value = std::copysign(MaxFiniteFp8Fp4(op->dtype), value); } diff --git a/src/tl_templates/cuda/common.h b/src/tl_templates/cuda/common.h index ac4028e636..9a930cc9a3 100644 --- a/src/tl_templates/cuda/common.h +++ b/src/tl_templates/cuda/common.h @@ -1486,14 +1486,16 @@ using tl::tfloat32_t; int width = warpSize) { \ RAW raw = reinterpret_cast(val); \ RAW ret = static_cast( \ - __shfl_down_sync(mask, static_cast(raw), delta, width)); \ + __shfl_down_sync(mask, static_cast(raw), \ + static_cast(delta), width)); \ return reinterpret_cast(ret); \ } \ TL_PATCH TL_DEVICE TYPE __shfl_up_sync(unsigned mask, TYPE val, int delta, \ int width = warpSize) { \ RAW raw = reinterpret_cast(val); \ - RAW ret = static_cast( \ - __shfl_up_sync(mask, static_cast(raw), delta, width)); \ + RAW ret = \ + static_cast(__shfl_up_sync(mask, static_cast(raw), \ + static_cast(delta), width)); \ return reinterpret_cast(ret); \ } @@ -1501,8 +1503,7 @@ TL_DEFINE_SHFL_SYNC_OVERLOADS(half_t, uint16_t) TL_DEFINE_SHFL_SYNC_OVERLOADS(bfloat16_t, uint16_t) TL_DEFINE_SHFL_SYNC_OVERLOADS(tl::float_e4m3_t, uint8_t) TL_DEFINE_SHFL_SYNC_OVERLOADS(tl::float_e5m2_t, uint8_t) - -#undef TL_DEFINE_SHFL_SYNC_OVERLOADS +// Kept defined: cuda_fp4.h also applies it to fp4_e2_t. // // Optimized type-punned warp shuffle helpers for 16-bit types diff --git a/src/tl_templates/cuda/cuda_fp4.h b/src/tl_templates/cuda/cuda_fp4.h index 0a62bb4039..1f1c4b25bc 100644 --- a/src/tl_templates/cuda/cuda_fp4.h +++ b/src/tl_templates/cuda/cuda_fp4.h @@ -44,6 +44,31 @@ struct fp4_e2_t { TL_DEVICE operator __half() const { return __half(float(*this)); } }; +// fp4_e2_t converts implicitly to several types, so shuffles and comparisons +// on it are ambiguous, and a bare max/min would compare raw encodings. Shuffle +// the raw bits and compare through float. +TL_DEFINE_SHFL_SYNC_OVERLOADS(fp4_e2_t, uint8_t) + +namespace tl { +template <> TL_DEVICE fp4_e2_t fast_max(fp4_e2_t a, fp4_e2_t b) { + return fp4_e2_t(::fmaxf(static_cast(a), static_cast(b))); +} + +template <> TL_DEVICE fp4_e2_t fast_min(fp4_e2_t a, fp4_e2_t b) { + return fp4_e2_t(::fminf(static_cast(a), static_cast(b))); +} +} // namespace tl + +// Global, not in tl: generated code calls max/min unqualified, and fp4_e2_t +// is a global type, so argument-dependent lookup would not find tl:: ones. +TL_DEVICE fp4_e2_t max(fp4_e2_t lhs, fp4_e2_t rhs) { + return tl::fast_max(lhs, rhs); +} + +TL_DEVICE fp4_e2_t min(fp4_e2_t lhs, fp4_e2_t rhs) { + return tl::fast_min(lhs, rhs); +} + // Tag for tcgen05 unpacked FP4 shared-memory layout. The hardware atom carries // 16 4-bit payload values in the low 64 bits of a 128-bit aligned region. // See: diff --git a/testing/python/language/test_tilelang_language_reduce.py b/testing/python/language/test_tilelang_language_reduce.py index b512f2b0f1..4855c08069 100644 --- a/testing/python/language/test_tilelang_language_reduce.py +++ b/testing/python/language/test_tilelang_language_reduce.py @@ -1166,6 +1166,59 @@ def test_reduce_fp8_e4m3_dst_runtime(reduce_fn, torch_op): assert torch.allclose(out.float(), ref, atol=1e-1) +def _make_fp4_reduce_kernel(reduce_fn, M, N): + # fp16 I/O holding exact fp4 values avoids packed-fp4 tensors on the host. + @T.prim_func + def kernel(A: T.Tensor((M, N), T.float16), B: T.Tensor((M,), T.float16)): + with T.Kernel(1, threads=128): + src = T.alloc_fragment((M, N), T.float4_e2m1fn) + dst = T.alloc_fragment((M,), T.float4_e2m1fn) + T.copy(A, src) + reduce_fn(src, dst, dim=1) + T.copy(dst, B) + + return kernel + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(10, 0) +@pytest.mark.parametrize(("reduce_fn", "torch_op"), [(T.reduce_max, "max"), (T.reduce_min, "min")], ids=["max", "min"]) +def test_reduce_fp4_dst_runtime(reduce_fn, torch_op): + """reduce_max/reduce_min into a float4_e2m1fn dst failed to compile + (ambiguous shuffle and comparison) and a bare max/min compared raw + encodings (GH-2998).""" + M, N = 32, 64 + k = _compile(_make_fp4_reduce_kernel(reduce_fn, M, N)) + + vals = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.float16) + vals = torch.cat([vals, -vals]) + a = vals[torch.randint(0, len(vals), (M, N))].cuda() + out = k(a) + ref = getattr(a, torch_op)(dim=1).values + torch.testing.assert_close(out, ref, atol=0, rtol=0) + + +@tilelang.testing.requires_cuda +def test_fp8_e8m0_non_finite_const_codegen(): + """e8m0fnu has no inf and no sign; non-finite constants must clamp to its + own range, not the fp4 limit of 6.0.""" + pos_inf = tvm.tirx.const(float("inf"), "float8_e8m0fnu") + neg_inf = tvm.tirx.const(float("-inf"), "float8_e8m0fnu") + + @T.prim_func + def kernel(B: T.Tensor((2,), T.float8_e8m0fnu)): + with T.Kernel(1, threads=32): + B[0] = pos_inf + B[1] = neg_inf + + target = {"kind": "cuda", "arch": "sm_100a"} + with tvm.target.Target(target): + src = tilelang.lower(kernel, target=target).kernel_source + assert "fp8_e8_t(0x1p+127f" in src + assert "fp8_e8_t(0x1p-127f" in src + assert "6.000000e+00" not in src + + @tilelang.testing.requires_cuda def test_reduce_packed_max_nan_propagate_uses_nan_intrinsics(): k = _compile(_make_nan_reduce_kernel(T.reduce_max, 128, 128, T.float16, threads=256, nan_propagate=True)) From ba145eb8f8a2f7086349b04726b37d2b30421873 Mon Sep 17 00:00:00 2001 From: LeiWang1999 Date: Thu, 22 Oct 2026 18:19:01 +0800 Subject: [PATCH 3/7] [Fix] Size FP4 reduction workspaces for scalar storage --- src/transform/lower_tile_op.cc | 6 ++++++ .../language/test_tilelang_language_reduce.py | 20 ++++++++++++++++--- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/src/transform/lower_tile_op.cc b/src/transform/lower_tile_op.cc index f75dba576f..f3a402e9af 100644 --- a/src/transform/lower_tile_op.cc +++ b/src/transform/lower_tile_op.cc @@ -1105,6 +1105,12 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { AddWorkspaceCallback add_workspace_callback = [this](int num_elem, DataType dtype) { + // Runtime helpers index scalar FP4 workspace slots as byte-sized + // fp4_e2_t objects. Reserve two packed elements per slot so buffer + // storage sizing and shared-memory merging account for the full byte. + if (dtype.is_float4_e2m1fn() && dtype.is_scalar()) { + num_elem *= 2; + } auto workspace = decl_buffer({PrimExpr(num_elem)}, dtype, "workspace", "shared.dyn"); // Record workspace under the innermost block scope so its lifetime diff --git a/testing/python/language/test_tilelang_language_reduce.py b/testing/python/language/test_tilelang_language_reduce.py index b9a8f621da..525445428f 100644 --- a/testing/python/language/test_tilelang_language_reduce.py +++ b/testing/python/language/test_tilelang_language_reduce.py @@ -1310,21 +1310,35 @@ def kernel(A: T.Tensor((M, N), T.float16), B: T.Tensor((M,), T.float16)): @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(10, 0) @pytest.mark.parametrize(("reduce_fn", "torch_op"), [(T.reduce_max, "max"), (T.reduce_min, "min")], ids=["max", "min"]) -def test_reduce_fp4_dst_runtime(reduce_fn, torch_op): +@pytest.mark.parametrize(("M", "N"), [(32, 64), (2, 512), (1, 1024)]) +def test_reduce_fp4_dst_runtime(reduce_fn, torch_op, M, N): """reduce_max/reduce_min into a float4_e2m1fn dst failed to compile (ambiguous shuffle and comparison) and a bare max/min compared raw encodings (GH-2998).""" - M, N = 32, 64 k = _compile(_make_fp4_reduce_kernel(reduce_fn, M, N)) - vals = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.float16) + vals = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4], dtype=torch.float16) vals = torch.cat([vals, -vals]) a = vals[torch.randint(0, len(vals), (M, N))].cuda() + # A unique extremum in the last column must reach the other warps. + a[:, -1] = 6 if torch_op == "max" else -6 out = k(a) ref = getattr(a, torch_op)(dim=1).values torch.testing.assert_close(out, ref, atol=0, rtol=0) +@tilelang.testing.requires_cuda +def test_reduce_fp4_cross_warp_workspace_codegen(): + target = {"kind": "cuda", "arch": "sm_100a"} + with tvm.target.Target(target): + artifact = tilelang.lower(_make_fp4_reduce_kernel(T.reduce_max, 1, 1024), target=target) + (device_func,) = artifact.device_mod.functions.values() + # AllReduce indexes one byte-sized fp4_e2_t per thread, even though + # ordinary FP4 buffers pack two elements into each byte. + assert "tl::AllReduce= 128 + + @tilelang.testing.requires_cuda def test_fp8_e8m0_non_finite_const_codegen(): """e8m0fnu has no inf and no sign; non-finite constants must clamp to its From d0aecddcfab9c56280ef73c5651664aa32c3decc Mon Sep 17 00:00:00 2001 From: LeiWang1999 Date: Thu, 22 Oct 2026 19:14:25 +0800 Subject: [PATCH 4/7] [CUDA] Delegate FP8/FP4 constants to conversion constructors --- src/backend/common/op/reduce.h | 3 + src/cuda/codegen/codegen_cuda.cc | 48 ++++------ src/op/reducer.cc | 3 + .../language/test_tilelang_language_reduce.py | 96 ++++++++++++++++--- 4 files changed, 105 insertions(+), 45 deletions(-) diff --git a/src/backend/common/op/reduce.h b/src/backend/common/op/reduce.h index af0e8f967d..e5d8ac685a 100644 --- a/src/backend/common/op/reduce.h +++ b/src/backend/common/op/reduce.h @@ -205,6 +205,9 @@ inline PrimExpr MakeInitValue(const ReduceOpNode &op, int vsize = 1) { scalar = make_const(op.dst->dtype, SignedMin(bits)); } else if (is_uint) { scalar = make_const(op.dst->dtype, 0); + } else if (dst_dtype.is_float8_e8m0fnu()) { + // E8M0 conversion takes the magnitude, so -Inf is not a lower bound. + scalar = make_const(dst_dtype, std::ldexp(1.0, -127)); } else { scalar = make_const(op.dst->dtype, -INFINITY); } diff --git a/src/cuda/codegen/codegen_cuda.cc b/src/cuda/codegen/codegen_cuda.cc index b761192f43..c80bdd23eb 100644 --- a/src/cuda/codegen/codegen_cuda.cc +++ b/src/cuda/codegen/codegen_cuda.cc @@ -6228,24 +6228,6 @@ void CodeGenTileLangCUDA::VisitExpr_(const BroadcastNode *op, os << ')'; } -namespace { -// Max finite magnitude for fp8/fp4 dtypes with no (or no non-finite) -// infinity encoding, e.g. the e4m3 family, e8m0fnu and float4_e2m1fn. -double MaxFiniteFp8Fp4(const DataType &dtype) { - if (dtype.is_float8_e4m3fnuz()) { - return 240.0; - } - if (dtype.is_float8_e4m3() || dtype.is_float8_e4m3fn()) { - return 448.0; - } - if (dtype.is_float8_e8m0fnu()) { - return std::ldexp(1.0, 127); - } - ICHECK(dtype.is_float4()) << "No max-finite value for " << dtype; - return 6.0; -} -} // namespace - inline void PrintConst(const FloatImmNode *op, std::ostream &os, CodeGenTileLangCUDA *p) { // NOLINT(*) // Type code is kBFloat/kFloat16 @@ -6308,21 +6290,23 @@ inline void PrintConst(const FloatImmNode *op, std::ostream &os, os << "::bitcast(0x7e)"; return; } - // No literal spelling for non-finite values here (e.g. -inff/nanf are - // invalid identifiers); clamp to the format's max-finite magnitude. - double value = op->value; - if (std::isnan(value)) { - value = MaxFiniteFp8Fp4(op->dtype); - } else if (std::isinf(value) && value < 0 && - op->dtype.is_float8_e8m0fnu()) { - // e8m0fnu is unsigned: its smallest value, 2^-127, is the lower bound. - value = std::ldexp(1.0, -127); - } else if (std::isinf(value)) { - value = std::copysign(MaxFiniteFp8Fp4(op->dtype), value); - } p->PrintType(op->dtype, os); - os << '(' << FlexibleHexFormat(value) << 'f'; - os << "/*" << std::scientific << value << "*/"; + os << '('; + // Use valid FP32 constants and let the conversion constructor handle + // saturation and NaNs according to the destination format. + if (std::isinf(op->value)) { + if (op->value < 0) { + os << '-'; + } + os << "CUDART_INF_F"; + p->need_math_constants_h_ = true; + } else if (std::isnan(op->value)) { + os << "CUDART_NAN_F"; + p->need_math_constants_h_ = true; + } else { + os << FlexibleHexFormat(op->value) << 'f'; + os << "/*" << std::scientific << op->value << "*/"; + } os << ')'; return; } diff --git a/src/op/reducer.cc b/src/op/reducer.cc index 463becd228..b60895114f 100644 --- a/src/op/reducer.cc +++ b/src/op/reducer.cc @@ -327,6 +327,9 @@ PrimExpr ReducerV2Identity(ReducerV2OpType op, DataType dtype) { return make_const(dtype, signed_min()); if (is_uint) return make_const(dtype, 0); + // E8M0 conversion takes the magnitude, so -Inf is not a lower bound. + if (dtype.is_float8_e8m0fnu()) + return make_const(dtype, std::ldexp(1.0, -127)); return make_const(dtype, -INFINITY); case ReducerV2OpType::kMin: if (is_int) diff --git a/testing/python/language/test_tilelang_language_reduce.py b/testing/python/language/test_tilelang_language_reduce.py index 525445428f..c5e2b2bc74 100644 --- a/testing/python/language/test_tilelang_language_reduce.py +++ b/testing/python/language/test_tilelang_language_reduce.py @@ -1339,25 +1339,95 @@ def test_reduce_fp4_cross_warp_workspace_codegen(): assert int(device_func.attrs["dyn_shared_memory_buf"]) >= 128 +def _make_non_finite_const_kernel(dtype): + pos_inf = tvm.tirx.const(float("inf"), dtype) + neg_inf = tvm.tirx.const(float("-inf"), dtype) + nan = tvm.tirx.const(float("nan"), dtype) + + @T.prim_func + def kernel(A: T.Tensor((3,), T.float32), B: T.Tensor((6,), T.float32)): + with T.Kernel(1, threads=1): + values = T.alloc_local((6,), dtype) + values[0] = pos_inf + values[1] = neg_inf + values[2] = nan + for i in T.serial(3): + values[i + 3] = T.cast(A[i], dtype) + for i in T.serial(6): + B[i] = T.cast(values[i], T.float32) + + return kernel + + @tilelang.testing.requires_cuda -def test_fp8_e8m0_non_finite_const_codegen(): - """e8m0fnu has no inf and no sign; non-finite constants must clamp to its - own range, not the fp4 limit of 6.0.""" - pos_inf = tvm.tirx.const(float("inf"), "float8_e8m0fnu") - neg_inf = tvm.tirx.const(float("-inf"), "float8_e8m0fnu") +@pytest.mark.parametrize( + ("dtype", "cuda_type"), + [(T.float8_e4m3, "fp8_e4_t"), (T.float8_e4m3fn, "fp8_e4_t"), (T.float8_e8m0fnu, "fp8_e8_t"), (T.float4_e2m1fn, "fp4_e2_t")], +) +def test_fp8_fp4_non_finite_const_codegen(dtype, cuda_type): + target = {"kind": "cuda", "arch": "sm_100a"} + with tvm.target.Target(target): + src = tilelang.lower(_make_non_finite_const_kernel(dtype), target=target).kernel_source + assert "#include " in src + for value in ("CUDART_INF_F", "-CUDART_INF_F", "CUDART_NAN_F"): + assert f"{cuda_type}({value})" in src + + +def _check_non_finite_const_matches_runtime_conversion(dtype): + kernel = _compile(_make_non_finite_const_kernel(dtype)) + values = torch.tensor([float("inf"), float("-inf"), float("nan")], device="cuda") + output = kernel(values) + torch.testing.assert_close(output[:3], output[3:], atol=0, rtol=0, equal_nan=True) + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(8, 9) +@pytest.mark.parametrize("dtype", [T.float8_e4m3, T.float8_e4m3fn]) +def test_fp8_non_finite_const_matches_runtime_conversion(dtype): + _check_non_finite_const_matches_runtime_conversion(dtype) + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(10, 0) +@pytest.mark.parametrize("dtype", [T.float8_e8m0fnu, T.float4_e2m1fn]) +def test_blackwell_non_finite_const_matches_runtime_conversion(dtype): + _check_non_finite_const_matches_runtime_conversion(dtype) + + +@tilelang.testing.requires_cuda +@pytest.mark.parametrize("op", ["max", "min"]) +@pytest.mark.parametrize("deferred", [False, True], ids=["local", "reducer"]) +def test_reduce_fp8_e8m0_identity_codegen(op, deferred): + reduce_fn = T.reduce_max if op == "max" else T.reduce_min + dtype = T.float8_e8m0fnu @T.prim_func - def kernel(B: T.Tensor((2,), T.float8_e8m0fnu)): - with T.Kernel(1, threads=32): - B[0] = pos_inf - B[1] = neg_inf + def kernel(A: T.Tensor((4,), T.float32), B: T.Tensor((1,), T.float32)): + with T.Kernel(1, threads=1): + if deferred: + acc = T.alloc_reducer((1,), dtype, op=op) + T.reducer_init(acc) + for i in T.Parallel(4): + T.reducer_update(acc[0], T.cast(A[i], dtype)) + result = T.alloc_fragment((1,), dtype) + T.finalize_reducer(acc, result) + B[0] = T.cast(result[0], T.float32) + else: + src = T.alloc_local((4,), dtype) + dst = T.alloc_local((1,), dtype) + for i in T.serial(4): + src[i] = T.cast(A[i], dtype) + reduce_fn(src, dst, dim=0) + B[0] = T.cast(dst[0], T.float32) target = {"kind": "cuda", "arch": "sm_100a"} with tvm.target.Target(target): - src = tilelang.lower(kernel, target=target).kernel_source - assert "fp8_e8_t(0x1p+127f" in src - assert "fp8_e8_t(0x1p-127f" in src - assert "6.000000e+00" not in src + source = tilelang.lower(kernel, target=target).kernel_source + if op == "max": + assert "fp8_e8_t(0x1p-127f" in source + assert "-CUDART_INF_F" not in source + else: + assert "fp8_e8_t(CUDART_INF_F)" in source @tilelang.testing.requires_cuda From 60e38175f99bc14d533275e403b320fc1b703a08 Mon Sep 17 00:00:00 2001 From: LeiWang1999 Date: Mon, 26 Oct 2026 13:14:01 +0800 Subject: [PATCH 5/7] [CUDA] Defer E8M0 reduction identities --- src/backend/common/op/reduce.h | 3 -- src/op/reducer.cc | 3 -- .../language/test_tilelang_language_reduce.py | 36 ------------------- 3 files changed, 42 deletions(-) diff --git a/src/backend/common/op/reduce.h b/src/backend/common/op/reduce.h index e5d8ac685a..af0e8f967d 100644 --- a/src/backend/common/op/reduce.h +++ b/src/backend/common/op/reduce.h @@ -205,9 +205,6 @@ inline PrimExpr MakeInitValue(const ReduceOpNode &op, int vsize = 1) { scalar = make_const(op.dst->dtype, SignedMin(bits)); } else if (is_uint) { scalar = make_const(op.dst->dtype, 0); - } else if (dst_dtype.is_float8_e8m0fnu()) { - // E8M0 conversion takes the magnitude, so -Inf is not a lower bound. - scalar = make_const(dst_dtype, std::ldexp(1.0, -127)); } else { scalar = make_const(op.dst->dtype, -INFINITY); } diff --git a/src/op/reducer.cc b/src/op/reducer.cc index b60895114f..463becd228 100644 --- a/src/op/reducer.cc +++ b/src/op/reducer.cc @@ -327,9 +327,6 @@ PrimExpr ReducerV2Identity(ReducerV2OpType op, DataType dtype) { return make_const(dtype, signed_min()); if (is_uint) return make_const(dtype, 0); - // E8M0 conversion takes the magnitude, so -Inf is not a lower bound. - if (dtype.is_float8_e8m0fnu()) - return make_const(dtype, std::ldexp(1.0, -127)); return make_const(dtype, -INFINITY); case ReducerV2OpType::kMin: if (is_int) diff --git a/testing/python/language/test_tilelang_language_reduce.py b/testing/python/language/test_tilelang_language_reduce.py index c5e2b2bc74..f1d0da21bc 100644 --- a/testing/python/language/test_tilelang_language_reduce.py +++ b/testing/python/language/test_tilelang_language_reduce.py @@ -1394,42 +1394,6 @@ def test_blackwell_non_finite_const_matches_runtime_conversion(dtype): _check_non_finite_const_matches_runtime_conversion(dtype) -@tilelang.testing.requires_cuda -@pytest.mark.parametrize("op", ["max", "min"]) -@pytest.mark.parametrize("deferred", [False, True], ids=["local", "reducer"]) -def test_reduce_fp8_e8m0_identity_codegen(op, deferred): - reduce_fn = T.reduce_max if op == "max" else T.reduce_min - dtype = T.float8_e8m0fnu - - @T.prim_func - def kernel(A: T.Tensor((4,), T.float32), B: T.Tensor((1,), T.float32)): - with T.Kernel(1, threads=1): - if deferred: - acc = T.alloc_reducer((1,), dtype, op=op) - T.reducer_init(acc) - for i in T.Parallel(4): - T.reducer_update(acc[0], T.cast(A[i], dtype)) - result = T.alloc_fragment((1,), dtype) - T.finalize_reducer(acc, result) - B[0] = T.cast(result[0], T.float32) - else: - src = T.alloc_local((4,), dtype) - dst = T.alloc_local((1,), dtype) - for i in T.serial(4): - src[i] = T.cast(A[i], dtype) - reduce_fn(src, dst, dim=0) - B[0] = T.cast(dst[0], T.float32) - - target = {"kind": "cuda", "arch": "sm_100a"} - with tvm.target.Target(target): - source = tilelang.lower(kernel, target=target).kernel_source - if op == "max": - assert "fp8_e8_t(0x1p-127f" in source - assert "-CUDART_INF_F" not in source - else: - assert "fp8_e8_t(CUDART_INF_F)" in source - - @tilelang.testing.requires_cuda def test_reduce_packed_max_nan_propagate_uses_nan_intrinsics(): k = _compile(_make_nan_reduce_kernel(T.reduce_max, 128, 128, T.float16, threads=256, nan_propagate=True)) From a3a030cc9da2f13621aa9a80d30e601a1d551208 Mon Sep 17 00:00:00 2001 From: LeiWang1999 Date: Mon, 26 Oct 2026 14:19:40 +0800 Subject: [PATCH 6/7] Reject sub-byte reduction accumulators and remove FP4 support patches --- src/op/reduce.cc | 8 ++ src/op/reduce.h | 3 + src/op/reducer.cc | 3 + src/tl_templates/cuda/common.h | 11 +-- src/tl_templates/cuda/cuda_fp4.h | 25 ----- src/transform/lower_tile_op.cc | 6 -- .../language/test_tilelang_language_reduce.py | 99 ++++++++++++++----- tilelang/language/allocate.py | 4 +- tilelang/language/reduce_op.py | 5 +- 9 files changed, 100 insertions(+), 64 deletions(-) diff --git a/src/op/reduce.cc b/src/op/reduce.cc index 4f896c22f8..03ef767078 100644 --- a/src/op/reduce.cc +++ b/src/op/reduce.cc @@ -68,6 +68,13 @@ void RegisterReduceImpl(ReduceImpl impl) { ReduceImplRegistry().push_back(impl); } +void CheckReductionDType(DataType dtype, const char *op) { + CHECK(dtype.bits() >= 8 || dtype.is_bool(), ValueError) + << op << " does not support sub-byte accumulator dtype " << dtype + << ". Use float16/float32 or an integer dtype of at least 8 bits for " + "the output/reducer buffer."; +} + ReduceOp::ReduceOp(Array args, Map annotations) { ObjectPtr node = make_object(); node->annotations = annotations; @@ -79,6 +86,7 @@ ReduceOp::ReduceOp(Array args, Map annotations) { node->SetAccessRegions({src_access, dst_access}); node->src = node->srcRegion_->buffer; node->dst = node->dstRegion_->buffer; + CheckReductionDType(node->dst->dtype, "reduce"); std::string reduce_type = args[2].as().value()->value; node->dim = args[3].as().value()->value; node->type = ReduceType(reduce_type); diff --git a/src/op/reduce.h b/src/op/reduce.h index 255da764a9..f01755ca4a 100644 --- a/src/op/reduce.h +++ b/src/op/reduce.h @@ -17,6 +17,9 @@ namespace tl { using namespace tirx; using namespace ffi; +/// Reject sub-byte numeric accumulators; bool reductions remain supported. +void CheckReductionDType(DataType dtype, const char *op); + /// Supported reduction operation types enum class ReduceTypeEnum : uint8_t { kSum, ///< Sum reduction diff --git a/src/op/reducer.cc b/src/op/reducer.cc index 463becd228..64ec23a338 100644 --- a/src/op/reducer.cc +++ b/src/op/reducer.cc @@ -20,6 +20,7 @@ #include "backend/common/op/reduce.h" #include "builtin_registry.h" +#include "reduce.h" #include "utils.h" namespace tvm { @@ -382,6 +383,7 @@ ReducerInitOp::ReducerInitOp(ffi::Array args, access.region = BufferRegion::FullRegion(access.region->buffer); access.access_mask = kAccessWrite; node->reducer = access.region->buffer; + CheckReductionDType(node->reducer->dtype, "reducer_init"); if (args.size() == 2) { ICHECK(args[1].dtype() == node->reducer->dtype) << "reducer_init init value dtype " << args[1].dtype() @@ -615,6 +617,7 @@ FinalizeReducerOp::FinalizeReducerOp( reducer_access.access_mask = kAccessReadWrite; node->reducer = reducer_access.region->buffer; node->SetAccessRegions({reducer_access}); + CheckReductionDType(node->reducer->dtype, "finalize_reducer"); node->op = (ReducerV2OpType)*as_const_int(args[1]); // Optional explicit collective plan (reducer v2 narrow plans): flattened // (reducing_threads, scale) pairs starting at args[2]. diff --git a/src/tl_templates/cuda/common.h b/src/tl_templates/cuda/common.h index a483ccae10..4a2c988d30 100644 --- a/src/tl_templates/cuda/common.h +++ b/src/tl_templates/cuda/common.h @@ -1680,16 +1680,14 @@ using tl::tfloat32_t; int width = warpSize) { \ RAW raw = reinterpret_cast(val); \ RAW ret = static_cast( \ - __shfl_down_sync(mask, static_cast(raw), \ - static_cast(delta), width)); \ + __shfl_down_sync(mask, static_cast(raw), delta, width)); \ return reinterpret_cast(ret); \ } \ TL_PATCH TL_DEVICE TYPE __shfl_up_sync(unsigned mask, TYPE val, int delta, \ int width = warpSize) { \ RAW raw = reinterpret_cast(val); \ - RAW ret = \ - static_cast(__shfl_up_sync(mask, static_cast(raw), \ - static_cast(delta), width)); \ + RAW ret = static_cast( \ + __shfl_up_sync(mask, static_cast(raw), delta, width)); \ return reinterpret_cast(ret); \ } @@ -1697,7 +1695,8 @@ TL_DEFINE_SHFL_SYNC_OVERLOADS(half_t, uint16_t) TL_DEFINE_SHFL_SYNC_OVERLOADS(bfloat16_t, uint16_t) TL_DEFINE_SHFL_SYNC_OVERLOADS(tl::float_e4m3_t, uint8_t) TL_DEFINE_SHFL_SYNC_OVERLOADS(tl::float_e5m2_t, uint8_t) -// Kept defined: cuda_fp4.h also applies it to fp4_e2_t. + +#undef TL_DEFINE_SHFL_SYNC_OVERLOADS // // Optimized type-punned warp shuffle helpers for 16-bit types diff --git a/src/tl_templates/cuda/cuda_fp4.h b/src/tl_templates/cuda/cuda_fp4.h index 12bca175a4..47156c2582 100644 --- a/src/tl_templates/cuda/cuda_fp4.h +++ b/src/tl_templates/cuda/cuda_fp4.h @@ -44,31 +44,6 @@ struct fp4_e2_t { TL_DEVICE operator __half() const { return __half(float(*this)); } }; -// fp4_e2_t converts implicitly to several types, so shuffles and comparisons -// on it are ambiguous, and a bare max/min would compare raw encodings. Shuffle -// the raw bits and compare through float. -TL_DEFINE_SHFL_SYNC_OVERLOADS(fp4_e2_t, uint8_t) - -namespace tl { -template <> TL_DEVICE fp4_e2_t fast_max(fp4_e2_t a, fp4_e2_t b) { - return fp4_e2_t(::fmaxf(static_cast(a), static_cast(b))); -} - -template <> TL_DEVICE fp4_e2_t fast_min(fp4_e2_t a, fp4_e2_t b) { - return fp4_e2_t(::fminf(static_cast(a), static_cast(b))); -} -} // namespace tl - -// Global, not in tl: generated code calls max/min unqualified, and fp4_e2_t -// is a global type, so argument-dependent lookup would not find tl:: ones. -TL_DEVICE fp4_e2_t max(fp4_e2_t lhs, fp4_e2_t rhs) { - return tl::fast_max(lhs, rhs); -} - -TL_DEVICE fp4_e2_t min(fp4_e2_t lhs, fp4_e2_t rhs) { - return tl::fast_min(lhs, rhs); -} - // Tag for tcgen05 unpacked FP4 shared-memory layout. The hardware atom carries // 16 4-bit payload values in the low 64 bits of a 128-bit aligned region. // See: diff --git a/src/transform/lower_tile_op.cc b/src/transform/lower_tile_op.cc index 2ea93374e8..112c5b55f1 100644 --- a/src/transform/lower_tile_op.cc +++ b/src/transform/lower_tile_op.cc @@ -1105,12 +1105,6 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { AddWorkspaceCallback add_workspace_callback = [this](int num_elem, DataType dtype) { - // Runtime helpers index scalar FP4 workspace slots as byte-sized - // fp4_e2_t objects. Reserve two packed elements per slot so buffer - // storage sizing and shared-memory merging account for the full byte. - if (dtype.is_float4_e2m1fn() && dtype.is_scalar()) { - num_elem *= 2; - } auto workspace = decl_buffer({PrimExpr(num_elem)}, dtype, "workspace", "shared.dyn"); // Record workspace under the innermost block scope so its lifetime diff --git a/testing/python/language/test_tilelang_language_reduce.py b/testing/python/language/test_tilelang_language_reduce.py index f1d0da21bc..c0ff013c4b 100644 --- a/testing/python/language/test_tilelang_language_reduce.py +++ b/testing/python/language/test_tilelang_language_reduce.py @@ -1293,13 +1293,74 @@ def test_reduce_fp8_e4m3_dst_runtime(reduce_fn, torch_op): assert torch.allclose(out.float(), ref, atol=1e-1) -def _make_fp4_reduce_kernel(reduce_fn, M, N): - # fp16 I/O holding exact fp4 values avoids packed-fp4 tensors on the host. +@tilelang.testing.requires_cuda +@pytest.mark.parametrize(("dtype", "op"), [(T.float4_e2m1fn, "sum"), (T.float4_e2m1fn, "max"), (T.int4, "min"), ("uint4", "bitand")]) +@pytest.mark.parametrize("scope", ["local", "local.fragment"]) +def test_reduce_rejects_subbyte_accumulator(dtype, op, scope): + @T.prim_func + def kernel(A: T.Tensor((1, 128), T.int32), B: T.Tensor((1,), T.int32)): + with T.Kernel(1, threads=128): + src = T.alloc_buffer((1, 128), T.int32, scope=scope) + dst = T.alloc_buffer((1,), dtype, scope=scope) + T.copy(A, src) + # Exercise the native contract, including TIR that bypasses Python. + T.evaluate( + T.call_intrin("handle", tvm.ir.Op.get("tl.tileop.reduce"), to_tile_region(src, "r"), to_tile_region(dst, "w"), op, 1, True) + ) + T.copy(dst, B) + + with pytest.raises(ValueError, match=f"sub-byte accumulator dtype {dtype}"): + tilelang.lower(kernel, target="cuda") + + +@tilelang.testing.requires_cuda +@pytest.mark.parametrize("dtype", [T.float4_e2m1fn, T.int4]) +@pytest.mark.parametrize("legacy", [False, True], ids=["v2", "legacy"]) +def test_finalize_reducer_rejects_subbyte_accumulator(dtype, legacy): + @T.prim_func + def kernel(A: T.Tensor((128,), T.float32), B: T.Tensor((1,), T.float32)): + with T.Kernel(1, threads=128): + acc = T.alloc_reducer((1,), dtype, op="sum", replication="all" if legacy else None) + if legacy: + T.clear(acc) + for i in T.Parallel(128): + acc[0] += T.cast(A[i], dtype) + T.finalize_reducer(acc) + T.copy(acc, B) + else: + T.reducer_init(acc) + for i in T.Parallel(128): + T.reducer_update(acc[0], A[i]) + dst = T.alloc_fragment((1,), dtype) + T.finalize_reducer(acc, dst) + T.copy(dst, B) + + with pytest.raises(ValueError, match=f"sub-byte accumulator dtype {dtype}"): + tilelang.lower(kernel, target="cuda") + + +def test_reduce_bool_accumulator_runtime(): @T.prim_func - def kernel(A: T.Tensor((M, N), T.float16), B: T.Tensor((M,), T.float16)): + def kernel(A: T.Tensor((2, 128), T.bool), B: T.Tensor((2,), T.bool)): + with T.Kernel(1, threads=128): + src = T.alloc_fragment((2, 128), T.bool) + dst = T.alloc_fragment((2,), T.bool) + T.copy(A, src) + T.reduce_bitor(src, dst, dim=1) + T.copy(dst, B) + + a = torch.zeros((2, 128), device="cuda", dtype=torch.bool) + a[1, -1] = True + torch.testing.assert_close(_compile(kernel)(a), a.any(dim=1)) + + +def _make_fp4_reduce_kernel(reduce_fn, M, N, dtype): + # FP16 input holding exact FP4 values avoids packed FP4 tensors on the host. + @T.prim_func + def kernel(A: T.Tensor((M, N), T.float16), B: T.Tensor((M,), dtype)): with T.Kernel(1, threads=128): src = T.alloc_fragment((M, N), T.float4_e2m1fn) - dst = T.alloc_fragment((M,), T.float4_e2m1fn) + dst = T.alloc_fragment((M,), dtype) T.copy(A, src) reduce_fn(src, dst, dim=1) T.copy(dst, B) @@ -1309,36 +1370,24 @@ def kernel(A: T.Tensor((M, N), T.float16), B: T.Tensor((M,), T.float16)): @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(10, 0) -@pytest.mark.parametrize(("reduce_fn", "torch_op"), [(T.reduce_max, "max"), (T.reduce_min, "min")], ids=["max", "min"]) -@pytest.mark.parametrize(("M", "N"), [(32, 64), (2, 512), (1, 1024)]) -def test_reduce_fp4_dst_runtime(reduce_fn, torch_op, M, N): - """reduce_max/reduce_min into a float4_e2m1fn dst failed to compile - (ambiguous shuffle and comparison) and a bare max/min compared raw - encodings (GH-2998).""" - k = _compile(_make_fp4_reduce_kernel(reduce_fn, M, N)) +@pytest.mark.parametrize( + ("reduce_fn", "torch_op"), [(T.reduce_sum, "sum"), (T.reduce_max, "amax"), (T.reduce_min, "amin")], ids=["sum", "max", "min"] +) +@pytest.mark.parametrize("dtype", [T.float16, T.float32]) +def test_reduce_fp4_input_wider_accumulator_runtime(reduce_fn, torch_op, dtype): + M, N = 1, 1024 + k = _compile(_make_fp4_reduce_kernel(reduce_fn, M, N, dtype)) vals = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4], dtype=torch.float16) vals = torch.cat([vals, -vals]) a = vals[torch.randint(0, len(vals), (M, N))].cuda() # A unique extremum in the last column must reach the other warps. - a[:, -1] = 6 if torch_op == "max" else -6 + a[:, -1] = -6 if torch_op == "amin" else 6 out = k(a) - ref = getattr(a, torch_op)(dim=1).values + ref = getattr(a.to(getattr(torch, dtype)), torch_op)(dim=1) torch.testing.assert_close(out, ref, atol=0, rtol=0) -@tilelang.testing.requires_cuda -def test_reduce_fp4_cross_warp_workspace_codegen(): - target = {"kind": "cuda", "arch": "sm_100a"} - with tvm.target.Target(target): - artifact = tilelang.lower(_make_fp4_reduce_kernel(T.reduce_max, 1, 1024), target=target) - (device_func,) = artifact.device_mod.functions.values() - # AllReduce indexes one byte-sized fp4_e2_t per thread, even though - # ordinary FP4 buffers pack two elements into each byte. - assert "tl::AllReduce= 128 - - def _make_non_finite_const_kernel(dtype): pos_inf = tvm.tirx.const(float("inf"), dtype) neg_inf = tvm.tirx.const(float("-inf"), dtype) diff --git a/tilelang/language/allocate.py b/tilelang/language/allocate.py index c84391bc0e..dd47ac884b 100644 --- a/tilelang/language/allocate.py +++ b/tilelang/language/allocate.py @@ -314,7 +314,9 @@ def alloc_reducer(shape: ShapeType, dtype: DType, op: ReducerOp = "sum", replica Args: shape (tuple): Logical shape of the reduction result. - dtype (str): Element data type (e.g., 'float32', 'int32'). + dtype (str): Accumulator data type (e.g., 'float32', 'int32'). Sub-byte + numeric dtypes (such as FP4/int4) are unsupported; use float16/float32 + or an integer dtype of at least 8 bits. Boolean reductions are allowed. op (str): Combine op: "sum", "max", "min", "bitand", "bitor" or "bitxor" (the bitwise ops require an integer dtype). replication (str | None): Deprecated legacy (v1) knob. Passing "all" diff --git a/tilelang/language/reduce_op.py b/tilelang/language/reduce_op.py index 5bed017658..406e3c8c98 100644 --- a/tilelang/language/reduce_op.py +++ b/tilelang/language/reduce_op.py @@ -35,7 +35,10 @@ def reduce( Args: buffer (tirx.Buffer): Input buffer to reduce - out (tirx.Buffer): Output buffer to store results + out (tirx.Buffer): Output buffer whose dtype also determines accumulation + precision. Sub-byte numeric dtypes (such as FP4/int4) are unsupported; + use float16/float32 or an integer dtype of at least 8 bits. Boolean + reductions are allowed. The input may use a narrower dtype. reduce_type (str): Type of reduction ('max', 'min', 'sum', 'abssum') dim (int): Dimension along which to perform reduction clear (bool): Whether to initialize the output buffer before reduction From 9f0d8111dd21af47da41ceb2ae26a45e6ee56edb Mon Sep 17 00:00:00 2001 From: LeiWang1999 Date: Mon, 26 Oct 2026 15:42:59 +0800 Subject: [PATCH 7/7] Remove boolean reduction test that forces CUDA on ROCm --- .../language/test_tilelang_language_reduce.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/testing/python/language/test_tilelang_language_reduce.py b/testing/python/language/test_tilelang_language_reduce.py index c0ff013c4b..8f1a7eb8bc 100644 --- a/testing/python/language/test_tilelang_language_reduce.py +++ b/testing/python/language/test_tilelang_language_reduce.py @@ -1339,21 +1339,6 @@ def kernel(A: T.Tensor((128,), T.float32), B: T.Tensor((1,), T.float32)): tilelang.lower(kernel, target="cuda") -def test_reduce_bool_accumulator_runtime(): - @T.prim_func - def kernel(A: T.Tensor((2, 128), T.bool), B: T.Tensor((2,), T.bool)): - with T.Kernel(1, threads=128): - src = T.alloc_fragment((2, 128), T.bool) - dst = T.alloc_fragment((2,), T.bool) - T.copy(A, src) - T.reduce_bitor(src, dst, dim=1) - T.copy(dst, B) - - a = torch.zeros((2, 128), device="cuda", dtype=torch.bool) - a[1, -1] = True - torch.testing.assert_close(_compile(kernel)(a), a.any(dim=1)) - - def _make_fp4_reduce_kernel(reduce_fn, M, N, dtype): # FP16 input holding exact FP4 values avoids packed FP4 tensors on the host. @T.prim_func