triton-windows 3.2.0.post17__cp312-cp312-win_amd64.whl → 3.2.0.post18__cp312-cp312-win_amd64.whl

triton/backends/nvidia/compiler.py

@@ -372,7 +372,9 @@ class CUDABackend(BaseBackend):
                 ptxas, *line_info, *fmad, '-v', *opt_level, f'--gpu-name=sm_{capability}{suffix}', fsrc.name, '-o', fbin
             ]
             try:
-                subprocess.run(ptxas_cmd, check=True, close_fds=True, stdout=flog, stderr=flog)
+                # close_fds=True on Windows and False on Linux, see https://github.com/triton-lang/triton/pull/4357
+                # On Windows, both stdout and stderr need to be redirected to flog
+                subprocess.run(ptxas_cmd, check=True, close_fds=True if os.name == 'nt' else False, stdout=flog, stderr=flog)
             except subprocess.CalledProcessError as e:
                 with open(flog.name) as log_file:
                     log = log_file.read()

triton/backends/nvidia/driver.py

@@ -69,6 +69,10 @@ def compile_module_from_src(src, name):
             so = _build(name, src_path, tmpdir, library_dirs(), include_dir, libraries)
             with open(so, "rb") as f:
                 cache_path = cache.put(f.read(), so_name, binary=True)
+
+    # Loading module with relative path may cause error
+    cache_path = os.path.abspath(cache_path)
+
     import importlib.util
     spec = importlib.util.spec_from_file_location(name, cache_path)
     mod = importlib.util.module_from_spec(spec)

triton/runtime/build.py

@@ -1,4 +1,5 @@
 import contextlib
+import functools
 import sys
 import io
 import sysconfig
@@ -21,8 +22,12 @@ def quiet():
         sys.stdout, sys.stderr = old_stdout, old_stderr
 
 
+@functools.cache
 def get_cc():
     cc = os.environ.get("CC")
+    if cc is None:
+        # Find and check MSVC and Windows SDK from environment variables set by Launch-VsDevShell.ps1 or VsDevCmd.bat
+        cc, _, _ = find_msvc_winsdk(env_only=True)
     if cc is None:
         # Bundled TinyCC
         cc = os.path.join(sysconfig.get_paths()["platlib"], "triton", "runtime", "tcc", "tcc.exe")
@@ -89,13 +94,19 @@ def _build(name, src, srcdir, library_dirs, include_dirs, libraries):
     if "python3" not in libraries:
         libraries += ["python3"]
     if is_msvc(cc):
-        msvc_winsdk_inc_dirs, msvc_winsdk_lib_dirs = find_msvc_winsdk()
+        _, msvc_winsdk_inc_dirs, msvc_winsdk_lib_dirs = find_msvc_winsdk()
         include_dirs += msvc_winsdk_inc_dirs
         library_dirs += msvc_winsdk_lib_dirs
     cc_cmd = _cc_cmd(cc, src, so, include_dirs, library_dirs, libraries)
-    ret = subprocess.check_call(cc_cmd)
-    if ret == 0:
-        return so
+
+    try:
+        ret = subprocess.check_call(cc_cmd)
+        if ret == 0:
+            return so
+    except Exception as e:
+        print("Failed to compile. cc_cmd:", cc_cmd)
+        raise e
+
     # fallback on setuptools
     extra_compile_args = []
     if is_msvc(cc):

triton/tools/compile.c

@@ -64,4 +64,5 @@ CUresult {kernel_name}(CUstream stream, {signature}) {{
     // TODO: shared memory
     if(gX * gY * gZ > 0)
       return cuLaunchKernel({kernel_name}_func, gX, gY, gZ, {num_warps} * 32, 1, 1, {shared}, stream, args, NULL);
+    return (CUresult)NULL;
 }}

triton/windows_utils.py

@@ -57,12 +57,31 @@ def check_msvc(msvc_base_path: Path, version: str) -> bool:
     return all(
         x.exists()
         for x in [
+            msvc_base_path / version / "bin" / "Hostx64" / "x64" / "cl.exe",
             msvc_base_path / version / "include" / "vcruntime.h",
             msvc_base_path / version / "lib" / "x64" / "vcruntime.lib",
         ]
     )
 
 
