Skip to content

Commit

Permalink
[inductor] Log precompilation time (pytorch#136395)
Browse files Browse the repository at this point in the history
This has been useful for diagnosing the long compile time issues I've seen in the Triton CPU backend.

Pull Request resolved: pytorch#136395
Approved by: https://github.com/eellison
  • Loading branch information
int3 authored and pytorchmergebot committed Sep 24, 2024
1 parent 802ba79 commit e4d2942
Showing 1 changed file with 9 additions and 1 deletion.
10 changes: 9 additions & 1 deletion torch/_inductor/select_algorithm.py
Original file line number Diff line number Diff line change
Expand Up @@ -1283,7 +1283,9 @@ def no_op(*args, **kwargs):

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)

Expand All @@ -1305,6 +1307,12 @@ def wait_on_futures():
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)

Expand Down

0 comments on commit e4d2942

Please sign in to comment.