[aoti] Initial Metal support (pytorch#153959)

An example generated file: P1816629015

Pull Request resolved: pytorch#153959
Approved by: https://github.com/malfet, https://github.com/desertfire
ghstack dependencies: pytorch#153964

Author: angelayi

test/inductor/test_mps_basic.py

@@ -181,6 +181,22 @@ def fn(x, y):
         )
 
 
+class MPSBasicTestsAOTI(TestCase):
+    def test_add_mps(self):
+        class M(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y
+
+        inp = (torch.ones(3, 3, device="mps"), torch.ones(3, 3, device="mps"))
+        m = M().to("mps")
+        res2 = m(*inp)
+        ep = torch.export.export(m, inp)
+        path = torch._inductor.aoti_compile_and_package(ep, "here.pt2")
+        m = torch._inductor.aoti_load_package(path)
+        res = m(*inp)
+        assert torch.allclose(res, res2)
+
+
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
 

torch/_export/__init__.py

@@ -165,7 +165,8 @@ def aot_load(so_path: str, device: str) -> Callable:
         runner = torch._C._aoti.AOTIModelContainerRunnerCuda(so_path, 1, device)  # type: ignore[assignment, call-arg]
     elif device == "xpu" or device.startswith("xpu:"):
         runner = torch._C._aoti.AOTIModelContainerRunnerXpu(so_path, 1, device)  # type: ignore[assignment, call-arg]
-
+    elif device == "mps" or device.startswith("mps:"):
+        runner = torch._C._aoti.AOTIModelContainerRunnerMps(so_path, 1)  # type: ignore[assignment, call-arg]
     else:
         raise RuntimeError("Unsupported device " + device)
 

torch/_inductor/codegen/common.py

@@ -447,6 +447,7 @@ def init_backend_registration() -> None:
     from .cpp_wrapper_cpu import CppWrapperCpu
     from .cpp_wrapper_cpu_array_ref import CppWrapperCpuArrayRef
     from .cpp_wrapper_gpu import CppWrapperGpu
+    from .cpp_wrapper_mps import CppWrapperMps
     from .cuda_combined_scheduling import CUDACombinedScheduling
     from .halide import HalideScheduling
     from .mps import MetalScheduling
@@ -494,7 +495,7 @@ def init_backend_registration() -> None:
             "mps",
             MetalScheduling,
             PythonWrapperCodegen,
-            CppWrapperGpu,
+            CppWrapperMps,
         )
 
     private_backend = torch._C._get_privateuse1_backend_name()

torch/_inductor/codegen/cpp_utils.py

@@ -80,6 +80,7 @@
     "cpu": "at::kCPU",
     "cuda": "at::kCUDA",
     "xpu": "at::kXPU",
