Add JNI wrapper for the cuFile API (GDS) [skip ci]

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - PR #6847 Add a cmake find module for cuFile in JNI code
 - PR #6902 Implement `DataFrame.quantile` for `datetime` and `timedelta` data types
+- PR #6940 Add JNI wrapper for the cuFile API (GDS)
 - PR #6814 Implement `cudf::reduce` for `decimal32` and `decimal64` (part 1)
 - PR #6929 Add `Index.set_names` api
 

java/README.md

@@ -119,3 +119,16 @@ then build the jar:
 cd src/cudf/java
 mvn clean install -DPER_THREAD_DEFAULT_STREAM=ON
 ```
+
+## GPUDirect Storage (GDS)
+
+The JNI code can be built with *GPUDirect Storage* (GDS) support, which enables direct copying 
+between GPU device buffers and supported filesystems (see
+https://docs.nvidia.com/gpudirect-storage/).
+
+To enable GDS support, first make sure GDS is installed (see
+https://docs.nvidia.com/gpudirect-storage/troubleshooting-guide/index.html), then run:  
+```shell script
+cd src/cudf/java
+mvn clean install -DUSE_GDS=ON
+```

java/pom.xml

@@ -147,6 +147,7 @@
         <CUDA_STATIC_RUNTIME>OFF</CUDA_STATIC_RUNTIME>
         <PER_THREAD_DEFAULT_STREAM>OFF</PER_THREAD_DEFAULT_STREAM>
         <RMM_LOGGING_LEVEL>INFO</RMM_LOGGING_LEVEL>
+        <USE_GDS>OFF</USE_GDS>
         <native.build.path>${project.build.directory}/cmake-build</native.build.path>
         <slf4j.version>1.7.30</slf4j.version>
     </properties>
@@ -158,6 +159,28 @@
                 <cxx.flags>-Wno-deprecated-declarations</cxx.flags>
             </properties>
         </profile>
+        <profile>
+            <id>no-cufile-tests</id>
+            <activation>
+                <property>
+                    <name>USE_GDS</name>
+                    <value>!ON</value>
+                </property>
+            </activation>
+            <build>
+                <plugins>
+                    <plugin>
+                        <groupId>org.apache.maven.plugins</groupId>
+                        <artifactId>maven-compiler-plugin</artifactId>
+                        <configuration>
+                            <excludes>
+                                <exclude>**/CuFileTest.java</exclude>
+                            </excludes>
+                        </configuration>
+                    </plugin>
+                </plugins>
+            </build>
+        </profile>
         <profile>
             <id>release</id>
             <distributionManagement>
@@ -349,6 +372,7 @@
                                     <arg value="-DCUDA_STATIC_RUNTIME=${CUDA_STATIC_RUNTIME}" />
                                     <arg value="-DPER_THREAD_DEFAULT_STREAM=${PER_THREAD_DEFAULT_STREAM}" />
                                     <arg value="-DRMM_LOGGING_LEVEL=${RMM_LOGGING_LEVEL}" />
+                                    <arg value="-DUSE_GDS=${USE_GDS}" />
                                     <arg value="-DCMAKE_CXX_FLAGS=${cxx.flags}"/>
                                     <arg value="-DCMAKE_EXPORT_COMPILE_COMMANDS=${CMAKE_EXPORT_COMPILE_COMMANDS}"/>
                                     <arg value="-DCUDF_CPP_BUILD_DIR=${CUDF_CPP_BUILD_DIR}"/>
@@ -479,6 +503,7 @@
                                     <directory>${native.build.path}</directory>
                                     <includes>
                                         <include>libcudfjni.so</include>
+                                        <include>libcufilejni.so</include>
                                     </includes>
                                 </resource>
                                 <resource>

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

@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package ai.rapids.cudf;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.File;
+
+/**
+ * JNI wrapper for accessing the cuFile API.
+ * <p>
+ * Using this wrapper requires GPUDirect Storage (GDS)/cuFile to be installed in the target
+ * environment, and the jar to be built with `USE_GDS=ON`. Otherwise it will throw an exception when
+ * loading.
+ * <p>
+ * The Java APIs are experimental and subject to change.
+ *
+ * @see <a href="https://docs.nvidia.com/gpudirect-storage/">GDS documentation</a>
+ */
+public class CuFile {
+  private static final Logger log = LoggerFactory.getLogger(CuFile.class);
+  private static boolean initialized = false;
+  private static long driverPointer = 0;
+
+  static {
+    initialize();
+  }
+
+  /**
+   * Load the native libraries needed for libcufilejni, if not loaded already; open the cuFile
+   * driver, and add a shutdown hook to close it.
+   */
+  private static synchronized void initialize() {
+    if (!initialized) {
+      try {
+        NativeDepsLoader.loadNativeDeps(new String[]{"cufilejni"});
+        driverPointer = createDriver();
+        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
+          destroyDriver(driverPointer);
+        }));
+        initialized = true;
+      } catch (Throwable t) {
+        log.error("Could not load cuFile jni library...", t);
+      }
+    }
+  }
+
+  private static native long createDriver();
+
+  private static native void destroyDriver(long pointer);
+
+  /**
+   * Copy a device buffer to a given file path synchronously.
+   * <p>
+   * This method is NOT thread safe if the path points to the same file on disk.
+   *
+   * @param path   The file path to copy to.
+   * @param buffer The device buffer to copy from.
+   * @param append Whether to append to the file.
+   * @return The file offset from which the buffer was appended.
+   */
+  public static long copyDeviceBufferToFile(File path, DeviceMemoryBuffer buffer, boolean append) {
+    return copyToFile(path.getAbsolutePath(), buffer.getAddress(), buffer.getLength(), append);
+  }
+
+  /**
+   * Copy a file into a device buffer synchronously.
+   * <p>
+   * This method is NOT thread safe if the path points to the same file on disk.
+   *
+   * @param buffer     The device buffer to copy into.
+   * @param path       The file path to copy from.
+   * @param fileOffset The file offset from which to copy the content.
+   */
+  public static void copyFileToDeviceBuffer(DeviceMemoryBuffer buffer, File path, long fileOffset) {
+    copyFromFile(buffer.getAddress(), buffer.getLength(), path.getAbsolutePath(), fileOffset);
+  }
+
+  private static native long copyToFile(String path, long address, long length, boolean append);
+
+  private static native void copyFromFile(long address, long length, String path, long fileOffset);
+}

java/src/main/native/CMakeLists.txt

@@ -289,8 +289,7 @@ include_directories("${THRUST_INCLUDE}"
                     "${JNI_INCLUDE_DIRS}"
                     "${CUDF_INCLUDE}"
                     "${RMM_INCLUDE}"
-                    "${ARROW_INCLUDE}"
-                    "$<$<BOOL:${cuFile_FOUND}>:${cuFile_INCLUDE_DIRS}>")
+                    "${ARROW_INCLUDE}")
 
 ###################################################################################################
 # - library paths ---------------------------------------------------------------------------------
@@ -321,6 +320,12 @@ add_library(cudfjni SHARED ${SOURCE_FILES})
 #Override RPATH for cudfjni
 SET_TARGET_PROPERTIES(cudfjni PROPERTIES BUILD_RPATH "\$ORIGIN")
 
+if(USE_GDS)
+    add_library(cufilejni SHARED "src/CuFileJni.cpp")
+    target_include_directories(cufilejni PRIVATE "${cuFile_INCLUDE_DIRS}")
+    target_link_libraries(cufilejni PRIVATE "${cuFile_LIBRARIES}")
+endif(USE_GDS)
+
 ###################################################################################################
 # - build options ---------------------------------------------------------------------------------