Skip to content
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.

Sign up for GitHub

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

Make rmm::prefetch fault tolerant #1649

Open
wants to merge 1 commit into
base: branch-24.10
Choose a base branch
from
+19 −3
Open

Make rmm::prefetch fault tolerant #1649

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions include/rmm/mr/device/prefetch_resource_adaptor.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,11 @@ namespace rmm::mr {
/**
* @brief Resource that prefetches all memory allocations.
*
* Note that prefetching is a no-op if the upstream resource is not using
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This adaptor has more information than rmm::prefetch does. It could (in its ctor) detect that the upstream MR is not managed memory, or that the device doesn't support managed memory. Then it could skip the prefetching, saving a call to cudaMemPrefetchAsync on every alloc.

If it doesn't do this, then there is overhead on alloc on systems that don't support prefetching, and so this isn't truly a no-op and the docs are technically inaccurate.

I think it would be better to make it truly a no-op. What do you think @bdice ?

Copy link
Contributor Author

@bdice bdice Aug 20, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree 100% and want to do that! I did not implement this because I did not know how to check if the upstream MR is using managed memory. Do you have advice on how to do that?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do a test allocation, try to prefetch it, check for errors, deallocate it.

* managed memory or the device does not support managed memory. In those
* cases, this adapter does nothing. Therefore it is only recommended to use
* this adapter with a managed memory resource on a supported system.
*
* @tparam Upstream Type of the upstream resource used for
* allocation/deallocation.
*/
Expand Down
12 changes: 9 additions & 3 deletions include/rmm/prefetch.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ namespace rmm {
* @brief Prefetch memory to the specified device on the specified stream.
*
* This function is a no-op if the pointer is not to CUDA managed memory.
* This function is a no-op if the CUDA device does not support managed memory.
*
* @throw rmm::cuda_error if the prefetch fails.
*
Expand All @@ -48,9 +49,14 @@ void prefetch(void const* ptr,
rmm::cuda_stream_view stream)
{
auto result = cudaMemPrefetchAsync(ptr, size, device.value(), stream.value());
// InvalidValue error is raised when non-managed memory is passed to cudaMemPrefetchAsync
// We should treat this as a no-op
if (result != cudaErrorInvalidValue && result != cudaSuccess) { RMM_CUDA_TRY(result); }
// cudaErrorInvalidValue is raised when non-managed memory is passed to
// cudaMemPrefetchAsync. cudaErrorInvalidDevice is raised when the device
// does not support managed memory. We treat both cases as a no-op, and
// document that this function does nothing in those cases.
if (result != cudaSuccess && result != cudaErrorInvalidValue &&
result != cudaErrorInvalidDevice) {
RMM_CUDA_TRY(result);
}
}

/**
Expand Down
5 changes: 5 additions & 0 deletions python/rmm/rmm/_lib/memory_resource.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1058,6 +1058,11 @@ cdef class PrefetchResourceAdaptor(UpstreamResourceAdaptor):
"""
Memory resource that prefetches all allocations.

Note that prefetching is a no-op if the upstream resource is not using
managed memory or the device does not support managed memory. In those
cases, this adapter does nothing. Therefore it is only recommended to
use this adapter with a managed memory resource on a supported system.

Parameters
----------
upstream : DeviceMemoryResource
Expand Down
Loading