diff --git a/tests/end_to_end/test_sanitizer.py b/tests/end_to_end/test_sanitizer.py index a039a906..542a1691 100644 --- a/tests/end_to_end/test_sanitizer.py +++ b/tests/end_to_end/test_sanitizer.py @@ -1004,6 +1004,61 @@ def test_tl_min_return_indices(): ), f"Sanitizer reported {len(reduce_indices_sanitizer.records)} error(s)" +# ======== Reduce Result Dtype Tests =========== + + +@tilelens.trace(client=SymbolicSanitizer()) +@triton.jit +def fp16_softmax_kernel(x_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + # tl.max promotes fp16 to fp32, so tl.exp gets an fp32 argument. + e = tl.exp(x - tl.max(x, 0)) + tl.store(out_ptr + offs, e / tl.sum(e, 0)) + + +def test_reduce_max_promotes_fp16_like_triton(): + x = torch.randn((8,), dtype=torch.float16) + out = torch.empty_like(x) + fp16_softmax_kernel[(1,)](x, out, BLOCK=x.numel()) + + +@tilelens.trace(client=SymbolicSanitizer()) +@triton.jit +def reduce_dtypes_kernel(i8_ptr, u8_ptr, f16_ptr, bf16_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + i8 = tl.load(i8_ptr + offs) + u8 = tl.load(u8_ptr + offs) + f16 = tl.load(f16_ptr + offs) + bf16 = tl.load(bf16_ptr + offs) + tl.static_assert(tl.sum(i8, 0).dtype == tl.int32) + tl.static_assert(tl.sum(u8, 0).dtype == tl.uint32) + tl.static_assert(tl.sum(f16, 0).dtype == tl.float16) + tl.static_assert(tl.sum(f16, 0, dtype=tl.float32).dtype == tl.float32) + tl.static_assert(tl.max(i8, 0).dtype == tl.int32) + tl.static_assert(tl.min(u8, 0).dtype == tl.int32) + tl.static_assert(tl.max(f16, 0).dtype == tl.float32) + tl.static_assert(tl.min(bf16, 0).dtype == tl.float32) + f16_max, f16_idx = tl.max(f16, 0, return_indices=True) + tl.static_assert(f16_max.dtype == tl.float16) + tl.static_assert(f16_idx.dtype == tl.int32) + bf16_min, _ = tl.min(bf16, 0, return_indices=True) + tl.static_assert(bf16_min.dtype == tl.float32) + + +def test_reduce_result_dtypes_match_triton(): + """tl.sum, tl.max, and tl.min promote small integer and float inputs as + triton.language.standard does.""" + n = 8 + reduce_dtypes_kernel[(1,)]( + torch.ones((n,), dtype=torch.int8), + torch.ones((n,), dtype=torch.uint8), + torch.ones((n,), dtype=torch.float16), + torch.ones((n,), dtype=torch.bfloat16), + BLOCK=n, + ) + + # ======== Reduce + Broadcast Tests =========== reduce_broadcast_sanitizer = SymbolicSanitizer() diff --git a/tilelens/clients/symbolic_engine.py b/tilelens/clients/symbolic_engine.py index af435c88..678133ba 100644 --- a/tilelens/clients/symbolic_engine.py +++ b/tilelens/clients/symbolic_engine.py @@ -87,9 +87,11 @@ TensorPointerStore, ) from ..core.symbolic_metadata import ( + BFLOAT16, FLOAT32, INT1, INT32, + UINT32, DTYPE_BY_NAME, SymbolicDType, SymbolicPointerDType, @@ -1613,7 +1615,14 @@ class ReduceSymbolicExpr(SymbolicExpr): keepdims: SymbolicExpr axis: SymbolicExpr | None - def __init__(self, op: str, input: Any, axis: Any = None, keepdims: bool = False): + def __init__( + self, + op: str, + input: Any, + axis: Any = None, + keepdims: bool = False, + dtype: Any = None, + ): if op not in self._SUPPORTED_OPS: raise NotImplementedError(f"Unsupported reduce op: {op}") super().__init__(op) @@ -1650,6 +1659,8 @@ def __init__(self, op: str, input: Any, axis: Any = None, keepdims: bool = False scalar_ty: SymbolicDType if op in ("argmax", "argmin"): scalar_ty = INT32 + elif dtype is not None: + scalar_ty = SymbolicExpr._NORMALIZE_VALUE(dtype) else: assert self.input.dtype is not None scalar_ty = self.input.dtype @@ -2698,10 +2709,39 @@ def _op_expand_dims_overrider(self, arg, axis): def _op_broadcast_overrider(self, arg, shape): return SymbolicExpr.create("broadcast", SymbolicExpr.from_value(arg), shape) - def _op_reduce_sum_overrider(self, input, axis=None, keep_dims=False, **kwargs): - return SymbolicExpr.create( - "sum", SymbolicExpr.from_value(input), axis, keep_dims - ) + # tl.sum, tl.max, and tl.min convert their input before reducing + # (triton.language.standard); these give the resulting dtype. + @staticmethod + def _reduce_sum_dtype(dtype: Any) -> Any: + if not isinstance(dtype, SymbolicScalarDType) or dtype.primitive_bitwidth >= 32: + return dtype + if dtype.name.startswith("int") and dtype != INT1: + return INT32 + if dtype.name.startswith("uint") or dtype == INT1: + return UINT32 + return dtype + + @staticmethod + def _reduce_minmax_dtype(dtype: Any, return_indices: bool) -> Any: + if dtype == BFLOAT16: + return FLOAT32 + if ( + return_indices + or not isinstance(dtype, SymbolicScalarDType) + or dtype.primitive_bitwidth >= 32 + ): + return dtype + if dtype.name.startswith(("fp", "bf", "float")): + return FLOAT32 + return INT32 + + def _op_reduce_sum_overrider( + self, input, axis=None, keep_dims=False, dtype=None, **kwargs + ): + input_sym = SymbolicExpr.from_value(input) + if dtype is None: + dtype = self._reduce_sum_dtype(input_sym.dtype) + return SymbolicExpr.create("sum", input_sym, axis, keep_dims, dtype) def _op_reduce_xor_overrider(self, input, axis=None, keep_dims=False, **kwargs): return SymbolicExpr.create( @@ -2715,16 +2755,20 @@ def _op_reduce_or_overrider(self, input, axis=None, keep_dims=False, **kwargs): def _op_reduce_max_overrider(self, input, axis=None, keep_dims=False, **kwargs): input_sym = SymbolicExpr.from_value(input) - val = SymbolicExpr.create("max", input_sym, axis, keep_dims) - if kwargs.get("return_indices", False): + return_indices = kwargs.get("return_indices", False) + dtype = self._reduce_minmax_dtype(input_sym.dtype, return_indices) + val = SymbolicExpr.create("max", input_sym, axis, keep_dims, dtype) + if return_indices: idx = SymbolicExpr.create("argmax", input_sym, axis, keep_dims) return (val, idx) return val def _op_reduce_min_overrider(self, input, axis=None, keep_dims=False, **kwargs): input_sym = SymbolicExpr.from_value(input) - val = SymbolicExpr.create("min", input_sym, axis, keep_dims) - if kwargs.get("return_indices", False): + return_indices = kwargs.get("return_indices", False) + dtype = self._reduce_minmax_dtype(input_sym.dtype, return_indices) + val = SymbolicExpr.create("min", input_sym, axis, keep_dims, dtype) + if return_indices: idx = SymbolicExpr.create("argmin", input_sym, axis, keep_dims) return (val, idx) return val