[AOTI] Fix a cubin file path issue

desertfire · desertfire · commit dc6d59c16a5e · 2024-11-05T19:54:46.000-08:00
Summary: When we use aoti_compile_and_package to package the AOTI compiled artifacts, cubin files will be included, and at the deploy time, we should setup the cubin file directory to the right path that contains unziped cubin files.

[ghstack-poisoned]
diff --git a/test/inductor/test_aot_inductor_package.py b/test/inductor/test_aot_inductor_package.py
@@ -9,6 +9,7 @@
 import torch
 from torch._inductor.package import AOTICompiledModel, load_package, package_aoti
 from torch._inductor.test_case import TestCase
+from torch._inductor.utils import fresh_inductor_cache
 from torch.export import Dim
 from torch.testing._internal.common_utils import IS_FBCODE
 from torch.testing._internal.triton_utils import HAS_CUDA
@@ -107,6 +108,44 @@ def forward(self, x, y):
         )
         self.check_model(Model(), example_inputs)
 
+    def test_remove_intermediate_files(self):
+        # For CUDA, generated cpp files contain absolute path to the generated cubin files.
+        # With the package artifact, that cubin path should be overriden at the run time,
+        # so removing those intermeidate files in this test to verify that.
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y
+
+        example_inputs = (
+            torch.randn(10, 10, device=self.device),
+            torch.randn(10, 10, device=self.device),
+        )
+        model = Model()
+        with torch.no_grad():
+            torch.manual_seed(0)
+            model = model.to(self.device)
+            ref_model = copy.deepcopy(model)
+            ref_inputs = copy.deepcopy(example_inputs)
+            expected = ref_model(*ref_inputs)
+
+            torch.manual_seed(0)
+            with tempfile.NamedTemporaryFile(suffix=".pt2") as f:
+                ep = torch.export.export(
+                    model,
+                    example_inputs,
+                )
+                with fresh_inductor_cache():
+                    # cubin files are removed when exiting this context
+                    package_path = torch._inductor.aoti_compile_and_package(
+                        ep,
+                        example_inputs,
+                        package_path=f.name,
+                    )  # type: ignore[arg-type]
+                loaded = torch._inductor.aoti_load_package(package_path)
+                actual = loaded(*example_inputs)
+
+            self.assertEqual(actual, expected)
+
     def test_linear(self):
         class Model(torch.nn.Module):
             def __init__(self) -> None:
diff --git a/torch/csrc/inductor/aoti_package/model_package_loader.cpp b/torch/csrc/inductor/aoti_package/model_package_loader.cpp
@@ -54,6 +54,13 @@ std::string create_temp_dir() {
   return temp_dir;
 #endif
 }
+
+#ifdef _WIN32
+const std::string k_separator = "\\";
+#else
+const std::string k_separator = "/";
+#endif
+
 } // namespace
 
 namespace torch::inductor {
@@ -286,6 +293,8 @@ AOTIModelPackageLoader::AOTIModelPackageLoader(
   std::string cpp_filename = "";
   std::string consts_filename = "";
   std::string found_filenames = ""; // Saving for bookkeeping
+  std::string model_directory =
+      "data" + k_separator + "aotinductor" + k_separator + model_name;
 
   for (uint32_t i = 0; i < zip_archive.m_total_files; i++) {
     uint32_t filename_len =
@@ -303,11 +312,10 @@ AOTIModelPackageLoader::AOTIModelPackageLoader(
     found_filenames += " ";
 
     // Only compile files in the specified model directory
-    std::string model_directory = "data/aotinductor/" + model_name;
     if (filename_str.length() >= model_directory.length() &&
         filename_str.substr(0, model_directory.length()) == model_directory) {
       std::string output_path_str = temp_dir;
-      output_path_str += "/";
+      output_path_str += k_separator;
       output_path_str += filename_str;
 
       // Create the parent directory if it doesn't exist
@@ -378,7 +386,8 @@ AOTIModelPackageLoader::AOTIModelPackageLoader(
     throw std::runtime_error("Unsupported device found: " + device);
   }
 
-  runner_ = registered_aoti_runner[device](so_path, 1, device, "");
+  std::string cubin_dir = temp_dir + k_separator + model_directory;
+  runner_ = registered_aoti_runner[device](so_path, 1, device, cubin_dir);
 
   std::remove(temp_dir.c_str());
 }
