-
Notifications
You must be signed in to change notification settings - Fork 74k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Crash fix cudaMallocAsync usage of TF_CUDA_MALLOC_ASYNC_SUPPORTED_PREALLOC. #50962
Crash fix cudaMallocAsync usage of TF_CUDA_MALLOC_ASYNC_SUPPORTED_PREALLOC. #50962
Conversation
Please make the PR description a bit more descriptive (what was the bug?). CC @hanbinyoon |
I updated the description. |
1cb58ca
to
17d3b28
Compare
tensorflow/core/common_runtime/gpu/gpu_cudamallocasync_allocator.cc
Outdated
Show resolved
Hide resolved
17d3b28
to
e41fc39
Compare
e41fc39
to
f95e862
Compare
I rebased this PR since the dependent PR is merged. |
Any update? |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Only a minor comment.
The CI seems to have failed:
Is this related to this PR? |
Some of the refactoring caused crash when using TF_CUDA_MALLOC_ASYNC_SUPPORTED_PREALLOC.
This PR fix them and add a test to be sure it continue to work.
This PR is on top of #50961 as otherwise there would be code conflict.
Now, GpuCudaMallocAsyncAllocator doesn't create its own compute stream. But it use one stream set after construction when GpuCudaMallocAsyncAllocator::SetStream is called. So it is now impossible to preallocate all memory during the object construction.
So this PR postpone the memory preallocation to when the stream is available, in SetStream.
This PR also check that the stream isn't set when SetStream is called.
@sanjoy