Update

[ghstack-poisoned]

test/inductor/test_padding.py

@@ -500,8 +500,13 @@ class PaddingTest(TestCaseBase):
         forward_wrapper = wrapper_codes[0]
 
         # make sure the load for softmax is aligned
+        if bias:
+            # addmm -> mm + bias and bias is fused with softmax
+            softmax_load_str = "tl.load(in_out_ptr0 + (r0_1 + 30528*x0)"
+        else:
+            softmax_load_str = "tl.load(in_ptr0 + (r0_1 + 30528*x0)"
         self.assertTrue(
-            "tl.load(in_ptr0 + (r0_1 + 30528*x0)" in forward_wrapper,
+            softmax_load_str in forward_wrapper,
             f"forward_wrapper: {forward_wrapper}",
         )
 

test/inductor/test_torchinductor.py

@@ -15280,7 +15280,7 @@ if RUN_GPU:
                 ),
                 (
                     fn3,
-                    "triton_poi_fused_native_layer_norm_relu",
+                    "triton_poi_fused_addmm_native_layer_norm",
                     (torch.randn(4, 4, device=GPU_TYPE),),
                 ),
             ]
@@ -15293,7 +15293,7 @@ if RUN_GPU:
                 ),
                 (
                     fn3,
-                    "triton_poi_fused_LayerNorm_ReLU",
+                    "triton_poi_fused_LayerNorm_Linear_ReLU",
                     (torch.randn(4, 4, device=GPU_TYPE),),
                 ),
             ]

torch/_dynamo/variables/higher_order_ops.py

@@ -1693,7 +1693,7 @@ class AssociativeScanHigherOrderVariable(TorchHigherOrderOperatorVariable):
         )
 
         from torch._higher_order_ops.utils import _maybe_fake_tracing
