[AUTOTUNER] adding simple report flag for autotuner runs

[bringlein]

For our application, we wanted to investigate when, how often, and for how long the triton autotuner is triggered (to asses the impact on the latency of the application).
Therefore, we implemented a simple report flag to the autotuner decorator:

@triton.autotune(
    configs=[...],  
    key=[...],
    report=True,  # new, optional
)
@triton.jit
def fused_add_rmsnorm_triton(...

If this flag is set to true, the autotuner will print the following statement, every time a new autotune run is triggered:

Autotuner for function JITFunction(__main__:fused_add_rmsnorm_triton) finished after 4.15s; best config selected: BLOCK_N_SIZE: 512, num_warps: 8, num_ctas: 1, num_stages: 3;

There are no prints if a cached configuration is used.

(We thought this could be also a helpful feature for others, therefore we created this PR directly and we can make/discuss requested changes here. However, if you think this should be discussed instead in an issue, let us know and sorry).

[Jokeren]

I'm OK with the change. Should it be a log or just print? Triton doesn't have a logging system though

cc @jlebar @ThomasRaoux

[jlebar]

I think this is a good idea!

[jlebar]

Rename the flag to print? "report" can be a noun or a verb, and that ambiguity is confusing here. (Do I pass the report in as a parameter? What is the report? Or is the autotuning report passed as an outparameter somehow?)

[bringlein]

ok, changed it to print_autotune_stats

[jlebar]

Nit: Reword to something like what you wrote in the commit message, which is more helpful. For example: "If print is true, Triton will print a log message each time it autotunes a function."

[bringlein]

Fair point, I updated it.

[ThomasRaoux]

I wonder why this can't be done by using a profiler?

[Jokeren]

I wonder why this can't be done by using a profiler?

Indeed this is a case I used the profiler a lot...as proton can rename kernels based on constants. The problem is probably because proton is not available yet...

[jlebar]

Personally I think something lightweight like this is nice to have even if we have heavier-weight like Proton. It's basically zero complexity overhead and can be really useful for quick-and-dirty debugging.

[Jokeren]

Personally I think something lightweight like this is nice to have even if we have heavier-weight like Proton. It's basically zero complexity overhead and can be really useful for quick-and-dirty debugging.

I think it also depends on whether you want to check tuning time on CPU + GPU or just GPU time. This PR seems to get end to end tuning time.

[Jokeren]

Can we use self.bench_time?

[bringlein]

Good point, the hook won't make such a difference. I changed it

[bringlein]

Thanks for your helpful comments!

I think it also depends on whether you want to check tuning time on CPU + GPU or just GPU time. This PR seems to get end to end tuning time.

Exactly, we were/are interested in the end-to-end time of the autotuner, so that we could tell easily if a variance in latency of our application was caused by the triton autotuner or smth else.

[bringlein]

Should it be a log or just print? Triton doesn't have a logging system though

I was also thinking if besides True/False maybe a streaming object could be passed to this argument, so to route where the log message should be printed to. But I guess this would assume more about an existing/coming logging infrastructure than just a simple print.

[Jokeren]

Thanks for your helpful comments!

I think it also depends on whether you want to check tuning time on CPU + GPU or just GPU time. This PR seems to get end to end tuning time.

Exactly, we were/are interested in the end-to-end time of the autotuner, so that we could tell easily if a variance in latency of our application was caused by the triton autotuner or smth else.

Hi @ThomasRaoux , since they are interested in end-to-end statistics. It might be fine to print some debugging information? What's your thought?

[ptillet]

Yeah I think this can be helpful to some people and really doesn't add much complexity since best_config is already here

[ThomasRaoux]

Having debug logs makes sense, do we want this to be a front end option or an env variable? In general debug features are controlled by env variables.

[jlebar]

I'm fine with either an envvar or the in-code flag. It sound like we'd have consensus if we went with the env var? I propose TRITON_PRINT_AUTOTUNING as a strawperson.

[bringlein]

I agree, controlling the logging via an environment variable serves the purpose better.
I changed the implementation to check for TRITON_PRINT_AUTOTUNING=1.

[jlebar]

Thanks!

[jlebar]

I can merge this once we make the final few remaining changes!

[jlebar]

Rebased and am trying to merge this.

[bringlein]

Thanks @jlebar! I was on a trip and didn't had the time to fix/react to your comments.