Support online build on Windows

Author: woct0rdho

python/setup.py

@@ -181,11 +181,13 @@ def get_json_package_info():
 def get_llvm_package_info():
     system = platform.system()
     try:
-        arch = {"x86_64": "x64", "arm64": "arm64", "aarch64": "arm64"}[platform.machine()]
+        arch = {"x86_64": "x64", "AMD64": "x64", "arm64": "arm64", "aarch64": "arm64"}[platform.machine()]
     except KeyError:
         arch = platform.machine()
     if system == "Darwin":
         system_suffix = f"macos-{arch}"
+    elif system == "Windows":
+        system_suffix = f"windows-{arch}"
     elif system == "Linux":
         if arch == 'arm64':
             system_suffix = 'ubuntu-arm64'
@@ -263,6 +265,17 @@ def update_symlink(link_path, source_path):
     link_path.symlink_to(source_path, target_is_directory=True)
 
 
+def download_and_extract_archive(url, extract_path):
+    with open_url(url) as response:
+        if url.endswith(".zip"):
+            file_bytes = BytesIO(response.read())
+            with zipfile.ZipFile(file_bytes, "r") as file:
+                file.extractall(path=extract_path)
+        else:
+            with tarfile.open(fileobj=response, mode="r|*") as file:
+                file.extractall(path=extract_path)
+
+
 def get_thirdparty_packages(packages: list):
     triton_cache_path = get_triton_cache_path()
     thirdparty_cmake_args = []
@@ -284,14 +297,7 @@ def get_thirdparty_packages(packages: list):
                 shutil.rmtree(package_root_dir)
             os.makedirs(package_root_dir, exist_ok=True)
             print(f'downloading and extracting {p.url} ...')
-            with open_url(p.url) as response:
-                if p.url.endswith(".zip"):
-                    file_bytes = BytesIO(response.read())
-                    with zipfile.ZipFile(file_bytes, "r") as file:
-                        file.extractall(path=package_root_dir)
-                else:
-                    with tarfile.open(fileobj=response, mode="r|*") as file:
-                        file.extractall(path=package_root_dir)
+            download_and_extract_archive(p.url, package_root_dir)
             # write version url to package_dir
             with open(os.path.join(package_dir, "version.txt"), "w") as f:
                 f.write(p.url)
@@ -316,9 +322,11 @@ def download_and_copy(name, src_func, dst_path, variable, version, url_func):
     system = platform.system()
     arch = platform.machine()
     # NOTE: This might be wrong for jetson if both grace chips and jetson chips return aarch64
-    arch = {"arm64": "sbsa", "aarch64": "sbsa"}.get(arch, arch)
-    supported = {"Linux": "linux", "Darwin": "linux"}
+    arch = {"AMD64": "x86_64", "arm64": "sbsa", "aarch64": "sbsa"}.get(arch, arch)
+    supported = {"Linux": "linux", "Darwin": "linux", "Windows": "windows"}
     url = url_func(supported[system], arch, version)
+    if system == "Windows":
+        url = url.replace(".tar.xz", ".zip")
     src_path = src_func(supported[system], arch, version)
     tmp_path = os.path.join(triton_cache_path, "nvidia", name)  # path to cache the download
     dst_path = os.path.join(base_dir, os.pardir, "third_party", "nvidia", "backend", dst_path)  # final binary path
@@ -331,8 +339,7 @@ def download_and_copy(name, src_func, dst_path, variable, version, url_func):
         download = download or curr_version.group(1) != version
     if download:
         print(f'downloading and extracting {url} ...')
-        file = tarfile.open(fileobj=open_url(url), mode="r|*")
-        file.extractall(path=tmp_path)
+        download_and_extract_archive(url, tmp_path)
     os.makedirs(os.path.split(dst_path)[0], exist_ok=True)
     print(f'copy {src_path} to {dst_path} ...')
     if os.path.isdir(src_path):