Changed condition to sysconfig call

Signed-off-by: Gregory Shimansky <[email protected]>

Author: gshimansky

python/triton/compiler/compiler.py

@@ -16,6 +16,7 @@
 import re
 import functools
 import os
+import sysconfig
 
 
 @dataclass
@@ -177,7 +178,7 @@ def triton_key():
 
     # backend
     libtriton_hash = hashlib.sha256()
-    ext = "so" if os.name != "nt" else "pyd"
+    ext = sysconfig.get_config_var("EXT_SUFFIX").split(".")[-1]
     with open(os.path.join(TRITON_PATH, "_C", "libtriton." + ext), "rb") as f:
         while True:
             chunk = f.read(1024**2)

third_party/intel/backend/driver.py

@@ -2,6 +2,7 @@
 import os
 import hashlib
 import shutil
+import sysconfig
 import tempfile
 from pathlib import Path
 from functools import cached_property
@@ -108,7 +109,7 @@ def include_dir(self) -> list[str]:
 def compile_module_from_src(src, name):
     key = hashlib.sha256(src.encode("utf-8")).hexdigest()
     cache = get_cache_manager(key)
-    file_name = name + (".pyd" if os.name == "nt" else ".so")
+    file_name = f"{name}.{sysconfig.get_config_var('EXT_SUFFIX').split('.')[-1]}"
     cache_path = cache.get_file(file_name)
     if cache_path is None:
         with tempfile.TemporaryDirectory() as tmpdir: