[GLUTEN-10933][VL] cuDF: Move lock to ShuffleReader

backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHSparkPlanExecApi.scala

@@ -454,7 +454,8 @@ class CHSparkPlanExecApi extends SparkPlanExecApi with Logging {
   override def createColumnarBatchSerializer(
       schema: StructType,
       metrics: Map[String, SQLMetric],
-      shuffleWriterType: ShuffleWriterType): Serializer = {
+      shuffleWriterType: ShuffleWriterType,
+      enableCudf: Boolean): Serializer = {
     val readBatchNumRows = metrics("avgReadBatchNumRows")
     val numOutputRows = metrics("numOutputRows")
     val dataSize = metrics("dataSize")

backends-velox/src-celeborn/main/scala/org/apache/spark/shuffle/VeloxCelebornColumnarBatchSerializer.scala

@@ -103,7 +103,8 @@ private class CelebornColumnarBatchSerializerInstance(
         batchSize,
         readerBufferSize,
         deserializerBufferSize,
-        shuffleWriterType.name
+        shuffleWriterType.name,
+        false
       )
     // Close shuffle reader instance as lately as the end of task processing,
     // since the native reader could hold a reference to memory pool that

backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala

@@ -633,7 +633,8 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi {
   override def createColumnarBatchSerializer(
       schema: StructType,
       metrics: Map[String, SQLMetric],
-      shuffleWriterType: ShuffleWriterType): Serializer = {
+      shuffleWriterType: ShuffleWriterType,
+      enableCudf: Boolean): Serializer = {
     val numOutputRows = metrics("numOutputRows")
     val deserializeTime = metrics("deserializeTime")
     val readBatchNumRows = metrics("avgReadBatchNumRows")
@@ -658,7 +659,8 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi {
           numOutputRows,
           deserializeTime,
           decompressTime,
-          shuffleWriterType)
+          shuffleWriterType,
+          enableCudf)
     }
   }
 

backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala

@@ -18,10 +18,11 @@ package org.apache.gluten.extension
 
 import org.apache.gluten.config.{GlutenConfig, VeloxConfig}
 import org.apache.gluten.cudf.VeloxCudfPlanValidatorJniWrapper
-import org.apache.gluten.execution.{CudfTag, LeafTransformSupport, TransformSupport, WholeStageTransformer}
+import org.apache.gluten.execution.{CudfTag, LeafTransformSupport, TransformSupport, VeloxResizeBatchesExec, WholeStageTransformer}
+import org.apache.gluten.extension.CudfNodeValidationRule.setTagForWholeStageTransformer
 
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.{ColumnarShuffleExchangeExec, GPUColumnarShuffleExchangeExec, SparkPlan}
 
 // Add the node name prefix 'Cudf' to GlutenPlan when can offload to cudf
 case class CudfNodeValidationRule(glutenConf: GlutenConfig) extends Rule[SparkPlan] {
@@ -31,33 +32,65 @@ case class CudfNodeValidationRule(glutenConf: GlutenConfig) extends Rule[SparkPl
       return plan
     }
     plan.transformUp {
+      case shuffle @ ColumnarShuffleExchangeExec(
+            _,
+            v @ VeloxResizeBatchesExec(w: WholeStageTransformer, _, _),
+            _,
+            _,
+            _) =>
+        setTagForWholeStageTransformer(w)
+        if (w.isCudf) {
+          log.info("VeloxResizeBatchesExec is not supported in GPU")
+        }
+        GPUColumnarShuffleExchangeExec(
+          shuffle.outputPartitioning,
+          w,
+          shuffle.shuffleOrigin,
+          shuffle.projectOutputAttributes,
+          shuffle.advisoryPartitionSize)
+
+      case shuffle @ ColumnarShuffleExchangeExec(_, w: WholeStageTransformer, _, _, _) =>
+        setTagForWholeStageTransformer(w)
+        GPUColumnarShuffleExchangeExec(
+          shuffle.outputPartitioning,
+          w,
+          shuffle.shuffleOrigin,
+          shuffle.projectOutputAttributes,
+          shuffle.advisoryPartitionSize)
+
       case transformer: WholeStageTransformer =>
-        if (!VeloxConfig.get.cudfEnableTableScan) {
-          // Spark3.2 does not have exists
-          val hasLeaf = transformer.find {
-            case _: LeafTransformSupport => true
-            case _ => false
-          }.isDefined
-          if (!hasLeaf && VeloxConfig.get.cudfEnableValidation) {
-            if (
-              VeloxCudfPlanValidatorJniWrapper.validate(
-                transformer.substraitPlan.toProtobuf.toByteArray)
-            ) {
-              transformer.foreach {
-                case _: LeafTransformSupport =>
-                case t: TransformSupport =>
-                  t.setTagValue(CudfTag.CudfTag, true)
-                case _ =>
-              }
-              transformer.setTagValue(CudfTag.CudfTag, true)
-            }
-          } else {
-            transformer.setTagValue(CudfTag.CudfTag, !hasLeaf)
+        setTagForWholeStageTransformer(transformer)
+        transformer
+    }
+  }
+}
+
+object CudfNodeValidationRule {
+  def setTagForWholeStageTransformer(transformer: WholeStageTransformer): Unit = {
+    if (!VeloxConfig.get.cudfEnableTableScan) {
+      // Spark3.2 does not have exists
+      val hasLeaf = transformer.find {
+        case _: LeafTransformSupport => true
+        case _ => false
+      }.isDefined
+      if (!hasLeaf && VeloxConfig.get.cudfEnableValidation) {
+        if (
+          VeloxCudfPlanValidatorJniWrapper.validate(
+            transformer.substraitPlan.toProtobuf.toByteArray)
+        ) {
+          transformer.foreach {
+            case _: LeafTransformSupport =>
+            case t: TransformSupport =>
+              t.setTagValue(CudfTag.CudfTag, true)
+            case _ =>
           }
-        } else {
           transformer.setTagValue(CudfTag.CudfTag, true)
         }
-        transformer
+      } else {
+        transformer.setTagValue(CudfTag.CudfTag, !hasLeaf)
+      }
+    } else {
+      transformer.setTagValue(CudfTag.CudfTag, true)
     }
   }
 }

backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala

@@ -51,7 +51,8 @@ class ColumnarBatchSerializer(
     numOutputRows: SQLMetric,
     deserializeTime: SQLMetric,
     decompressTime: SQLMetric,
-    shuffleWriterType: ShuffleWriterType)
+    shuffleWriterType: ShuffleWriterType,
+    enableCudf: Boolean)
   extends Serializer
   with Serializable {
 
@@ -63,7 +64,8 @@ class ColumnarBatchSerializer(
       numOutputRows,
       deserializeTime,
       decompressTime,
-      shuffleWriterType)
+      shuffleWriterType,
+      enableCudf)
   }
 
   override def supportsRelocationOfSerializedObjects: Boolean = true
