From 0fcf65a1c843b55ca64e1919e77c145de70032fd Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Mon, 16 Mar 2026 18:59:28 +0000 Subject: [PATCH] [FEAT] Add if/else branch condition tracking to sanitizer Eliminate false OOB alarms for memory accesses inside guarded branches within loops. The sanitizer now tracks branch conditions symbolically and merges them into Z3 constraints for deferred and immediate checks. Architecture mirrors the existing for-loop patching pattern: - AST transformer patches visit_If on the Triton interpreter - IfFrame tracks predicate, support constraints, and concrete result - Branch constraints are merged into PendingCheck.constraints - Concrete materialization substitutes loop idx and PID vars Known limitations: - Compound boolean expressions (and/or/not) lose symbolic info - Data-dependent conditions (tl.load results) fall back to defaults --- tests/end_to_end/test_sanitizer.py | 78 +++++++++++++ tests/unit/test_sanitizer.py | 35 ++++++ triton_viz/clients/sanitizer/sanitizer.py | 10 ++ triton_viz/clients/symbolic_engine.py | 122 ++++++++++++++++++++- triton_viz/core/callbacks.py | 8 ++ triton_viz/core/client.py | 63 ++++++++++- triton_viz/core/patch.py | 51 ++++++++- triton_viz/transformers/if_else_patcher.py | 103 +++++++++++++++++ 8 files changed, 466 insertions(+), 4 deletions(-) create mode 100644 triton_viz/transformers/if_else_patcher.py diff --git a/tests/end_to_end/test_sanitizer.py b/tests/end_to_end/test_sanitizer.py index ebf63630..7c22d3e0 100644 --- a/tests/end_to_end/test_sanitizer.py +++ b/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() diff --git a/tests/unit/test_sanitizer.py b/tests/unit/test_sanitizer.py index e6793466..405bc134 100644 --- a/tests/unit/test_sanitizer.py +++ b/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). diff --git a/triton_viz/clients/sanitizer/sanitizer.py b/triton_viz/clients/sanitizer/sanitizer.py index d297f36e..97340161 100644 --- a/triton_viz/clients/sanitizer/sanitizer.py +++ b/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 [] diff --git a/triton_viz/clients/symbolic_engine.py b/triton_viz/clients/symbolic_engine.py index 0ac17016..1981413f 100644 --- a/triton_viz/clients/symbolic_engine.py +++ b/triton_viz/clients/symbolic_engine.py @@ -20,11 +20,15 @@ Sum, And, Or, + Not, simplify, + substitute, Int2BV, BV2Int, BitVecRef, BoolVal, + is_true, + is_false, ) from z3.z3 import BoolRef, ArithRef, IntNumRef, ExprRef, Tactic, Probe @@ -32,7 +36,7 @@ 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. @@ -1746,6 +1767,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) ) @@ -2078,6 +2100,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 @@ -2095,3 +2118,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), + ) diff --git a/triton_viz/core/callbacks.py b/triton_viz/core/callbacks.py index 37cd9e64..67221ac6 100644 --- a/triton_viz/core/callbacks.py +++ b/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) diff --git a/triton_viz/core/client.py b/triton_viz/core/client.py index 2fc5130d..6d7432df 100644 --- a/triton_viz/core/client.py +++ b/triton_viz/core/client.py @@ -11,11 +11,13 @@ unpatch_op, patch_for_loop, unpatch_for_loop, + patch_if_else, + unpatch_if_else, patch_calls, LoopIter, ) from functools import wraps -from .callbacks import OpCallbacks, ForLoopCallbacks +from .callbacks import OpCallbacks, ForLoopCallbacks, IfElseCallbacks from .patch import patch_lang, unpatch_lang, OPERATION_REGISTRY from .config import config as cfg @@ -84,6 +86,9 @@ def register_op_callback( def register_for_loop_callback(self) -> ForLoopCallbacks: ... + def register_if_else_callback(self) -> IfElseCallbacks: + return IfElseCallbacks() + @abstractmethod def finalize(self) -> list: ... @@ -122,6 +127,8 @@ def __init__(self, clients: list[Client] | None = None): self.launch = Launch() self._lock = threading.Lock() self._clear_loop_hooks() + self._clear_if_hooks() + self._if_patching_active: bool = False def _lock_context(self): if cfg.num_sms > 1: @@ -174,15 +181,26 @@ def patch_run(self, fn, backend: str): with patch_calls(backend): # Collect all for-loop callbacks from clients all_loop_callbacks = [] + all_if_callbacks = [] for client in self.clients.values(): for namespace, attrs in namespaces.items(): # patch ops for attr, op in attrs.items(): callbacks = client.register_op_callback(op) patch_op(namespace, attr, callbacks, backend=backend) all_loop_callbacks.append(client.register_for_loop_callback()) + all_if_callbacks.append(client.register_if_else_callback()) self._populate_loop_hooks(all_loop_callbacks) patch_for_loop() + + # Only enable if/else patching if at least one client provides + # an eval_condition_callback. Without a provider, the AST rewrite + # would have no way to determine the condition's truth value. + self._populate_if_hooks(all_if_callbacks) + self._if_patching_active = self._eval_condition_callback is not None + if self._if_patching_active: + patch_if_else() + patch_lang(fn, backend, client_manager=self) try: yield @@ -192,6 +210,10 @@ def patch_run(self, fn, backend: str): unpatch_op(namespace, attr, backend) unpatch_for_loop() self._clear_loop_hooks() + if self._if_patching_active: + unpatch_if_else() + self._if_patching_active = False + self._clear_if_hooks() unpatch_lang(backend) def pre_run_callback(self, fn: Callable) -> bool: @@ -307,3 +329,42 @@ def loop_iter_wrapper( else: iterable = iterable_callable(*args, **kwargs) return LoopIter(self, iterable, lineno, range_type) + + # --- If/else callback management --- + + def _clear_if_hooks(self) -> None: + self._pre_if_hooks: list[Callable] = [] + self._eval_condition_callback: Callable | None = None + self._flip_condition_hooks: list[Callable] = [] + self._post_if_hooks: list[Callable] = [] + + def _populate_if_hooks(self, callbacks_list: list[IfElseCallbacks]) -> None: + self._clear_if_hooks() + for cb in callbacks_list: + if cb.pre_if_callback is not None: + self._pre_if_hooks.append(cb.pre_if_callback) + if cb.eval_condition_callback is not None: + if self._eval_condition_callback is not None: + raise RuntimeError("Only one eval_condition_callback allowed") + self._eval_condition_callback = cb.eval_condition_callback + if cb.flip_condition_callback is not None: + self._flip_condition_hooks.append(cb.flip_condition_callback) + if cb.post_if_callback is not None: + self._post_if_hooks.append(cb.post_if_callback) + + def pre_if(self, condition: Any, lineno: int) -> None: + for hook in self._pre_if_hooks: + hook(condition, lineno) + + def eval_condition(self, lineno: int) -> bool: + if self._eval_condition_callback is not None: + return self._eval_condition_callback(lineno) + return True + + def flip_condition(self, lineno: int) -> None: + for hook in self._flip_condition_hooks: + hook(lineno) + + def post_if(self, lineno: int) -> None: + for hook in self._post_if_hooks: + hook(lineno) diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py index 17239045..c2730628 100644 --- a/triton_viz/core/patch.py +++ b/triton_viz/core/patch.py @@ -29,6 +29,7 @@ from triton.tools.tensor_descriptor import TensorDescriptor from triton.runtime import JITFunction from ..transformers.for_loop_patcher import _visit_For as triton_viz_visit_For +from ..transformers.if_else_patcher import _visit_If as triton_viz_visit_If from .flip_patch import patch_flip from ..frontends.base import AdapterResult, OPERATION_REGISTRY @@ -294,19 +295,25 @@ class _LoopPatcher: def __init__(self): self._orig_visit_for: Callable | None = None + self._had_original: bool = False self._patched: bool = False def patch(self) -> None: if not self._patched: - self._orig_visit_for = getattr(_OrigASTTransformer, "visit_For", None) + orig = getattr(_OrigASTTransformer, "visit_For", None) + self._orig_visit_for = orig + self._had_original = orig is not None _OrigASTTransformer.visit_For = triton_viz_visit_For # type: ignore[assignment] self._patched = True def unpatch(self) -> None: if not self._patched: return - if self._orig_visit_for is not None: + if self._had_original: _OrigASTTransformer.visit_For = self._orig_visit_for + else: + if hasattr(_OrigASTTransformer, "visit_For"): + delattr(_OrigASTTransformer, "visit_For") self._patched = False @@ -321,6 +328,44 @@ def unpatch_for_loop(): _loop_patcher.unpatch() +class _IfElsePatcher: + """Manages AST patching for if/else branch interception.""" + + def __init__(self): + self._orig_visit_if: Callable | None = None + self._had_original: bool = False + self._patched: bool = False + + def patch(self) -> None: + if not self._patched: + orig = getattr(_OrigASTTransformer, "visit_If", None) + self._orig_visit_if = orig + self._had_original = orig is not None + _OrigASTTransformer.visit_If = triton_viz_visit_If # type: ignore[assignment] + self._patched = True + + def unpatch(self) -> None: + if not self._patched: + return + if self._had_original: + _OrigASTTransformer.visit_If = self._orig_visit_if + else: + if hasattr(_OrigASTTransformer, "visit_If"): + delattr(_OrigASTTransformer, "visit_If") + self._patched = False + + +_if_else_patcher = _IfElsePatcher() + + +def patch_if_else(): + _if_else_patcher.patch() + + +def unpatch_if_else(): + _if_else_patcher.unpatch() + + def patch_lang(fn, backend, client_manager=None): if backend == "triton": scope = _triton_snapshot_scope(fn) @@ -344,6 +389,8 @@ def patch_lang(fn, backend, client_manager=None): if client_manager is not None: fn.__globals__["_triton_viz_loop_patcher"] = client_manager + if getattr(client_manager, "_if_patching_active", False): + fn.__globals__["_triton_viz_if_patcher"] = client_manager patch_flip(scope, lambda: _current_client_manager) diff --git a/triton_viz/transformers/if_else_patcher.py b/triton_viz/transformers/if_else_patcher.py new file mode 100644 index 00000000..bd6adb19 --- /dev/null +++ b/triton_viz/transformers/if_else_patcher.py @@ -0,0 +1,103 @@ +import ast + + +def _visit_If(self, node: ast.If): # type: ignore[override] + """ + Transform if/elif/else to inject branch tracking hooks. + + if cond: _triton_viz_if_patcher.pre_if(cond, LINENO) + body1 ==> try: + else: if _triton_viz_if_patcher.eval_condition(LINENO): + body2 body1 + else: + _triton_viz_if_patcher.flip_condition(LINENO) + body2 + finally: + _triton_viz_if_patcher.post_if(LINENO) + + elif chains work automatically because generic_visit transforms the inner + If node (which sits in node.orelse) first. + """ + # Handle nested if statements (including elif chains) first + self.generic_visit(node) + + lineno_const = ast.Constant(value=node.lineno) + patcher_name = ast.Name(id="_triton_viz_if_patcher", ctx=ast.Load()) + + # _triton_viz_if_patcher.pre_if(cond, LINENO) + pre_if_call = ast.Expr( + value=ast.Call( + func=ast.Attribute( + value=patcher_name, + attr="pre_if", + ctx=ast.Load(), + ), + args=[node.test, lineno_const], + keywords=[], + ) + ) + + # _triton_viz_if_patcher.eval_condition(LINENO) + eval_condition_call = ast.Call( + func=ast.Attribute( + value=patcher_name, + attr="eval_condition", + ctx=ast.Load(), + ), + args=[lineno_const], + keywords=[], + ) + + # _triton_viz_if_patcher.flip_condition(LINENO) + flip_condition_stmt = ast.Expr( + value=ast.Call( + func=ast.Attribute( + value=patcher_name, + attr="flip_condition", + ctx=ast.Load(), + ), + args=[lineno_const], + keywords=[], + ) + ) + + # _triton_viz_if_patcher.post_if(LINENO) + post_if_call = ast.Expr( + value=ast.Call( + func=ast.Attribute( + value=patcher_name, + attr="post_if", + ctx=ast.Load(), + ), + args=[lineno_const], + keywords=[], + ) + ) + + # Build the else body: flip_condition + original else body + if node.orelse: + new_orelse = [flip_condition_stmt] + node.orelse + else: + new_orelse = [flip_condition_stmt, ast.Pass()] + + # New if: if eval_condition(LINENO): body else: flip + orelse + new_if = ast.If( + test=eval_condition_call, + body=node.body, + orelse=new_orelse, + ) + + # Wrap in try/finally to guarantee post_if runs + try_node = ast.Try( + body=[new_if], + handlers=[], + orelse=[], + finalbody=[post_if_call], + ) + + # Return list: [pre_if_call, try_node] + # NodeTransformer splices lists into parent body + result = [pre_if_call, try_node] + for n in result: + ast.fix_missing_locations(n) + return result