remove extra sync, and make sure copyNext is always synchronous with …

…the cuda stream
NVIDIA · Nov 22, 2024 · b48b801 · b48b801
1 parent 2619cab
commit b48b801
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/spill/SpillFramework.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/spill/SpillFramework.scala
@@ -205,9 +205,8 @@ object SpillableHostBufferHandle extends Logging {
         while (chunkedPacker.hasNext) {
           withResource(chunkedPacker.next(bb)) { n =>
             builder.copyNext(n, Cuda.DEFAULT_STREAM)
-            // we are calling chunked packer on `bb` again each time, we need
-            // to synchronize before we ask for the next chunk
-            Cuda.DEFAULT_STREAM.sync()
+            // copyNext is synchronous w.r.t. the cuda stream passed,
+            // no need to synchronize here.
           }
         }
       }
@@ -221,7 +220,6 @@ object SpillableHostBufferHandle extends Logging {
     withResource(
       SpillFramework.stores.hostStore.makeBuilder(handle)) { builder =>
       builder.copyNext(buff, Cuda.DEFAULT_STREAM)
-      Cuda.DEFAULT_STREAM.sync()
       builder.build
     }
   }
@@ -1163,7 +1161,7 @@ class SpillableHostStore(val maxSize: Option[Long] = None)
 
     override def copyNext(mb: DeviceMemoryBuffer, stream: Cuda.Stream): Unit = {
       GpuTaskMetrics.get.spillToHostTime {
-        singleShotBuffer.copyFromMemoryBufferAsync(
+        singleShotBuffer.copyFromMemoryBuffer(
           copied,
           mb,
           0,