[aoti] Print output name for sympy.Expr as well (#138524)

henrylhtsang · pytorchmergebot · commit 3f9f6048dabe · 2024-10-29T16:02:45.000Z
To avoid ``` NotImplementedError: unsupported type of output=s0*s1 ``` It seems like this was caused by the use of `_scaled_dot_product_flash_attention`. Fallback kernek: ``` FallbackKernel( python_kernel_name='torch.ops.aten._scaled_dot_product_flash_attention.default', name=buf55, layout=MultiOutputLayout(device=device(type='cuda', index=0)), inputs=[ComputedBuffer(name='buf52', layout=FixedLayout('cuda', torch.bfloat16, size=[1, 6, s0*s1, 64], stride=[384*s0*s1, 64*s0*s1, 64, 1]), data=Pointwise(device=device(type='cuda', index=0), dtype=torch.bfloat16, inner_fn=<function BaseView.make_loader.<locals>.loader at 0x7fcd7f99da20>, ranges=[1, 6, s0*s1, 64])), ComputedBuffer(name='buf53', layout=FixedLayout('cuda', torch.bfloat16, size=[1, 6, s0*s1, 64], stride=[384*s0*s1, 64*s0*s1, 64, 1]), data=Pointwise(device=device(type='cuda', index=0), dtype=torch.bfloat16, inner_fn=<function BaseView.make_loader.<locals>.loader at 0x7fcd7f99d480>, ranges=[1, 6, s0*s1, 64])), ComputedBuffer(name='buf54', layout=FixedLayout('cuda', torch.bfloat16, size=[1, 6, s0*s1, 64], stride=[384*s0*s1, 64*s0*s1, 64, 1]), data=Pointwise(device=device(type='cuda', index=0), dtype=torch.bfloat16, inner_fn=<function BaseView.make_loader.<locals>.loader at 0x7fcd7f99c430>, ranges=[1, 6, s0*s1, 64]))], constant_args=(0.125,), kwargs={'scale': 0.125}, output_view=None, python_kernel_name=torch.ops.aten._scaled_dot_product_flash_attention.default, cpp_kernel_name=at::_ops::_scaled_dot_product_flash_attention::call, ordered_kwargs_for_cpp_kernel=['scale'], op_overload=aten._scaled_dot_product_flash_attention.default, arg_properties=[{'name': 'query', 'type': Tensor, 'default_value': None}, {'name': 'key', 'type': Tensor, 'default_value': None}, {'name': 'value', 'type': Tensor, 'default_value': None}, {'name': 'dropout_p', 'type': float, 'default_value': 0.0}, {'name': 'is_causal', 'type': bool, 'default_value': False}, {'name': 'return_debug_mask', 'type': bool, 'default_value': False}], kwarg_properties=None, unbacked_bindings=None, mutation_outputs=[], origin_node=None, origins=OrderedSet([_scaled_dot_product_flash_attention]) ) ``` codegen with this pr ``` // Topologically Sorted Source Nodes: [scaled_dot_product_attention], Original ATen: [aten._scaled_dot_product_flash_attention] double var_147 = 0.125; AtenTensorHandle buf56_handle; AtenTensorHandle buf57_handle; auto buf55_4 = s0*s1; auto buf55_5 = s0*s1; AtenTensorHandle buf58_handle; AtenTensorHandle buf59_handle; AtenTensorHandle buf60_handle; AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_cuda__scaled_dot_product_flash_attention(convert_arrayref_tensor_to_tensor(buf52), convert_arrayref_tensor_to_tensor(buf53), convert_arrayref_tensor_to_tensor(buf54), 0.0, 0, 0, &var_147, &buf56_handle, &buf57_handle, nullptr, nullptr, &buf55_4, &buf55_5, &buf58_handle, &buf59_handle, &buf60_handle)); RAIIAtenTensorHandle buf56(buf56_handle); RAIIAtenTensorHandle buf57(buf57_handle); RAIIAtenTensorHandle buf58(buf58_handle); RAIIAtenTensorHandle buf59(buf59_handle); RAIIAtenTensorHandle buf60(buf60_handle); ``` Differential Revision: D64724460 Pull Request resolved: #138524 Approved by: https://github.com/chenyang78
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
@@ -1104,6 +1104,56 @@ def forward(self, x, y):
         )
         self.check_model(Repro(), example_inputs)
 
+    def test_fallback_kernel_with_symexpr_output(self):
+        if self.device != "cuda":
+            raise unittest.SkipTest("requires CUDA")
+
+        class Module(torch.nn.Module):
+            def forward(self, q, k, v):
+                q = q.reshape(
+                    q.shape[0],
+                    2,
+                    q.shape[2] * q.shape[3],
+                    q.shape[1] // 2,
+                )
+                k = k.reshape(
+                    k.shape[0],
+                    2,
+                    k.shape[2] * k.shape[3],
+                    k.shape[1] // 2,
+                )
+                v = v.reshape(
+                    v.shape[0],
+                    2,
+                    v.shape[2] * v.shape[3],
+                    v.shape[1] // 2,
+                )
+
+                res = torch.ops.aten._scaled_dot_product_flash_attention.default(
+                    q,
+                    k,
+                    v,
+                )
+                return res[0]
+
+        m = Module().to(device=self.device)
+        tensor_shape = (4, 32, 4, 4)
+        inputs = (
+            torch.randn(tensor_shape, dtype=torch.float16, device=self.device),
+            torch.randn(tensor_shape, dtype=torch.float16, device=self.device),
+            torch.randn(tensor_shape, dtype=torch.float16, device=self.device),
+        )
+
+        dynamic_shapes = {
+            "q": {2: Dim.DYNAMIC, 3: Dim.DYNAMIC},
+            "k": {2: Dim.DYNAMIC, 3: Dim.DYNAMIC},
+            "v": {2: Dim.DYNAMIC, 3: Dim.DYNAMIC},
+        }
+        ep = torch.export.export(m, inputs, dynamic_shapes=dynamic_shapes, strict=False)
+        path = torch._inductor.aot_compile(ep.module(), inputs)
+        aot_model = torch._export.aot_load(path, device=self.device)
+        torch.testing.assert_close(m(*inputs), aot_model(*inputs))
+
     def test_large_grid(self):
         if self.device != "cuda":
             raise unittest.SkipTest("requires CUDA")
diff --git a/torch/_inductor/codegen/cpp_wrapper_cpu.py b/torch/_inductor/codegen/cpp_wrapper_cpu.py
@@ -1055,9 +1055,9 @@ def generate_c_shim_fallback_kernel(self, fallback_kernel, args):
                 output_name = f"{output_name_base}_{idx}"
                 self.writeline(f"int64_t {output_name} = {output};")
                 output_args.append(f"&{output_name}")
-            elif isinstance(output, sympy.Symbol):
+            elif isinstance(output, sympy.Expr):
                 output_name = f"{output_name_base}_{idx}"
-                self.writeline(f"auto {output_name} = {output};")
+                self.writeline(f"auto {output_name} = {self.expr_printer(output)};")
                 output_args.append(f"&{output_name}")
             elif output is None:
                 output_args.append("nullptr")
