github/ROCm
1
1
Fork 0
mirror of https://github.com/ROCm/ROCm.git synced 2026-04-05 03:01:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

[FRONTEND] Add input dtypes to autotuning key (#2534)

Browse Source
This commit is contained in:
Adnan Akhundov
2023-10-24 05:29:30 +02:00
committed by GitHub
parent 5e6071254c
commit 50add54334
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/triton/runtime/autotuner.py
View File

@@ -96,7 +96,11 @@ class Autotuner(KernelInterface):
for name in self.arg_names:
if name in all_args:
_args.append(all_args[name])
key = tuple(_args[i] for i in self.key_idx)
key = [_args[i] for i in self.key_idx]
for arg in _args:
if hasattr(arg, "dtype"):
key.append(str(arg.dtype))
key = tuple(key)
if key not in self.cache:
# prune configs
pruned_configs = self.prune_configs(kwargs)