@@ -75,7 +77,8 @@ private class ColumnarBatchSerializerInstanceImpl(
     numOutputRows: SQLMetric,
     deserializeTime: SQLMetric,
     decompressTime: SQLMetric,
-    shuffleWriterType: ShuffleWriterType)
+    shuffleWriterType: ShuffleWriterType,
+    enableCudf: Boolean)
   extends ColumnarBatchSerializerInstance
   with Logging {
 
@@ -111,7 +114,8 @@ private class ColumnarBatchSerializerInstanceImpl(
       batchSize,
       readerBufferSize,
       deserializerBufferSize,
-      shuffleWriterType.name)
+      shuffleWriterType.name,
+      enableCudf)
     // Close shuffle reader instance as lately as the end of task processing,
     // since the native reader could hold a reference to memory pool that
     // was used to create all buffers read from shuffle reader. The pool

cpp/core/jni/JniWrapper.cc

@@ -1083,7 +1083,8 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleReaderJniWrappe
     jint batchSize,
     jlong readerBufferSize,
     jlong deserializerBufferSize,
-    jstring shuffleWriterType) {
+    jstring shuffleWriterType,
+    jboolean enableCudf) {
   JNI_METHOD_START
   auto ctx = getRuntime(env, wrapper);
 
@@ -1095,6 +1096,7 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleReaderJniWrappe
   options.batchSize = batchSize;
   options.readerBufferSize = readerBufferSize;
   options.deserializerBufferSize = deserializerBufferSize;
+  options.enableCudf = enableCudf;
 
   options.shuffleWriterType = ShuffleWriter::stringToType(jStringToCString(env, shuffleWriterType));
   std::shared_ptr<arrow::Schema> schema =

cpp/core/shuffle/Options.h

@@ -62,6 +62,11 @@ struct ShuffleReaderOptions {
 
   // Buffer size when deserializing rows into columnar batches. Only used for sort-based shuffle.
   int64_t deserializerBufferSize = kDefaultDeserializerBufferSize;
+
+  // When true, convert the buffers to cudf table.
+  // Add a lock after reader produces the Vector, the next operator should be CudfFromVelox.
+  // After move the shuffle read operation to gpu, move the lock to start read.
+  bool enableCudf = false;
 };
 
 struct ShuffleWriterOptions {

cpp/velox/CMakeLists.txt

@@ -199,7 +199,8 @@ if(ENABLE_S3)
 endif()
 
 if(ENABLE_GPU)
-  list(APPEND VELOX_SRCS cudf/CudfPlanValidator.cc)
+  list(APPEND VELOX_SRCS cudf/CudfPlanValidator.cc cudf/GpuLock.cc
+       shuffle/GpuShuffleReader.cc)
 endif()
 
 if(ENABLE_ENHANCED_FEATURES)

cpp/velox/compute/VeloxRuntime.cc

@@ -300,7 +300,8 @@ std::shared_ptr<ShuffleReader> VeloxRuntime::createShuffleReader(
       options.readerBufferSize,
       options.deserializerBufferSize,
       memoryManager(),
-      options.shuffleWriterType);
+      options.shuffleWriterType,
+      options.enableCudf);
 
   return std::make_shared<VeloxShuffleReader>(std::move(deserializerFactory));
 }

