Explicitly set legacy or per-thread default stream in JNI (#6690)

Fixes an issue where nvcomp could use the wrong stream when cudf is compiled with PTDS

CHANGELOG.md

@@ -34,6 +34,7 @@
 - PR #6540 Add dictionary support to `cudf::unary_operation`
 - PR #6537 Refactor ORC timezone
 - PR #6527 Refactor DeviceColumnViewAccess to avoid JNI returning an array
+- PR #6690 Explicitly set legacy or per-thread default stream in JNI
 - PR #6545 Pin cmake policies to cmake 3.17 version
 - PR #6556 Add dictionary support to `cudf::inner_join`, `cudf::left_join` and `cudf::full_join`
 - PR #6557 Support nullable timestamp columns in time range window functions

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

@@ -19,15 +19,20 @@
 import org.slf4j.LoggerFactory;
 
 public class Cuda {
+  // This needs to happen first before calling any native methods.
+  static {
+    NativeDepsLoader.loadNativeDeps();
+  }
+
   // Defined in driver_types.h in cuda library.
   static final int CPU_DEVICE_ID = -1;
+  static final long CUDA_STREAM_DEFAULT = 0;
+  static final long CUDA_STREAM_LEGACY = 1;
+  static final long CUDA_STREAM_PER_THREAD = 2;
+  private final static long DEFAULT_STREAM_ID = isPtdsEnabled() ? CUDA_STREAM_PER_THREAD : CUDA_STREAM_LEGACY;
   private static final Logger log = LoggerFactory.getLogger(Cuda.class);
   private static Boolean isCompat = null;
 
-  static {
-    NativeDepsLoader.loadNativeDeps();
-  }
-
   private static class StreamCleaner extends MemoryCleaner.Cleaner {
     private long stream;
 
@@ -39,7 +44,9 @@ private static class StreamCleaner extends MemoryCleaner.Cleaner {
     protected boolean cleanImpl(boolean logErrorIfNotClean) {
       boolean neededCleanup = false;
       long origAddress = stream;
-      if (stream != 0) {
+      if (stream != CUDA_STREAM_DEFAULT &&
+          stream != CUDA_STREAM_LEGACY &&
+          stream != CUDA_STREAM_PER_THREAD) {
         destroyStream(stream);
         stream = 0;
         neededCleanup = true;
@@ -89,7 +96,7 @@ public void waitOn(Event event) {
     }
 
     public long getStream() {
-      return cleaner == null ? 0 : cleaner.stream;
+      return cleaner == null ? DEFAULT_STREAM_ID : cleaner.stream;
     }
 
     /**
@@ -206,7 +213,7 @@ public void record(Stream stream) {
     }
 
     /**
-     * Captures the contents of stream 0 at the time of this call.
+     * Captures the contents of the default stream at the time of this call.
      */
     public void record() {
       record(DEFAULT_STREAM);