Make HostMemoryBuffer call into the DefaultHostMemoryAllocator

Signed-off-by: Robert (Bobby) Evans <[email protected]>
rapidsai · Oct 29, 2024 · ca35305 · ca35305
1 parent 63b773e
commit ca35305
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 12 deletions.
diff --git a/java/src/main/java/ai/rapids/cudf/DefaultHostMemoryAllocator.java b/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 new allocator to use.
    */
   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 new HostMemoryBuffer(UnsafeMemoryAccessor.allocate(bytes), bytes);
   }
 
   @Override
   public HostMemoryBuffer allocate(long bytes) {
-    return HostMemoryBuffer.allocate(bytes);
+    return allocate(bytes, HostMemoryBuffer.defaultPreferPinned);
   }
 }
diff --git a/java/src/main/java/ai/rapids/cudf/HostMemoryBuffer.java b/java/src/main/java/ai/rapids/cudf/HostMemoryBuffer.java
@@ -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);
   }
 
   /**