[FEAT] Add if/else branch condition tracking to sanitizer

tests/end_to_end/test_sanitizer.py

@@ -308,6 +308,84 @@ def test_loop_deferred_checks_simplify():
     assert load_index_checker.observed_offsets == []
 
 
+# ======== Branch Condition Tracking Tests ===========
+
+
+branch_checker: SymbolicSanitizer = SymbolicSanitizer(abort_on_error=False)
+
+
+@triton_viz.trace(client=branch_checker)
+@triton.jit
+def branch_false_positive_kernel(y_ptr, N: tl.constexpr):
+    """Load inside if branch that guards against OOB.
+
+    Without branch tracking, sanitizer would report OOB because
+    the symbolic loop var t could be 0, making (t - 1) = -1 which is OOB.
+    With branch tracking, the constraint t > 0 restricts t to [1], so
+    (t - 1) = 0 which is in bounds for a length-1 scalar tensor.
+    """
+    for t in tl.static_range(0, 2):
+        if t > 0:
+            tl.load(y_ptr + t - 1)  # t=1 → offset 0, safe for scalar
+
+
+def test_branch_false_positive_elimination():
+    """Sanitizer should NOT report OOB for loads guarded by branch conditions."""
+    branch_checker.records.clear()
+    y = torch.ones(1, dtype=torch.float32)
+    branch_false_positive_kernel[(1,)](y, N=1)
+    assert len(branch_checker.records) == 0, (
+        f"Expected 0 OOB records (branch should prevent false alarm), "
+        f"got {len(branch_checker.records)}"
+    )
+
+
+@triton_viz.trace(client=branch_checker)
+@triton.jit
+def branch_different_signatures_kernel(ptr, N: tl.constexpr):
+    """Same addr expr under different branches should get distinct signatures.
+
+    With t > 0 (t=1): load ptr + t = ptr + 1, OOB for length-1 tensor
+    With t == 0 (t=0): load ptr + t = ptr + 0, safe for length-1 tensor
+    """
+    for t in tl.static_range(0, 2):
+        if t > 0:
+            tl.load(ptr + t)  # t=1 → OOB for scalar
+        else:
+            tl.load(ptr + t)  # t=0 → safe for scalar
+
+
+def test_branch_different_signatures():
+    """Same addr expr under different branch conditions should not dedup."""
+    branch_checker.records.clear()
+    x = torch.ones(1, dtype=torch.float32)
+    branch_different_signatures_kernel[(1,)](x, N=1)
+    # The load at t=1 with offset 1 exceeds length-1 tensor
+    assert (
+        len(branch_checker.records) > 0
+    ), "Expected OOB for t=1 branch (offset exceeds tensor)"
+
+
+@triton_viz.trace(client=branch_checker)
+@triton.jit
+def branch_pid_immediate_kernel(ptr):
+    """Load guarded by pid == 0, tested with grid > 1."""
+    pid = tl.program_id(0)
+    if pid == 0:
+        tl.load(ptr + pid)
+
+
+def test_branch_pid_immediate():
+    """Sanitizer should NOT report OOB for pid-guarded immediate check."""
+    branch_checker.records.clear()
+    x = torch.ones(1, dtype=torch.float32)
+    # grid=2 means pid=0 and pid=1; load only happens for pid=0
+    branch_pid_immediate_kernel[(2,)](x)
+    assert (
+        len(branch_checker.records) == 0
+    ), f"Expected 0 OOB records (pid==0 guard), got {len(branch_checker.records)}"
+
+
 # Dedicated sanitizer for nested loop regression test
 nested_loop_checker = SymbolicSanitizer()
 

tests/unit/test_sanitizer.py

@@ -540,6 +540,41 @@ def test_load_dtype_block_of_pointers():
     assert load.dtype == tl.float32, f"Expected tl.float32, got {load.dtype}"
 
 
