[FIX] Concretize reductions in data-dependent addresses

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)
 
 

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