[inductor] support _scaled_dot_product_flash_attention fallback

Yang Chen · Yang Chen · commit 49d63a70d4af · 2023-09-25T22:55:16.000Z
This PR supports _scaled_dot_product_flash_attention fallback kernel. Note that in the abi_compatible mode, we retrieve outputs by passing output argument pointers rather than relying on std::get. It also fixes an issue related to dynamic shapes, where we wrongfully query undefined dynamic symbols. ghstack-source-id: 3c51dab Pull Request resolved: #110003
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
@@ -538,14 +538,49 @@ def forward(self, x, y):
             constraints=constraints,
         )
 
+    # scaled_dot_product_flash_attention
+    def test_sdpa(self):
+        class Repro(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, q, k, v):
+                return torch.nn.functional.scaled_dot_product_attention(q, k, v)[0]
+
+        example_inputs = (
+            torch.randn(1, 48, 64, 64, dtype=torch.bfloat16, device="cuda"),
+            torch.randn(1, 48, 64, 64, dtype=torch.bfloat16, device="cuda"),
+            torch.randn(1, 48, 64, 64, dtype=torch.bfloat16, device="cuda"),
+        )
+        self.check_model(Repro(), example_inputs)
+
+    def test_sdpa_2(self):
+        class Repro(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, q, k, v, x):
+                t = torch.nn.functional.scaled_dot_product_attention(
+                    q, k, v, is_causal=True
+                )[0]
+                return x + t
+
+        example_inputs = (
+            torch.randn(1, 48, 64, 64, dtype=torch.bfloat16, device="cuda"),
+            torch.randn(1, 48, 64, 64, dtype=torch.bfloat16, device="cuda"),
+            torch.randn(1, 48, 64, 64, dtype=torch.bfloat16, device="cuda"),
+            torch.randn(1, 48, 64, 64, dtype=torch.bfloat16, device="cuda"),
+        )
+        self.check_model(Repro(), example_inputs)
+
 
-class AOTInductorTestABICompatibile(TestCase):
+class AOTInductorTestABICompatible(TestCase):
     abi_compatible = True
     check_model = check_model
     check_model_with_multiple_inputs = check_model_with_multiple_inputs
 
 
-copy_tests(AOTInductorTestsTemplate, AOTInductorTestABICompatibile, "abi_compatible")
+copy_tests(AOTInductorTestsTemplate, AOTInductorTestABICompatible, "abi_compatible")
 
 
 class AOTInductorTestNonABICompatible(TestCase):
diff --git a/torch/_inductor/codegen/wrapper.py b/torch/_inductor/codegen/wrapper.py
@@ -290,6 +290,11 @@ def __init__(self):
         self.first_device_guard = True
         self.supports_intermediate_hooks = True
         self.expr_printer = pexpr
+        # Not all the dynamic symbols will be used in the generated code. This
+        # set contains those actually being defined by something like
+        # "{self.declare_shape} s0 = ...". It ensures that we are not going to
+        # emit queries for undefined symbols.
+        self.defined_symbols = set()
 
         self.write_header()
         self.write_prefix()
@@ -581,6 +586,7 @@ def is_expr(x):
         for name, shape in graph_inputs_expr:
             shape = V.graph.sizevars.simplify(shape)
             if shape in needed:
+                self.defined_symbols.add(shape)
                 needed.remove(shape)
                 code.writeline(f"{self.declare}{shape} = {name}{self.ending}")
 
@@ -589,6 +595,7 @@ def is_expr(x):
             for dim, shape in enumerate(shapes):
                 shape = V.graph.sizevars.simplify(shape)
                 if shape in needed:
