tree 4858b1020d5e426d0c8f2ac48923e5fe6acce8b3
parent 5ca2294439457d7311c46db6e09f296986ac5037
author Christian Sigg <csigg@google.com> 1641379652 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> 1641380071 -0800

Make `tfrt_gpu.blas/dnn/solver.handle` only refer to context instead of stream.

Change the ops that previously only took a handle to also take a stream. In the kernel implementations, set the handle's stream before executing the actual work. It is the lowering's responsibility to serialize ops that operate on the same handle.

This is a preparation step to wrap the handle creation into a tfrt.once. Those calls should not take a stream, because that stream might have been created internally (instead of passed in) and would therefore not match previous executions.

We want to wrap handle creation into a tfrt.once because handles are quite expensive to create and destroy (~1ms/8.5MB for cuDNN, ~0.4ms/8.5MB for cuBLAS). With this change, we will be able to amortize the cycles and eventually (more work required) allocate the bytes before the main loop.

PiperOrigin-RevId: 419786266
Change-Id: I70a09879f2b228b88fdd64d9bd125082b8a68e51