-        from torch._inductor.utils import is_pointwise_use
+        from torch._inductor.utils import has_only_pointwise_uses
 
         with tx.fake_mode:
             sub_args_fake = [
@@ -1712,9 +1712,7 @@ class AssociativeScanHigherOrderVariable(TorchHigherOrderOperatorVariable):
 
             for node in fx.graph.nodes:
                 # Check that the combine_fn is pointwise, if combine_mode='pointwise'
-                if not all(
-                    is_pointwise_use(use) or use.op == "output" for use in node.users
-                ):
+                if not has_only_pointwise_uses(node, select_output=True):
                     raise RuntimeError(
                         "For combine_mode='pointwise', the combine_fn needs to be pointwise"
                     )

torch/_inductor/fx_passes/post_grad.py

@@ -51,8 +51,8 @@ from ..utils import (
     decode_device,
     get_all_devices,
     get_gpu_type,
+    has_uses_tagged_as,
     is_gpu,
-    is_pointwise_use,
     OPTIMUS_EXCLUDE_POST_GRAD,
 )
 from ..virtualized import V
@@ -1510,8 +1510,10 @@ def should_prefer_unfused_addmm(match):
     if not is_gpu(inp.meta["val"].device.type):
         return False
 
-    output = match.output_node()
-    return all(is_pointwise_use(use) for use in output.users)
+    return has_uses_tagged_as(
+        match.output_node(),
+        (torch.Tag.pointwise, torch.Tag.reduction),
+    )
 
 
 @register_graph_pattern(

torch/_inductor/lowering.py

@@ -80,9 +80,9 @@ from .ir import (
 from .utils import (
     ceildiv,
     decode_device,
+    has_only_pointwise_uses,
     is_dynamic,
     is_gpu,
-    is_pointwise_use,
     is_view,
     needs_fallback_due_to_atomic_add_limitations,
     pad_listlike,
@@ -1850,10 +1850,7 @@ def cat(inputs, dim=0):
         (len(inputs) <= config.max_pointwise_cat_inputs)
         and all(op_count(t) <= MAX_SIMPLE_OP_COUNT for t in inputs)
     ):
-        pointwise_uses = all(
-            is_pointwise_use(use, additional_pointwise_ops)
-            for use in V.current_node.users
-        )
+        pointwise_uses = has_only_pointwise_uses(V.current_node)
         # fuse in case we will be used in a pointwise node, and there are any inputs we
         # we can prevent materialization of.
         fuse_pointwise_use = (

torch/_inductor/utils.py

@@ -525,28 +525,107 @@ def is_view(op: torch._ops.OpOverload) -> bool:
     return any(a.alias_info is not None for a in op._schema.arguments)
 
 
-def is_pointwise_use(
-    use: Node,
-    is_pointwise_fn: Callable[[torch._ops.OpOverload], bool] = lambda _: False,
+class LogicalConnective(enum.Enum):
+    OR = enum.auto()
+    AND = enum.auto()
+
+
+def has_uses(
+    target: Node,
+    use_selector_fn: Callable[[torch._ops.OpOverload], bool] = lambda _: False,
+    use_aggregate_type: LogicalConnective = LogicalConnective.OR,
+    *,
+    select_output: bool = False,
 ) -> bool:
     """
-    Do all uses of this op have torch.Tag.pointwise or return True for optional `is_pointwise_fn`
+    Given a target, explore the uses of `target` by applying `use_selector_fn`
+    on them, and then aggregate these booleans with the `use_aggregate_type`
+    logical connective.
+
+    Uses in view ops will follow the views uses.
+    """
+
+    def get_use_aggregate_fn(
+        use_aggregate_type: LogicalConnective,
+    ) -> Callable[[Iterator[Any]], bool]:
+        match use_aggregate_type:
+            case LogicalConnective.AND:
+                return all
+            case LogicalConnective.OR:
+                return any
+            case _:
+                return any
+
+    use_aggregate_fn = get_use_aggregate_fn(use_aggregate_type)
+
+    def has_uses_impl(use: Node) -> bool:
+        if select_output and use.op == "output":
+            return True
+        if use.op != "call_function":
+            return False
+        if not (
+            isinstance(use.target, torch._ops.OpOverload)
+            or use.target is operator.getitem
+        ):
+            return False
+
+        target = cast(torch._ops.OpOverload, use.target)
+        # Process getitem and view
+        if target is operator.getitem or is_view(target):
+            return use_aggregate_fn(has_uses_impl(user) for user in use.users)
+
+        return use_selector_fn(target)
+
+    return use_aggregate_fn(has_uses_impl(user) for user in target.users)
+
+
+def has_only_uses(
+    target: Node,
+    use_selector_fn: Callable[[torch._ops.OpOverload], bool] = lambda _: False,
+    *,
+    select_output: bool = False,
+) -> bool:
+    return has_uses(
+        target, use_selector_fn, LogicalConnective.AND, select_output=select_output
+    )
+
+
+def has_uses_tagged_as(
+    target: Node,
+    use_tags: Collection[torch.Tag],
+    use_aggregate_type: LogicalConnective = LogicalConnective.OR,
+    *,
+    select_output: bool = False,
+) -> bool:
+    """
+    Is there a use with given tags?
+    """
+
+    return has_uses(
+        target,
+        lambda use: any(tag in use_tags for tag in use.tags),
+        use_aggregate_type,
+        select_output=select_output,
+    )
+
+
+def has_only_pointwise_uses(
+    target: Node,
+    *,
+    select_output: bool = False,
+) -> bool:
+    """
+    Do all uses of target have torch.Tag.pointwise?
 
     Uses in views ops will follow the views uses
     """
 
-    if use.op != "call_function":
-        return False
-    if not (
-        isinstance(use.target, torch._ops.OpOverload) or use.target is operator.getitem
-    ):
-        return False
-
-    target = cast(torch._ops.OpOverload, use.target)
-    if target is operator.getitem or is_view(target):
-        return all(is_pointwise_use(u, is_pointwise_fn) for u in use.users)
-
-    return torch.Tag.pointwise in target.tags or is_pointwise_fn(target)
+    return has_uses_tagged_as(
+        target,
+        use_tags=(torch.Tag.pointwise,),
+        use_aggregate_type=LogicalConnective.AND,
+        select_output=select_output,
+    )
 
 
 def gen_gm_and_inputs(