Stream synchronize before deallocating SAM

[rongou]

Description

While investigating cuml benchmarks, I found an issue with the current system_memory_resource that causes segfault. Roughly it's in code like this:

void foo(...) {
  rmm::device_uvector<T> tmp(bufferSize, stream);
  // launch cuda kernels making use of tmp
}

When the function returns, the device_uvector would go out of scope and get deleted, while the cuda kernel might still be in flight. With cudaFree, the CUDA runtime would perform implicit synchronization to make sure the kernel finishes before actually freeing the memory, but with SAM we don't have that guarantee, thus causing use-after-free errors.

This is a rather simple fix. In the future we may want to use CUDA events to make this less blocking.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

Closes #1656

[rongou]

This is similar to what we are trying to do in cuPy on the python side: cupy/cupy#8442

[harrism]

This is now a synchronous deallocation, as in the cuda::mr::memory_resource concept (but not the cuda::mr::async_memory_resource concept). As we continue refactoring toward those concepts, there will be two functions: deallocate_async which takes a stream, and deallocate which does not.

[rongou]

Hmm I think it's the other way around. The async version assumes everything is stream ordered, but since we are dealing with malloc/free which doesn't understand cuda streams, we have to synchronize here. The synchronous version would leave the synchronization to the caller, so we don't need to sync here.

[harrism]

I get what you are saying. I'm talking about the fact that deallocate_async is named "async" but will have to synchronize. So that should be documented for users.

[rongou]

I updated the doc, is there anything else you want me to do here?

[harrism]

Waiting on second approval.

[harrism]

Reminder, we ask that every PR be associated with an issue. The PR template says this.

[rongou]

Reminder, we ask that every PR be associated with an issue. The PR template says this.

Done.

[harrism]

Please mention in the docstring of this function that it synchronizes stream.

[rongou]

Done.

[harrism]

Speaking to Vivek, he advised that we not rely on cudaHostFree being synchronised either (though it is). So we should duplicate this PR for some other host memory MRs.

[wence-]

This makes sense to me, thanks.

[rongou]

Test failure seems to be unrelated, can't reproduce it on my machine.

[rongou]

Can this be merged? Thanks!

[harrism]

Yeah, I reran the failed test last night. It must have been a flaky GPU?

[harrism]

/merge