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/_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/utils.py

@@ -549,6 +549,70 @@ def is_pointwise_use(
     return torch.Tag.pointwise in target.tags or is_pointwise_fn(target)
 
 
+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,
+) -> bool:
+    """
+    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 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_uses_tagged_as(
+    target: Node,
+    use_tags: Collection[torch.Tag],
+    use_aggregate_type: LogicalConnective = LogicalConnective.OR,
+) -> 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
+    )
+
+
 def gen_gm_and_inputs(
     target: Any, args: list[Any], kwargs: dict[str, Any]
 ) -> tuple[GraphModule, list[torch.Tensor]]: