ai.rapids.cudf

Class Rmm

java.lang.Object

ai.rapids.cudf.Rmm

public class Rmm
extends Object

This is the binding class for rmm lib.

Nested Class Summary

Nested Classes

Modifier and Type	Class and Description
static class	Rmm.LogConf What to send RMM alloc and free logs to.

Constructor Summary

Constructors

Constructor and Description
Rmm()

Method Summary

Modifier and Type	Method and Description
static 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 void	clearEventHandler() Clears the active RMM event handler if one is set.
static long	getTotalBytesAllocated() Return the amount of RMM memory allocated in bytes.
static void	initialize(int allocationMode, boolean enableLogging, long poolSize) Initialize memory manager state and storage.
static void	initialize(int allocationMode, boolean enableLogging, long poolSize, long maxPoolSize) Initialize memory manager state and storage.
static void	initialize(int allocationMode, Rmm.LogConf logConf, long poolSize) Initialize memory manager state and storage.
static void	initialize(int allocationMode, Rmm.LogConf logConf, long poolSize, long maxPoolSize) Initialize memory manager state and storage.
static void	initialize(int allocationMode, Rmm.LogConf logConf, long poolSize, long maxPoolSize, long allocationAlignment, long alignmentThreshold) Initialize memory manager state and storage.
static boolean	isInitialized() Check if RMM has been initialized already or not.
static Rmm.LogConf	logTo(File location) Create a config that will write alloc/free logs to a file.
static Rmm.LogConf	logToStderr() Create a config that will write alloc/free logs to stderr.
static Rmm.LogConf	logToStdout() Create a config that will write alloc/free logs to stdout.
static void	setEventHandler(RmmEventHandler handler) Sets the event handler to be called on RMM events (e.g.: allocation failure).
static void	shutdown() Shut down any initialized RMM instance.
static void	shutdown(long forceGCInterval, long maxWaitTime, TimeUnit units) Shut down any initialized RMM instance.

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Detail

Rmm

public Rmm()

Method Detail

logTo

public static Rmm.LogConf logTo(File location)

Create a config that will write alloc/free logs to a file.

logToStdout

public static Rmm.LogConf logToStdout()

Create a config that will write alloc/free logs to stdout.

logToStderr

public static Rmm.LogConf logToStderr()

Create a config that will write alloc/free logs to stderr.

initialize

public static void initialize(int allocationMode,
                              boolean enableLogging,
                              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 using RmmAllocationMode.CUDA_DEFAULT, RmmAllocationMode.POOL, RmmAllocationMode.ARENA and RmmAllocationMode.CUDA_MANAGED_MEMORY

enableLogging - Enable logging memory manager events

poolSize - The initial pool size in bytes

Throws:

IllegalStateException - if RMM has already been initialized

RmmException

initialize

public static void initialize(int allocationMode,
                              boolean enableLogging,
                              long poolSize,
                              long maxPoolSize)
                       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 using RmmAllocationMode.CUDA_DEFAULT, RmmAllocationMode.POOL, RmmAllocationMode.ARENA and RmmAllocationMode.CUDA_MANAGED_MEMORY

enableLogging - Enable logging memory manager events

poolSize - The initial pool size in bytes

maxPoolSize - The maximum size the pool is allowed to grow. If the specified value is <= 0 then the maximum pool size will not be artificially limited.

Throws:

IllegalStateException - if RMM has already been initialized

RmmException

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 using RmmAllocationMode.CUDA_DEFAULT, RmmAllocationMode.POOL, RmmAllocationMode.ARENA and RmmAllocationMode.CUDA_MANAGED_MEMORY

logConf - How to do logging or null if you don't want to

poolSize - The initial pool size in bytes

Throws:

IllegalStateException - if RMM has already been initialized

RmmException

initialize

public static void initialize(int allocationMode,
                              Rmm.LogConf logConf,
                              long poolSize,
                              long maxPoolSize)
                       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 using RmmAllocationMode.CUDA_DEFAULT, RmmAllocationMode.POOL, RmmAllocationMode.ARENA and RmmAllocationMode.CUDA_MANAGED_MEMORY

logConf - How to do logging or null if you don't want to

poolSize - The initial pool size in bytes

maxPoolSize - The maximum size the pool is allowed to grow. If the specified value is <= 0 then the pool size will not be artificially limited.

Throws:

IllegalStateException - if RMM has already been initialized

IllegalArgumentException - if a max pool size is specified but the allocation mode is not RmmAllocationMode.POOL or RmmAllocationMode.ARENA, or the maximum pool size is below the initial size.

RmmException

initialize

public static void initialize(int allocationMode,
                              Rmm.LogConf logConf,
                              long poolSize,
                              long maxPoolSize,
                              long allocationAlignment,
                              long alignmentThreshold)
                       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 using RmmAllocationMode.CUDA_DEFAULT, RmmAllocationMode.POOL, RmmAllocationMode.ARENA and RmmAllocationMode.CUDA_MANAGED_MEMORY

logConf - How to do logging or null if you don't want to

poolSize - The initial pool size in bytes

maxPoolSize - The maximum size the pool is allowed to grow. If the specified value is <= 0 then the pool size will not be artificially limited.

allocationAlignment - The size to which allocations are aligned.

alignmentThreshold - Only allocations with size larger than or equal to this threshold are aligned with `allocationAlignment`.

Throws:

IllegalStateException - if RMM has already been initialized

IllegalArgumentException - if a max pool size is specified but the allocation mode is not RmmAllocationMode.POOL or RmmAllocationMode.ARENA, or the maximum pool size is below the initial size.

RmmException

isInitialized

public static boolean isInitialized()
                             throws RmmException

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.

setEventHandler

public static void setEventHandler(RmmEventHandler handler)
                            throws RmmException

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

clearEventHandler

public static void clearEventHandler()
                              throws RmmException

Clears the active RMM event handler if one is set.

Throws:

RmmException

shutdown

public static void shutdown()
                     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. 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

public static DeviceMemoryBuffer alloc(long size)

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

public static DeviceMemoryBuffer alloc(long size,
                                       Cuda.Stream stream)

Allocate device memory and return a pointer to device memory.

Parameters:

size - The size in bytes of the allocated memory region

stream - The stream in which to synchronize this command.

Returns:

Returned pointer to the allocated memory