+def find_msvc_env() -> tuple[Optional[Path], Optional[str]]:
+    msvc_base_path = os.getenv("VCINSTALLDIR")
+    if msvc_base_path is None:
+        return None, None
+    msvc_base_path = Path(msvc_base_path) / "Tools" / "MSVC"
+
+    version = os.getenv("VCToolsVersion")
+    if not check_msvc(msvc_base_path, version):
+        warnings.warn(
+            f"Environment variables VCINSTALLDIR = {os.getenv('VCINSTALLDIR')}, "
+            f"VCToolsVersion = {os.getenv('VCToolsVersion')} are set, "
+            "but this MSVC installation is incomplete."
+        )
+        return None, None
+
+    return msvc_base_path, version
+
+
 def find_msvc_vswhere() -> tuple[Optional[Path], Optional[str]]:
     vswhere_path = find_in_program_files(
         r"Microsoft Visual Studio\Installer\vswhere.exe"
@@ -144,20 +163,28 @@ def find_msvc_hardcoded() -> tuple[Optional[Path], Optional[str]]:
     return None, None
 
 
-def find_msvc() -> tuple[list[str], list[str]]:
-    msvc_base_path, version = find_msvc_vswhere()
-    if msvc_base_path is None:
-        msvc_base_path, version = find_msvc_envpath()
-    if msvc_base_path is None:
-        msvc_base_path, version = find_msvc_hardcoded()
-    if msvc_base_path is None:
+def find_msvc(env_only: bool) -> tuple[Optional[str], list[str], list[str]]:
+    if env_only:
+        fs = [find_msvc_env]
+    else:
+        fs = [
+            find_msvc_env,
+            find_msvc_vswhere,
+            find_msvc_envpath,
+            find_msvc_hardcoded,
+        ]
+    for f in fs:
+        msvc_base_path, version = f()
+        if msvc_base_path:
+            return (
+                str(msvc_base_path / version / "bin" / "Hostx64" / "x64" / "cl.exe"),
+                [str(msvc_base_path / version / "include")],
+                [str(msvc_base_path / version / "lib" / "x64")],
+            )
+
+    if not env_only:
         warnings.warn("Failed to find MSVC.")
-        return [], []
-
-    return (
-        [str(msvc_base_path / version / "include")],
-        [str(msvc_base_path / version / "lib" / "x64")],
-    )
+    return None, [], []
 
 
 def check_winsdk(winsdk_base_path: Path, version: str) -> bool:
@@ -170,6 +197,26 @@ def check_winsdk(winsdk_base_path: Path, version: str) -> bool:
     )
 
 
+def find_winsdk_env() -> tuple[Optional[Path], Optional[str]]:
+    winsdk_base_path = os.getenv("WindowsSdkDir")
+    if winsdk_base_path is None:
+        return None, None
+    winsdk_base_path = Path(winsdk_base_path)
+
+    version = os.getenv("WindowsSDKVersion")
+    if version:
+        version = version.rstrip("\\")
+    if not check_winsdk(winsdk_base_path, version):
+        warnings.warn(
+            f"Environment variables WindowsSdkDir = {os.getenv('WindowsSdkDir')}, "
+            f"WindowsSDKVersion = {os.getenv('WindowsSDKVersion')} are set, "
+            "but this Windows SDK installation is incomplete."
+        )
+        return None, None
+
+    return winsdk_base_path, version
+
+
 def find_winsdk_registry() -> tuple[Optional[Path], Optional[str]]:
     try:
         reg = winreg.ConnectRegistry(None, winreg.HKEY_LOCAL_MACHINE)
@@ -212,32 +259,46 @@ def find_winsdk_hardcoded() -> tuple[Optional[Path], Optional[str]]:
     return winsdk_base_path, version
 
 
-def find_winsdk() -> tuple[list[str], list[str]]:
-    winsdk_base_path, version = find_winsdk_registry()
-    if winsdk_base_path is None:
-        winsdk_base_path, version = find_winsdk_hardcoded()
-    if winsdk_base_path is None:
+def find_winsdk(env_only: bool) -> tuple[list[str], list[str]]:
+    if env_only:
+        fs = [find_winsdk_env]
+    else:
+        fs = [
+            find_winsdk_env,
+            find_winsdk_registry,
+            find_winsdk_hardcoded,
+        ]
+    for f in fs:
+        winsdk_base_path, version = f()
+        if winsdk_base_path:
+            return (
+                [
+                    str(winsdk_base_path / "Include" / version / "shared"),
+                    str(winsdk_base_path / "Include" / version / "ucrt"),
+                    str(winsdk_base_path / "Include" / version / "um"),
+                ],
+                [
+                    str(winsdk_base_path / "Lib" / version / "ucrt" / "x64"),
+                    str(winsdk_base_path / "Lib" / version / "um" / "x64"),
+                ],
+            )
+
+    if not env_only:
         warnings.warn("Failed to find Windows SDK.")
-        return [], []
-
-    return (
-        [
-            str(winsdk_base_path / "Include" / version / "shared"),
-            str(winsdk_base_path / "Include" / version / "ucrt"),
-            str(winsdk_base_path / "Include" / version / "um"),
-        ],
-        [
-            str(winsdk_base_path / "Lib" / version / "ucrt" / "x64"),
-            str(winsdk_base_path / "Lib" / version / "um" / "x64"),
-        ],
-    )
+    return [], []
 
 
 @functools.cache
