experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[KERNEL_GEN] Cache cuda stream to avoid OOM.

Browse Source 
PiperOrigin-RevId: 336719808
Change-Id: I6422a64e897b75f4259eb4c0d6dad0f1f50affc9
This commit is contained in:
Alexander Belyaev 2020-10-12 12:32:21 -07:00 committed by TensorFlower Gardener
parent c32906c588
commit e2aa1ff751
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/compiler/mlir/tools/kernel_gen/tf_cuda_runtime_wrappers.cc
View File

@ -64,9 +64,14 @@ extern "C" void mgpuLaunchKernel(CUfunction function, intptr_t gridX,
} }
extern "C" CUstream mgpuStreamCreate() { extern "C" CUstream mgpuStreamCreate() {
static CUstream stream = []() {
// TODO(b/170649852): This is neither thread-safe nor handles
// creation/descruction of one stream per context.
CUstream stream = nullptr; CUstream stream = nullptr;
CUDA_REPORT_IF_ERROR(cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING)); CUDA_REPORT_IF_ERROR(cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING));
return stream; return stream;
}();
return stream;
} }
extern "C" void mgpuStreamSynchronize(CUstream stream) { extern "C" void mgpuStreamSynchronize(CUstream stream) {