[inductor] Simplify remove_kernel_local_buffers (pytorch#139452)

I plan to reuse `can_buffer_be_removed_through_fusion` in some heuristics.

Pull Request resolved: pytorch#139452
Approved by: https://github.com/shunting314
ghstack dependencies: pytorch#139364, pytorch#139365, pytorch#139370

Author: jansel

torch/_inductor/codegen/common.py

@@ -2528,27 +2528,16 @@ def remove_kernel_local_buffers(self) -> None:
             for buf in self.store_buffer_names
             if buf in scheduler.name_to_buf
         )
-        names_to_remove = []
-        for out_buf in self.store_buffer_names:
-            if out_buf not in scheduler.name_to_buf:
-                # Aux buffers created during kernel codegen
-                names_to_remove.append(out_buf)
-                continue
-            users = scheduler.name_to_buf[out_buf].users
-            assert users is not None
-            users = OrderedSet(user.get_name() for user in users if not user.is_weak)
-            if users.issubset(fused_node_names):
-                names_to_remove.append(out_buf)
-
-        def remove_filter(n: str) -> bool:
-            return (
-                n not in self.must_keep_buffers
-                and n not in self.args.input_buffers
-                and n not in scheduler.mutation_renames
-                and n not in scheduler.mutation_real_name
-            )
-
-        names_to_remove = [*filter(remove_filter, names_to_remove)]
+        names_to_remove: OrderedSet[str] = OrderedSet()
+        for name in self.store_buffer_names:
+            if (
+                name not in self.must_keep_buffers
+                and name not in self.args.input_buffers
+                and scheduler.can_buffer_be_removed_through_fusion(
+                    name, fused_node_names
+                )
+            ):
+                names_to_remove.add(name)
 
         for name in names_to_remove:
             if name in self.args.inplace_buffers:

torch/_inductor/scheduler.py

@@ -3399,6 +3399,19 @@ def get_order(n: torch.fx.Node) -> int:
             _, last = max(origins, key=operator.itemgetter(0))
             V.graph.wrapper_code.enter_context(last)
 
+    def can_buffer_be_removed_through_fusion(
+        self, name: str, fused_node_names: OrderedSet[str]
+    ) -> bool:
+        try:
+            users = self.name_to_buf[name].users
+        except KeyError:
+            return False
+        return (
+            all(user.is_weak or user.get_name() in fused_node_names for user in users)
+            and name not in self.mutation_renames
+            and name not in self.mutation_real_name
+        )
+
     def codegen(self) -> None:
         with dynamo_timed("Scheduler.codegen"):
             return self._codegen()