Optimization pass to remove unneeded masking (#109)

Author: jansel

helion/_compiler/compile_environment.py

@@ -333,6 +333,11 @@ def size_hint(self) -> int:
         assert isinstance(size, (int, torch.SymInt))
         return CompileEnvironment.current().size_hint(size)
 
+    def size_matches(self, numel: sympy.Expr | None) -> bool:
+        if numel is None or not isinstance(self.size, (int, torch.SymInt)):
+            return False
+        return numel == self.numel
+
     def mark_alternate_size(self, size: torch.SymInt | int | None) -> None:
         """If a block size is used with a different size, we need to clear the hint to enable masking."""
         if isinstance(self.size, AutoSize):

helion/_compiler/indexing_strategy.py

@@ -13,6 +13,7 @@
 from .ast_extension import expr_from_string
 from .compile_environment import CompileEnvironment
 from .host_function import HostFunction
+from .tile_strategy import DeviceLoopState
 from .tile_strategy import TileStrategy
 from .variable_origin import BlockSizeOrigin
 
@@ -203,7 +204,7 @@ def compute_shape(
         tensor: torch.Tensor, index: list[object]
     ) -> list[int | torch.SymInt]:
         assert isinstance(tensor, torch.Tensor)
-        assert isinstance(index, (list, tuple))
+        assert isinstance(index, (list, tuple)), index
         input_size = collections.deque(tensor.size())
         output_size = []
         for k in index:
@@ -455,8 +456,9 @@ def is_supported(
         index: list[object],
         extra_mask: ast.AST | None,
     ) -> bool:
+        # TODO(jansel): TensorDescriptor has some extra restrictions that are not captured here.
         if extra_mask is not None:
-            # TODO(jansel): block_ptr with extra_mask
+            # TODO(jansel): support block_ptr with extra_mask
             return False
         for k in index:
             if isinstance(k, torch.SymInt):
@@ -465,10 +467,27 @@ def is_supported(
                 if isinstance(symbol, sympy.Symbol):
                     origin = HostFunction.current().expr_to_origin.get(symbol)
                 if origin and isinstance(origin.origin, BlockSizeOrigin):
+                    block_index = origin.origin.block_size_idx
                     try:
-                        state.codegen.offset_var(origin.origin.block_size_idx)
+                        state.codegen.offset_var(block_index)
                     except NotImplementedError:
                         return False
+                    loop_state = state.codegen.active_device_loops[block_index][-1]
+                    if isinstance(loop_state, DeviceLoopState):
+                        """
+                        Check for a corner case where the loop size does not match the tensor size.
+                        In this case, the block masking will be incorrect.  So we check if the
+                        masking is needed and bail if it is.
+                        """
+                        end = loop_state.end_bounds[block_index]
+                        if (
+                            not CompileEnvironment.current()
+                            .block_sizes[block_index]
+                            .size_matches(end)
+                        ):
+                            assert state.fx_node is not None
+                            if "masked_value" in state.fx_node.meta:
+                                return False
             if isinstance(k, torch.Tensor):
                 # indirect loads don't work with block_ptr
                 return False

helion/_compiler/inductor_lowering.py

@@ -13,6 +13,7 @@
 import torch
 from torch._dynamo.convert_frame import compile_lock
 from torch._inductor import config as inductor_config
+from torch._inductor import ir
 from torch._inductor.codegen.simd import SIMDKernelFeatures
 from torch._inductor.codegen.simd import constant_repr
 from torch._inductor.codegen.triton import TritonKernel
@@ -43,6 +44,9 @@
 from .ast_extension import expr_from_string
 from .ast_extension import statement_from_string
 from .compile_environment import CompileEnvironment
+from .node_masking import apply_masking
+from .node_masking import cached_masked_value
+from .node_masking import mask_node_inputs
 from .tile_strategy import TileStrategy
 
 if TYPE_CHECKING:
@@ -185,7 +189,9 @@ def convert_arg(arg: Node) -> TensorBox:
                 )
             ),
         )