-def find_msvc_winsdk() -> tuple[list[str], list[str]]:
-    msvc_inc_dirs, msvc_lib_dirs = find_msvc()
-    winsdk_inc_dirs, winsdk_lib_dirs = find_winsdk()
-    return msvc_inc_dirs + winsdk_inc_dirs, msvc_lib_dirs + winsdk_lib_dirs
+def find_msvc_winsdk(
+    env_only: bool = False,
+) -> tuple[Optional[str], list[str], list[str]]:
+    msvc_bin_path, msvc_inc_dirs, msvc_lib_dirs = find_msvc(env_only)
+    winsdk_inc_dirs, winsdk_lib_dirs = find_winsdk(env_only)
+    return (
+        msvc_bin_path,
+        msvc_inc_dirs + winsdk_inc_dirs,
+        msvc_lib_dirs + winsdk_lib_dirs,
+    )
 
 
 @functools.cache

{triton_windows-3.2.0.post17.dist-info → triton_windows-3.2.0.post18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: triton-windows
-Version: 3.2.0.post17
+Version: 3.2.0.post18
 Summary: A language and compiler for custom Deep Learning operations
 Home-page: https://github.com/woct0rdho/triton-windows
 Author: Philippe Tillet, Dian Wu

{triton_windows-3.2.0.post17.dist-info → triton_windows-3.2.0.post18.dist-info}/RECORD

@@ -2,8 +2,8 @@ triton/__init__.py,sha256=buC26NIOMT6cwmoYkc5jk46jb3fploY6_bOZYKwzdKQ,1347
 triton/_internal_testing.py,sha256=4pzyc_36u05khUveZ1TWL3MQ-7EVMJ1C2c1dRd8BMsw,4269
 triton/errors.py,sha256=8WfnuRKLG578mgY6cBA3ECruVMf9ULEKFNgRcJ6IhWM,89
 triton/testing.py,sha256=fX3pn9bjC3Z-z5qzSKW56C_2WF8h3mHLy5RJqpZ-HsA,19382
-triton/windows_utils.py,sha256=aQMItmuZNXaki8zSB7HMvhy1RsXlmz1GcO-lORm1IIk,10852
-triton/_C/libtriton.pyd,sha256=G30VIWeW8rJzfziZGm8Wq8sMs4uVvfQtfGd5uwvWrlU,87254528
+triton/windows_utils.py,sha256=YUl-1QbLINQRaAAMNYPjLiTFZlVsCNi9mTFck5aemwk,12778
+triton/_C/libtriton.pyd,sha256=ZesLdcWMEpclb3Z5E1TkoFpCx94VAydMNNx5jP6qYPI,87260672
 triton/backends/__init__.py,sha256=opAo_vgEMt3tLO_bYFrYGksnIu0qohbmyuu_s3-rNAs,1595
 triton/backends/compiler.py,sha256=JZiiEbB9Wws3tjU6KXrydKtlOQI7Suk-mTYPlafa0Qk,11388
 triton/backends/driver.py,sha256=QX_6P1Go9ajdlHZi4Hv3nCtdHyDA6o8_lM3NMnlH1mk,1386
@@ -105,9 +105,9 @@ triton/backends/amd/include/roctracer/ext/prof_protocol.h,sha256=6FAcvVD-dNM7uul
 triton/backends/amd/lib/ockl.bc,sha256=wQKCzkKukIHbu0lyjKUYlhndc7S27xto6L54J0Bn-C0,246124
 triton/backends/amd/lib/ocml.bc,sha256=UPNTXW0gCXUNB-c6orSYwb-mz9_mjUc7zny_vfFza44,205964
 triton/backends/nvidia/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-triton/backends/nvidia/compiler.py,sha256=7-WtsBYdNn_RIHj6-IhkkDwh9fJvqkfjarj3Q757Z3s,16420
+triton/backends/nvidia/compiler.py,sha256=cKKtU92WsvwgU30m902d8OI7MnVGI_2RriZdi1ujNDQ,16649
 triton/backends/nvidia/driver.c,sha256=8m1JB-t52GN1LoAFNlpLmM9DHuBxNIHL2VHFLTyJp8g,18906
-triton/backends/nvidia/driver.py,sha256=daDEQB2ExjTZ2F0-k-iv5DZBRSvKlVnt-IuvsDZtZHs,17593
+triton/backends/nvidia/driver.py,sha256=B_d89ULS6F63CgtUA04TRUdpws3e6ZVFI3YbG5-j7S0,17696
 triton/backends/nvidia/bin/ptxas.exe,sha256=2m-Z7N6rSw_zUngI0Zza4eMzhqglEEQQ6sMkev8cAr8,19580416
 triton/backends/nvidia/include/cuda.h,sha256=6SN2FpBhqz9-wdKnWN4_R3q18-auO67gdcBp-cEiSpE,1072817
 triton/backends/nvidia/lib/libdevice.10.bc,sha256=XC-uN8huaMOjhgWpX1EtfRLV89uYYxC-R_VzBKpype4,473728
@@ -134,7 +134,7 @@ triton/language/extra/hip/__init__.py,sha256=ieSER4LeX9_0horChGUUVwpuKAprkuka8uG
 triton/language/extra/hip/libdevice.py,sha256=EVraUfeXzQmN3F5Lleg2mohVcbFWOWlLaAH1nkbqtV4,16841
 triton/runtime/__init__.py,sha256=mKL5cqIBDUw2WO80NRCh4s1G8KYaqgM59TTAbTkPPjQ,621
 triton/runtime/autotuner.py,sha256=BJe69v9MSMSzdkvYSUDrvXrAFeLZ1x6A-7aUmpz2Le0,17271
-triton/runtime/build.py,sha256=elg0nv2lfaNzj-kfDv73-hAGir477OVQa8-0z-oQMjk,4428
+triton/runtime/build.py,sha256=rZdId897p7aZmu8Jg_LhORkFgiD6wiE0JYiG6C7TdpU,4771
 triton/runtime/cache.py,sha256=xMLfnhhgGoG2H5u7zQTydao4H0nenvAy0Or7btKWLiM,10591
 triton/runtime/driver.py,sha256=VZ-883Xri71R72lHB6usIpLo3gGLbZJkAlLP3ewWSpc,1509
 triton/runtime/errors.py,sha256=oj73dn34qJbLhOjakakAuZPSv-laZyIYylJiJwREA8Y,787
@@ -232,13 +232,13 @@ triton/runtime/tcc/lib/python3.def,sha256=o8MijZkdWOnHNPU4fmJrT1qw3fWSlaegC969H_
 triton/runtime/tcc/lib/user32.def,sha256=EcYohyyDgmz9fLBoOR-vszLeJ2YkBUoNGvSnuXrkum0,10439
 triton/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 triton/tools/build_extern.py,sha256=jCr-2hu3nLGBIJhCGUQ1jAyzLttughjkiPGEwRFjLR0,13673
-triton/tools/compile.c,sha256=rjuAQ8b-2DTtbj29SgK1NxJI5BSU2P9ccp9wa5p8Iyc,2090
+triton/tools/compile.c,sha256=Kb62pyFgeabw_d8ciNV3O2dLN2_cb6KlAlrFaRVUVVc,2117
 triton/tools/compile.h,sha256=n9QKIFZTL4RSsiXtAxBP9XGSnxjyaevQQ9bBpwDsvAg,332
 triton/tools/compile.py,sha256=b3yNnVgoBk8WzOs87JrZPDIyasdSgAslOWmxse1J6yM,6761
 triton/tools/disasm.py,sha256=BBO4bALdLcWgWDLhQdYHLlTx3oo8g_d8maeE_Uu-FmU,5088
 triton/tools/experimental_descriptor.py,sha256=0Wqy96Cc6YLh9o0eTknW-Lfvha6lfRSfe8bswkcPHMs,1260
 triton/tools/link.py,sha256=u7qtfZRLriZkAMEGNvj8YF-k1cthmLL7BwHYqBgT63E,11871
-triton_windows-3.2.0.post17.dist-info/METADATA,sha256=eBQszKwOYSWNxTzpro6SuqUlVl2LOHDz6Qxd2P4EsxM,1514
-triton_windows-3.2.0.post17.dist-info/WHEEL,sha256=9_mdKl-_Yr--f6EOmQ6F9FLhZL1o0okl1kr9b_edmVI,101
-triton_windows-3.2.0.post17.dist-info/top_level.txt,sha256=iIyUoyO6Ld8jQHWY9plUyco-emBXKTqVmuAXJAGvmEk,211
-triton_windows-3.2.0.post17.dist-info/RECORD,,
+triton_windows-3.2.0.post18.dist-info/METADATA,sha256=lPIYRbAkCh0zKhcxyfDBKA654-t_DfpozxMbXpdH9ho,1514
+triton_windows-3.2.0.post18.dist-info/WHEEL,sha256=ovhA9_Ei_7ok2fAych90j-feDV4goiAxbO7REePtvw0,101
+triton_windows-3.2.0.post18.dist-info/top_level.txt,sha256=iIyUoyO6Ld8jQHWY9plUyco-emBXKTqVmuAXJAGvmEk,211
+triton_windows-3.2.0.post18.dist-info/RECORD,,

{triton_windows-3.2.0.post17.dist-info → triton_windows-3.2.0.post18.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (77.0.1)
+Generator: setuptools (78.1.0)
 Root-Is-Purelib: false
 Tag: cp312-cp312-win_amd64