Support the Equivalent of profiler.key_averages(group_by_stack_n=n)

[james47]

What is your question?

The pytorch profiler can aggregate runtimes by the operation and its traceback via key_averages(group_by_stack_n=n): https://docs.pytorch.org/tutorials/beginner/profiler.html#print-profiler-result

Identifying the top operators / kernels and their source user Python code is helpful to find the main bottlenecks and perform targeted optimizations. Does HTA support something equivalent, or could you shed some light on implementing one? Thanks!

I also found get_aten_op_kernels_and_delay(), but seems like it traces down from aten ops only instead of source user training code.

Code

No response

What have you tried?

No response

Environment

No response

[fengxizhou]

It's possible but requires writing some code. The basic idea is to construct a call stack graph and then traverse that graph to print the aggregation.