[inductor] [cpp] fix gemm_output_name conflict (#136419)

chunyuan-w · pytorchmergebot · commit c2637a7b2656 · 2024-09-27T12:23:17.000Z
Fixes the max-autotune failure of `soft_actor_critic` of Torchbench in FP32 single thread dynamic shape case: ```log File "/home/user/inductor/pytorch/torch/_inductor/codegen/cpp_micro_gemm.py", line 136, in codegen_call C_ptr = f"&({kernel.index(C, [0, 0])})" File "/home/user/inductor/pytorch/torch/_inductor/codegen/cpp_template_kernel.py", line 135, in index else self.args.input(node.get_name()) File "/home/user/inductor/pytorch/torch/_inductor/codegen/common.py", line 1251, in input assert name not in V.graph.removed_buffers, name AssertionError: buf_GemmOut ``` The 1st and 2nd linear does not need to use local buffer while the 3rd linear needs to use local buffer. The 3rd linear which uses local buffer will add its global buffer (named as `buf_GemmOut`) into `V.graph.removed_buffers`. When scheduling the nodes, the 1st linear (won't use local buffer) will get its output buffer (also named as `buf_GemmOut`) from the input and found that it's in the `V.graph.removed_buffers` and raise AssertionError. The issue is that the output buffer of all these linears are all names with `buf_GemmOut`, which have a conflict. Rename these buffers by adding the name of the `template_buffer` as the prefix. Pull Request resolved: #136419 Approved by: https://github.com/leslie-fang-intel, https://github.com/jgong5 ghstack dependencies: #136418, #136518
diff --git a/test/inductor/test_cpu_select_algorithm.py b/test/inductor/test_cpu_select_algorithm.py
@@ -739,6 +739,60 @@ def forward(self, x):
             self.common(mod, (v,), atol=atol, rtol=rtol)
         self.assertEqual(counters["inductor"]["select_algorithm_autotune"], 1)
 
+    @inductor_config.patch({"freezing": True})
+    @patches
+    @torch.no_grad
+    @unittest.skipIf(not TEST_MKL, "Test requires MKL")
+    @set_num_threads(1)
+    @dynamo_config.patch({"dynamic_shapes": True, "assume_static_by_default": False})
+    @parametrize("batch_size", (256,))
+    @parametrize("in_features", (3,))
+    @parametrize("out_features", (1024,))
+    @parametrize("out_features2", (2,))
+    @parametrize("bias", (True, False))
+    @dtypes(torch.float)
+    def test_linear_local_and_global_buffer_dynamic_shapes(
+        self, batch_size, in_features, out_features, out_features2, bias, dtype
+    ):
+        # Reproducer from soft_actor_critic
+        class M(torch.nn.Module):
+            def __init__(self, bias):
+                super().__init__()
+                self.linear = torch.nn.Linear(in_features, out_features, bias)
+                self.linear1 = torch.nn.Linear(out_features, out_features, bias)
+                self.linear2 = torch.nn.Linear(out_features, out_features2, bias)
+
+            def forward(self, arg7_1):
+                addmm_3 = self.linear(arg7_1)
+                relu_2 = torch.ops.aten.relu.default(addmm_3)
+
+                addmm_4 = self.linear1(relu_2)
+                relu_3 = torch.ops.aten.relu.default(addmm_4)
+
+                addmm_5 = self.linear2(relu_3)
+
+                split_1 = torch.ops.aten.split.Tensor(addmm_5, 1, 1)
+                getitem_2 = split_1[0]
+                getitem_3 = split_1[1]
+
+                tanh_1 = torch.ops.aten.tanh.default(getitem_3)
+
+                add_62 = torch.ops.aten.add.Tensor(tanh_1, 1)
+
+                mul_36 = torch.ops.aten.mul.Tensor(add_62, 6.0)
+                add_69 = torch.ops.aten.add.Tensor(mul_36, -10.0)
+
+                exp_1 = torch.ops.aten.exp.default(add_69)
+                return (getitem_2, exp_1)
+
+        counters.clear()
+        v = torch.randn(batch_size, in_features).to(dtype=dtype)
+        mod = M(bias=bias).to(dtype=dtype).eval()
+        with verify(dtype) as (atol, rtol):
+            self.common(mod, (v,), atol=atol, rtol=rtol)
+        self.assertEqual(counters["inductor"]["select_algorithm_autotune"], 3)
+        self.assertEqual(counters["inductor"]["cpp_epilogue_fusion_counter"], 2)
+
     @inductor_config.patch({"freezing": True})
     @patches
     @torch.no_grad
diff --git a/torch/_inductor/codegen/cpp_gemm_template.py b/torch/_inductor/codegen/cpp_gemm_template.py
@@ -909,14 +909,14 @@ def copy_inner(index):
         #     --> zero or more out-of-template epilogues (`epilogue_nodes`) -->
         #   Y
         if epilogue_creators:
-            gemm_output_name = "buf_GemmOut"
+            gemm_output_name = f"{template_buffer.get_name()}_GemmOut"
             gemm_output_buffer = ir.Buffer(gemm_output_name, template_buffer.layout)
             current_input_buffer = gemm_output_buffer
             for i, creator in enumerate(epilogue_creators):
                 if i == len(epilogue_creators) - 1:
                     buffer_name = template_buffer.get_name()
                 else:
-                    buffer_name = f"buf_GemmOut_epilogue_{i}"
+                    buffer_name = f"{gemm_output_name}_epilogue_{i}"
                 epilogues.append(
                     ir.ComputedBuffer(
                         name=buffer_name,
