[AOTI] Fix aot_inductor_package test errors (#148279)

Summary: Fix fbcode test failures introduced by https://github.com/pytorch/pytorch/pull/147975. Make sure script.ld is copied to the build-time directory.

Differential Revision: D70454149

Pull Request resolved: https://github.com/pytorch/pytorch/pull/148279
Approved by: https://github.com/zoranzhao

setup.py

@@ -1363,6 +1363,7 @@ def main():
         "include/sleef.h",
         "_inductor/codegen/*.h",
         "_inductor/codegen/aoti_runtime/*.cpp",
+        "_inductor/script.ld",
         "_export/serde/*.yaml",
         "_export/serde/*.thrift",
         "share/cmake/ATen/*.cmake",

torch/_inductor/codecache.py

@@ -54,6 +54,7 @@ from torch._inductor.codegen.rocm.compile_command import (
     rocm_compiler,
 )
 from torch._inductor.cpp_builder import (
+    _LINKER_SCRIPT,
     _set_gpu_runtime_env,
     _TORCH_PATH,
     _transform_cuda_paths,
@@ -1714,6 +1715,7 @@ class AotCodeCompiler:
                 )
                 so_build_options.save_flags_to_json(linker_flags)
                 generated_files.append(linker_flags)
+                generated_files.append(_LINKER_SCRIPT)
 
                 # If we only want to package the cpp, then we need to save the
                 # weights separately into a bin, and we also need to prevent compiling the so

torch/_inductor/cpp_builder.py

@@ -723,10 +723,7 @@ def _setup_standard_sys_libs(
         include_dirs.append(build_paths.linux_kernel_include)
         include_dirs.append("include")
 
-        if aot_mode and not use_relative_path:
-            linker_script = _LINKER_SCRIPT
-        else:
-            linker_script = os.path.basename(_LINKER_SCRIPT)
+        linker_script = os.path.basename(_LINKER_SCRIPT)
 
         if _is_clang(cpp_compiler):
             passthrough_args.append(" --rtlib=compiler-rt")

torch/_inductor/package/package.py

@@ -2,8 +2,6 @@ import io
 import json
 import logging
 import os
-import shlex
-import subprocess
 import tempfile
 import zipfile
 from pathlib import Path
@@ -13,7 +11,7 @@ from typing_extensions import Self
 import torch
 import torch._inductor
 import torch.utils._pytree as pytree
-from torch._inductor import exc
+from torch._inductor import config
 from torch._inductor.cpp_builder import BuildOptionsBase, CppBuilder
 from torch.export._tree_utils import reorder_kwargs
 from torch.types import FileLike
@@ -92,14 +90,6 @@ class PT2ArchiveReader:
         return self.archive_file.namelist()
 
 
-def _run_command_and_check(cmd: str) -> None:
-    cmd = shlex.split(cmd)
-    try:
-        subprocess.run(cmd, check=True)
-    except subprocess.CalledProcessError as e:
-        raise exc.CppCompileError(cmd, e.output) from e
-
-
 def compile_so(aoti_dir: str, aoti_files: list[str], so_path: str) -> str:
     def get_aoti_file_with_suffix(suffix: str) -> str:
         for file in aoti_files:
@@ -117,32 +107,32 @@ def compile_so(aoti_dir: str, aoti_files: list[str], so_path: str) -> str:
     with open(file_name + "_compile_flags.json") as f:
         compile_flags = json.load(f)
 
-    compile_options = BuildOptionsBase(**compile_flags)
+    compile_options = BuildOptionsBase(
+        **compile_flags, use_relative_path=config.is_fbcode()
+    )
     object_builder = CppBuilder(
         name=file_name,
         sources=cpp_file,
         BuildOption=compile_options,
     )
-    compile_cmd = object_builder.get_command_line()
     output_o = object_builder.get_target_file_path()
-
-    _run_command_and_check(compile_cmd)
+    object_builder.build()
 
     # Parse linker flags and build the .so file
     with open(file_name + "_linker_flags.json") as f:
         linker_flags = json.load(f)
 
-    linker_options = BuildOptionsBase(**linker_flags)
+    linker_options = BuildOptionsBase(
+        **linker_flags, use_relative_path=config.is_fbcode()
+    )
     so_builder = CppBuilder(
         name=os.path.split(so_path)[-1],
         sources=[output_o, consts_o],
         BuildOption=linker_options,
         output_dir=so_path,
     )
-    link_cmd = so_builder.get_command_line()
     output_so = so_builder.get_target_file_path()
-
-    _run_command_and_check(link_cmd)
+    so_builder.build()
 
     # mmapped weights
     serialized_weights_filename = file_name + "_serialized_weights.bin"

torch/csrc/inductor/aoti_package/model_package_loader.cpp

@@ -98,6 +98,11 @@ std::tuple<std::string, std::string> get_cpp_compile_command(
 
   std::string file_ext = compile_only ? ".o" : ".so";
   std::string target_file = output_dir + filename + file_ext;
+  std::string target_dir = output_dir;
+  if (target_dir.empty()) {
+    size_t parent_path_idx = filename.find_last_of(k_separator);
+    target_dir = filename.substr(0, parent_path_idx);
+  }
 
   std::string cflags_args;
   for (auto& arg : compile_options["cflags"]) {
@@ -131,7 +136,15 @@ std::tuple<std::string, std::string> get_cpp_compile_command(
 
   std::string passthrough_parameters_args;
   for (auto& arg : compile_options["passthrough_args"]) {
-    passthrough_parameters_args += arg.get<std::string>() + " ";
+    std::string arg_str = arg.get<std::string>();
+    std::string target = "script.ld";
+    std::string replacement = target_dir;
+    replacement.append(k_separator).append(target);
+    size_t pos = arg_str.find(target);
+    if (pos != std::string::npos) {
+      arg_str.replace(pos, target.length(), replacement);
+    }
+    passthrough_parameters_args += arg_str + " ";
   }
 
   std::string compile_only_arg = compile_only ? "-c" : "";