diff --git a/java/src/main/java/ai/rapids/cudf/DefaultHostMemoryAllocator.java b/java/src/main/java/ai/rapids/cudf/DefaultHostMemoryAllocator.java index 469c8d16fe4..3178a3ea309 100644 --- 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,7 +32,7 @@ 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; @@ -40,11 +40,17 @@ public static void set(HostMemoryAllocator 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 index bfb959b12c1..7251ce643d4 100644 --- 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); } /**