Exclude nop nodes from foreach

mlazos · mlazos · commit f94a6e5efd6f · 2023-10-17T23:33:01.000Z
diff --git a/test/inductor/test_foreach.py b/test/inductor/test_foreach.py
@@ -519,6 +519,23 @@ def fn(x1, x2, x3, w1, w2, w3):
 
         self.assertEqual(torch._inductor.metrics.generated_kernel_count, 2)
 
+    @requires_cuda()
+    def test_zero_elems(self):
+        def fn(a0, a1, b0, b1):
+            return torch._foreach_add([a0, a1], [b0, b1])
+
+        self.check_model_cuda(
+            fn,
+            (
+                torch.rand(0, device="cuda:0"),
+                torch.rand(10, 10, device="cuda:0"),
+                torch.rand(0, device="cuda:0"),
+                torch.rand(10, 10, device="cuda:0"),
+            ),
+        )
+
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 1)
+
     @requires_cuda()
     @bin_ops
     def test_2d_blocking(self, op):
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -3086,7 +3086,6 @@ def create(cls, inputs, dim):
                 input_unwrapped.is_input_buffer()
                 and inputs[i].get_device().type == "cuda"
                 and not is_dynamic(input_buffer)
-                and not input_buffer.is_no_op()
             ):
                 buffer_names.append(input_buffer.get_name())
 
diff --git a/torch/_inductor/scheduler.py b/torch/_inductor/scheduler.py
@@ -1206,14 +1206,19 @@ def create_foreach_nodes(self):
         kept_node_names = self.name_to_fused_node.keys()
 
         for names in V.graph.lists.values():
-            removed_node_names.update(names)
-
-            names = [name for name in names if name in kept_node_names]
+            names = [
+                name
+                for name in names
+                if name in kept_node_names
+                and not isinstance(self.name_to_node[name], NopKernelSchedulerNode)
+            ]
             if not names:
                 # All nodes eliminated
                 continue
 
+            removed_node_names.update(names)
             snodes = [self.name_to_node[name] for name in names]
+
             fe_node = ForeachKernelSchedulerNode(self, snodes)
 
             fe_nodes.append(fe_node)
