[AOTInductor] Store loaded kernels in the model (#110554)

khabinov · pytorchmergebot · commit cf1b494afd0d · 2023-10-05T10:17:05.000Z
Defining kernels as static vars is problematic for subsequent model loading on non-default CUDA devices. Assuming those kernels were loaded in context of the device #0, so, they are not nullptr anymore, therefore kernels won't work on devices other than the device #0. This change makes devices remembered at model level in AOT mode. Pull Request resolved: #110554 Approved by: https://github.com/chenyang78, https://github.com/desertfire
diff --git a/.ci/pytorch/test.sh b/.ci/pytorch/test.sh
@@ -282,6 +282,7 @@ test_inductor_distributed() {
   echo "Testing multi_gpu tests in test_torchinductor"
   pytest test/inductor/test_torchinductor.py -k test_multi_gpu
   pytest test/inductor/test_aot_inductor.py -k test_non_default_cuda_device
+  pytest test/inductor/test_aot_inductor.py -k test_replicate_on_devices
 
   # this runs on both single-gpu and multi-gpu instance. It should be smart about skipping tests that aren't supported
   # with if required # gpus aren't available
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
@@ -54,15 +54,18 @@
 
 class AOTInductorModelRunner:
     @classmethod
-    def load(cls, model, example_inputs, options=None, constraints=None):
+    def compile(cls, model, example_inputs, options=None, constraints=None):
         # The exact API is subject to change
         so_path, exported = torch._export.aot_compile(
             model,
             example_inputs,
             options=options,
             constraints=constraints,
         )
+        return so_path, exported
 
+    @classmethod
+    def load(cls, so_path, example_inputs):
         is_cpu = all(x.device.type == "cpu" for x in example_inputs)
         if IS_FBCODE:
             from .fb import test_aot_inductor_model_runner_pybind
@@ -85,7 +88,7 @@ def load(cls, model, example_inputs, options=None, constraints=None):
                 with_cuda=not is_cpu,
             ).run
 
-        return optimized, exported
+        return optimized
 
     @classmethod
     def run_compiled(cls, optimized, exported, example_inputs):
@@ -97,9 +100,10 @@ def run_compiled(cls, optimized, exported, example_inputs):
 
     @classmethod
     def run(cls, model, example_inputs, options=None, constraints=None):
-        optimized, exported = AOTInductorModelRunner.load(
+        so_path, exported = AOTInductorModelRunner.compile(
             model, example_inputs, options=options, constraints=constraints
         )
+        optimized = AOTInductorModelRunner.load(so_path, example_inputs)
         return AOTInductorModelRunner.run_compiled(optimized, exported, example_inputs)
 
     @classmethod
@@ -110,12 +114,13 @@ def run_multiple(
         options=None,
         constraints=None,
     ):
-        optimized, exported = AOTInductorModelRunner.load(
+        so_path, exported = AOTInductorModelRunner.compile(
             model,
             list_example_inputs[0],
             options=options,
             constraints=constraints,
         )
+        optimized = AOTInductorModelRunner.load(so_path, list_example_inputs[0])
         list_output_tensors = []
         for example_inputs in list_example_inputs:
             list_output_tensors.append(
@@ -702,9 +707,49 @@ def forward(self, x1, x2):
         example_inputs = (a, b)
         self.check_model(Model(), example_inputs, constraints=constraints)
 
+    @requires_multigpu()
+    def test_replicate_on_devices(self):
+        if self.device != "cuda":
+            raise unittest.SkipTest("requires CUDA")
+
+        class Model(torch.nn.Module):
+            def __init__(self, w1, w2):
+                super().__init__()
+                self.w1 = w1
+                self.w2 = w2
+
+            def forward(self, x, y):
+                a = x * self.w1
+                b = y * self.w2
+                return a + b
+
+        w1 = torch.randn(10, 10)
+        w2 = torch.randn(10, 10)
+        inputs = (torch.randn(10, 10), torch.randn(10, 10))
+        result_cpu = Model(w1, w2)(*inputs)
+
+        # Compile model with AOTInductor
+        with torch.cuda.device(0), config.patch(
+            "aot_inductor.abi_compatible", self.abi_compatible
+        ):
+            so_path, exported = AOTInductorModelRunner.compile(
+                model=Model(w1.cuda(0), w2.cuda(0)),
+                example_inputs=tuple(t.cuda(0) for t in inputs),
+            )
+
+        # Run model on cuda:N
+        for i in range(torch.cuda.device_count()):
+            with torch.cuda.device(i):
+                example_inputs = tuple(t.cuda(i) for t in inputs)
+                optimized = AOTInductorModelRunner.load(so_path, example_inputs)
+                result_cuda = AOTInductorModelRunner.run_compiled(
+                    optimized, exported, example_inputs
+                )
+            self.assertTrue(same(result_cpu, result_cuda.cpu()))
+
     @requires_multigpu()
     def test_non_default_cuda_device(self):
-        if self.device == "cpu":
+        if self.device != "cuda":
             raise unittest.SkipTest("requires CUDA")
 
         class Model(torch.nn.Module):
diff --git a/torch/_inductor/codegen/wrapper.py b/torch/_inductor/codegen/wrapper.py
@@ -1184,6 +1184,9 @@ def write_wrapper_decl(self):
 
             if V.graph.aot_mode:
                 self.prefix.writeline("inputs.clear();")
+                self.prefix.writeline(
+                    "auto& kernels = *dynamic_cast<AOTInductorModelKernels*>(this->kernels_.get());"
+                )
 
     def codegen_input_size_var_decl(self, code: IndentedBuffer, name):
         if config.aot_inductor.abi_compatible:
@@ -1203,6 +1206,17 @@ def codegen_input_stride_var_decl(self, code: IndentedBuffer, name):
         else:
             super().codegen_input_stride_var_decl(code, name)
 
+    def codegen_model_kernels(self):
+        self.prefix.writeline("namespace {")
+        self.prefix.writeline(
+            "class AOTInductorModelKernels : public AOTInductorModelKernelsBase {"
+        )
+        self.prefix.writeline("  public:")
+        for kernel in self.src_to_kernel.values():
+            self.prefix.writeline(f"    CUfunction {kernel}{{nullptr}};")
+        self.prefix.writeline("};")
+        self.prefix.writeline("}  // namespace")
+
     def codegen_model_constructor(self):
         """
         // Generated code example
@@ -1270,10 +1284,15 @@ def codegen_model_constructor(self):
                 name = f"output{idx}"
                 self.write_input_output_info("outputs_info_", idx, name)
 
+            self.prefix.writeline(
+                "this->kernels_ = std::make_unique<AOTInductorModelKernels>();"
+            )
+
         self.prefix.writeline("}")
 
     def generate(self):
         if V.graph.aot_mode:
+            self.codegen_model_kernels()
             self.codegen_model_constructor()
         self.write_wrapper_decl()
         return super().generate()
@@ -1929,9 +1948,16 @@ def write_header(self):
             }  // anonymous namespace
 
             static inline CUfunction loadKernel(
-                    const std::string &filePath,
+                    std::string filePath,
                     const std::string &funcName,
-                    uint32_t sharedMemBytes) {
+                    uint32_t sharedMemBytes,
+                    const std::optional<std::string> &cubinDir = std::nullopt) {
+                if (cubinDir) {
+                    std::filesystem::path p1{*cubinDir};
+                    std::filesystem::path p2{filePath};
+                    filePath = (p1 / p2.filename()).string();
+                }
+
                 CUmodule mod;
                 CUfunction func;
                 CUDA_DRIVER_CHECK(cuModuleLoad(&mod, filePath.c_str()));
@@ -1977,9 +2003,10 @@ def define_kernel(
 
     def generate(self):
         self.prefix.writeline("\n")
-        for kernel in self.src_to_kernel.values():
-            self.prefix.writeline(f"static CUfunction {kernel} = nullptr;")
-        self.prefix.writeline("\n")
+        if not V.graph.aot_mode:
+            for kernel in self.src_to_kernel.values():
+                self.prefix.writeline(f"static CUfunction {kernel} = nullptr;")
+            self.prefix.writeline("\n")
         return super().generate()
 
     def generate_load_kernel(self, name, params):
@@ -1991,29 +2018,19 @@ def generate_load_kernel(self, name, params):
         ), f"cubin file should already exist at this moment: {cubin_path}"
 
         shared_mem = params.get("shared_mem", 0)
-        self.writeline(f"if ({name} == nullptr) {{")
+
         if V.graph.aot_mode:
-            runtime_cubin_path_var = f"var_{next(self.arg_var_id)}"
-            self.writeline("    if (this->cubin_dir_.has_value()) {")
-            self.writeline(
-                f"""        std::filesystem::path {runtime_cubin_path_var}{{this->cubin_dir_.value()}};"""
-            )
-            self.writeline(
-                f"""        {runtime_cubin_path_var} /= "{os.path.basename(cubin_path)}";"""
-            )
+            self.writeline(f"if (kernels.{name} == nullptr) {{")
             self.writeline(
-                f"""        {name} = loadKernel({runtime_cubin_path_var}.c_str(), "{mangled_name}", {shared_mem});"""
+                f"""    kernels.{name} = loadKernel("{cubin_path}", "{mangled_name}", {shared_mem}, this->cubin_dir_);"""
             )
-            self.writeline("    } else {")
-            self.writeline(
-                f"""        {name} = loadKernel("{cubin_path}", "{mangled_name}", {shared_mem});"""
-            )
-            self.writeline("    }")
+            self.writeline("}")
         else:
+            self.writeline(f"if ({name} == nullptr) {{")
             self.writeline(
                 f"""    {name} = loadKernel("{cubin_path}", "{mangled_name}", {shared_mem});"""
             )
-        self.writeline("}")
+            self.writeline("}")
 
     def generate_args_decl(self, call_args):
         dynamic_symbols = V.graph.sizevars.free_symbols()
@@ -2105,9 +2122,10 @@ def generate_kernel_call(
 
         if grid_has_unbacked_symbols:
             self.writeline(f"if ({grid_name}.is_non_zero()) {{")
+        kernel_var_name = f"kernels.{name}" if V.graph.aot_mode else name
         self.writeline(
             "launchKernel({}, {}, {}, {}, {}, {}, {}, {});".format(
-                name,
+                kernel_var_name,
                 f"{grid_name}.grid_x",
                 f"{grid_name}.grid_y",
                 f"{grid_name}.grid_z",
diff --git a/torch/csrc/inductor/aoti_runtime/model.h b/torch/csrc/inductor/aoti_runtime/model.h
@@ -283,6 +283,12 @@ class AOTInductorModelBase {
   int device_idx_;
 };
 
+// Codegen-ed classes can derive from this to keep pointers to loaded kernels.
+class AOTInductorModelKernelsBase {
+ public:
+  virtual ~AOTInductorModelKernelsBase() = default;
+};
+
 class AOTInductorModel : public AOTInductorModelBase<AOTInductorModel> {
  public:
   AOTInductorModel(std::shared_ptr<ConstantMap>, std::optional<std::string>);
@@ -303,6 +309,9 @@ class AOTInductorModel : public AOTInductorModelBase<AOTInductorModel> {
       std::optional<std::string> cubin_dir) {
     return std::make_unique<AOTInductorModel>(constants, cubin_dir);
   }
+
+ private:
+  std::unique_ptr<AOTInductorModelKernelsBase> kernels_;
 };
 
 #ifdef USE_CUDA