+                    self.defined_symbols.add(shape)
                     needed.remove(shape)
                     code.writeline(
                         f"{self.declare}{shape} = {sizeof(name)}[{dim}]{self.ending}"
@@ -599,6 +606,7 @@ def is_expr(x):
             for dim, shape in enumerate(shapes):
                 shape = V.graph.sizevars.simplify(shape)
                 if shape in needed:
+                    self.defined_symbols.add(shape)
                     needed.remove(shape)
                     code.writeline(
                         f"{self.declare}{shape} = {strideof(name)}[{dim}]{self.ending}"
@@ -617,7 +625,7 @@ def codegen_python_sizevar(self, x: Expr) -> str:
     def codegen_sizevar(self, x: Expr) -> str:
         return self.codegen_python_sizevar(x)
 
-    def codegen_tuple_access(self, basename: str, index: str) -> str:
+    def codegen_tuple_access(self, basename: str, name: str, index: str) -> str:
         return f"{basename}[{index}]"
 
     def codegen_python_shape_tuple(self, shape: Tuple[Expr, ...]) -> str:
@@ -637,6 +645,9 @@ def codegen_reinterpret_view(self, name, size, stride, offset, writer) -> str:
         offset = self.codegen_sizevar(offset)
         return f"reinterpret_tensor({name}, {size}, {stride}, {offset})"
 
+    def codegen_multi_output(self, name, value):
+        self.writeline(f"{self.declare}{name} = {value}{self.ending}")
+
     def benchmark_compiled_module(self, output):
         def add_fake_input(name, shape, stride, device, dtype):
             output.writeline(
@@ -1182,7 +1193,11 @@ def write_wrapper_decl(self):
 
             if V.graph.aot_mode:
                 self.prefix.writeline("inputs.clear();")
-                dynamic_symbols = V.graph.sizevars.free_symbols()
+                dynamic_symbols = [
+                    s
+                    for s in V.graph.sizevars.free_symbols()
+                    if s in self.defined_symbols
+                ]
                 for dim in dynamic_symbols:
                     self.prefix.writeline(
                         f'auto dim_{dim} = find_dynamic_dim("{dim}");'
@@ -1402,21 +1417,38 @@ def generate_c_shim_extern_kernel_call(self, kernel, args):
         kernel = "aoti_torch_" + kernel.split("::")[-1]
         self.writeline(f"AOTI_TORCH_ERROR_CODE_CHECK({kernel}({', '.join(args)}));")
 
+    def generate_c_shim_extern_kernel_alloc_call(self, extern_kernel, args):
+        output_args = []
+        output_raii_handles = []
+        output_name_base = extern_kernel.get_name()
+        for idx, output in enumerate(extern_kernel.outputs):
+            if isinstance(output, ir.MultiOutput):
+                name = f"{output.get_name()}"
+                output_handle_name = f"{name}_handle"
+                assert (
+                    output.indices[0][1] == idx
+                ), f"expected {output.indices[1]=} == {idx=} for {output_name_base=}"
+                self.writeline(f"AtenTensorHandle {output_handle_name};")
+                output_args.append(f"&{output_handle_name}")
+                output_raii_handles.append(
+                    f"RAIIAtenTensorHandle {name}({output_handle_name});"
+                )
+            elif isinstance(output, int):
+                output_name = f"{output_name_base}_{idx}"
+                self.writeline(f"int64_t {output_name} = {output};")
+                output_args.append(f"&{output_name}")
+            elif output is None:
+                output_args.append("nullptr")
+            else:
+                raise NotImplementedError("unsupported type of {output=}")
+        args = args + output_args
+        self.generate_c_shim_extern_kernel_call(extern_kernel.kernel, args)
+        for raii_handle in output_raii_handles:
+            self.writeline(raii_handle)
+
     def generate_extern_kernel_alloc(self, extern_kernel, args):
         if V.graph.aot_mode and config.aot_inductor.abi_compatible:
-            output_name = extern_kernel.get_name()
-            self.writeline(f"AtenTensorHandle {output_name};")
-            kernel = extern_kernel.kernel
-            size = self.codegen_shape_tuple(tuple(extern_kernel.get_size()))
-            stride = self.codegen_shape_tuple(tuple(extern_kernel.get_stride()))
-            args = [
-                f"&{output_name}",
-                str(len(extern_kernel.get_size())),  # ndim
-                self.codegen_int_array_var(size),
-                self.codegen_int_array_var(stride),
-            ] + args
-            # TODO: support extern kernel that allocates
-            self.generate_c_shim_extern_kernel_call(kernel, args)
+            self.generate_c_shim_extern_kernel_alloc_call(extern_kernel, args)
         else:
             super().generate_extern_kernel_alloc(extern_kernel, args)
 
@@ -1461,8 +1493,12 @@ def add_benchmark_harness(self, output):
     def codegen_sizevar(self, x: Expr) -> str:
         return self.expr_printer(V.graph.sizevars.simplify(x))
 
-    def codegen_tuple_access(self, basename: str, index: str) -> str:
-        return f"std::get<{index}>({basename})"
+    def codegen_tuple_access(self, basename: str, name: str, index: str) -> str:
+        if V.graph.aot_mode and config.aot_inductor.abi_compatible:
+            # in the abi_compatible mode, outputs are returned via arguments
+            return name
+        else:
+            return f"std::get<{index}>({basename})"
 
     def codegen_shape_tuple(self, shape: Tuple[Expr, ...]) -> str:
         parts = list(map(self.codegen_sizevar, shape))
@@ -1584,6 +1620,11 @@ def codegen_reinterpret_view(self, name, size, stride, offset, writer) -> str:
             args = [name, size, stride, offset]
             return f"reinterpret_tensor({', '.join(args)})"
 
+    def codegen_multi_output(self, name, value):
+        # if V.graph.aot_mode and name in set(V.graph.get_output_names()):
+        if not config.aot_inductor.abi_compatible:
+            super().codegen_multi_output(name, value)
+
     def generate_extern_kernel_args_decl_if_needed(
         self, op_overload, raw_args, output_args
     ):
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -3627,6 +3627,10 @@ def __init__(
             tuple(tensor_args),
             tuple(nontensor_args),
         )
+        # We need output buffers for generating kernel arguments in the
+        # abi-compatible mode, where we retrieve outputs by pass each individual
+        # output through the abi-compatible interface.
+        self.outputs = []
         self.use_cpp_op_schema = False
 
         self.op_overload = kernel
@@ -3879,7 +3883,8 @@ def generate_output(output, indices):
                 assert output is None, "FallbackKernel output type is not supported"
                 return None
 
-        return generate_output(example_output, [])
+        packed.outputs = generate_output(example_output, [])
+        return packed.outputs
 
     def apply_constraint(self):
         return super().apply_constraint()
@@ -3899,7 +3904,7 @@ def codegen_list_tuple_access(self, basename, indices):
             elif itype == tuple:
                 # cpp wrapper code needs to use std::get<> to access a tuple
                 tuple_access = V.graph.wrapper_code.codegen_tuple_access(
-                    basename, str(i)
+                    basename, self.get_name(), str(i)
                 )
                 return self.codegen_list_tuple_access(tuple_access, indices[1:])
             else:
@@ -3908,10 +3913,10 @@ def codegen_list_tuple_access(self, basename, indices):
             return basename
 
     def codegen(self, wrapper):
-        line = V.graph.wrapper_code.declare
-        line += f"{self.get_name()} = {self.codegen_list_tuple_access(self.inputs[0].get_name(), self.indices)}"
-        line += V.graph.wrapper_code.ending
-        V.graph.wrapper_code.writeline(line)
+        V.graph.wrapper_code.codegen_multi_output(
+            self.get_name(),
+            self.codegen_list_tuple_access(self.inputs[0].get_name(), self.indices),
+        )
         self.codegen_size_asserts(V.graph.wrapper_code)
 
     def __init__(self, layout, input, indices: List[Tuple[Any, ...]]):
diff --git a/torch/csrc/inductor/aoti_torch/c/shim.h b/torch/csrc/inductor/aoti_torch/c/shim.h
@@ -152,6 +152,25 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_tensor_from_blob(
     AtenTensorHandle* ret // returns new reference
 );
 
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch__scaled_dot_product_flash_attention(
+    AtenTensorHandle query,
+    AtenTensorHandle key,
+    AtenTensorHandle value,
+    double dropout_p,
+    bool is_causal,
+    bool return_debug_mask,
+    double scale,
+    AtenTensorHandle* ret0, // returns new reference
+    AtenTensorHandle* ret1, // returns new reference
+    AtenTensorHandle* ret2, // returns new reference
+    AtenTensorHandle* ret3, // returns new reference
+    int64_t* ret4,
+    int64_t* ret5,
+    AtenTensorHandle* ret6, // returns new reference
+    AtenTensorHandle* ret7, // returns new reference
+    AtenTensorHandle* ret8 // returns new reference
+);
+
 AOTI_TORCH_EXPORT AOTITorchError
 aoti_torch_tensor_copy_(AtenTensorHandle src, AtenTensorHandle dst);
 
diff --git a/torch/csrc/inductor/aoti_torch/shim_common.cpp b/torch/csrc/inductor/aoti_torch/shim_common.cpp
@@ -6,16 +6,19 @@
 #include <torch/csrc/inductor/aoti_torch/tensor_converter.h>
 #include <torch/csrc/inductor/aoti_torch/utils.h>
 #include <torch/csrc/inductor/inductor_ops.h>
+#include <cstdarg>
 #include <cstdint>
 #include <cstdio>
 #include <iostream>
 #include <memory>
+#include <tuple>
 
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/Functions.h>
 #else
 
 #include <ATen/ops/_addmm_activation.h>
+#include <ATen/ops/_scaled_dot_product_flash_attention.h>
 #include <ATen/ops/addmm.h>
 #include <ATen/ops/as_strided.h>
 #include <ATen/ops/bmm.h>
@@ -182,6 +185,62 @@ AOTITorchError aoti_torch_create_tensor_from_blob(
   });
 }
 
+AOTITorchError aoti_torch__scaled_dot_product_flash_attention(
+    AtenTensorHandle query,
+    AtenTensorHandle key,
+    AtenTensorHandle value,
+    double dropout_p,
+    bool is_causal,
+    bool return_debug_mask,
+    double scale,
+    AtenTensorHandle* ret0, // returns new reference
+    AtenTensorHandle* ret1, // returns new reference
+    AtenTensorHandle* ret2, // returns new reference
+    AtenTensorHandle* ret3, // returns new reference
+    int64_t* ret4,
+    int64_t* ret5,
+    AtenTensorHandle* ret6, // returns new reference
+    AtenTensorHandle* ret7, // returns new reference
+    AtenTensorHandle* ret8 // returns new reference
+) {
+  AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
+    at::Tensor* query_tensor = tensor_handle_to_tensor_pointer(query);
+    at::Tensor* key_tensor = tensor_handle_to_tensor_pointer(key);
+    at::Tensor* value_tensor = tensor_handle_to_tensor_pointer(value);
+    auto [r0, r1, r2, r3, r4, r5, r6, r7, r8] =
+        at::_scaled_dot_product_flash_attention(
+            *query_tensor,
+            *key_tensor,
+            *value_tensor,
+            dropout_p,
+            is_causal,
+            return_debug_mask,
+            scale);
+
+    at::Tensor* ret0_tensor = new at::Tensor(std::move(r0));
+    *ret0 = tensor_pointer_to_tensor_handle(ret0_tensor);
+    at::Tensor* ret1_tensor = new at::Tensor(std::move(r1));
+    *ret1 = tensor_pointer_to_tensor_handle(ret1_tensor);
+    // ret2 and ret3 may be null
+    if (ret2) {
+      at::Tensor* ret2_tensor = new at::Tensor(std::move(r2));
+      *ret2 = tensor_pointer_to_tensor_handle(ret2_tensor);
+    }
+    if (ret3) {
+      at::Tensor* ret3_tensor = new at::Tensor(std::move(r3));
+      *ret3 = tensor_pointer_to_tensor_handle(ret3_tensor);
+    }
+    *ret4 = r4;
+    *ret5 = r5;
+    at::Tensor* ret6_tensor = new at::Tensor(std::move(r6));
+    *ret6 = tensor_pointer_to_tensor_handle(ret6_tensor);
+    at::Tensor* ret7_tensor = new at::Tensor(std::move(r7));
+    *ret7 = tensor_pointer_to_tensor_handle(ret7_tensor);
+    at::Tensor* ret8_tensor = new at::Tensor(std::move(r8));
+    *ret8 = tensor_pointer_to_tensor_handle(ret8_tensor);
+  });
+}
+
 // TODO: implement a more efficient version instead of calling into aten
 AOTITorchError aoti_torch_tensor_copy_(
     AtenTensorHandle src,
