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

[RUNTIME] Make entry point cache key depend on triton version hash (#765)

Browse Source
This commit is contained in:
Philippe Tillet
2022-10-11 13:24:30 -07:00
committed by GitHub
parent 09cc2d454b
commit af76c989eb
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/triton/compiler.py
View File

@@ -1197,10 +1197,10 @@ def _build(name, src, srcdir):
return so return so
def make_so_cache_key(signature, constants): def make_so_cache_key(version_hash, signature, constants):
# Get unique key for the compiled code # Get unique key for the compiled code
signature = {k: 'ptr' if v[0] == '*' else v for k, v in signature.items()} signature = {k: 'ptr' if v[0] == '*' else v for k, v in signature.items()}
key = f"{''.join(signature.values())}{constants}" key = f"{version_hash}-{''.join(signature.values())}{constants}"
key = hashlib.md5(key.encode("utf-8")).hexdigest() key = hashlib.md5(key.encode("utf-8")).hexdigest()
return key return key
@@ -1221,7 +1221,7 @@ def compile(fn, signature: str, device: int = -1, constants=dict(), num_warps: i
# cache manager # cache manager
name = fn.__name__ name = fn.__name__
# name of files that are cached # name of files that are cached
so_cache_key = make_so_cache_key(signature, constants) so_cache_key = make_so_cache_key(triton.runtime.jit.version_key(), signature, constants)
so_cache_manager = CacheManager(so_cache_key) so_cache_manager = CacheManager(so_cache_key)
so_name = f"{name}.so" so_name = f"{name}.so"
# retrieve stub from cache if it exists # retrieve stub from cache if it exists