[CUDA][Reduce] Fix non-finite constants and reject sub-byte accumulators

src/cuda/codegen/codegen_cuda.cc

@@ -6261,8 +6261,22 @@ inline void PrintConst(const FloatImmNode *op, std::ostream &os,
       return;
     }
     p->PrintType(op->dtype, os);
-    os << '(' << FlexibleHexFormat(op->value) << 'f';
-    os << "/*" << std::scientific << op->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;
   }

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<PrimExpr> args, Map<String, ObjectRef> annotations) {
   ObjectPtr<ReduceOpNode> node = make_object<ReduceOpNode>();
   node->annotations = annotations;
@@ -79,6 +86,7 @@ ReduceOp::ReduceOp(Array<PrimExpr> args, Map<String, ObjectRef> 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<StringImm>().value()->value;
   node->dim = args[3].as<IntImm>().value()->value;
   node->type = ReduceType(reduce_type);

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

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<PrimExpr> 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].

testing/python/language/test_tilelang_language_reduce.py

@@ -1261,6 +1261,173 @@ 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
+@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 _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,), dtype)
+            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_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 == "amin" else 6
+    out = k(a)
+    ref = getattr(a.to(getattr(torch, dtype)), torch_op)(dim=1)
+    torch.testing.assert_close(out, ref, atol=0, rtol=0)
+
+
+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
+@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 <math_constants.h>" 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
 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))

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"

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