cpp/velox/compute/WholeStageResultIterator.cc

@@ -23,10 +23,10 @@
 #include "velox/exec/PlanNodeStats.h"
 #ifdef GLUTEN_ENABLE_GPU
 #include <cudf/io/types.hpp>
-#include <mutex>
 #include "velox/experimental/cudf/CudfConfig.h"
 #include "velox/experimental/cudf/connectors/hive/CudfHiveConnectorSplit.h"
 #include "velox/experimental/cudf/exec/ToCudf.h"
+#include "cudf/GpuLock.h"
 #endif
 
 using namespace facebook;
@@ -75,11 +75,11 @@ WholeStageResultIterator::WholeStageResultIterator(
     : memoryManager_(memoryManager),
       veloxCfg_(
           std::make_shared<facebook::velox::config::ConfigBase>(std::unordered_map<std::string, std::string>(confMap))),
-      taskInfo_(taskInfo),
-      veloxPlan_(planNode),
 #ifdef GLUTEN_ENABLE_GPU
-      lock_(mutex_, std::defer_lock),
+      enableCudf_(veloxCfg_->get<bool>(kCudfEnabled, kCudfEnabledDefault)),
 #endif
+      taskInfo_(taskInfo),
+      veloxPlan_(planNode),
       scanNodeIds_(scanNodeIds),
       scanInfos_(scanInfos),
       streamIds_(streamIds) {
@@ -90,13 +90,6 @@ WholeStageResultIterator::WholeStageResultIterator(
   }
   getOrderedNodeIds(veloxPlan_, orderedNodeIds_);
 
-#ifdef GLUTEN_ENABLE_GPU
-  enableCudf_ = veloxCfg_->get<bool>(kCudfEnabled, kCudfEnabledDefault);
-  if (enableCudf_) {
-    lock_.lock();
-  }
-#endif
-
   auto fileSystem = velox::filesystems::getFileSystem(spillDir, nullptr);
   GLUTEN_CHECK(fileSystem != nullptr, "File System for spilling is null!");
   fileSystem->mkdir(spillDir);
@@ -213,10 +206,6 @@ WholeStageResultIterator::WholeStageResultIterator(
   }
 }
 
-#ifdef GLUTEN_ENABLE_GPU
-std::mutex WholeStageResultIterator::mutex_;
-#endif
-
 std::shared_ptr<velox::core::QueryCtx> WholeStageResultIterator::createNewVeloxQueryCtx() {
   std::unordered_map<std::string, std::shared_ptr<velox::config::ConfigBase>> connectorConfigs;
   connectorConfigs[kHiveConnectorId] = createConnectorConfig();
@@ -236,17 +225,6 @@ std::shared_ptr<velox::core::QueryCtx> WholeStageResultIterator::createNewVeloxQ
 }
 
 std::shared_ptr<ColumnarBatch> WholeStageResultIterator::next() {
-  auto result = nextInternal();
-#ifdef GLUTEN_ENABLE_GPU
-  if (result == nullptr && enableCudf_) {
-    lock_.unlock();
-  }
-#endif
-
-  return result;
-}
-
-std::shared_ptr<ColumnarBatch> WholeStageResultIterator::nextInternal() {
   tryAddSplitsToTask();
   if (task_->isFinished()) {
     return nullptr;

cpp/velox/compute/WholeStageResultIterator.h

@@ -27,6 +27,9 @@
 #include "velox/connectors/hive/iceberg/IcebergSplit.h"
 #include "velox/core/PlanNode.h"
 #include "velox/exec/Task.h"
+#ifdef GLUTEN_ENABLE_GPU
+#include "cudf/GpuLock.h"
+#endif
 
 namespace gluten {
 
@@ -48,8 +51,8 @@ class WholeStageResultIterator : public ColumnarBatchIterator {
       task_->requestCancel().wait();
     }
 #ifdef GLUTEN_ENABLE_GPU
-    if (enableCudf_ && lock_.owns_lock()) {
-      lock_.unlock();
+    if (enableCudf_) {
+      unlockGpu();
     }
 #endif
   }
@@ -75,8 +78,6 @@ class WholeStageResultIterator : public ColumnarBatchIterator {
   }
 
  private:
-  std::shared_ptr<ColumnarBatch> nextInternal();
-
   /// Get the Spark confs to Velox query context.
   std::unordered_map<std::string, std::string> getQueryContextConf();
 
@@ -113,6 +114,9 @@ class WholeStageResultIterator : public ColumnarBatchIterator {
 
   /// Config, task and plan.
   std::shared_ptr<config::ConfigBase> veloxCfg_;
+#ifdef GLUTEN_ENABLE_GPU
+  const bool enableCudf_;
+#endif
   const SparkTaskInfo taskInfo_;
   std::shared_ptr<facebook::velox::exec::Task> task_;
   std::shared_ptr<const facebook::velox::core::PlanNode> veloxPlan_;
@@ -124,13 +128,6 @@ class WholeStageResultIterator : public ColumnarBatchIterator {
   /// Metrics
   std::unique_ptr<Metrics> metrics_{};
 
-#ifdef GLUTEN_ENABLE_GPU
-  // Mutex for thread safety.
-  static std::mutex mutex_;
-  std::unique_lock<std::mutex> lock_;
-  bool enableCudf_;
-#endif
-
   /// All the children plan node ids with postorder traversal.
   std::vector<facebook::velox::core::PlanNodeId> orderedNodeIds_;