message: Adding GetGpuStream routine (same as GetCudaStream, but with the generic *Gpu* name)
diff: @@ -82,6 +82,19 @@ inline const char* GpuGetErrorString(hipError_t error) {
  }
  #endif
  +// Exact copy from GetCudaStream() in gpu_launch_config.h
  +// Returns a raw reference to the current cuda stream. Required by a
  +// number of kernel calls (for which StreamInterface* does not work),
  +// i.e. CUB and certain cublas primitives.
  +inline const gpuStream_t& GetGpuStream(OpKernelContext* context) {
  + const gpuStream_t* ptr = CHECK_NOTNULL(
  + reinterpret_cast<const gpuStream_t*>(context->op_device_context()
  + ->stream()
  + ->implementation()
  + ->GpuStreamMemberHack()));
  + return *ptr;
  +}
  +
  // Launches a GPU kernel through cudaLaunchKernel in CUDA environment, or
  // hipLaunchKernel in ROCm environment with the given arguments.
  //
  