-        new_node.meta["lowering"] = lowering_cls(buffer, used_input_names)
+        new_node.meta["lowering"] = lowering = lowering_cls(buffer, used_input_names)
+        if isinstance(lowering, ReductionLowering):
+            lowering.add_input_mask(new_node)
         nodes.append(new_node)
         extra_input_names.append(buffer.get_name())
 
@@ -269,6 +275,10 @@ class Lowering:
     def codegen(self, ctx: GraphInterpreter, node: torch.fx.Node) -> object:
         raise NotImplementedError
 
+    def get_masked_value(self, node: torch.fx.Node) -> float | bool | None:
+        """Get the masked value for this node."""
+        return None
+
 
 @dataclasses.dataclass
 class InductorLowering(Lowering):
@@ -361,6 +371,11 @@ def codegen(self, ctx: GraphInterpreter, node: torch.fx.Node) -> object:
             output_name = _unpack_opsvalue(self.buffer.data.inner_fn(indices))
             return expr_from_string(output_name)
 
+    def get_masked_value(self, node: torch.fx.Node) -> float | bool | None:
+        """Get the masked value for this node."""
+        # TODO(jansel): use valueranges to determine masked value
+        return None
+
 
 @dataclasses.dataclass
 class ReductionLowering(InductorLowering):