+def test_if_else_ast_transformer_structure():
+    """Verify the AST transformer produces the expected try/finally structure."""
+    import ast
+    from triton_viz.transformers.if_else_patcher import _visit_If
+
+    source = "if cond:\n    x = 1\nelse:\n    x = 2\n"
+    tree = ast.parse(source)
+    if_node = tree.body[0]
+
+    class FakeTransformer:
+        def generic_visit(self, node):
+            return node
+
+    transformer = FakeTransformer()
+    result = _visit_If(transformer, if_node)
+
+    # Should return a list of [pre_if Expr, Try node]
+    assert isinstance(result, list)
+    assert len(result) == 2
+    assert isinstance(result[0], ast.Expr)  # pre_if call
+    assert isinstance(result[1], ast.Try)  # try/finally
+
+    try_node = result[1]
+    assert len(try_node.finalbody) == 1  # post_if in finally
+    assert isinstance(try_node.finalbody[0], ast.Expr)
+
+    # The try body should be a single If node
+    assert len(try_node.body) == 1
+    assert isinstance(try_node.body[0], ast.If)
+
+    inner_if = try_node.body[0]
+    # The else branch should start with flip_condition
+    assert len(inner_if.orelse) >= 2  # flip_condition + original else body
+
+
 def test_store_dtype_block_of_pointers():
     """tl.store on a block of pointers should not derive a dtype (store returns None).
 

triton_viz/clients/sanitizer/sanitizer.py

@@ -45,6 +45,7 @@
     PendingCheck,
     Z3Expr,
     ConstraintConjunction,
+    _and_constraints,
 )
 from .data import OutOfBoundsRecordZ3
 from ...utils.traceback_utils import (
@@ -322,6 +323,12 @@ def _handle_access_check(self, expr: SymbolicExpr) -> None:
         """
         # check memory access using z3
         z3_addr, z3_constraints = expr.eval()
+
+        # Merge current branch path constraints into z3_constraints
+        if self.branch_stack:
+            path_parts = [f.path_constraint() for f in self.branch_stack]
+            z3_constraints = _and_constraints(z3_constraints, *path_parts)
+
         if not self.loop_stack:
             self._check_range_satisfiable(z3_addr, z3_constraints, expr)
             return
@@ -643,6 +650,9 @@ def _loop_hook_after(self, lineno: int) -> None:
     def register_for_loop_callback(self) -> ForLoopCallbacks:
         return SymbolicClient.register_for_loop_callback(self)
 
+    def register_if_else_callback(self):
+        return SymbolicClient.register_if_else_callback(self)
+
     def finalize(self) -> list:
         return []
 

triton_viz/clients/symbolic_engine.py

@@ -20,19 +20,23 @@
     Sum,
     And,
     Or,
+    Not,
     simplify,
+    substitute,
     Int2BV,
     BV2Int,
     BitVecRef,
     BoolVal,
+    is_true,
+    is_false,
 )
 from z3.z3 import BoolRef, ArithRef, IntNumRef, ExprRef, Tactic, Probe
 
 import triton.language as tl
 from triton.runtime.interpreter import TensorHandle, _get_np_dtype
 
 from ..core.client import Client
-from ..core.callbacks import OpCallbacks, ForLoopCallbacks
+from ..core.callbacks import OpCallbacks, ForLoopCallbacks, IfElseCallbacks
 from ..core.data import (
     Op,
     UnaryOp,
@@ -130,10 +134,23 @@ class LoopContext:
     start: int = 0
     stop: int = 0
     step: int = 1
+    current_concrete_idx: int | None = None
     signature_cache: dict[int, int] = field(default_factory=dict)
     pending_checks: list[PendingCheck] = field(default_factory=list)
 
 
+@dataclass
+class IfFrame:
+    lineno: int
+    predicate: BoolRef  # The condition Z3 expr (negated for else)
+    support: BoolRef | None  # Supporting constraints from expr.eval()
+    concrete_result: bool  # Concrete evaluation result
+
+    def path_constraint(self) -> ConstraintConjunction:
+        """Return combined predicate + support constraints for this branch."""
+        return _and_constraints(self.support, self.predicate)
+
+
 class SymbolicExprDataWrapper:
     """
     This wrapper is used as a workaround of triton interpreter legacy code.
@@ -580,6 +597,10 @@ def has_op(self, op_name: str) -> bool:
         self._has_op_cache[op_name] = False
         return False
 
+    def has_op_in(self, op_names: frozenset[str]) -> bool:
+        """Return True when the subtree contains any op in the given set."""
+        return any(self.has_op(name) for name in op_names)
+
     def to_tree_str(self) -> str:
         """
         Render the AST as an ASCII tree using anytree.RenderTree.
@@ -1755,6 +1776,7 @@ class SymbolicClient(Client):
     def __init__(self):
         super().__init__()
         self.loop_stack: list[LoopContext] = []
+        self.branch_stack: list[IfFrame] = []
         SymbolicExpr.set_loop_ctx_provider(
             lambda *_args, **_kwargs: (self.loop_stack[-1] if self.loop_stack else None)
         )
@@ -2095,6 +2117,7 @@ def _loop_hook_iter_overrider(self, lineno, idx):
         if self._should_skip_loop_hooks():
             return idx
         if self.loop_stack and self.loop_stack[-1].lineno == lineno:
+            self.loop_stack[-1].current_concrete_idx = idx
             return self.loop_stack[-1].idx
         return idx
 
@@ -2112,3 +2135,100 @@ def register_for_loop_callback(self) -> ForLoopCallbacks:
             loop_iter_overrider=self.lock_fn(self._loop_hook_iter_overrider),
             after_loop_callback=self.lock_fn(self._loop_hook_after),
         )
