From 31d6f0ede6932ebf57395c1ed552629e6e32c44f Mon Sep 17 00:00:00 2001 From: Rong Ou Date: Thu, 1 Apr 2021 15:01:12 -0700 Subject: [PATCH] add an allocate method with stream in java DeviceMemoryBuffer --- .../main/java/ai/rapids/cudf/DeviceMemoryBuffer.java | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/java/src/main/java/ai/rapids/cudf/DeviceMemoryBuffer.java b/java/src/main/java/ai/rapids/cudf/DeviceMemoryBuffer.java index 5753ecea74d..fa888625d47 100644 --- a/java/src/main/java/ai/rapids/cudf/DeviceMemoryBuffer.java +++ b/java/src/main/java/ai/rapids/cudf/DeviceMemoryBuffer.java @@ -122,7 +122,17 @@ private DeviceMemoryBuffer(long address, long lengthInBytes, DeviceMemoryBuffer * @return the buffer */ public static DeviceMemoryBuffer allocate(long bytes) { - return Rmm.alloc(bytes); + return allocate(bytes, Cuda.DEFAULT_STREAM); + } + + /** + * Allocate memory for use on the GPU. You must close it when done. + * @param bytes size in bytes to allocate + * @param stream The stream in which to synchronize this command + * @return the buffer + */ + public static DeviceMemoryBuffer allocate(long bytes, Cuda.Stream stream) { + return Rmm.alloc(bytes, stream); } /**