[reland][inductor] switch AotCodeCompiler to new cpp_builder (pytorch#130127)

Changes:
1. Switch `AotCodeCompiler` to new cpp_builder.
2. Only use `deprecated_cpp_compile_command` for `fb_code`, due to I can't debug anymore on no Meta internal environment access.
3. Add `TODO` comments for further some Meta employee help on contine to do this work.
4. Due to item 3, we only remaining `deprecated_cpp_compile_command` for `fb_code` to be fix, let's remove `validate_new_cpp_commands`.

Pull Request resolved: pytorch#130127
Approved by: https://github.com/jgong5, https://github.com/jansel

Author: xuhancn

test/inductor/test_torchinductor.py

@@ -7121,11 +7121,6 @@ def fn(x):
 
         self.common(fn, [torch.randn(64, 64)])
 
-    def test_new_cpp_build_logical(self):
-        from torch._inductor.codecache import validate_new_cpp_commands
-
-        validate_new_cpp_commands()
-
     def test_as_strided(self):
         def fn(x):
             return (

torch/_inductor/codecache.py

@@ -1508,7 +1508,7 @@ def get_include_and_linking_paths(
     return ipaths, lpaths_str, libs_str, macros, build_arch_flags
 
 
-def cpp_compile_command(
+def deprecated_cpp_compile_command(
     input: Union[str, List[str]],
     output: str,
     warning_all: bool = True,
@@ -1522,6 +1522,13 @@ def cpp_compile_command(
     use_mmap_weights: bool = False,
     extra_flags: Sequence[str] = (),
 ) -> str:
+    """
+    Please don't use this function in new development code.
+    It was planed to remove after we switched to new cpp_builder, but I can't access to Meta
+    internal environment to fix AotCodeCompiler fb_code.
+    TODO: need some Meta employee help on fix AotCodeCompiler fb_code, and then delete this
+    deprecated function.
+    """
     ipaths, lpaths, libs, macros, build_arch_flags = get_include_and_linking_paths(
         include_pytorch, vec_isa, cuda, aot_mode
     )
@@ -1576,6 +1583,20 @@ def cpp_compile_command(
     ).strip()
 
 
+def _temp_validate_new_and_old_command(new_cmd: List[str], old_cmd: List[str]) -> None:
+    """
+    TODO: Will remove the temp code after switch to new cpp_builder
+    """
+    new_diff: List[str] = [x for x in new_cmd if x not in old_cmd]
+    old_diff: List[str] = [y for y in old_cmd if y not in new_cmd]
+    if new_diff or old_diff:
+        print("!!! new_cmd: ", new_cmd)
+        print("!!! old_cmd: ", old_cmd)
+        print("!!! new_diff: ", new_diff)
+        print("!!! old_diff: ", old_diff)
+        raise RuntimeError("Error in new and old command different.")
+
+
 def run_command_and_check(cmd: str) -> None:
     cmd = shlex.split(cmd)
     try:
@@ -1823,7 +1844,6 @@ def _compile_consts_darwin(consts: bytes) -> str:
                 else os.path.splitext(input_path)[0] + ".so"
             )
 
-            output_o = os.path.splitext(input_path)[0] + ".o"
             consts_size = sum(
                 torch.ops.mkldnn._nbytes(tensor)
                 if tensor.is_mkldnn
@@ -1836,43 +1856,50 @@ def _compile_consts_darwin(consts: bytes) -> str:
             if config.aot_inductor.force_mmap_weights:
                 use_mmap_weights = True
 
-            if config.aot_inductor.package:
-                (
-                    object_output_name,
-                    object_output_dir,
-                ) = get_name_and_dir_from_output_file_path(input_path)
-                object_build_options = CppTorchCudaOptions(
-                    vec_isa=picked_vec_isa,
-                    cuda=cuda,
-                    aot_mode=graph.aot_mode,
-                    compile_only=True,
-                    use_absolute_path=use_absolute_path,
-                    use_mmap_weights=use_mmap_weights,
-                )
-                object_builder = CppBuilder(
-                    name=object_output_name,
-                    sources=input_path,
-                    output_dir=object_output_dir,
-                    BuildOption=object_build_options,
-                )
-                compile_cmd = object_builder.get_command_line()
-                output_o = object_builder.get_target_file_path()
+            (
+                object_output_name,
+                object_output_dir,
+            ) = get_name_and_dir_from_output_file_path(input_path)
+            object_build_options = CppTorchCudaOptions(
+                vec_isa=picked_vec_isa,
+                cuda=cuda,
+                aot_mode=graph.aot_mode,
+                compile_only=True,
+                use_absolute_path=use_absolute_path,
+                use_mmap_weights=use_mmap_weights,
+            )
+            object_builder = CppBuilder(
+                name=object_output_name,
+                sources=input_path,
+                output_dir=object_output_dir,
+                BuildOption=object_build_options,
+            )
+            compile_cmd = object_builder.get_command_line()
+            output_o = object_builder.get_target_file_path()
 
+            if config.aot_inductor.package:
                 compile_flags = os.path.splitext(input_path)[0] + "_compile_flags.json"
                 object_build_options.save_flags_to_file(compile_flags)
-
             else:
-                # TODO: replace this with using the CppBuilder above
-                compile_cmd = cpp_compile_command(
-                    input=input_path,
-                    output=output_o,
-                    vec_isa=picked_vec_isa,
-                    cuda=cuda,
-                    aot_mode=graph.aot_mode,
-                    compile_only=True,
-                    use_absolute_path=use_absolute_path,
-                    use_mmap_weights=use_mmap_weights,
-                )
+                if config.is_fbcode():
+                    # TODO: enable AotCodeCompiler fb_code, and remove deprecated_cpp_compile_command.
+                    compile_cmd_old = deprecated_cpp_compile_command(
+                        input=input_path,
+                        output=output_o,
+                        vec_isa=picked_vec_isa,
+                        cuda=cuda,
+                        aot_mode=graph.aot_mode,
+                        compile_only=True,
+                        use_absolute_path=use_absolute_path,
+                        use_mmap_weights=use_mmap_weights,
+                    )
+                    # TODO: Enable below code to debug in fb_code.
+                    """
+                    _temp_validate_new_and_old_command(
+                        compile_cmd.split(" "), compile_cmd_old.split(" ")
+                    )
+                    """
+                    compile_cmd = compile_cmd_old
 
                 log.debug("aot compilation command: %s", compile_cmd)
                 if fbcode_aot_cpu_re:
@@ -1935,28 +1962,25 @@ def _pad_to_alignment(raw_bytes: bytes) -> bytes:
                 "darwin": _compile_consts_darwin,
             }[sys.platform](aot_constants)
 
-            if config.aot_inductor.package:
-                output_name, output_dir = get_name_and_dir_from_output_file_path(
-                    output_so
-                )
-                so_build_options = CppTorchCudaOptions(
-                    vec_isa=picked_vec_isa,
-                    cuda=cuda,
-                    aot_mode=graph.aot_mode,
-                    use_absolute_path=use_absolute_path,
-                )
-                so_builder = CppBuilder(
-                    name=output_name,
-                    sources=[output_o, consts_o],
-                    output_dir=output_dir,
-                    BuildOption=so_build_options,
-                )
-                link_cmd = so_builder.get_command_line()
-                output_so = so_builder.get_target_file_path()
+            output_name, output_dir = get_name_and_dir_from_output_file_path(output_so)
+            so_build_options = CppTorchCudaOptions(
+                vec_isa=picked_vec_isa,
+                cuda=cuda,
+                aot_mode=graph.aot_mode,
+                use_absolute_path=use_absolute_path,
+            )
+            so_builder = CppBuilder(
+                name=output_name,
+                sources=[output_o, consts_o],
+                output_dir=output_dir,
+                BuildOption=so_build_options,
+            )
+            link_cmd = so_builder.get_command_line()
+            output_so = so_builder.get_target_file_path()
 
+            if config.aot_inductor.package:
                 linker_flags = os.path.splitext(input_path)[0] + "_linker_flags.json"
                 so_build_options.save_flags_to_file(linker_flags)
-
                 from torch._inductor.package import package_aoti
 
                 if use_mmap_weights:
@@ -1969,37 +1993,45 @@ def _pad_to_alignment(raw_bytes: bytes) -> bytes:
 
                 archive_path = package_aoti(os.path.split(input_path)[0])
                 return archive_path
+            else:
+                if config.is_fbcode():
+                    # TODO: enable AotCodeCompiler fb_code, and remove deprecated_cpp_compile_command.
+                    link_cmd_old = deprecated_cpp_compile_command(
+                        input=[output_o, consts_o],
+                        output=output_so,
+                        vec_isa=picked_vec_isa,
+                        cuda=cuda,
+                        aot_mode=graph.aot_mode,
+                        use_absolute_path=use_absolute_path,
+                    )
+                    # TODO: Enable below code to debug in fb_code.
+                    """
+                    _temp_validate_new_and_old_command(
+                        link_cmd.split(" "), link_cmd_old.split(" ")
+                    )
+                    """
+                    link_cmd = link_cmd_old
 
-            # TODO: replace this with using the CppBuilder above
-            link_cmd = cpp_compile_command(
-                input=[output_o, consts_o],
-                output=output_so,
-                vec_isa=picked_vec_isa,
-                cuda=cuda,
-                aot_mode=graph.aot_mode,
-                use_absolute_path=use_absolute_path,
-            )
+                log.debug("aot linkage command: %s", link_cmd)
+                if fbcode_aot_cpu_re:
+                    compile_file([output_o, consts_o], output_so, link_cmd.split())
+                    os.chmod(output_so, 0o755)
+                else:
+                    run_command_and_check(link_cmd)
 
-            log.debug("aot linkage command: %s", link_cmd)
-            if fbcode_aot_cpu_re:
-                compile_file([output_o, consts_o], output_so, link_cmd.split())
-                os.chmod(output_so, 0o755)
-            else:
-                run_command_and_check(link_cmd)
-
-            if use_mmap_weights:
-                with open(output_so, "a+b") as f_so:
-                    so_size = f_so.tell()
-                    # Page align the weights
-                    f_so.write(b" " * (16384 - so_size % 16384))
-                    f_so.write(serialized_weights)
-                    f_so.write(struct.pack("q", magic_number))
-
-            # Append cmds to the end of codegen-ed wrapper file
-            with open(input_path, "a") as f:
-                f.write("\n")
-                f.write(f"// Compile cmd\n// {compile_cmd}\n")
-                f.write(f"// Link cmd\n// {link_cmd}\n")
+                if use_mmap_weights:
+                    with open(output_so, "a+b") as f_so:
+                        so_size = f_so.tell()
+                        # Page align the weights
+                        f_so.write(b" " * (16384 - so_size % 16384))
+                        f_so.write(serialized_weights)
+                        f_so.write(struct.pack("q", magic_number))
+
+                # Append cmds to the end of codegen-ed wrapper file
+                with open(input_path, "a") as f:
+                    f.write("\n")
+                    f.write(f"// Compile cmd\n// {compile_cmd}\n")
+                    f.write(f"// Link cmd\n// {link_cmd}\n")
 
         return output_so
 
@@ -2499,120 +2531,6 @@ class CppWrapperCodeCache(CppPythonBindingsCodeCache):
     )
 
 
-# TODO: Will remove the temp code after switch to new cpp_builder
-def _temp_validate_new_and_old_command(new_cmd: List[str], old_cmd: List[str]) -> None:
-    new_diff: List[str] = [x for x in new_cmd if x not in old_cmd]
-    old_diff: List[str] = [y for y in old_cmd if y not in new_cmd]
-
-    if new_diff or old_diff:
-        print("!!! new_cmd: ", new_cmd)
-        print("!!! old_cmd: ", old_cmd)
-        print("!!! new_diff: ", new_diff)
-        print("!!! old_diff: ", old_diff)
-        raise RuntimeError("Error in new and old command different.")
-
-
-def _do_validate_cpp_commands(
-    include_pytorch: bool,
-    cuda: bool,
-    compile_only: bool,
-    mmap_weights: bool,
-    use_absolute_path: bool,
-    aot_mode: bool,
-) -> None:
-    # PreCI will failed if test machine can't run cuda.
-    temp_dir = tempfile.TemporaryDirectory()
-    test_dir_path = temp_dir.name
-    test_cuda = torch.cuda.is_available() and cuda
-    input_path = os.path.join(test_dir_path, "dummy_file.cpp")
-    output_path = os.path.join(test_dir_path, "dummy_file.so")
-    extra_flags = ["-D TEST_EXTRA_FLAGS"]
-    if compile_only:
-        output_path = os.path.join(test_dir_path, "dummy_file.o")
-    picked_isa = pick_vec_isa()
-
-    # Simulate fb_code env:
-    if not (aot_mode and not use_absolute_path):
-        input_path = os.path.basename(input_path)
-        output_path = os.path.basename(output_path)
-
-    # Fix test_new_cpp_build_logical failed on MacOS
-    if sys.platform != "linux":
-        aot_mode = False
-
-    old_cmd = cpp_compile_command(
-        input=input_path,
-        output=output_path,
-        include_pytorch=include_pytorch,
-        vec_isa=picked_isa,
-        cuda=test_cuda,
-        aot_mode=aot_mode,
-        compile_only=compile_only,
-        use_absolute_path=use_absolute_path,
-        use_mmap_weights=mmap_weights,
-        extra_flags=extra_flags,
-    ).split(" ")
-
-    name, dir = get_name_and_dir_from_output_file_path(input_path)
-
-    dummy_build_option = CppTorchCudaOptions(
-        vec_isa=picked_isa,
-        include_pytorch=include_pytorch,
-        cuda=test_cuda,
-        aot_mode=aot_mode,
-        compile_only=compile_only,
-        use_absolute_path=use_absolute_path,
-        use_mmap_weights=mmap_weights,
-        extra_flags=extra_flags,
-    )
-
-    dummy_builder = CppBuilder(
-        name=name,
-        sources=input_path,
-        output_dir=dir,
-        BuildOption=dummy_build_option,
-    )
-    new_cmd = dummy_builder.get_command_line().split(" ")
-
-    _temp_validate_new_and_old_command(new_cmd, old_cmd)
-
-    temp_dir.cleanup()
-
-
-# TODO: Will remove the temp code after switch to new cpp_builder
-# It could help on sync new cpp_builder generate same command line as the old one.
-def validate_new_cpp_commands() -> None:
-    cuda = [True, False]
-    use_mmap_weights = [True, False]
-    compile_only = [True, False]
-    include_pytorch = [True, False]
-    use_absolute_path = [True, False]
-    aot_mode = [False, True]
-
-    # Try to pass it in fb_code.
-    if config.is_fbcode():
-        return
-
-    for x in cuda:
-        for y in use_mmap_weights:
-            for z in compile_only:
-                for m in include_pytorch:
-                    for n in use_absolute_path:
-                        for o in aot_mode:
-                            print(
-                                f"!!! cuda:{x}, use_mmap_weights:{y}, compile_only:{z}, include_pytorch:{m},"
-                                f" use_absolute_path:{n}, aot_mode:{o}"
-                            )
-                            _do_validate_cpp_commands(
-                                include_pytorch=m,
-                                cuda=x,
-                                mmap_weights=y,
-                                compile_only=z,
-                                use_absolute_path=n,
-                                aot_mode=o,
-                            )
-
-
 @clear_on_fresh_inductor_cache
 class HalideCodeCache(CppPythonBindingsCodeCache):
     cache: Dict[str, Callable[[], Union[ModuleType, CDLL]]] = {}