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

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

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