From 1bb3aac2a918e95483fd52b1523855014bc03a34 Mon Sep 17 00:00:00 2001
From: Alessandro Bellina <abellina@nvidia.com>
Date: Thu, 12 May 2022 15:43:07 -0500
Subject: [PATCH] Adds the JNI call for Cuda.deviceSynchronize (#10839)

Running tests locally, but putting this up as WIP for now.

Discussing with @jlowe a solution to https://github.com/NVIDIA/spark-rapids/issues/4818 could involve `cudaDeviceSynchronize.` I noticed that's not in our JNI exposed calls, so I am adding it here.

Authors:
  - Alessandro Bellina (https://github.com/abellina)

Approvers:
  - Jason Lowe (https://github.com/jlowe)

URL: https://github.com/rapidsai/cudf/pull/10839
---
 java/src/main/java/ai/rapids/cudf/Cuda.java | 6 ++++++
 java/src/main/native/src/CudaJni.cpp        | 8 ++++++++
 2 files changed, 14 insertions(+)

diff --git a/java/src/main/java/ai/rapids/cudf/Cuda.java b/java/src/main/java/ai/rapids/cudf/Cuda.java
index 21843527fc2..56a754279fc 100755
--- a/java/src/main/java/ai/rapids/cudf/Cuda.java
+++ b/java/src/main/java/ai/rapids/cudf/Cuda.java
@@ -596,4 +596,10 @@ public static void multiBufferCopyAsync(long [] destAddrs,
    * no effect.
    */
   public static native void profilerStop();
+
+  /**
+   * Synchronizes the whole device using cudaDeviceSynchronize.
+   * @note this is very expensive and should almost never be used
+   */
+  public static native void deviceSynchronize();
 }
diff --git a/java/src/main/native/src/CudaJni.cpp b/java/src/main/native/src/CudaJni.cpp
index 926521c55f9..ce1ad1b1671 100644
--- a/java/src/main/native/src/CudaJni.cpp
+++ b/java/src/main/native/src/CudaJni.cpp
@@ -390,4 +390,12 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cuda_profilerStop(JNIEnv *env, jclass
   CATCH_STD(env, );
 }
 
+JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cuda_deviceSynchronize(JNIEnv *env, jclass clazz) {
+  try {
+    cudf::jni::auto_set_device(env);
+    CUDF_CUDA_TRY(cudaDeviceSynchronize());
+  }
+  CATCH_STD(env, );
+}
+
 } // extern "C"