diff --git a/tests/end_to_end/test_sanitizer.py b/tests/end_to_end/test_sanitizer.py index a039a906..a5af8cb2 100644 --- a/tests/end_to_end/test_sanitizer.py +++ b/tests/end_to_end/test_sanitizer.py @@ -1209,6 +1209,45 @@ def test_sanitizer_supports_data_dependent_cumsum_index(): assert len(cumsum_sanitizer.records) == 0 +reduce_index_sanitizer = SymbolicSanitizer(abort_on_error=False) + + +@tilelens.trace(client=reduce_index_sanitizer) +@triton.jit +def reduce_offset_kernel(idx_ptr, src_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + idx = tl.load(idx_ptr + offs) + # Each offset is a reduction of loaded values, so the sanitizer must + # concretize the reduction. + tl.store(out_ptr + 0, tl.load(src_ptr + tl.max(idx, 0))) + tl.store(out_ptr + 1, tl.load(src_ptr + tl.min(idx, 0))) + tl.store(out_ptr + 2, tl.load(src_ptr + tl.argmax(idx, 0))) + tl.store(out_ptr + 3, tl.load(src_ptr + tl.argmin(idx, 0))) + tl.store(out_ptr + 4, tl.load(src_ptr + tl.sum(idx, 0))) + tl.store(out_ptr + 5, tl.load(src_ptr + tl.xor_sum(idx, 0))) + tl.store(out_ptr + 6, tl.load(src_ptr + tl.reduce_or(idx, 0))) + + +@pytest.mark.parametrize( + "src_numel, oob_offsets", [(16, []), (7, [7]), (5, [5, 5, 5, 7])] +) +def test_sanitizer_concretizes_reductions_in_addresses(src_numel, oob_offsets): + reduce_index_sanitizer.records.clear() + + # max 5, min 0, argmax 1, argmin 3, sum 7, xor 5, or 5 + idx = torch.tensor([1, 5, 1, 0], dtype=torch.int32) + src = torch.zeros((src_numel,), dtype=torch.float32) + out = torch.empty((7,), dtype=torch.float32) + + reduce_offset_kernel[(1,)](idx, src, out, BLOCK=idx.numel()) + + offsets = sorted( + (record.violation_address - src.data_ptr()) // src.element_size() + for record in reduce_index_sanitizer.records + ) + assert offsets == oob_offsets + + ashr_index_sanitizer = SymbolicSanitizer(abort_on_error=False) diff --git a/tilelens/clients/symbolic_engine.py b/tilelens/clients/symbolic_engine.py index af435c88..d72dc8f7 100644 --- a/tilelens/clients/symbolic_engine.py +++ b/tilelens/clients/symbolic_engine.py @@ -1662,6 +1662,33 @@ def _to_z3_impl(self) -> tuple[Z3Expr, ConstraintConjunction]: raise NotImplementedError(f"Eval for op {self.op} is not implemented") return handler(self) + def concretize(self) -> Any: + input_concrete = self.input.concretize() + data = input_concrete.data + axis = self.axis.to_py() if self.axis is not None else None + axis = int(axis) if axis is not None else None + keepdims = bool(self.keepdims.to_py()) + out_np = dtype_to_numpy(self.dtype) + is_float = np.issubdtype(data.dtype, np.floating) + if self.op == "sum": + out = np.sum(data.astype(out_np), axis=axis, keepdims=keepdims) + elif self.op in ("max", "min"): + # Triton's max and min ignore NaNs (maxnumf/minnumf), like fmax/fmin. + if self.op == "max": + ufunc = np.fmax if is_float else np.maximum + else: + ufunc = np.fmin if is_float else np.minimum + out = ufunc.reduce(data.astype(out_np), axis=axis, keepdims=keepdims) + elif self.op == "argmax": + out = np.argmax(data, axis=axis, keepdims=keepdims) + elif self.op == "argmin": + out = np.argmin(data, axis=axis, keepdims=keepdims) + elif self.op == "xor_sum": + out = np.bitwise_xor.reduce(data, axis=axis, keepdims=keepdims) + else: + out = np.bitwise_or.reduce(data, axis=axis, keepdims=keepdims) + return SymbolicTensorValue(np.asarray(out).astype(out_np), self.dtype) + def _reduce_sum(self) -> tuple[Z3Expr, ConstraintConjunction]: arr, constraints = self.input._to_z3() return Sum(arr), constraints