PR #11172: GpuTimer: use delay kernel to improve accuracy
PR #11172: GpuTimer: use delay kernel to improve accuracy

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

Originally imported from GitHub PR https://github.com/openxla/xla/pull/9757.

This ensures that all the device operations to be timed are queued to
the relevant stream before any of them are executed, resulting in a more
accurate measurement. This is skipped if CUDA_LAUNCH_BLOCKING=1 or
unified addressing is not available.

Only launch the delay kernel if it has been explicitly flagged that a
warm-up run, which might have triggered lazy kernel loading, has been
executed.

Reproducing information from https://github.com/openxla/xla/pull/9757 below:

Quoting my comment in the code:
```c++
// When a timer is created it launches a delay kernel into the given stream and
// queues a start event immediately afterwards. This delay kernel blocks
// execution on the stream until GetElapsedDuration() is called, at which point
// an end event is queued and the delay kernel exits. This allows the device
// execution time of the tasks queued to the stream while the timer is active
// to be measured more accurately.
```
this should improve the accuracy of the measurements that are used to make
auto-tuning decisions, especially for small kernels.

With the example HLO:
```
  parameter_0 = bf16[128,12,128]{2,1,0} parameter(0)
  parameter_1 = bf16[12,128,128]{2,1,0} parameter(1)
  ROOT dot.26 = bf16[12,128,128]{2,1,0} dot(parameter_0, parameter_1),
    lhs_batch_dims={1}, lhs_contracting_dims={0}, rhs_batch_dims={0},
    rhs_contracting_dims={2}
```
a cuBLAS kernel was selected with a CUPTI/NSys-measured runtime of
3.9µs and before this change, GpuTimer measured a runtime of 17.5µs.
With this change, GpuTimer measures a runtime of 8.1µs.

With the example HLO:
```
  p0 = bf16[128,1536]{0,1} parameter(0)
  p1 = s8[1536,12288]{0,1} parameter(1)
  c = bf16[1536,12288]{0,1} convert(p1)
  ROOT d = bf16[128,12288]{1,0} dot(p0, c), lhs_contracting_dims={1},
    rhs_contracting_dims={0}
```
five cuDNN plans were auto-tuned, with CUPTI/Nsys-measured runtimes of
{34.9, 18.5, 30.9, 32.4, 31.5}µs. Before this change, GpuTimer measured
runtimes of {76.6, 40.0, 52.5, 55.7, 51.8}µs. With this change, GpuTimer
measures runtimes of {39.6, 23.0, 35.4, 35.1, 35.1}µs.

In summary, with this change, GpuTimer gives results that are much
closer to the CUPTI/Nsys measurements, with a ~uniform offset of ~4µs. A
constant offset doesn't matter for auto-tuning.

There are a couple of edge cases that have special treatment:
- if `CUDA_LAUNCH_BLOCKING=1` then the delay kernel will not achieve anything,
  so we don't launch it
- if any code in the timed region synchronises the device then the host will
  wait for the delay kernel to time out before actually launching the kernels
  to be measured. This will lead to a poor quality measurement, and an error
  message is printed. This condition can be met if one of the kernels being
  measured is lazily loaded, as lazy loading can trigger synchronisation. Best
  practice is to execute a warmup run (without timing enabled) before the timed
  execution.

The `GpuTimer::Create*` signatures that were deprecated in https://github.com/openxla/xla/pull/9841 do not benefit
from these accuracy improvements.
Copybara import of the project:

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

GpuTimer: use delay kernel to improve accuracy

Originally imported from GitHub PR #9757.

This ensures that all the device operations to be timed are queued to
the relevant stream before any of them are executed, resulting in a more
accurate measurement. This is skipped if CUDA_LAUNCH_BLOCKING=1 or
unified addressing is not available.

Quoting my comment in the code:
// When a timer is created it launches a delay kernel into the given stream and
// queues a start event immediately afterwards. This delay kernel blocks
// execution on the stream until GetElapsedDuration() is called, at which point
// an end event is queued and the delay kernel exits. This allows the device
// execution time of the tasks queued to the stream while the timer is active
// to be measured more accurately.
this should improve the accuracy of the measurements that are used to make
auto-tuning decisions, especially for small kernels.

With the example HLO:
  parameter_0 = bf16[128,12,128]{2,1,0} parameter(0)
  parameter_1 = bf16[12,128,128]{2,1,0} parameter(1)
  ROOT dot.26 = bf16[12,128,128]{2,1,0} dot(parameter_0, parameter_1),
    lhs_batch_dims={1}, lhs_contracting_dims={0}, rhs_batch_dims={0},
    rhs_contracting_dims={2}
a cuBLAS kernel was selected with a CUPTI/NSys-measured runtime of
3.9µs and before this change, GpuTimer measured a runtime of 17.5µs.
With this change, GpuTimer measures a runtime of 8.1µs.

With the example HLO:
  p0 = bf16[128,1536]{0,1} parameter(0)
  p1 = s8[1536,12288]{0,1} parameter(1)
  c = bf16[1536,12288]{0,1} convert(p1)
  ROOT d = bf16[128,12288]{1,0} dot(p0, c), lhs_contracting_dims={1},
    rhs_contracting_dims={0}
five cuDNN plans were auto-tuned, with CUPTI/Nsys-measured runtimes of
{34.9, 18.5, 30.9, 32.4, 31.5}µs. Before this change, GpuTimer measured
runtimes of {76.6, 40.0, 52.5, 55.7, 51.8}µs. With this change, GpuTimer
measures runtimes of {39.6, 23.0, 35.4, 35.1, 35.1}µs.

In summary, with this change, GpuTimer gives results that are much
closer to the CUPTI/Nsys measurements, with a ~uniform offset of ~4µs. A
constant offset doesn't matter for auto-tuning.

There are a couple of edge cases that have special treatment:
- if `CUDA_LAUNCH_BLOCKING=1` then the delay kernel will not achieve anything,
  so we don't launch it
- if any code in the timed region synchronises the device then the host will
  wait for the delay kernel to time out before actually launching the kernels
  to be measured. This will lead to a poor quality measurement, and an error
  message is printed. This condition can be met if one of the kernels being
  measured is lazily loaded, as lazy loading can trigger synchronisation. Best
  practice is to execute a warmup run (without timing enabled) before the timed
  execution.

The `GpuTimer::Create*` signatures that were deprecated in #9841 do not benefit
from these accuracy improvements.

--
684275034a8bea55b3cbb9a1b40104feb1addad1 by Olli Lupton <olupton@nvidia.com>:

GpuTimer: use delay kernel on request only.

Only launch the delay kernel if it has been explicitly flagged that a
warm-up run, which might have triggered lazy kernel loading, has been
executed.

Merging this change closes #11172

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/11172 from olupton:re-land-delay-kernel 684275034a8bea55b3cbb9a1b40104feb1addad1