@@ -383,6 +398,25 @@ def __init__(
         assert block_index is not None
         self.block_index: int = block_index
 
+    @property
+    def reduction_type(self) -> str:
+        reduction = self.buffer.data
+        assert isinstance(reduction, Reduction)
+        return reduction.reduction_type
+
+    def add_input_mask(self, node: torch.fx.Node) -> None:
+        """Modify the node to apply masking for the reduction if needed."""
+        reduction_type = self.reduction_type
+        input_dtype = None
+        for inp in node.all_input_nodes:
+            if isinstance(inp.meta["val"], torch.Tensor):
+                input_dtype = inp.meta["val"].dtype
+                break
+        assert input_dtype is not None
+        default = ir.Reduction.default_accumulator(reduction_type, input_dtype)
+        assert isinstance(default, (float, int, bool))
+        mask_node_inputs(node, default)
+
     def codegen(self, ctx: GraphInterpreter, node: torch.fx.Node) -> object:
         reduction = self.buffer.data
         assert isinstance(reduction, Reduction)
@@ -463,6 +497,11 @@ def normalize_args_kwargs(
         node.args = (*bound.arguments.values(),)
         node.kwargs = {}
 
+    def get_masked_value(self, node: torch.fx.Node) -> float | bool | None:
+        if self.api_func._get_masked_value is not None:
+            return self.api_func._get_masked_value(node)
+        return None
+
 
 @dataclasses.dataclass
 class SympyExprLowering(Lowering):
@@ -471,31 +510,61 @@ class SympyExprLowering(Lowering):
     def codegen(self, ctx: GraphInterpreter, node: torch.fx.Node) -> object:
         return expr_from_string(ctx.cg.device_function.user_sympy_expr(self.expr))
 
+    def get_masked_value(self, node: torch.fx.Node) -> float | bool | None:
+        if isinstance(self.expr, sympy.Integer):
+            return int(self.expr)
+        if isinstance(self.expr, sympy.Float):
+            return float(self.expr)
+        return None
+
 
 @dataclasses.dataclass
 class LambdaLowering(Lowering):
     fn: Callable[..., object]
+    masked_value_fn: Callable[[torch.fx.Node], float | bool | None] | None = None
 
     def codegen(self, ctx: GraphInterpreter, node: torch.fx.Node) -> object:
         return self.fn(ctx, node)
 
+    def get_masked_value(self, node: torch.fx.Node) -> float | bool | None:
+        if self.masked_value_fn is not None:
+            return self.masked_value_fn(node)
+        return None
+
+
+def passthrough_masked_value(node: torch.fx.Node) -> float | bool | None:
+    for input_node in node.all_input_nodes:
+        if isinstance(input_node.meta["val"], torch.Tensor):
+            return cached_masked_value(input_node)
+    return None
+
 
 aten_lowering_dispatch: dict[object, Callable[[torch.fx.Node], Lowering]] = {}
 
 
-def default_make_lowering(handler: CodegenHandler, node: torch.fx.Node) -> Lowering:
-    return LambdaLowering(handler)
+def default_make_lowering(
+    handler: CodegenHandler,
+    node: torch.fx.Node,
+    masked_value_fn: Callable[[torch.fx.Node], float | bool | None] | None = None,
+) -> Lowering:
+    return LambdaLowering(handler, masked_value_fn=masked_value_fn)
 
 
 def register_lowering(
     fn: object,
     make_lowering: Callable[
         [CodegenHandler, torch.fx.Node], Lowering
     ] = default_make_lowering,
+    masked_value_fn: Callable[[torch.fx.Node], float | bool | None] | None = None,
 ) -> Callable[[CodegenHandler], CodegenHandler]:
     def decorator(handler: CodegenHandler) -> CodegenHandler:
         assert fn not in aten_lowering_dispatch, f"Lowering for {fn} already registered"
-        aten_lowering_dispatch[fn] = lambda node: make_lowering(handler, node)
+        # pyre-ignore[28]
+        aten_lowering_dispatch[fn] = lambda node: make_lowering(
+            handler,
+            node,
+            masked_value_fn=masked_value_fn,
+        )
         return handler
 
     return decorator
@@ -521,7 +590,12 @@ def codegen_getitem(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
 
 
 # pyre-fixme[56]
-@register_lowering(torch.ops.aten.full.default)
+@register_lowering(
+    torch.ops.aten.full.default,
+    masked_value_fn=lambda n: (
+        n.args[1] if isinstance(n.args[1], (int, float, bool)) else None
+    ),
+)
 def codegen_full(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
     env = CompileEnvironment.current()
     size, fill_value = map_arg(node.args, lambda n: n.meta["val"])
@@ -539,7 +613,9 @@ def codegen_full(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
 
 
 # pyre-fixme[56]
-@register_lowering(torch.ops.aten.unsqueeze.default)
+@register_lowering(
+    torch.ops.aten.unsqueeze.default, masked_value_fn=passthrough_masked_value
+)
 def codegen_unsqueeze(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
     assert not node.kwargs, "getitem kwargs not supported"
     tensor, dim = map_arg(node.args, lambda arg: ctx.env[arg])
@@ -557,10 +633,14 @@ def codegen_unsqueeze(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
     )
 
 
-@register_lowering(torch.ops.aten.squeeze.dim)
-@register_lowering(torch.ops.aten.view.default)
+@register_lowering(torch.ops.aten.squeeze.dim, masked_value_fn=passthrough_masked_value)
+@register_lowering(
+    torch.ops.aten.view.default, masked_value_fn=passthrough_masked_value
+)
 # pyre-fixme[56]
-@register_lowering(torch.ops.aten.reshape.default)
+@register_lowering(
+    torch.ops.aten.reshape.default, masked_value_fn=passthrough_masked_value
+)
 def codegen_view(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
     assert not node.kwargs, "view kwargs not supported"
     tensor = map_arg(node.args[0], lambda arg: ctx.env[arg])
@@ -572,7 +652,9 @@ def codegen_view(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
 
 
 # pyre-fixme[56]
-@register_lowering(torch.ops.aten.permute.default)
+@register_lowering(
+    torch.ops.aten.permute.default, masked_value_fn=passthrough_masked_value
+)
 def codegen_permute(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
     assert not node.kwargs, "getitem kwargs not supported"
     tensor, dims = map_arg(node.args, lambda arg: ctx.env[arg])
@@ -586,7 +668,9 @@ def codegen_permute(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
 
 
 # pyre-fixme[56]
-@register_lowering(torch.ops.aten.expand.default)
+@register_lowering(
+    torch.ops.aten.expand.default, masked_value_fn=passthrough_masked_value
+)
 def codegen_expand(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
     assert not node.kwargs, "getitem kwargs not supported"
     tensor, _ = map_arg(node.args, lambda arg: ctx.env[arg])
@@ -606,7 +690,11 @@ def codegen_expand(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
     )
 
 
-def apply_dot_requirements(handler: CodegenHandler, node: torch.fx.Node) -> Lowering:
+def apply_dot_requirements(
+    handler: CodegenHandler,
+    node: torch.fx.Node,
+    masked_value_fn: Callable[[torch.fx.Node], float | bool | None] | None = None,
+) -> Lowering:
     """Apply min_dot_size requirements to the config_spec"""
     assert not node.kwargs, "dot kwargs not supported"
     assert len(node.args) in (2, 3)
@@ -625,7 +713,14 @@ def apply_dot_requirements(handler: CodegenHandler, node: torch.fx.Node) -> Lowe
         block_idx = TileStrategy.get_block_index(shape)
         if block_idx is not None:
             env.block_sizes[block_idx].update_min_block(min_size, allow_flattened=True)
-    return LambdaLowering(handler)
+    # inputs to the dot operation must be zero-masked
+    *maybe_acc, lnode, rnode = node.args
+    assert isinstance(lnode, torch.fx.Node)
+    assert isinstance(rnode, torch.fx.Node)
+    lnode = apply_masking(lnode, base_node=node, other=0)
+    rnode = apply_masking(rnode, base_node=node, other=0)
+    node.args = (*maybe_acc, lnode, rnode)
+    return LambdaLowering(handler, masked_value_fn=masked_value_fn)
 
 
 @register_lowering(torch.ops.aten.bmm.default, apply_dot_requirements)

helion/_compiler/node_masking.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import functools
+
+import torch.fx
+from torch.fx.experimental import proxy_tensor
+
+from helion.language._tracing_ops import _mask_to
+
+
+def mask_node_inputs(
+    node: torch.fx.Node,
+    other: float | bool = 0,
+) -> None:
+    """Inplace update the node's args and kwargs to apply masking if needed."""
+    apply = functools.partial(apply_masking, other=other, base_node=node)
+    node.args = torch.fx.map_arg(node.args, apply)
+    node.kwargs = torch.fx.map_arg(node.kwargs, apply)
+
+
+def apply_masking(
+    node: torch.fx.Node,
+    *,
+    base_node: torch.fx.Node,
+    other: float | bool = 0,
+) -> torch.fx.Node:
+    """Analyze the node and apply masking if needed."""
+    current_mask = cached_masked_value(node)
+    if current_mask == other:
+        return node  # already masked, no need to change it
+    for user in node.users:
+        if user.op == "call_function" and user.target == _mask_to:
+            if user.args[1] == other:
+                assert user.args[0] is node
+                return user  # reuse existing mask_to node
+    from helion._compiler.inductor_lowering import APIFuncLowering
+
+    # If we reach here, we need to create a new mask_to node
+    with node.graph.inserting_before(base_node):
+        new_node = node.graph.call_function(_mask_to, (node, other), {})
+    new_node.meta.update(base_node.meta)
+    with proxy_tensor.disable_proxy_modes_tracing():
+        new_node.meta["val"] = node.meta["val"].clone()
+    # pyre-ignore[6]
+    new_node.meta["lowering"] = APIFuncLowering(_mask_to)
+    new_node.meta["masked_value"] = other
+    return new_node
+
+
+def cached_masked_value(
+    node: torch.fx.Node,
+) -> float | bool | None:
+    """Determine the current masked value for the node."""
+    if "masked_value" in node.meta:
+        return node.meta["masked_value"]
+    if node.op != "call_function":
+        return None
+    node.meta["masked_value"] = result = node.meta["lowering"].get_masked_value(node)
+    return result