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

support CUDA async memory resource in JNI #9201

Merged
merged 5 commits into from
Sep 13, 2021
Merged

support CUDA async memory resource in JNI #9201

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0531784
support cuda async memory resource in jni
rongou Sep 9, 2021
445656f
Merge remote-tracking branch 'upstream/branch-21.10' into cuda-async-mr
rongou Sep 9, 2021
75ac47b
review feedback
rongou Sep 9, 2021
989a43b
add hard limit
rongou Sep 10, 2021
5d91944
Merge remote-tracking branch 'upstream/branch-21.10' into cuda-async-mr
rongou Sep 10, 2021
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
31 changes: 20 additions & 11 deletions java/src/main/java/ai/rapids/cudf/Rmm.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,8 @@ public static LogConf logToStderr() {
* @param allocationMode Allocation strategy to use. Bit set using
* {@link RmmAllocationMode#CUDA_DEFAULT},
* {@link RmmAllocationMode#POOL},
* {@link RmmAllocationMode#ARENA} and
* {@link RmmAllocationMode#ARENA},
* {@link RmmAllocationMode#CUDA_ASYNC} and
* {@link RmmAllocationMode#CUDA_MANAGED_MEMORY}
* @param enableLogging Enable logging memory manager events
* @param poolSize The initial pool size in bytes
Expand All @@ -106,7 +107,8 @@ public static void initialize(int allocationMode, boolean enableLogging, long po
* @param allocationMode Allocation strategy to use. Bit set using
* {@link RmmAllocationMode#CUDA_DEFAULT},
* {@link RmmAllocationMode#POOL},
* {@link RmmAllocationMode#ARENA} and
* {@link RmmAllocationMode#ARENA},
* {@link RmmAllocationMode#CUDA_ASYNC} and
* {@link RmmAllocationMode#CUDA_MANAGED_MEMORY}
* @param enableLogging Enable logging memory manager events
* @param poolSize The initial pool size in bytes
Expand Down Expand Up @@ -138,7 +140,8 @@ public static void initialize(int allocationMode, boolean enableLogging, long po
* @param allocationMode Allocation strategy to use. Bit set using
* {@link RmmAllocationMode#CUDA_DEFAULT},
* {@link RmmAllocationMode#POOL},
* {@link RmmAllocationMode#ARENA} and
* {@link RmmAllocationMode#ARENA},
* {@link RmmAllocationMode#CUDA_ASYNC} and
* {@link RmmAllocationMode#CUDA_MANAGED_MEMORY}
* @param logConf How to do logging or null if you don't want to
* @param poolSize The initial pool size in bytes
Expand All @@ -159,7 +162,8 @@ public static synchronized void initialize(int allocationMode, LogConf logConf,
* @param allocationMode Allocation strategy to use. Bit set using
* {@link RmmAllocationMode#CUDA_DEFAULT},
* {@link RmmAllocationMode#POOL},
* {@link RmmAllocationMode#ARENA} and
* {@link RmmAllocationMode#ARENA},
* {@link RmmAllocationMode#CUDA_ASYNC} and
* {@link RmmAllocationMode#CUDA_MANAGED_MEMORY}
* @param logConf How to do logging or null if you don't want to
* @param poolSize The initial pool size in bytes
Expand All @@ -168,8 +172,9 @@ public static synchronized void initialize(int allocationMode, LogConf logConf,
* @throws IllegalStateException if RMM has already been initialized
* @throws IllegalArgumentException if a max pool size is specified but the allocation mode
* is not {@link RmmAllocationMode#POOL} or
* {@link RmmAllocationMode#ARENA}, or the maximum pool size is
* below the initial size.
* {@link RmmAllocationMode#ARENA} or
* {@link RmmAllocationMode#CUDA_ASYNC}, or the maximum pool
* size is below the initial size.
*/
public static synchronized void initialize(int allocationMode, LogConf logConf, long poolSize,
long maxPoolSize) throws RmmException {
Expand All @@ -186,7 +191,8 @@ public static synchronized void initialize(int allocationMode, LogConf logConf,
* @param allocationMode Allocation strategy to use. Bit set using
* {@link RmmAllocationMode#CUDA_DEFAULT},
* {@link RmmAllocationMode#POOL},
* {@link RmmAllocationMode#ARENA} and
* {@link RmmAllocationMode#ARENA},
* {@link RmmAllocationMode#CUDA_ASYNC} and
* {@link RmmAllocationMode#CUDA_MANAGED_MEMORY}
* @param logConf How to do logging or null if you don't want to
* @param poolSize The initial pool size in bytes
Expand All @@ -198,18 +204,21 @@ public static synchronized void initialize(int allocationMode, LogConf logConf,
* @throws IllegalStateException if RMM has already been initialized
* @throws IllegalArgumentException if a max pool size is specified but the allocation mode
* is not {@link RmmAllocationMode#POOL} or
* {@link RmmAllocationMode#ARENA}, or the maximum pool size is
* below the initial size.
* {@link RmmAllocationMode#ARENA} or
* {@link RmmAllocationMode#CUDA_ASYNC}, or the maximum pool
* size is below the initial size.
*/
public static synchronized void initialize(int allocationMode, LogConf logConf, long poolSize,
long maxPoolSize, long allocationAlignment, long alignmentThreshold) throws RmmException {
if (initialized) {
throw new IllegalStateException("RMM is already initialized");
}
if (maxPoolSize > 0) {
if (allocationMode != RmmAllocationMode.POOL && allocationMode != RmmAllocationMode.ARENA) {
if (allocationMode != RmmAllocationMode.POOL &&
allocationMode != RmmAllocationMode.ARENA &&
allocationMode != RmmAllocationMode.CUDA_ASYNC) {
jlowe marked this conversation as resolved.
Show resolved Hide resolved
throw new IllegalArgumentException(
"Pool limit only supported in POOL or ARENA allocation mode");
"Pool limit only supported in POOL or ARENA or CUDA_ASYNC allocation mode");
}
if (maxPoolSize < poolSize) {
throw new IllegalArgumentException("Pool limit of " + maxPoolSize
Expand Down
4 changes: 4 additions & 0 deletions java/src/main/java/ai/rapids/cudf/RmmAllocationMode.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,4 +32,8 @@ public class RmmAllocationMode {
* Use arena suballocation strategy
*/
public static final int ARENA = 0x00000004;
/**
* Use CUDA async suballocation strategy
*/
public static final int CUDA_ASYNC = 0x00000008;
}
7 changes: 7 additions & 0 deletions java/src/main/native/src/RmmJni.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@

#include <rmm/mr/device/aligned_resource_adaptor.hpp>
#include <rmm/mr/device/arena_memory_resource.hpp>
#include <rmm/mr/device/cuda_async_memory_resource.hpp>
#include <rmm/mr/device/cuda_memory_resource.hpp>
#include <rmm/mr/device/logging_resource_adaptor.hpp>
#include <rmm/mr/device/managed_memory_resource.hpp>
Expand Down Expand Up @@ -344,6 +345,7 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_initializeInternal(
bool use_pool_alloc = allocation_mode & 1;
bool use_managed_mem = allocation_mode & 2;
bool use_arena_alloc = allocation_mode & 4;
bool use_cuda_async_alloc = allocation_mode & 8;
if (use_pool_alloc) {
auto pool_limit = (max_pool_size > 0) ?
thrust::optional<std::size_t>{static_cast<std::size_t>(max_pool_size)} :
Expand All @@ -365,6 +367,11 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_initializeInternal(
Initialized_resource = rmm::mr::make_owning_wrapper<rmm::mr::arena_memory_resource>(
std::make_shared<rmm::mr::cuda_memory_resource>(), pool_size, pool_limit);
}
} else if (use_cuda_async_alloc) {
auto pool_limit = (max_pool_size > 0) ? thrust::optional<std::size_t>{max_pool_size} :
thrust::optional<std::size_t>{};
Initialized_resource = std::make_shared<rmm::mr::cuda_async_memory_resource>(
thrust::optional<std::size_t>{pool_size}, pool_limit);
jlowe marked this conversation as resolved.
Show resolved Hide resolved
jlowe marked this conversation as resolved.
Show resolved Hide resolved
} else if (use_managed_mem) {
Initialized_resource = std::make_shared<rmm::mr::managed_memory_resource>();
} else {
Expand Down