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"