[Build] Support windows in setup.py

Fix windows build issues in setup.py

1. Build executable name and dynamic lib name based on platform.
2. Set the src directory based on build config for windows.
3. Avoid using os.geteuid() on windows.

For pytorch#4661

Author: python3kgae

setup.py

@@ -48,6 +48,7 @@
 
 import contextlib
 import os
+import platform
 import re
 import sys
 
@@ -162,6 +163,31 @@ def write_to_python_file(cls, path: str) -> None:
             fp.write("\n".join(lines) + "\n")
 
 
+# The build type is determined by the DEBUG environment variable. If DEBUG is
+# set to a non-empty value, the build type is Debug. Otherwise, the build type
+# is Release.
+def get_build_type(is_debug=None) -> str:
+    debug = int(os.environ.get("DEBUG", 0)) if is_debug is None else is_debug
+    cfg = "Debug" if debug else "Release"
+    return cfg
+
+
+def get_dynamic_lib_name(name: str) -> str:
+    if platform.system() == "Windows":
+        return name + ".dll"
+    elif platform.system() == "Darwin":
+        return "lib" + name + ".dylib"
+    else:
+        return "lib" + name + ".so"
+
+
+def get_executable_name(name: str) -> str:
+    if platform.system() == "Windows":
+        return name + ".exe"
+    else:
+        return name
+
+
 class _BaseExtension(Extension):
     """A base class that maps an abstract source to an abstract destination."""
 
@@ -189,9 +215,17 @@ def src_path(self, installer: "InstallerBuildExt") -> Path:
             installer: The InstallerBuildExt instance that is installing the
                 file.
         """
-        # TODO(dbort): share the cmake-out location with CustomBuild. Can get a
-        # handle with installer.get_finalized_command('build')
-        cmake_cache_dir: Path = Path().cwd() / installer.build_temp / "cmake-out"
+        # Share the cmake-out location with CustomBuild.
+        cmake_cache_dir = Path(installer.get_finalized_command("build").cmake_cache_dir)
+
+        cfg = get_build_type(installer.debug)
+
+        if os.name == "nt":
+            # Replace %BUILD_TYPE% with the current build type.
+            self.src = self.src.replace("%BUILD_TYPE%", cfg)
+        else:
+            # Remove %BUILD_TYPE% from the path.
+            self.src = self.src.replace("/%BUILD_TYPE%", "")
 
         # Construct the full source path, resolving globs. If there are no glob
         # pattern characters, this will just ensure that the source file exists.
@@ -212,17 +246,39 @@ class BuiltFile(_BaseExtension):
     `ext_modules`.
     """
 
-    def __init__(self, src: str, dst: str):
+    def __init__(
+        self,
+        src_dir: str,
+        src_name: str,
+        dst: str,
+        is_executable: bool = False,
+        is_dynamic_lib: bool = False,
+    ):
         """Initializes a BuiltFile.
 
         Args:
-            src: The path to the file to install, relative to the cmake-out
-                directory. May be an fnmatch-style glob that matches exactly one
-                file.
+            src_dir: The directory of the file to install, relative to the cmake-out
+                directory. A placeholder %BUILD_TYPE% will be replaced with the build
+                type for multi-config generators (like Visual Studio) where the build
+                output is in a subdirectory named after the build type. For single-
+                config generators (like Makefile Generators or Ninja), this placeholder
+                will be removed.
+            src_name: The name of the file to install
             dst: The path to install to, relative to the root of the pip
                 package. If dst ends in "/", it is treated as a directory.
                 Otherwise it is treated as a filename.
+            is_executable: If True, the file is an executable. This is used to
+                determine the destination filename for executable.
+            is_dynamic_lib: If True, the file is a dynamic library. This is used
+                to determine the destination filename for dynamic library.
         """
