diff --git a/tests/end_to_end/test_sanitizer.py b/tests/end_to_end/test_sanitizer.py index a039a906..8aa28e88 100644 --- a/tests/end_to_end/test_sanitizer.py +++ b/tests/end_to_end/test_sanitizer.py @@ -1209,6 +1209,32 @@ def test_sanitizer_supports_data_dependent_cumsum_index(): assert len(cumsum_sanitizer.records) == 0 +bool_op_mask_sanitizer = SymbolicSanitizer(abort_on_error=False) + + +@tilelens.trace(client=bool_op_mask_sanitizer) +@triton.jit +def bool_op_mask_kernel(x_ptr, out_ptr, lo, hi, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + # Compiled Triton computes a logical and/or here; Python's `and` would + # keep only `offs < hi`, and the load would read before x. + in_range = (offs >= lo) and (offs < hi) + outside = (offs < lo) or (offs >= hi) + tl.store(out_ptr + offs, tl.load(x_ptr + offs - lo, mask=in_range), mask=in_range) + tl.store(out_ptr + offs, -1.0, mask=outside) + + +def test_sanitizer_and_or_on_tensors_follow_compiled_triton(): + bool_op_mask_sanitizer.records.clear() + + x = torch.arange(4, dtype=torch.float32) + out = torch.empty((8,), dtype=torch.float32) + + bool_op_mask_kernel[(1,)](x, out, 2, 6, BLOCK=8) + + assert len(bool_op_mask_sanitizer.records) == 0 + + ashr_index_sanitizer = SymbolicSanitizer(abort_on_error=False) diff --git a/tests/end_to_end/test_tracer.py b/tests/end_to_end/test_tracer.py index 9eef0015..188f7803 100644 --- a/tests/end_to_end/test_tracer.py +++ b/tests/end_to_end/test_tracer.py @@ -1,3 +1,4 @@ +import pytest import torch import triton import triton.language as tl @@ -223,3 +224,51 @@ def autotune_add_kernel_cache_on( assert ( bench_fn is not None and bench_fn.__name__ == "dummy_benchmarker" ), f"Expected dummy_benchmarker, got: {bench_fn}" + + +@triton.jit +def logical_ops_kernel(X, Y, SCALAR: tl.constexpr, MODE: tl.constexpr): + offsets = tl.program_id(0) if SCALAR else tl.arange(0, 16) + x = tl.load(X + offsets) + if MODE == "and": + result = (x >= 0) and (x <= 8) + elif MODE == "or": + result = (x < 0) or (x > 8) + elif MODE == "long_and": + result = (x > -2) and (x <= 8) and (x != 3) + elif MODE == "long_or": + result = (x == -4) or (x == 3) or (x > 8) + elif MODE == "nested": + result = (x > 0) and (x < 3) or (x > 8) and (x != 10) + elif MODE == "constexpr_false": + result = False and (x > 0) + else: + result = True and (x > 0) + tl.store(Y + offsets, result) + + +@pytest.mark.parametrize("scalar", [False, True]) +@pytest.mark.parametrize( + "mode", + ["and", "or", "long_and", "long_or", "nested", "constexpr_false", "constexpr_true"], +) +def test_and_or_on_tensors_follow_compiled_triton(scalar, mode): + """`and`/`or` on tensors compute logical_and/logical_or, as in compiled + Triton; only a compile-time operand short-circuits.""" + x = torch.arange(-4, 12, dtype=torch.int32) + y = torch.zeros_like(x, dtype=torch.bool) + + tilelens.trace("tracer")(logical_ops_kernel)[(16 if scalar else 1,)]( + x, y, scalar, mode + ) + + expected = { + "and": (x >= 0) & (x <= 8), + "or": (x < 0) | (x > 8), + "long_and": (x > -2) & (x <= 8) & (x != 3), + "long_or": (x == -4) | (x == 3) | (x > 8), + "nested": (x > 0) & (x < 3) | (x > 8) & (x != 10), + "constexpr_false": torch.zeros_like(y), + "constexpr_true": x > 0, + }[mode] + torch.testing.assert_close(y, expected) diff --git a/tilelens/core/frontend/triton.py b/tilelens/core/frontend/triton.py index 30b4846a..c697d468 100644 --- a/tilelens/core/frontend/triton.py +++ b/tilelens/core/frontend/triton.py @@ -203,6 +203,83 @@ } +class _BoolOpChain: + """One operand of a kernel's `and`/`or`, as rewritten by `_BoolOpRewriter`. + + Triton's interpreter runs `and`/`or` with Python semantics and treats any + tensor with more than one element as true, so `a and b` returns `b`. + Compiled Triton short-circuits only on compile-time values and combines + runtime operands with logical_and/logical_or; the rewrite chains these + objects with `<` so that operands are still evaluated once, left to right, + and the chain stops only at a compile-time value that decides the result. + """ + + def __init__(self, is_and: bool, value: Any = None) -> None: + self.is_and = is_and + self.value = value + self.short_circuit = False + self.tensor_values: list[Any] = [] + + def __lt__(self, rhs: "_BoolOpChain") -> "_BoolOpChain": + rhs.tensor_values = self.tensor_values + if isinstance(rhs.value, tl.core.base_value) and not isinstance( + rhs.value, (tl.constexpr, tl.tuple) + ): + rhs.tensor_values.append(rhs.value) + else: + rhs.short_circuit = bool(rhs.value) != self.is_and + return rhs + + def __bool__(self) -> bool: + return not self.short_circuit + + def get_result(self) -> Any: + if self.short_circuit or not self.tensor_values: + return self.value + method = "logical_and" if self.is_and else "logical_or" + result = self.tensor_values[-1] + for value in reversed(self.tensor_values[:-1]): + result = getattr(value, method)(result) + return result + + +class _BoolOpRewriter(ast.NodeTransformer): + """Rewrite `a and b` as the native chain `op(True) < op(True, a) < + op(True, b)`, where `op` names `_BoolOpChain`, so that operands keep their + scope and their lazy, left-to-right evaluation.""" + + def __init__(self, helper_name: str) -> None: + self.helper_name = helper_name + + def visit_BoolOp(self, node: ast.BoolOp) -> ast.expr: + node = cast(ast.BoolOp, self.generic_visit(node)) + is_and = isinstance(node.op, ast.And) + + def operand(value: ast.expr | None = None) -> ast.expr: + args: list[ast.expr] = [ast.Constant(value=is_and)] + if value is not None: + args.append(value) + return ast.Call( + func=ast.Name(id=self.helper_name, ctx=ast.Load()), + args=args, + keywords=[], + ) + + chain = ast.Compare( + left=operand(), + ops=[ast.Lt() for _ in node.values], + comparators=[operand(value) for value in node.values], + ) + return ast.copy_location( + ast.Call( + func=ast.Attribute(value=chain, attr="get_result", ctx=ast.Load()), + args=[], + keywords=[], + ), + node, + ) + + @dataclass(frozen=True) class FakeTensor: _data_ptr: int @@ -251,6 +328,7 @@ def __init__(self): self._thread_local_interpreter_state.grid_idx = None self._current_client_manager = None self._loop_wrapper_arg = "_tilelens_loop_iter_wrapper" + self._bool_op_arg = "_tilelens_bool_op" self._loop_ast_methods: dict[str, Callable | object] = {} self._loop_ast_patched = False self._patch_calls_scope = 0 @@ -354,22 +432,32 @@ def _visit_triton_function_def( self, transformer: ast.NodeTransformer, node: ast.FunctionDef ): node = cast(ast.FunctionDef, transformer.generic_visit(node)) + # Run after Triton's own rewrite, which does not visit assigned values + # and must not rewrite the chains again. A Triton whose interpreter + # rewrites `and`/`or` itself keeps its own. + if "visit_BoolOp" not in vars(_OrigASTTransformer): + node = cast(ast.FunctionDef, _BoolOpRewriter(self._bool_op_arg).visit(node)) # Triton compiles rewritten kernels with a local namespace, but loop # bodies resolve names from the function defaults/globals. Add the - # wrapper as a defaulted kw-only arg so the rewritten loop can call it - # without mutating the user's global dictionary. - if self._loop_wrapper_arg in { + # helpers as defaulted kw-only args so the rewritten loops and boolean + # operators can call them without mutating the user's global dictionary. + existing = { arg.arg for arg in (node.args.posonlyargs + node.args.args + node.args.kwonlyargs) - }: - return node - - node.args.kwonlyargs.append(ast.arg(arg=self._loop_wrapper_arg)) - node.args.kw_defaults.append( - ast.Name(id=self._loop_wrapper_arg, ctx=ast.Load()) - ) + } + for name in self._hidden_kernel_args(): + if name in existing: + continue + node.args.kwonlyargs.append(ast.arg(arg=name)) + node.args.kw_defaults.append(ast.Name(id=name, ctx=ast.Load())) return ast.fix_missing_locations(node) + def _hidden_kernel_args(self) -> dict[str, Any]: + return { + self._loop_wrapper_arg: self.loop_iter_wrapper, + self._bool_op_arg: _BoolOpChain, + } + def loop_iter_wrapper( self, iterable_callable: Callable, @@ -939,9 +1027,10 @@ def _prepare_grid_launch(self, grid_executor, args_dev, kwargs): *args_hst, **kwargs_hst, ) - # The hidden loop-wrapper default is an implementation detail of the - # rewritten function; clients should not see it as a kernel argument. - args.pop(self._loop_wrapper_arg, None) + # The hidden helper defaults are implementation details of the + # rewritten function; clients should not see them as kernel arguments. + for name in self._hidden_kernel_args(): + args.pop(name, None) call_args = {} for name, arg in args.items(): if name in grid_executor.constexprs: @@ -1126,13 +1215,10 @@ def jit_function_call(jit_function, *args, **kwargs): def compile_and_exec(rewriter, transformed_ast): # Triton's FunctionRewriter executes code with `rewriter.kwargs` - # as the local namespace. Temporarily add the hidden loop wrapper - # there so rewritten function defaults can bind it. + # as the local namespace. Temporarily add the hidden helpers + # there so rewritten function defaults can bind them. old_kwargs = rewriter.kwargs - rewriter.kwargs = { - **old_kwargs, - self._loop_wrapper_arg: self.loop_iter_wrapper, - } + rewriter.kwargs = {**old_kwargs, **self._hidden_kernel_args()} try: return old_compile_and_exec(rewriter, transformed_ast) finally: