Change the java API so a global default host allocator can be set.

java/src/main/java/ai/rapids/cudf/DefaultHostMemoryAllocator.java

@@ -1,6 +1,6 @@
 /*
  *
- *  Copyright (c) 2023, NVIDIA CORPORATION.
+ *  Copyright (c) 2023-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.
@@ -32,19 +32,25 @@ public static HostMemoryAllocator get() {
 
   /**
    * Sets a new default host memory allocator implementation by default.
-   * @param hostMemoryAllocator
+   * @param hostMemoryAllocator the allocator to use from here on
    */
   public static void set(HostMemoryAllocator hostMemoryAllocator) {
     instance = hostMemoryAllocator;
   }
 
   @Override
   public HostMemoryBuffer allocate(long bytes, boolean preferPinned) {
-    return HostMemoryBuffer.allocate(bytes, preferPinned);
+    if (preferPinned) {
+      HostMemoryBuffer pinnedBuffer = PinnedMemoryPool.tryAllocate(bytes);
+      if (pinnedBuffer != null) {
+        return pinnedBuffer;
+      }
+    }
+    return HostMemoryBuffer.allocateRaw(bytes);
   }
 
   @Override
   public HostMemoryBuffer allocate(long bytes) {
-    return HostMemoryBuffer.allocate(bytes);
+    return allocate(bytes, HostMemoryBuffer.defaultPreferPinned);
   }
 }

java/src/main/java/ai/rapids/cudf/HostMemoryBuffer.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.
@@ -41,7 +41,7 @@
  * Be aware that the off heap memory limits set by Java do not apply to these buffers.
  */
 public class HostMemoryBuffer extends MemoryBuffer {
-  private static final boolean defaultPreferPinned;
+  static final boolean defaultPreferPinned;
   private static final Logger log = LoggerFactory.getLogger(HostMemoryBuffer.class);
 
   static {
@@ -135,13 +135,7 @@ public boolean isClean() {
    * @return the newly created buffer
    */
   public static HostMemoryBuffer allocate(long bytes, boolean preferPinned) {
-    if (preferPinned) {
-      HostMemoryBuffer pinnedBuffer = PinnedMemoryPool.tryAllocate(bytes);
-      if (pinnedBuffer != null) {
-        return pinnedBuffer;
-      }
-    }
-    return new HostMemoryBuffer(UnsafeMemoryAccessor.allocate(bytes), bytes);
+    return DefaultHostMemoryAllocator.get().allocate(bytes, preferPinned);
   }
 
   /**
@@ -155,6 +149,10 @@ public static HostMemoryBuffer allocate(long bytes) {
     return allocate(bytes, defaultPreferPinned);
   }
 
+  public static HostMemoryBuffer allocateRaw(long bytes) {
+    return new HostMemoryBuffer(UnsafeMemoryAccessor.allocate(bytes), bytes);
+  }
+
   /**
    * Create a host buffer that is memory-mapped to a file.
    * @param path path to the file to map into host memory

java/src/main/java/ai/rapids/cudf/MultiBufferDataSource.java

@@ -1,6 +1,6 @@
 /*
  *
- *  Copyright (c) 2023, NVIDIA CORPORATION.
+ *  Copyright (c) 2023-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.
@@ -26,7 +26,6 @@ public class MultiBufferDataSource extends DataSource {
   private final long sizeInBytes;
   private final HostMemoryBuffer[] hostBuffers;
   private final long[] startOffsets;
-  private final HostMemoryAllocator allocator;
 
   // Metrics
   private long hostReads = 0;
@@ -39,15 +38,6 @@ public class MultiBufferDataSource extends DataSource {
    * @param buffers the buffers that will back the data source.
    */
   public MultiBufferDataSource(HostMemoryBuffer ... buffers) {
-    this(DefaultHostMemoryAllocator.get(), buffers);
-  }
-
-  /**
-   * Create a new data source backed by multiple buffers.
-   * @param allocator the allocator to use for host buffers, if needed.
-   * @param buffers the buffers that will back the data source.
-   */
-  public MultiBufferDataSource(HostMemoryAllocator allocator, HostMemoryBuffer ... buffers) {
     int numBuffers = buffers.length;
     hostBuffers = new HostMemoryBuffer[numBuffers];
     startOffsets = new long[numBuffers];
@@ -61,7 +51,6 @@ public MultiBufferDataSource(HostMemoryAllocator allocator, HostMemoryBuffer ...
       currentOffset += hmb.getLength();
     }
     sizeInBytes = currentOffset;
-    this.allocator = allocator;
   }
 
   @Override
@@ -159,7 +148,7 @@ public HostMemoryBuffer hostRead(long offset, long amount) {
     } else {
       // We will have to allocate a new buffer and copy data into it.
       boolean success = false;
-      HostMemoryBuffer ret = allocator.allocate(realAmount, true);
+      HostMemoryBuffer ret = HostMemoryBuffer.allocate(realAmount, true);
       try {
         long amountRead = read(offset, ret, HostMemoryBuffer::copyFromHostBuffer);
         assert(amountRead == realAmount);

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.
@@ -165,14 +165,6 @@ 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);
-  }
-
   /**
    * Initialize the pool.
    *
@@ -235,47 +227,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.
-   *
-   * @param bytes size in bytes to allocate
-   * @return newly created buffer
-   */
-  public static HostMemoryBuffer allocate(long bytes, HostMemoryAllocator hostMemoryAllocator) {
-    HostMemoryBuffer result = tryAllocate(bytes);
-    if (result == null) {
-      result = hostMemoryAllocator.allocate(bytes, false);
-    }
-    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.
-   *
-   * @param bytes size in bytes to allocate
-   * @return newly created buffer
-   */
-  public static HostMemoryBuffer allocate(long bytes) {
-    return allocate(bytes, DefaultHostMemoryAllocator.get());
-  }
-
   /**
    * Get the number of bytes free in the pinned memory pool.
    *
@@ -319,8 +270,7 @@ public void close() {
   }
 
   /**
-   * 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.
+   * Pads a length of bytes to the alignment the CPU wants in the worst case.
    * @param bytes the size in bytes
    * @return the new padded size in bytes.
    */
@@ -374,62 +324,6 @@ private synchronized HostMemoryBuffer tryAllocateInternal(long 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;
@@ -446,12 +340,6 @@ private synchronized void free(MemorySection section) {
     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;
   }