+        if is_executable and is_dynamic_lib:
+            raise ValueError("is_executable and is_dynamic_lib cannot be both True.")
+        if is_executable:
+            src_name = get_executable_name(src_name)
+        elif is_dynamic_lib:
+            src_name = get_dynamic_lib_name(src_name)
+        src = os.path.join(src_dir, src_name)
         # This is not a real extension, so use a unique name that doesn't look
         # like a module path. Some of setuptools's autodiscovery will look for
         # extension names with prefixes that match certain module paths.
@@ -397,7 +453,7 @@ def __init__(self):
         self.saved_env = {}
 
     def __enter__(self):
-        if os.geteuid() == 0 and "HOME" in os.environ:
+        if os.name != "nt" and os.geteuid() == 0 and "HOME" in os.environ:
             log.info("temporarily unsetting HOME while running as root")
             self.saved_env["HOME"] = os.environ.pop("HOME")
         return self
@@ -432,8 +488,7 @@ def initialize_options(self):
     def run(self):
         self.dump_options()
 
-        debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug
-        cfg = "Debug" if debug else "Release"
+        cfg = get_build_type(self.debug)
 
         # get_python_lib() typically returns the path to site-packages, where
         # all pip packages in the environment are installed.
@@ -508,6 +563,14 @@ def run(self):
                 item for item in os.environ["CMAKE_BUILD_ARGS"].split(" ") if item
             ]
 
+        # CMAKE_BUILD_TYPE variable specifies the build type (configuration) for
+        # single-configuration generators (e.g., Makefile Generators or Ninja).
+        # For multi-config generators (like Visual Studio), CMAKE_BUILD_TYPE
+        # isn’t directly applicable.
+        # During the build step, --config specifies the configuration to build
+        # for multi-config generators.
+        build_args += ["--config", cfg]
+
         # Put the cmake cache under the temp directory, like
         # "pip-out/temp.<plat>/cmake-out".
         cmake_cache_dir = os.path.join(repo_root, self.build_temp, "cmake-out")
@@ -545,18 +608,24 @@ def run(self):
             "build/pip_data_bin_init.py.in",
             os.path.join(bin_dir, "__init__.py"),
         )
+        # Share the cmake-out location with _BaseExtension.
+        self.cmake_cache_dir = cmake_cache_dir
 
         # Finally, run the underlying subcommands like build_py, build_ext.
         build.run(self)
 
 
 def get_ext_modules() -> List[Extension]:
     """Returns the set of extension modules to build."""
-
     ext_modules = []
     if ShouldBuild.flatc():
         ext_modules.append(
-            BuiltFile("third-party/flatbuffers/flatc", "executorch/data/bin/")
+            BuiltFile(
+                src_dir="third-party/flatbuffers/%BUILD_TYPE%/",
+                src_name="flatc",
+                dst="executorch/data/bin/",
+                is_executable=True,
+            )
         )
 
     if ShouldBuild.pybindings():
@@ -570,17 +639,20 @@ def get_ext_modules() -> List[Extension]:
         )
     if ShouldBuild.llama_custom_ops():
         ext_modules.append(
-            # Install the prebuilt library for custom ops used in llama.
             BuiltFile(
-                "extension/llm/custom_ops/libcustom_ops_aot_lib.*",
-                "executorch/extension/llm/custom_ops/",
+                src_dir="extension/llm/custom_ops/%BUILD_TYPE%/",
+                src_name="custom_ops_aot_lib",
+                dst="executorch/extension/llm/custom_ops",
+                is_dynamic_lib=True,
             )
         )
         ext_modules.append(
             # Install the prebuilt library for quantized ops required by custom ops.
             BuiltFile(
-                "kernels/quantized/libquantized_ops_aot_lib.*",
-                "executorch/kernels/quantized/",
+                src_dir="kernels/quantized/%BUILD_TYPE%/",
+                src_name="quantized_ops_aot_lib",
+                dst="executorch/kernels/quantized/",
+                is_dynamic_lib=True,
             )
         )