Zero size allocation in StackDeviceMemory (#5100)

When the remaining space is empty, the zero size allocation will
get the pointer `end_`. However, `cudafree` the pointer `end_` will
cause to crash.

Fix: #5078 #4770 #4412 #4340 #3646

Signed-off-by: shengjun.li <shengjun.li@zilliz.com>
pull/5104/head
shengjun.li 2021-04-30 12:21:02 +08:00 committed by GitHub
parent ad1c889824
commit b46ae44087
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 0 deletions

View File

@ -65,6 +65,10 @@ StackDeviceMemory::Stack::getSizeAvailable() const {
char*
StackDeviceMemory::Stack::getAlloc(size_t size,
cudaStream_t stream) {
if (size == 0) {
return nullptr;
}
if (size > (end_ - head_)) {
// Too large for our stack
DeviceScope s(device_);
@ -133,6 +137,10 @@ void
StackDeviceMemory::Stack::returnAlloc(char* p,
size_t size,
cudaStream_t stream) {
if (size == 0) {
return;
}
if (p < start_ || p >= end_) {
// This is not on our stack; it was a one-off allocation
DeviceScope s(device_);