gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FRONTEND] Ensure version_key is called at most once (#519)

Browse Source 
Co-authored-by: hauntsaninja <>
This commit is contained in:
Shantanu
2022-05-23 13:40:08 -07:00
committed by GitHub
parent 205a493b10
commit 80f6a2698b
1 changed files with 33 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/triton/code_gen.py
View File

@@ -11,6 +11,7 @@ import subprocess
import sys import sys
import tempfile import tempfile
import textwrap import textwrap
import threading
import time import time
import warnings import warnings
from typing import Dict, Set, Tuple, Union from typing import Dict, Set, Tuple, Union
@@ -1058,8 +1059,20 @@ class Autotuner:
return self.kernel(*args, num_warps=config.num_warps, num_stages=config.num_stages, **kwargs, **config.kwargs) return self.kernel(*args, num_warps=config.num_warps, num_stages=config.num_stages, **kwargs, **config.kwargs)
@functools.lru_cache() _version_key_lock = threading.Lock()
_version_key = None
def version_key(): def version_key():
global _version_key
if _version_key is not None:
return _version_key
with _version_key_lock:
if _version_key is not None:
return _version_key
import pkgutil import pkgutil
contents = [] contents = []
# frontend # frontend
@@ -1078,7 +1091,8 @@ def version_key():
ptxas_version = hashlib.md5(subprocess.check_output(["ptxas", "--version"])).hexdigest() ptxas_version = hashlib.md5(subprocess.check_output(["ptxas", "--version"])).hexdigest()
except Exception: except Exception:
ptxas_version = '' ptxas_version = ''
return '-'.join(triton.__version__) + '-' + ptxas_version + '-' + '-'.join(contents) _version_key = '-'.join(triton.__version__) + '-' + ptxas_version + '-' + '-'.join(contents)
return _version_key
class DependenciesFinder(ast.NodeVisitor): class DependenciesFinder(ast.NodeVisitor):