Package ai.rapids.cudf

Class PinnedMemoryPool

java.lang.Object

ai.rapids.cudf.PinnedMemoryPool

All Implemented Interfaces:

AutoCloseable

public final class PinnedMemoryPool
extends Object
implements AutoCloseable

This is the JNI interface to a rmm::pool_memory_resource<rmm::pinned_host_memory_resource>.

Method Summary

Modifier and Type	Method	Description
static HostMemoryBuffer	allocate(long bytes)	Factory method to create a host buffer but preferably pointing to pinned memory.
static HostMemoryBuffer	allocate(long bytes, HostMemoryAllocator hostMemoryAllocator)	Factory method to create a host buffer but preferably pointing to pinned memory.
void	close()
static boolean	configureDefaultCudfPinnedPoolSize(long size)	Sets the size of the cuDF default pinned pool.
static long	getTotalPoolSizeBytes()	Get the number of bytes that the pinned memory pool was allocated with.
static void	initialize(long poolSize)	Initialize the pool.
static void	initialize(long poolSize, int gpuId)	Initialize the pool.
static void	initialize(long poolSize, int gpuId, boolean setCudfPinnedPoolMemoryResource)	Initialize the pool.
static boolean	isInitialized()	Check if the pool has been initialized or not.
static void	shutdown()	Shut down the RMM pool_memory_resource, nulling out our reference.
static HostMemoryBuffer	tryAllocate(long bytes)	Factory method to create a pinned host memory buffer.

Methods inherited from class java.lang.Object

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

Method Details

initialize

public static void initialize(long poolSize)

Initialize the pool.

Parameters:

poolSize - size of the pool to initialize.

initialize

public static void initialize(long poolSize,
 int gpuId)

Initialize the pool.

Parameters:

poolSize - size of the pool to initialize.

gpuId - gpu id to set to get memory pool from, -1 means to use default

initialize

public static void initialize(long poolSize,
 int gpuId,
 boolean setCudfPinnedPoolMemoryResource)

Initialize the pool.

Parameters:

poolSize - size of the pool to initialize.

gpuId - gpu id to set to get memory pool from, -1 means to use default

setCudfPinnedPoolMemoryResource - true if this pinned pool should be used by cuDF for pinned memory

isInitialized

public static boolean isInitialized()

Check if the pool has been initialized or not.

shutdown

public static void shutdown()

Shut down the RMM pool_memory_resource, nulling out our reference. Any allocation or free that is in flight will fail after this.

tryAllocate

public static HostMemoryBuffer tryAllocate(long bytes)

Factory method to create a pinned host memory buffer.

Parameters:

bytes - size in bytes to allocate

Returns:

newly created buffer or null if insufficient pinned memory

allocate

public static HostMemoryBuffer allocate(long bytes,
 HostMemoryAllocator hostMemoryAllocator)

Factory method to create a host buffer but preferably pointing to pinned memory. It is not guaranteed that the returned buffer will be pointer to pinned memory.

Parameters:

bytes - size in bytes to allocate

Returns:

newly created buffer

allocate

public static HostMemoryBuffer allocate(long bytes)

Factory method to create a host buffer but preferably pointing to pinned memory. It is not guaranteed that the returned buffer will be pointer to pinned memory.

Parameters:

bytes - size in bytes to allocate

Returns:

newly created buffer

getTotalPoolSizeBytes

public static long getTotalPoolSizeBytes()

Get the number of bytes that the pinned memory pool was allocated with.

close

public void close()

Specified by:

close in interface AutoCloseable

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.