Class Rmm
-
Nested Class Summary
Nested ClassesModifier and TypeClassDescriptionstatic class
What to send RMM alloc and free logs to. -
Constructor Summary
Constructors -
Method Summary
Modifier and TypeMethodDescriptionstatic DeviceMemoryBuffer
alloc
(long size) Allocate device memory and return a pointer to device memory, using stream 0.static DeviceMemoryBuffer
alloc
(long size, Cuda.Stream stream) Allocate device memory and return a pointer to device memory.static CudaMemoryBuffer
allocCuda
(long size, Cuda.Stream stream) Allocate device memory using `cudaMalloc` and return a pointer to device memory.static long
allocFromFallbackPinnedPool
(long size) static long
allocFromPinnedPool
(long poolPtr, long size) static void
static void
Clears the active RMM event handler if one is set.static boolean
configureDefaultCudfPinnedPoolSize
(long size) Sets the size of the cuDF default pinned pool.static void
freeFromFallbackPinnedPool
(long ptr, long size) static void
freeFromPinnedPool
(long poolPtr, long ptr, long size) static RmmDeviceMemoryResource
Get the RmmDeviceMemoryResource that was last set through the java APIs.static long
Returns the maximum amount of RMM memory (Bytes) outstanding during the lifetime of the process.static long
Get the most recently set pool size or -1 if RMM has not been initialized or pooling is not enabled.static long
Returns the maximum amount of RMM memory (Bytes) outstanding since the last `resetScopedMaximumOutstanding` call was issued (it is "scoped" because it's the maximum amount seen since the last reset).static long
Return the amount of RMM memory allocated in bytes.Get the currently set RmmTrackingResourceAdaptor that is set.static void
static void
initialize
(int allocationMode, Rmm.LogConf logConf, long poolSize) Initialize memory manager state and storage.static boolean
Check if RMM has been initialized already or not.static boolean
Return true if rmm is initialized and pooling has been enabled, else false.static Rmm.LogConf
Create a config that will write alloc/free logs to a file.static Rmm.LogConf
Create a config that will write alloc/free logs to stderr.static Rmm.LogConf
Create a config that will write alloc/free logs to stdout.static long
newPinnedPoolMemoryResource
(long initSize, long maxSize) static void
releasePinnedPoolMemoryResource
(long poolPtr) static void
Resets a scoped maximum counter of RMM memory used to keep track of usage between code sections while debugging.static void
resetScopedMaximumBytesAllocated
(long initialValue) Resets a scoped maximum counter of RMM memory used to keep track of usage between code sections while debugging.static long
setCudfPinnedPoolMemoryResource
(long poolPtr) static RmmDeviceMemoryResource
setCurrentDeviceResource
(RmmDeviceMemoryResource newResource, RmmDeviceMemoryResource expectedResource, boolean forceChange) Set the current device resource that RMM should use for all allocations and de-allocations.static void
setEventHandler
(RmmEventHandler handler) Sets the event handler to be called on RMM events (e.g.: allocation failure).static void
setEventHandler
(RmmEventHandler handler, boolean enableDebug) Sets the event handler to be called on RMM events (e.g.: allocation failure) and optionally enable debug mode (callbacks on every allocate and deallocate)static void
shutdown()
Shut down any initialized RMM instance.static void
Shut down any initialized RMM instance.
-
Constructor Details
-
Rmm
public Rmm()
-
-
Method Details
-
logTo
Create a config that will write alloc/free logs to a file. -
logToStdout
Create a config that will write alloc/free logs to stdout. -
logToStderr
Create a config that will write alloc/free logs to stderr. -
getCurrentDeviceResource
Get the RmmDeviceMemoryResource that was last set through the java APIs. This will not return the correct value if the resource was not set using the java APIs. It will return a null if the resource was never set through the java APIs. -
getTracker
Get the currently set RmmTrackingResourceAdaptor that is set. This might return null if RMM has nto been initialized. -
setCurrentDeviceResource
public static RmmDeviceMemoryResource setCurrentDeviceResource(RmmDeviceMemoryResource newResource, RmmDeviceMemoryResource expectedResource, boolean forceChange) Set the current device resource that RMM should use for all allocations and de-allocations. This should only be done if you feel comfortable that the current device resource has no pending allocations. Note that the caller of this is responsible for closing the current RmmDeviceMemoryResource that is returned by this. Assuming that it was not used to create the newResource. Please use the `shutdown` API to clear the resource as it does best effort clean up before shutting it down. If `newResource` is not null this will initialize the CUDA context for the calling thread if it is not already set. The caller is responsible for setting the desired CUDA device prior to this call if a specific device is already set.NOTE: All cudf methods will set the chosen CUDA device in the CUDA context of the calling thread after this returns and `newResource` was not null.
If `newResource` is null this will unset the default CUDA device and mark RMM as not initialized.
Be aware that for many of these APIs to work the RmmDeviceMemoryResource will need an `RmmTrackingResourceAdaptor`. If one is not found and `newResource` is not null it will be added to `newResource`.
Also be very careful with how you set this up. It is possible to set up an RmmDeviceMemoryResource that is just bad, like multiple pools or pools on top of an RmmAsyncMemoryResource, that does pooling already. Unless you know what you are doing it is best to just use the `initialize` API instead.
- Parameters:
newResource
- the new resource to set. If it is null an RmmCudaMemoryResource will be used, and RMM will be set as not initialized.expectedResource
- the resource that we expect to be set. This is to let us avoid race conditions with multiple things trying to set this at once. It should never happen, but just to be careful.forceChange
- if true then the expectedResource check is not done.
-
initialize
public static void initialize(int allocationMode, Rmm.LogConf logConf, long poolSize) throws RmmException Initialize memory manager state and storage. This will always initialize the CUDA context for the calling thread if it is not already set. The caller is responsible for setting the desired CUDA device prior to this call if a specific device is already set.NOTE: All cudf methods will set the chosen CUDA device in the CUDA context of the calling thread after this returns.
- Parameters:
allocationMode
- Allocation strategy to use. Bit set usingRmmAllocationMode.CUDA_DEFAULT
,RmmAllocationMode.POOL
,RmmAllocationMode.ARENA
,RmmAllocationMode.CUDA_ASYNC
,RmmAllocationMode.CUDA_ASYNC_FABRIC
andRmmAllocationMode.CUDA_MANAGED_MEMORY
logConf
- How to do logging or null if you don't want topoolSize
- The initial pool size in bytes- Throws:
IllegalStateException
- if RMM has already been initializedRmmException
-
configureDefaultCudfPinnedPoolSize
public static boolean configureDefaultCudfPinnedPoolSize(long size) Sets the size of the cuDF default pinned pool.- Parameters:
size
- initial and maximum size for the cuDF default pinned pool. Pass size=0 to disable the default pool.- Returns:
- true if we were able to setup the default resource, false if there was a resource already set.
-
getPoolSize
public static long getPoolSize()Get the most recently set pool size or -1 if RMM has not been initialized or pooling is not enabled. -
isPoolingEnabled
public static boolean isPoolingEnabled()Return true if rmm is initialized and pooling has been enabled, else false. -
isInitialized
Check if RMM has been initialized already or not.- Throws:
RmmException
-
getTotalBytesAllocated
public static long getTotalBytesAllocated()Return the amount of RMM memory allocated in bytes. Note that the result may be less than the actual amount of allocated memory if underlying RMM allocator decides to return more memory than what was requested. However, the result will always be a lower bound on the amount allocated. -
getMaximumTotalBytesAllocated
public static long getMaximumTotalBytesAllocated()Returns the maximum amount of RMM memory (Bytes) outstanding during the lifetime of the process. -
resetScopedMaximumBytesAllocated
public static void resetScopedMaximumBytesAllocated(long initialValue) Resets a scoped maximum counter of RMM memory used to keep track of usage between code sections while debugging.- Parameters:
initialValue
- an initial value (in Bytes) to use for this scoped counter
-
resetScopedMaximumBytesAllocated
public static void resetScopedMaximumBytesAllocated()Resets a scoped maximum counter of RMM memory used to keep track of usage between code sections while debugging. This resets the counter to 0 Bytes. -
getScopedMaximumBytesAllocated
public static long getScopedMaximumBytesAllocated()Returns the maximum amount of RMM memory (Bytes) outstanding since the last `resetScopedMaximumOutstanding` call was issued (it is "scoped" because it's the maximum amount seen since the last reset).If the memory used is net negative (for example if only frees happened since reset, and we reset to 0), then result will be 0.
If `resetScopedMaximumBytesAllocated` is never called, the scope is the whole program and is equivalent to `getMaximumTotalBytesAllocated`.
- Returns:
- the scoped maximum bytes allocated
-
setEventHandler
Sets the event handler to be called on RMM events (e.g.: allocation failure).- Parameters:
handler
- event handler to invoke on RMM events or null to clear an existing handler- Throws:
RmmException
- if an active handler is already set
-
setEventHandler
public static void setEventHandler(RmmEventHandler handler, boolean enableDebug) throws RmmException Sets the event handler to be called on RMM events (e.g.: allocation failure) and optionally enable debug mode (callbacks on every allocate and deallocate)NOTE: Only enable debug mode when necessary, as code will run much slower!
- Parameters:
handler
- event handler to invoke on RMM events or null to clear an existing handlerenableDebug
- if true enable debug callbacks in RmmEventHandler (onAllocated, onDeallocated)- Throws:
RmmException
- if an active handler is already set
-
clearEventHandler
Clears the active RMM event handler if one is set.- Throws:
RmmException
-
initDefaultCudaDevice
public static void initDefaultCudaDevice() -
cleanupDefaultCudaDevice
public static void cleanupDefaultCudaDevice() -
shutdown
Shut down any initialized RMM instance. This should be used very rarely. It does not need to be used when shutting down your process because CUDA will handle releasing all of the resources when your process exits. This really should only be used if you want to turn off the memory pool for some reasons. As such we make an effort to be sure no resources have been leaked before shutting down. This may involve forcing a JVM GC to collect any leaked java objects that still point to CUDA memory. By default this will do a gc every 2 seconds and wait for up to 4 seconds before throwing an RmmException if not all of the resources are freed.- Throws:
RmmException
- on any error. This includes if there are outstanding allocations that could not be collected.
-
shutdown
public static void shutdown(long forceGCInterval, long maxWaitTime, TimeUnit units) throws RmmException Shut down any initialized RMM instance. This should be used very rarely. It does not need to be used when shutting down your process because CUDA will handle releasing all of the resources when your process exits. This really should only be used if you want to turn off the memory pool for some reasons. As such we make an effort to be sure no resources have been leaked before shutting down. This may involve forcing a JVM GC to collect any leaked java objects that still point to CUDA memory.- Parameters:
forceGCInterval
- how frequently should we force a JVM GC. This is just a recommendation to the JVM to do a gc.maxWaitTime
- the maximum amount of time to wait for all objects to be collected before throwing an exception.units
- the units for forceGcInterval and maxWaitTime.- Throws:
RmmException
- on any error. This includes if there are outstanding allocations that could not be collected before maxWaitTime.
-
alloc
Allocate device memory and return a pointer to device memory, using stream 0.- Parameters:
size
- The size in bytes of the allocated memory region- Returns:
- Returned pointer to the allocated memory
-
alloc
Allocate device memory and return a pointer to device memory.- Parameters:
size
- The size in bytes of the allocated memory regionstream
- The stream in which to synchronize this command.- Returns:
- Returned pointer to the allocated memory
-
allocCuda
Allocate device memory using `cudaMalloc` and return a pointer to device memory.- Parameters:
size
- The size in bytes of the allocated memory regionstream
- The stream in which to synchronize this command.- Returns:
- Returned pointer to the allocated memory
-
newPinnedPoolMemoryResource
public static long newPinnedPoolMemoryResource(long initSize, long maxSize) -
setCudfPinnedPoolMemoryResource
public static long setCudfPinnedPoolMemoryResource(long poolPtr) -
releasePinnedPoolMemoryResource
public static void releasePinnedPoolMemoryResource(long poolPtr) -
allocFromPinnedPool
public static long allocFromPinnedPool(long poolPtr, long size) -
freeFromPinnedPool
public static void freeFromPinnedPool(long poolPtr, long ptr, long size) -
allocFromFallbackPinnedPool
public static long allocFromFallbackPinnedPool(long size) -
freeFromFallbackPinnedPool
public static void freeFromFallbackPinnedPool(long ptr, long size)
-