+
+    # ── If/else branch tracking ──────────────────────────────────
+
+    _LOAD_LIKE_OPS: ClassVar[frozenset[str]] = frozenset(
+        {"load", "tensor_pointer_load", "store", "tensor_pointer_store"}
+    )
+
+    def _if_pre_if(self, condition: Any, lineno: int) -> None:
+        # Extract SymbolicExpr from condition
+        sym_expr: SymbolicExpr | None = None
+        if isinstance(condition, SymbolicExpr):
+            sym_expr = condition
+        elif hasattr(condition, "handle") and isinstance(
+            condition.handle, SymbolicExpr
+        ):
+            sym_expr = condition.handle
+
+        # Fallback for unsupported conditions
+        if sym_expr is None or not isinstance(sym_expr, SymbolicExpr):
+            self.branch_stack.append(
+                IfFrame(lineno, BoolVal(bool(condition)), None, bool(condition))
+            )
+            return
+
+        # Check for load-like ops in the expression tree
+        if sym_expr.has_op_in(self._LOAD_LIKE_OPS):
+            self.branch_stack.append(
+                IfFrame(lineno, BoolVal(bool(condition)), None, bool(condition))
+            )
+            return
+
+        try:
+            pred_raw, support_z3 = sym_expr.eval()
+        except Exception:
+            self.branch_stack.append(
+                IfFrame(lineno, BoolVal(bool(condition)), None, bool(condition))
+            )
+            return
+
+        # Non-scalar Z3 expr: fallback
+        if isinstance(pred_raw, list):
+            self.branch_stack.append(
+                IfFrame(lineno, BoolVal(bool(condition)), None, bool(condition))
+            )
+            return
+
+        # Normalize predicate to BoolRef
+        pred_bool = _constraint_to_bool(pred_raw)
+
+        # Concrete materialization: substitute loop idx and PID vars
+        substitutions: list[tuple[ExprRef, ExprRef]] = []
+        for ctx in self.loop_stack:
+            if ctx.current_concrete_idx is not None:
+                substitutions.append((ctx.idx_z3, IntVal(ctx.current_concrete_idx)))
+        if self.grid_idx is not None:
+            substitutions.append((SymbolicExpr.PID0, IntVal(self.grid_idx[0])))
+            substitutions.append((SymbolicExpr.PID1, IntVal(self.grid_idx[1])))
+            substitutions.append((SymbolicExpr.PID2, IntVal(self.grid_idx[2])))
+
+        try:
+            concrete_pred = pred_bool
+            if substitutions:
+                concrete_pred = substitute(pred_bool, substitutions)
+            concrete_pred = simplify(concrete_pred)
+            if is_true(concrete_pred):
+                concrete_bool = True
+            elif is_false(concrete_pred):
+                concrete_bool = False
+            else:
+                # Cannot determine concretely; default to True (current behavior)
+                concrete_bool = True
+        except Exception:
+            concrete_bool = bool(condition)
+
+        self.branch_stack.append(IfFrame(lineno, pred_bool, support_z3, concrete_bool))
+
+    def _if_eval_condition(self, lineno: int) -> bool:
+        assert self.branch_stack and self.branch_stack[-1].lineno == lineno
+        return self.branch_stack[-1].concrete_result
+
+    def _if_flip_condition(self, lineno: int) -> None:
+        assert self.branch_stack and self.branch_stack[-1].lineno == lineno
+        frame = self.branch_stack[-1]
+        frame.predicate = Not(frame.predicate)
+        frame.concrete_result = not frame.concrete_result
+
+    def _if_post_if(self, lineno: int) -> None:
+        assert self.branch_stack and self.branch_stack[-1].lineno == lineno
+        self.branch_stack.pop()
+
+    def register_if_else_callback(self) -> IfElseCallbacks:
+        return IfElseCallbacks(
+            pre_if_callback=self.lock_fn(self._if_pre_if),
+            eval_condition_callback=self.lock_fn(self._if_eval_condition),
+            flip_condition_callback=self.lock_fn(self._if_flip_condition),
+            post_if_callback=self.lock_fn(self._if_post_if),
+        )

triton_viz/core/callbacks.py

@@ -19,3 +19,11 @@ class ForLoopCallbacks:
     loop_iter_overrider: Callable | None = None
     loop_iter_listener: Callable | None = None
     after_loop_callback: Callable | None = None
+
+
+@dataclass
+class IfElseCallbacks:
+    pre_if_callback: Callable | None = None  # (condition, lineno)
+    eval_condition_callback: Callable | None = None  # (lineno) -> bool
+    flip_condition_callback: Callable | None = None  # (lineno)
+    post_if_callback: Callable | None = None  # (lineno)