+    "mps": "at::kMPS",
 }
 
 LAYOUT_TO_ATEN = {

torch/_inductor/codegen/cpp_wrapper_cpu.py

@@ -131,7 +131,9 @@ def _generate_kernel_call_helper(
                 Only valid when cuda == True.
         """
         assert arg_types is not None and len(call_args) == len(arg_types), (
-            "Mismatch call_args and arg_types in generate_kernel_call"
+            "Mismatch call_args and arg_types in generate_kernel_call:\n"
+            f"call_args: {call_args}\n"
+            f"arg_types: {arg_types}"
         )
         new_args = []
         for idx, arg in enumerate(call_args):

torch/_inductor/codegen/cpp_wrapper_mps.py

@@ -0,0 +1,83 @@
+from typing import Any, Optional
+
+from ..ir import GraphPartitionSignature
+from ..virtualized import V
+from .cpp_wrapper_gpu import CppWrapperGpu
+from .wrapper import PythonWrapperCodegen
+
+
+class CppWrapperMps(CppWrapperGpu):
+    @staticmethod
+    def create(
+        is_subgraph: bool,
+        subgraph_name: Optional[str],
+        parent_wrapper: Optional[PythonWrapperCodegen],
+        partition_signatures: Optional[GraphPartitionSignature] = None,
+    ) -> "CppWrapperMps":
+        return CppWrapperMps()
+
+    def _generate_kernel_call_helper(
+        self,
+        kernel_name: str,
+        call_args: list[str],
+        **kwargs: dict[str, Any],
+    ) -> None:
+        """
+        Generates MPS kernel call code. It should look something like:
+        ```
+        auto mps_lib_0_func = mps_lib_0.getKernelFunction("generated_kernel");
+        auto mps_lib_0_func_handle = AOTIMetalKernelFunctionHandle(mps_lib_0_func.get());
+        mps_lib_0_func->runCommandBlock([&] {
+            mps_lib_0_func->startEncoding();
+            aoti_torch_mps_set_arg(mps_lib_0_func_handle, 0, buf0);
+            aoti_torch_mps_set_arg(mps_lib_0_func_handle, 1, arg0_1);
+            ...
+            mps_lib_0_func->dispatch(9);
+        });
+        ```
+        """
+        new_args = []
+        for idx, arg in enumerate(call_args[:-2]):
+            new_args.append(
+                f"aoti_torch_mps_set_arg({kernel_name}_handle, {idx}, {arg});\n"
+            )
+
+        threads, group_size = call_args[-2], call_args[-1]
+        if threads is None:
+            raise NotImplementedError("No threads or group_size provided")
+        elif group_size is None:
+            new_args.append(f"{kernel_name}->dispatch({threads});\n")
+        else:
+            new_args.append(f"{kernel_name}->dispatch({threads}, {group_size});\n")
+
+        # debug printer related logic for cpp kernel type.
+        debug_printer_manager = V.graph.wrapper_code.debug_printer
+        debug_printer_manager.set_printer_args(
+            call_args[:-2],
+            kernel_name,
+            None,
+            None,
+            "cpp",
+        )
+        with debug_printer_manager:
+            self.writeline(self.wrap_kernel_call(kernel_name, new_args))
+
+    def wrap_kernel_call(self, name: str, call_args: list[str]) -> str:
+        lib_name = name[: -len("_func")]
+        calling_args = "        ".join(call_args)
+        return f"""
+    auto {name} = {lib_name}.getKernelFunction("generated_kernel");
+    auto {name}_handle = AOTIMetalKernelFunctionHandle({name}.get());
+    {name}->runCommandBlock([&] {{
+        {name}->startEncoding();
+        {calling_args}
+    }});
+        """
+
+    @staticmethod
+    def get_device_include_path(device: str) -> str:
+        assert V.graph.aot_mode
+        return (
+            "#include <torch/csrc/inductor/aoti_include/mps.h>\n"
+            "#include <torch/csrc/inductor/aoti_torch/c/shim_mps.h>"
+        )

torch/_inductor/codegen/mps.py

@@ -775,11 +775,17 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:
         """Called at the end to generate a final kernel string"""
         self.codegen_body()
         code = IndentedBuffer()
-        code.writeline("compile_mps_shader('''")
+
+        if V.graph.cpp_wrapper:
+            code.writeline('(R"MTL(')
+        else:
+            code.writeline("compile_mps_shader('''")
+
         idx_vars = self.active_range_trees()
         with code.indent():
-            for header in self.headers:
-                code.writeline(f"#include <c10/metal/{header}.h>")
+            if not V.graph.cpp_wrapper:
+                for header in self.headers:
+                    code.writeline(f"#include <c10/metal/{header}.h>")
             if self.inside_reduction:
                 total_reduction_size = math.prod(
                     t.numel for t in self.range_trees if t.is_reduction
@@ -833,7 +839,11 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:
                 code.splice(self.indexing_code)
                 code.splice(self.body)
             code.writeline("}")
-        code.writeline("''')")
+
+        if V.graph.cpp_wrapper:
+            code.writeline(')MTL");')
+        else:
+            code.writeline("''')")
 
         return code.getvalue()
 
@@ -858,15 +868,31 @@ def call_kernel(self, name: str, node: Any = None) -> None:
                 )
                 for v in self.active_range_trees()
             ]
-            args += [f"threads=[{', '.join(threads)}]"]
+
+            if V.graph.cpp_wrapper:
+                args += [f"{', '.join(threads)}"]
+            else:
+                args += [f"threads=[{', '.join(threads)}]"]
+        else:
+            if V.graph.cpp_wrapper:
+                raise RuntimeError("We should always have threads?")
+
         if self.inside_reduction:
             threads = [
                 self.pexpr(sympy.Min(v.numel, self.max_threadgroup_size))  # type: ignore[misc]
                 if v.is_reduction
                 else "1"
                 for v in self.active_range_trees()
             ]
-            args += [f"group_size=[{', '.join(threads)}]"]
+            if V.graph.cpp_wrapper:
+                args += [f"{{{', '.join(threads)}}}"]
+            else:
+                args += [f"group_size=[{', '.join(threads)}]"]
+        else:
+            if V.graph.cpp_wrapper:
+                # Add a None so that we always have a group_size in the
+                # arguments. We won't use it if the value is None.
+                args += [None]  # type: ignore[list-item]
 
         wrapper.generate_kernel_call(
             name,
@@ -900,9 +926,10 @@ def __init__(self, scheduler: Optional[Scheduler]) -> None:
         super().__init__(scheduler)
         wrapper = V.graph.wrapper_code
         if wrapper is not None:
-            wrapper.header.splice(
-                "from torch._inductor.runtime.runtime_utils import compile_mps_shader"
-            )
+            if not V.graph.cpp_wrapper:
+                wrapper.header.splice(
+                    "from torch._inductor.runtime.runtime_utils import compile_mps_shader"
+                )
 
     def define_kernel(
         self, src_code: str, node_schedule: list[SchedulerNode], kernel: MetalKernel
@@ -914,10 +941,19 @@ def define_kernel(
             # TODO: Merge multiple kernels into a single library
             # Either using MultiKernel concept or overriding SIMDScheduling.codegen_node_scheduling
             mps_lib_name = f"mps_lib_{wrapper.next_kernel_suffix()}"
-            kernel_name = f"{mps_lib_name}.generated_kernel"
+
+            if V.graph.cpp_wrapper:
+                src_code = (
+                    f"at::native::mps::DynamicMetalShaderLibrary {mps_lib_name}"
+                    + src_code
+                )
+                kernel_name = f"{mps_lib_name}_func"
+            else:
+                kernel_name = f"{mps_lib_name}.generated_kernel"
+
             wrapper.src_to_kernel[src_code] = kernel_name
             origins, detailed_origins = get_kernel_metadata(node_schedule, wrapper)
             metadata_comment = f"{origins}\n{detailed_origins}"
-            wrapper.define_kernel(mps_lib_name, src_code, metadata_comment)
+            wrapper.define_kernel(mps_lib_name, src_code, metadata_comment, gpu=False)
 
         return kernel_name

torch/_inductor/codegen/mps_device_op_overrides.py

@@ -12,5 +12,13 @@ def set_device(self, device_idx: int) -> str:
         assert device_idx == 0
         return "pass  # MPS set device"
 
+    def kernel_driver(self) -> str:
+        return """
+            #include <ATen/native/mps/MetalShaderLibrary.h>
+        """
+
+    def cpp_kernel_type(self) -> str:
+        return "MTLFunction_t"
+
 
 register_device_op_overrides("mps", MPSDeviceOpOverrides())

torch/_inductor/utils.py

@@ -2475,7 +2475,7 @@ def is_gpu(device: Optional[str]) -> bool:
 
 
 def device_need_guard(device: str) -> bool:
-    return is_gpu(device)
+    return device != "mps" and is_gpu(device)  # TODO: MPS does not expose streams now
 
 
 def needs_fallback_due_to_atomic_add_limitations(dtype: torch.dtype) -> bool:

torch/csrc/inductor/aoti_runner/model_container_runner_mps.h

@@ -1,4 +1,4 @@
-#if !defined(C10_MOBILE) && !defined(ANDROID)
+#if defined(__APPLE__)
 #pragma once
 
 #include <torch/csrc/inductor/aoti_runner/model_container_runner.h>