saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-14 20:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[inductor] Log precompilation time (#136395)

Browse source 
This has been useful for diagnosing the long compile time issues I've seen in the Triton CPU backend.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/136395
Approved by: https://github.com/eellison
This commit is contained in:
Jez Ng 2024-09-21 18:35:05 -07:00 committed by PyTorch MergeBot
parent 802ba79121
commit e4d294221b
1 changed files with 9 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
torch/_inductor/select_algorithm.py
View file

@ -1283,7 +1283,9 @@ class AlgorithmSelectorCache(PersistentCache):
def precompile_with_captured_stdout(choice):
with restore_stdout_stderr(initial_stdout, initial_stderr):
return choice.precompile()
start_time = time.time()
choice.precompile()
return time.time() - start_time
executor = ThreadPoolExecutor(max_workers=num_workers)
@ -1305,6 +1307,12 @@ class AlgorithmSelectorCache(PersistentCache):
log.error(
"Exception %s for benchmark choice %s", e, futures[future]
)
else:
log.info(
"Precompiling benchmark choice %s took %.02fs",
futures[future],
future.result(),
)
executor.shutdown(wait=True)