From 7a8af31fcf91d9c434b701601d8f83f7205f6c70 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Mon, 4 Mar 2024 03:58:43 +0000 Subject: [PATCH 1/4] JNI rmm based pinned pool Signed-off-by: Alessandro Bellina --- .../java/ai/rapids/cudf/PinnedMemoryPool.java | 283 ++---------------- java/src/main/java/ai/rapids/cudf/Rmm.java | 10 +- java/src/main/native/src/RmmJni.cpp | 53 ++++ .../ai/rapids/cudf/PinnedMemoryPoolTest.java | 11 +- 4 files changed, 93 insertions(+), 264 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java index 9ce72ba237e..4f0288c9bfd 100644 --- a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java +++ b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java @@ -1,6 +1,6 @@ /* * - * Copyright (c) 2019-2023, NVIDIA CORPORATION. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -22,93 +22,30 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.Comparator; -import java.util.Iterator; import java.util.Objects; -import java.util.Optional; -import java.util.SortedSet; -import java.util.TreeSet; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; /** - * This provides a pool of pinned memory similar to what RMM does for device memory. + * This is the JNI interface to a rmm::pool_memory_resource. */ public final class PinnedMemoryPool implements AutoCloseable { private static final Logger log = LoggerFactory.getLogger(PinnedMemoryPool.class); - private static final long ALIGNMENT = ColumnView.hostPaddingSizeInBytes(); // These static fields should only ever be accessed when class-synchronized. // Do NOT use singleton_ directly! Use the getSingleton accessor instead. private static volatile PinnedMemoryPool singleton_ = null; private static Future initFuture = null; - - private final long totalPoolSize; - private final long pinnedPoolBase; - private final SortedSet freeHeap = new TreeSet<>(new SortedByAddress()); - private int numAllocatedSections = 0; - private long availableBytes; - - private static class SortedBySize implements Comparator { - @Override - public int compare(MemorySection s0, MemorySection s1) { - return Long.compare(s0.size, s1.size); - } - } - - private static class SortedByAddress implements Comparator { - @Override - public int compare(MemorySection s0, MemorySection s1) { - return Long.compare(s0.baseAddress, s1.baseAddress); - } - } - - private static class MemorySection { - private long baseAddress; - private long size; - - MemorySection(long baseAddress, long size) { - this.baseAddress = baseAddress; - this.size = size; - } - - boolean canCombine(MemorySection other) { - boolean ret = (other.baseAddress + other.size) == baseAddress || - (baseAddress + size) == other.baseAddress; - log.trace("CAN {} COMBINE WITH {} ? {}", this, other, ret); - return ret; - } - - void combineWith(MemorySection other) { - assert canCombine(other); - log.trace("COMBINING {} AND {}", this, other); - this.baseAddress = Math.min(baseAddress, other.baseAddress); - this.size = other.size + this.size; - log.trace("COMBINED TO {}\n", this); - } - - MemorySection splitOff(long newSize) { - assert this.size > newSize; - MemorySection ret = new MemorySection(baseAddress, newSize); - this.baseAddress += newSize; - this.size -= newSize; - return ret; - } - - @Override - public String toString() { - return "PINNED: " + size + " bytes (0x" + Long.toHexString(baseAddress) - + " to 0x" + Long.toHexString(baseAddress + size) + ")"; - } - } + private long poolHandle; + private long poolSize; private static final class PinnedHostBufferCleaner extends MemoryBuffer.MemoryBufferCleaner { - private MemorySection section; + private long address; private final long origLength; - PinnedHostBufferCleaner(MemorySection section, long length) { - this.section = section; + PinnedHostBufferCleaner(long address, long length) { + this.address = address; origLength = length; } @@ -116,15 +53,15 @@ private static final class PinnedHostBufferCleaner extends MemoryBuffer.MemoryBu protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) { boolean neededCleanup = false; long origAddress = 0; - if (section != null) { - origAddress = section.baseAddress; + if (address != -1) { + origAddress = address; try { - PinnedMemoryPool.freeInternal(section); + PinnedMemoryPool.freeInternal(address, origLength); } finally { // Always mark the resource as freed even if an exception is thrown. // We cannot know how far it progressed before the exception, and // therefore it is unsafe to retry. - section = null; + address = -1; } neededCleanup = true; } @@ -137,7 +74,7 @@ protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) { @Override public boolean isClean() { - return section == null; + return address == -1; } } @@ -161,16 +98,8 @@ private static PinnedMemoryPool getSingleton() { return singleton_; } - private static void freeInternal(MemorySection section) { - Objects.requireNonNull(getSingleton()).free(section); - } - - /** - * Used to indicate that memory was allocated from a reservation. This primarily is for - * keeping track of outstanding allocations. - */ - private static void reserveAllocInternal(MemorySection section) { - Objects.requireNonNull(getSingleton()).reserveAllocHappened(section); + private static void freeInternal(long address, long origLength) { + Objects.requireNonNull(getSingleton()).free(address, origLength); } /** @@ -209,12 +138,14 @@ public static boolean isInitialized() { } /** - * Shut down the pool of memory. If there are outstanding allocations this may fail. + * Shut down the RMM pool_memory_resource, nulling out our reference. Any allocation + * or free that is in flight will fail after this. */ public static synchronized void shutdown() { PinnedMemoryPool pool = getSingleton(); if (pool != null) { pool.close(); + pool = null; } initFuture = null; singleton_ = null; @@ -235,21 +166,6 @@ public static HostMemoryBuffer tryAllocate(long bytes) { return result; } - /** - * Factory method to create a pinned host memory reservation. - * - * @param bytes size in bytes to reserve - * @return newly created reservation or null if insufficient pinned memory to cover it. - */ - public static HostMemoryReservation tryReserve(long bytes) { - HostMemoryReservation result = null; - PinnedMemoryPool pool = getSingleton(); - if (pool != null) { - result = pool.tryReserveInternal(bytes); - } - return result; - } - /** * 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. @@ -276,26 +192,13 @@ public static HostMemoryBuffer allocate(long bytes) { return allocate(bytes, DefaultHostMemoryAllocator.get()); } - /** - * Get the number of bytes free in the pinned memory pool. - * - * @return amount of free memory in bytes or 0 if the pool is not initialized - */ - public static long getAvailableBytes() { - PinnedMemoryPool pool = getSingleton(); - if (pool != null) { - return pool.getAvailableBytesInternal(); - } - return 0; - } - /** * Get the number of bytes that the pinned memory pool was allocated with. */ public static long getTotalPoolSizeBytes() { PinnedMemoryPool pool = getSingleton(); if (pool != null) { - return pool.getTotalPoolSizeInternal(); + return pool.poolSize; } return 0; } @@ -306,157 +209,31 @@ private PinnedMemoryPool(long poolSize, int gpuId) { Cuda.setDevice(gpuId); Cuda.freeZero(); } - this.totalPoolSize = poolSize; - this.pinnedPoolBase = Cuda.hostAllocPinned(poolSize); - freeHeap.add(new MemorySection(pinnedPoolBase, poolSize)); - this.availableBytes = poolSize; + this.poolHandle = Rmm.newPinnedPoolMemoryResource(poolSize, poolSize); + this.poolSize = poolSize; } @Override public void close() { - assert numAllocatedSections == 0 : "Leaked " + numAllocatedSections + " pinned allocations"; - Cuda.freePinned(pinnedPoolBase); + Rmm.releasePinnedPoolMemoryResource(this.poolHandle); + this.poolHandle = -1; } /** - * Pads a length of bytes to the alignment the CPU wants in the worst case. This helps to - * calculate the size needed for a reservation if there are multiple buffers. - * @param bytes the size in bytes - * @return the new padded size in bytes. + * This makes an attempt to allocate pinned memory, and if the pinned memory allocation fails + * it will return null, instead of throw. */ - public static long padToCpuAlignment(long bytes) { - return ((bytes + ALIGNMENT - 1) / ALIGNMENT) * ALIGNMENT; - } - - private synchronized MemorySection tryGetInternal(long bytes, String what) { - if (freeHeap.isEmpty()) { - log.debug("No free pinned memory left"); - return null; - } - // Align the allocation - long alignedBytes = padToCpuAlignment(bytes); - Optional firstFit = freeHeap.stream() - .filter(section -> section.size >= alignedBytes) - .findFirst(); - if (!firstFit.isPresent()) { - if (log.isDebugEnabled()) { - MemorySection largest = freeHeap.stream() - .max(new SortedBySize()) - .orElse(new MemorySection(0, 0)); - log.debug("Insufficient pinned memory. {} needed, {} found", alignedBytes, largest.size); - } - return null; - } - MemorySection first = firstFit.get(); - log.debug("{} {}/{} bytes pinned from {} FREE COUNT {} OUTSTANDING COUNT {}", - what, bytes, alignedBytes, first, freeHeap.size(), numAllocatedSections); - freeHeap.remove(first); - MemorySection allocated; - if (first.size == alignedBytes) { - allocated = first; - } else { - allocated = first.splitOff(alignedBytes); - freeHeap.add(first); - } - numAllocatedSections++; - availableBytes -= allocated.size; - log.debug("{} {} free {} outstanding {}", what, allocated, freeHeap, numAllocatedSections); - return allocated; - } - private synchronized HostMemoryBuffer tryAllocateInternal(long bytes) { - MemorySection allocated = tryGetInternal(bytes, "allocate"); - if (allocated == null) { + long allocated = Rmm.allocFromPinnedPool(this.poolHandle, bytes); + if (allocated == -1) { return null; } else { - return new HostMemoryBuffer(allocated.baseAddress, bytes, + return new HostMemoryBuffer(allocated, bytes, new PinnedHostBufferCleaner(allocated, bytes)); } } - private class PinnedReservation implements HostMemoryReservation { - private MemorySection section = null; - - public PinnedReservation(MemorySection section) { - this.section = section; - } - - @Override - public synchronized HostMemoryBuffer allocate(long bytes, boolean preferPinned) { - return this.allocate(bytes); - } - - @Override - public synchronized HostMemoryBuffer allocate(long bytes) { - if (section == null || section.size < bytes) { - throw new OutOfMemoryError("Reservation didn't have enough space " + bytes + " / " + - (section == null ? 0 : section.size)); - } - long alignedSize = padToCpuAlignment(bytes); - MemorySection allocated; - if (section.size >= bytes && section.size <= alignedSize) { - allocated = section; - section = null; - // No need for reserveAllocInternal because the original section is already tracked - } else { - allocated = section.splitOff(alignedSize); - PinnedMemoryPool.reserveAllocInternal(allocated); - } - return new HostMemoryBuffer(allocated.baseAddress, bytes, - new PinnedHostBufferCleaner(allocated, bytes)); - } - - @Override - public synchronized void close() throws Exception { - if (section != null) { - try { - PinnedMemoryPool.freeInternal(section); - } finally { - // Always mark the resource as freed even if an exception is thrown. - // We cannot know how far it progressed before the exception, and - // therefore it is unsafe to retry. - section = null; - } - } - } - } - - private HostMemoryReservation tryReserveInternal(long bytes) { - MemorySection allocated = tryGetInternal(bytes, "allocate"); - if (allocated == null) { - return null; - } else { - return new PinnedReservation(allocated); - } - } - - private synchronized void free(MemorySection section) { - log.debug("Freeing {} with {} outstanding {}", section, freeHeap, numAllocatedSections); - availableBytes += section.size; - Iterator it = freeHeap.iterator(); - while(it.hasNext()) { - MemorySection current = it.next(); - if (section.canCombine(current)) { - it.remove(); - section.combineWith(current); - } - } - freeHeap.add(section); - numAllocatedSections--; - log.debug("After freeing {} outstanding {}", freeHeap, numAllocatedSections); - } - - private synchronized void reserveAllocHappened(MemorySection section) { - if (section != null && section.size > 0) { - numAllocatedSections++; - } - } - - private synchronized long getAvailableBytesInternal() { - return this.availableBytes; - } - - private long getTotalPoolSizeInternal() { - return this.totalPoolSize; + private synchronized void free(long address, long size) { + Rmm.freeFromPinnedPool(this.poolHandle, address, size); } -} +} \ No newline at end of file diff --git a/java/src/main/java/ai/rapids/cudf/Rmm.java b/java/src/main/java/ai/rapids/cudf/Rmm.java index 66c053f15b2..552da62382a 100755 --- a/java/src/main/java/ai/rapids/cudf/Rmm.java +++ b/java/src/main/java/ai/rapids/cudf/Rmm.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2023, NVIDIA CORPORATION. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -581,4 +581,12 @@ static native long newEventHandlerResourceAdaptor(long handle, long trackerHandl static native long releaseEventHandlerResourceAdaptor(long handle, boolean debug); private static native void setCurrentDeviceResourceInternal(long newHandle); + + public static native long newPinnedPoolMemoryResource(long initSize, long maxSize); + + public static native void releasePinnedPoolMemoryResource(long poolPtr); + + public static native long allocFromPinnedPool(long poolPtr, long size); + + public static native void freeFromPinnedPool(long poolPtr, long ptr, long size); } diff --git a/java/src/main/native/src/RmmJni.cpp b/java/src/main/native/src/RmmJni.cpp index 81b8241bab0..40fdf45d196 100644 --- a/java/src/main/native/src/RmmJni.cpp +++ b/java/src/main/native/src/RmmJni.cpp @@ -31,11 +31,13 @@ #include #include #include +#include #include "cudf_jni_apis.hpp" using rmm::mr::device_memory_resource; using rmm::mr::logging_resource_adaptor; +using rmm_pinned_pool_t = rmm::mr::pool_memory_resource; namespace { @@ -746,4 +748,55 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_setCurrentDeviceResourceInternal( } CATCH_STD(env, ) } + +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_newPinnedPoolMemoryResource(JNIEnv *env, + jclass clazz, + jlong init, + jlong max) { + try { + cudf::jni::auto_set_device(env); + auto pool = new rmm_pinned_pool_t(new rmm::mr::pinned_host_memory_resource(), init, max); + return reinterpret_cast(pool); + } + CATCH_STD(env, 0) +} + +JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_releasePinnedPoolMemoryResource(JNIEnv *env, + jclass clazz, + jlong pool_ptr) { + try { + cudf::jni::auto_set_device(env); + delete reinterpret_cast(pool_ptr); + } + CATCH_STD(env, ) +} + +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_allocFromPinnedPool(JNIEnv *env, + jclass clazz, + jlong pool_ptr, + jlong size) { + try { + cudf::jni::auto_set_device(env); + auto pool = reinterpret_cast(pool_ptr); + void *ret = pool->allocate(size); + return reinterpret_cast(ret); + } + catch (const std::exception& unused) { + return -1; + } +} + +JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_freeFromPinnedPool(JNIEnv *env, jclass clazz, + jlong pool_ptr, + jlong ptr, + jlong size) { + try { + cudf::jni::auto_set_device(env); + auto pool = reinterpret_cast(pool_ptr); + void *cptr = reinterpret_cast(ptr); + pool->deallocate(cptr, size); + } + CATCH_STD(env, ) +} + } diff --git a/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java b/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java index 16628d7be36..ef37e78dece 100644 --- a/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java +++ b/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java @@ -1,6 +1,6 @@ /* * - * Copyright (c) 2019, NVIDIA CORPORATION. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -85,30 +85,22 @@ void testFragmentationAndExhaustion() { try { buffers[0] = PinnedMemoryPool.tryAllocate(1024); assertNotNull(buffers[0]); - assertEquals(14*1024L, PinnedMemoryPool.getAvailableBytes()); buffers[1] = PinnedMemoryPool.tryAllocate(2048); assertNotNull(buffers[1]); - assertEquals(12*1024L, PinnedMemoryPool.getAvailableBytes()); buffers[2] = PinnedMemoryPool.tryAllocate(4096); assertNotNull(buffers[2]); - assertEquals(8*1024L, PinnedMemoryPool.getAvailableBytes()); buffers[1].close(); - assertEquals(10*1024L, PinnedMemoryPool.getAvailableBytes()); buffers[1] = null; buffers[1] = PinnedMemoryPool.tryAllocate(8192); assertNotNull(buffers[1]); - assertEquals(2*1024L, PinnedMemoryPool.getAvailableBytes()); buffers[3] = PinnedMemoryPool.tryAllocate(2048); assertNotNull(buffers[3]); - assertEquals(0L, PinnedMemoryPool.getAvailableBytes()); buffers[4] = PinnedMemoryPool.tryAllocate(64); assertNull(buffers[4]); buffers[0].close(); - assertEquals(1024L, PinnedMemoryPool.getAvailableBytes()); buffers[0] = null; buffers[4] = PinnedMemoryPool.tryAllocate(64); assertNotNull(buffers[4]); - assertEquals(1024L - 64, PinnedMemoryPool.getAvailableBytes()); } finally { for (HostMemoryBuffer buffer : buffers) { if (buffer != null) { @@ -116,7 +108,6 @@ void testFragmentationAndExhaustion() { } } } - assertEquals(poolSize, PinnedMemoryPool.getAvailableBytes()); } @Test From a5cadc0870ad4e5b0a6466d2f0ecf345f3905cc2 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Mon, 4 Mar 2024 16:13:53 +0000 Subject: [PATCH 2/4] Run linter --- java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java | 4 ++-- java/src/main/native/src/RmmJni.cpp | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java index 4f0288c9bfd..17f05a9baf6 100644 --- a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java +++ b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java @@ -138,7 +138,7 @@ public static boolean isInitialized() { } /** - * Shut down the RMM pool_memory_resource, nulling out our reference. Any allocation + * Shut down the RMM pool_memory_resource, nulling out our reference. Any allocation * or free that is in flight will fail after this. */ public static synchronized void shutdown() { @@ -236,4 +236,4 @@ private synchronized HostMemoryBuffer tryAllocateInternal(long bytes) { private synchronized void free(long address, long size) { Rmm.freeFromPinnedPool(this.poolHandle, address, size); } -} \ No newline at end of file +} diff --git a/java/src/main/native/src/RmmJni.cpp b/java/src/main/native/src/RmmJni.cpp index 40fdf45d196..03f62774bb0 100644 --- a/java/src/main/native/src/RmmJni.cpp +++ b/java/src/main/native/src/RmmJni.cpp @@ -751,7 +751,7 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_setCurrentDeviceResourceInternal( JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_newPinnedPoolMemoryResource(JNIEnv *env, jclass clazz, - jlong init, + jlong init, jlong max) { try { cudf::jni::auto_set_device(env); @@ -771,7 +771,7 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_releasePinnedPoolMemoryResource(J CATCH_STD(env, ) } -JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_allocFromPinnedPool(JNIEnv *env, +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_allocFromPinnedPool(JNIEnv *env, jclass clazz, jlong pool_ptr, jlong size) { From 86c73b628bfaab1276c049046cfb7463d6eaf3f8 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Mon, 4 Mar 2024 23:13:19 +0000 Subject: [PATCH 3/4] Touch pinned memory to improve tests, and bug fixes --- .../ai/rapids/cudf/HostMemoryBufferTest.java | 4 +-- .../ai/rapids/cudf/PinnedMemoryPoolTest.java | 26 ++++++++++++++++--- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/java/src/test/java/ai/rapids/cudf/HostMemoryBufferTest.java b/java/src/test/java/ai/rapids/cudf/HostMemoryBufferTest.java index e848d4a89bf..b7fde511c38 100644 --- a/java/src/test/java/ai/rapids/cudf/HostMemoryBufferTest.java +++ b/java/src/test/java/ai/rapids/cudf/HostMemoryBufferTest.java @@ -1,6 +1,6 @@ /* * - * Copyright (c) 2019-2020, NVIDIA CORPORATION. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -187,7 +187,7 @@ public void testFilemap() throws Exception { } public static void initPinnedPoolIfNeeded(long size) { - long available = PinnedMemoryPool.getAvailableBytes(); + long available = PinnedMemoryPool.getTotalPoolSizeBytes(); if (available < size) { if (PinnedMemoryPool.isInitialized()) { PinnedMemoryPool.shutdown(); diff --git a/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java b/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java index ef37e78dece..dfe947ac44a 100644 --- a/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java +++ b/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java @@ -18,6 +18,7 @@ package ai.rapids.cudf; +import java.nio.ByteBuffer; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.slf4j.Logger; @@ -80,7 +81,7 @@ void allocate() { void testFragmentationAndExhaustion() { final long poolSize = 15 * 1024L; PinnedMemoryPool.initialize(poolSize); - assertEquals(poolSize, PinnedMemoryPool.getAvailableBytes()); + assertEquals(poolSize, PinnedMemoryPool.getTotalPoolSizeBytes()); HostMemoryBuffer[] buffers = new HostMemoryBuffer[5]; try { buffers[0] = PinnedMemoryPool.tryAllocate(1024); @@ -110,16 +111,33 @@ void testFragmentationAndExhaustion() { } } + @Test + void testTouchPinnedMemory() { + final long poolSize = 15 * 1024L; + PinnedMemoryPool.initialize(poolSize); + int bufLength = 256; + try(HostMemoryBuffer hmb = PinnedMemoryPool.allocate(bufLength); + HostMemoryBuffer hmb2 = PinnedMemoryPool.allocate(bufLength)) { + ByteBuffer bb = hmb.asByteBuffer(0, bufLength); + for (int i = 0; i < bufLength; i++) { + bb.put(i, (byte)i); + } + hmb2.copyFromHostBuffer(0, hmb, 0, bufLength); + ByteBuffer bb2 = hmb2.asByteBuffer(0, bufLength); + for (int i = 0; i < bufLength; i++) { + assertEquals(bb.get(i), bb2.get(i)); + } + } + } + @Test void testZeroSizedAllocation() { final long poolSize = 4 * 1024L; PinnedMemoryPool.initialize(poolSize); - assertEquals(poolSize, PinnedMemoryPool.getAvailableBytes()); + assertEquals(poolSize, PinnedMemoryPool.getTotalPoolSizeBytes()); try (HostMemoryBuffer buffer = PinnedMemoryPool.tryAllocate(0)) { assertNotNull(buffer); assertEquals(0, buffer.getLength()); - assertEquals(poolSize, PinnedMemoryPool.getAvailableBytes()); } - assertEquals(poolSize, PinnedMemoryPool.getAvailableBytes()); } } From 1949d5916cc8ce34d80c4696303251b3ba9ee076 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Tue, 5 Mar 2024 14:31:43 +0000 Subject: [PATCH 4/4] fix code styles --- java/src/main/native/src/RmmJni.cpp | 24 +++++++------------ .../ai/rapids/cudf/PinnedMemoryPoolTest.java | 2 +- 2 files changed, 9 insertions(+), 17 deletions(-) diff --git a/java/src/main/native/src/RmmJni.cpp b/java/src/main/native/src/RmmJni.cpp index 03f62774bb0..7b81b5ff4de 100644 --- a/java/src/main/native/src/RmmJni.cpp +++ b/java/src/main/native/src/RmmJni.cpp @@ -751,8 +751,7 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_setCurrentDeviceResourceInternal( JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_newPinnedPoolMemoryResource(JNIEnv *env, jclass clazz, - jlong init, - jlong max) { + jlong init, jlong max) { try { cudf::jni::auto_set_device(env); auto pool = new rmm_pinned_pool_t(new rmm::mr::pinned_host_memory_resource(), init, max); @@ -766,37 +765,30 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_releasePinnedPoolMemoryResource(J jlong pool_ptr) { try { cudf::jni::auto_set_device(env); - delete reinterpret_cast(pool_ptr); + delete reinterpret_cast(pool_ptr); } CATCH_STD(env, ) } -JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_allocFromPinnedPool(JNIEnv *env, - jclass clazz, - jlong pool_ptr, - jlong size) { +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_allocFromPinnedPool(JNIEnv *env, jclass clazz, + jlong pool_ptr, jlong size) { try { cudf::jni::auto_set_device(env); - auto pool = reinterpret_cast(pool_ptr); + auto pool = reinterpret_cast(pool_ptr); void *ret = pool->allocate(size); return reinterpret_cast(ret); - } - catch (const std::exception& unused) { - return -1; - } + } catch (const std::exception &unused) { return -1; } } JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_freeFromPinnedPool(JNIEnv *env, jclass clazz, - jlong pool_ptr, - jlong ptr, + jlong pool_ptr, jlong ptr, jlong size) { try { cudf::jni::auto_set_device(env); - auto pool = reinterpret_cast(pool_ptr); + auto pool = reinterpret_cast(pool_ptr); void *cptr = reinterpret_cast(ptr); pool->deallocate(cptr, size); } CATCH_STD(env, ) } - } diff --git a/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java b/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java index dfe947ac44a..8c6e29dbd0c 100644 --- a/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java +++ b/java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java @@ -116,7 +116,7 @@ void testTouchPinnedMemory() { final long poolSize = 15 * 1024L; PinnedMemoryPool.initialize(poolSize); int bufLength = 256; - try(HostMemoryBuffer hmb = PinnedMemoryPool.allocate(bufLength); + try(HostMemoryBuffer hmb = PinnedMemoryPool.allocate(bufLength); HostMemoryBuffer hmb2 = PinnedMemoryPool.allocate(bufLength)) { ByteBuffer bb = hmb.asByteBuffer(0, bufLength); for (int i = 0; i < bufLength; i++) {