PR #11141: Add warmup iterations introduced in #9757
PR #11141: Add warmup iterations introduced in #9757

Imported from GitHub PR https://github.com/openxla/xla/pull/11141

#9757 has been rolled back and re-landed multiple times.
This commit includes just those parts that should not cause problems, shrinking the remaining diff.
Copybara import of the project:

--
6981be3191a94040c97e42a231166f7db07f8a5f by Olli Lupton <olupton@nvidia.com>:

Add warmup iterations introduced in #9757

That PR has been rolled back and relanded multiple times. This commit
includes just those parts that should not cause problems with Google
production workflows.

Merging this change closes #11141

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/11141 from olupton:add-warmup-iterations 6981be3191a94040c97e42a231166f7db07f8a5f
