[fix] rewrite do_bench with fixed warmup/rep runs mode

[ConvolutedDog]

PR Category

Benchmark

Type of Change

Bug Fix

Description

In the original implementation of triton's do_bench function, when users specify the warmup and rep parameters, the behavior may not align with their expectations. This is because:

• In triton.testing.do_bench, the warmup and rep parameters are in milliseconds instead of reps, so that triton dynamically adjusts the number of warmup and repetition runs. And this is inconsistent with torch's measurement.

Therefore, this PR rewrites the triton.testing.do_bench function to support a new fixed_warmup_rep_runs parameter, which allows users to choose between two modes:

1. Dynamic Mode (the original implementation of triton.testing.do_bench): Estimates the runtime of the kernel and dynamically adjusts the number of warmup and repetition runs based on the provided warmup and rep times (in milliseconds).
2. Fixed Mode (default in this rewritten version, and consistent with torch's testing): Uses the provided warmup and rep values directly as the number of warmup and repetition runs.

Issue

None

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

No need

[StrongSpoon]

thanks for your idea and contribution! I recommend reformat the code according to CONTRIBUTING.md to pass the checks. we will review it soon.

[ConvolutedDog]

@StrongSpoon Thank you for your feedback! I've reformatted the code according to CONTRIBUTING.md. Please let me know if there's anything else that needs to be adjusted.