[WIP][VL] Support setting stage execution mode

backends-clickhouse/src/main/scala/org/apache/gluten/extension/LazyAggregateExpandRule.scala

@@ -69,6 +69,8 @@ case class LazyAggregateExpandRule(session: SparkSession) extends Rule[SparkPlan
               ExpandExecTransformer(projections, output, child)),
             _,
             _,
+            _,
+            _,
             _
           ) =>
         logDebug(s"xxx match plan:$shuffle")
@@ -116,6 +118,8 @@ case class LazyAggregateExpandRule(session: SparkSession) extends Rule[SparkPlan
               FilterExecTransformer(_, ExpandExecTransformer(projections, output, child))),
             _,
             _,
+            _,
+            _,
             _
           ) =>
         val partialAggregate = shuffle.child.asInstanceOf[CHHashAggregateExecTransformer]

backends-velox/src/main/java/org/apache/gluten/utils/GpuBufferColumnarBatchResizer.java

@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.utils;
+
+import org.apache.gluten.backendsapi.BackendsApiManager;
+import org.apache.gluten.runtime.Runtime;
+import org.apache.gluten.runtime.Runtimes;
+import org.apache.gluten.vectorized.ColumnarBatchInIterator;
+import org.apache.gluten.vectorized.ColumnarBatchOutIterator;
+
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+import java.util.Iterator;
+
+public final class GpuBufferColumnarBatchResizer {
+  private static class ResizeBatchOutIterator extends ColumnarBatchOutIterator {
+    GpuBufferBatchResizerJniWrapper jniWrapper;
+
+    ResizeBatchOutIterator(
+        GpuBufferBatchResizerJniWrapper jniWrapper, Runtime runtime, long iterHandle) {
+      super(runtime, iterHandle);
+      this.jniWrapper = jniWrapper;
+    }
+  }
+
+  public static ColumnarBatchOutIterator create(
+      int minOutputBatchSize, long maxPrefetchBatchBytes, Iterator<ColumnarBatch> in) {
+    final Runtime runtime =
+        Runtimes.contextInstance(
+            BackendsApiManager.getBackendName(), "GpuBufferColumnarBatchResizer");
+    GpuBufferBatchResizerJniWrapper jniWrapper = GpuBufferBatchResizerJniWrapper.create(runtime);
+    long outHandle =
+        jniWrapper.create(
+            minOutputBatchSize,
+            maxPrefetchBatchBytes,
+            new ColumnarBatchInIterator(BackendsApiManager.getBackendName(), in));
+    return new ResizeBatchOutIterator(jniWrapper, runtime, outHandle);
+  }
+}

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

@@ -112,7 +112,6 @@ object VeloxRuleApi {
 
     // Legacy: Post-transform rules.
     injector.injectPostTransform(c => AppendBatchResizeForShuffleInputAndOutput(c.caller.isAqe()))
-    injector.injectPostTransform(_ => GpuBufferBatchResizeForShuffleInputOutput())
     injector.injectPostTransform(_ => UnionTransformerRule())
     injector.injectPostTransform(_ => PartialFallbackRules())
     injector.injectPostTransform(_ => RemoveNativeWriteFilesSortAndProject())
@@ -147,6 +146,8 @@ object VeloxRuleApi {
       c => PreventBatchTypeMismatchInTableCache(c.caller.isCache(), Set(VeloxBatchType)))
     injector.injectFinal(
       c => GlutenAutoAdjustStageResourceProfile(new GlutenConfig(c.sqlConf), c.session))
+    injector.injectFinal(
+      c => AdjustStageExecutionMode(new GlutenConfig(c.sqlConf), c.session, c.caller.isAqe()))
     injector.injectFinal(c => GlutenFallbackReporter(new GlutenConfig(c.sqlConf), c.session))
     injector.injectFinal(_ => RemoveFallbackTagRule())
   }

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

@@ -36,8 +36,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.memory.SparkMemoryUtil
 import org.apache.spark.rdd.RDD
 import org.apache.spark.serializer.Serializer
-import org.apache.spark.shuffle.{GenShuffleReaderParameters, GenShuffleWriterParameters, GlutenShuffleReaderWrapper, GlutenShuffleWriterWrapper}
-import org.apache.spark.shuffle.utils.ShuffleUtil
+import org.apache.spark.shuffle.{GenShuffleReaderParameters, GenShuffleWriterParameters, GlutenShuffleReaderWrapper, GlutenShuffleWriterWrapper, VeloxShuffleUtils}
 import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec
 import org.apache.spark.sql.catalyst.expressions._
@@ -644,12 +643,12 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
    */
   override def genColumnarShuffleWriter[K, V](
       parameters: GenShuffleWriterParameters[K, V]): GlutenShuffleWriterWrapper[K, V] = {
-    ShuffleUtil.genColumnarShuffleWriter(parameters)
+    VeloxShuffleUtils.genColumnarShuffleWriter(parameters)
   }
 
   override def genColumnarShuffleReader[K, C](
       parameters: GenShuffleReaderParameters[K, C]): GlutenShuffleReaderWrapper[K, C] = {
-    ShuffleUtil.genColumnarShuffleReader(parameters)
+    VeloxShuffleUtils.genColumnarShuffleReader(parameters)
   }
 
   override def createColumnarWriteFilesExec(

backends-velox/src/main/scala/org/apache/gluten/execution/VeloxResizeBatchesExec.scala

@@ -20,7 +20,7 @@ import org.apache.gluten.backendsapi.velox.VeloxBatchType
 import org.apache.gluten.config.VeloxConfig
 import org.apache.gluten.extension.columnar.transition.Convention
 import org.apache.gluten.iterator.ClosableIterator
-import org.apache.gluten.utils.VeloxBatchResizer
+import org.apache.gluten.utils.{GpuBufferColumnarBatchResizer, VeloxBatchResizer}
 
 import org.apache.spark.sql.catalyst.expressions.SortOrder
 import org.apache.spark.sql.catalyst.plans.physical.Partitioning
@@ -35,20 +35,30 @@ import scala.collection.JavaConverters._
  */
 case class VeloxResizeBatchesExec(
     override val child: SparkPlan,
-    minOutputBatchSize: Int,
-    maxOutputBatchSize: Int,
-    preferredBatchBytes: Long)
+    executionMode: Option[StageExecutionMode] = None)
   extends ColumnarToColumnarExec(child) {
 
   override protected def mapIterator(in: Iterator[ColumnarBatch]): Iterator[ColumnarBatch] = {
-    VeloxBatchResizer
-      .create(
-        minOutputBatchSize,
-        maxOutputBatchSize,
-        preferredBatchBytes,
-        VeloxConfig.get.enableVeloxResizeBatchesCopyRanges,
-        in.asJava)
-      .asScala
+    val veloxConfig = VeloxConfig.get
+    executionMode match {
+      case Some(GPUStageMode) =>
+        GpuBufferColumnarBatchResizer
+          .create(
+            veloxConfig.cudfBatchSize,
+            veloxConfig.cudfShuffleMaxPrefetchBytes,
+            in.asJava)
+          .asScala
+      case _ =>
+        val range = veloxConfig.veloxResizeBatchesShuffleInputOutputRange
+        VeloxBatchResizer
+          .create(
+            range.min,
+            range.max,
+            veloxConfig.veloxPreferredBatchBytes,
+            veloxConfig.enableVeloxResizeBatchesCopyRanges,
+            in.asJava)
+          .asScala
+    }
   }
 
   override protected def closeIterator(out: Iterator[ColumnarBatch]): Unit = {

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

@@ -30,27 +30,20 @@ import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, ShuffleQuery
 case class AppendBatchResizeForShuffleInputAndOutput(isAdaptiveContext: Boolean)
   extends Rule[SparkPlan] {
   override def apply(plan: SparkPlan): SparkPlan = {
-    if (VeloxConfig.get.enableColumnarCudf) {
-      return plan
-    }
-
     val resizeBatchesShuffleInputEnabled = VeloxConfig.get.veloxResizeBatchesShuffleInput
     val resizeBatchesShuffleOutputEnabled = VeloxConfig.get.veloxResizeBatchesShuffleOutput
     if (!resizeBatchesShuffleInputEnabled && !resizeBatchesShuffleOutputEnabled) {
       return plan
     }
 
-    val range = VeloxConfig.get.veloxResizeBatchesShuffleInputOutputRange
-    val preferredBatchBytes = VeloxConfig.get.veloxPreferredBatchBytes
-
     val newPlan = if (resizeBatchesShuffleInputEnabled) {
-      addResizeBatchesForShuffleInput(plan, range.min, range.max, preferredBatchBytes)
+      addResizeBatchesForShuffleInput(plan)
     } else {
       plan
     }
 
     val resultPlan = if (isAdaptiveContext && resizeBatchesShuffleOutputEnabled) {
-      addResizeBatchesForShuffleOutput(newPlan, range.min, range.max, preferredBatchBytes)
+      addResizeBatchesForShuffleOutput(newPlan)
     } else {
       newPlan
     }
@@ -59,32 +52,26 @@ case class AppendBatchResizeForShuffleInputAndOutput(isAdaptiveContext: Boolean)
   }
 
   private def addResizeBatchesForShuffleInput(
-      plan: SparkPlan,
-      min: Int,
-      max: Int,
-      preferredBatchBytes: Long): SparkPlan = {
+      plan: SparkPlan): SparkPlan = {
     plan.transformUp {
       case shuffle: ColumnarShuffleExchangeExec
           if shuffle.shuffleWriterType.requiresResizingShuffleInput =>
         val appendBatches =
-          VeloxResizeBatchesExec(shuffle.child, min, max, preferredBatchBytes)
+          VeloxResizeBatchesExec(shuffle.child)
         shuffle.withNewChildren(Seq(appendBatches))
     }
   }
 
   private def addResizeBatchesForShuffleOutput(
-      plan: SparkPlan,
-      min: Int,
-      max: Int,
-      preferredBatchBytes: Long): SparkPlan = {
+      plan: SparkPlan): SparkPlan = {
     plan match {
       case s: ShuffleQueryStageExec if requiresResizingShuffleOutput(s) =>
-        VeloxResizeBatchesExec(s, min, max, preferredBatchBytes)
+        VeloxResizeBatchesExec(s)
       case a @ AQEShuffleReadExec(s @ ShuffleQueryStageExec(_, _, _), _)
           if requiresResizingShuffleOutput(s) =>
-        VeloxResizeBatchesExec(a, min, max, preferredBatchBytes)
+        VeloxResizeBatchesExec(a)
       case other =>
-        other.mapChildren(addResizeBatchesForShuffleOutput(_, min, max, preferredBatchBytes))
+        other.mapChildren(addResizeBatchesForShuffleOutput)
     }
   }
 

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

@@ -18,12 +18,13 @@ package org.apache.gluten.extension
 
 import org.apache.gluten.config.{GlutenConfig, VeloxConfig}
 import org.apache.gluten.cudf.VeloxCudfPlanValidatorJniWrapper
-import org.apache.gluten.exception.GlutenNotSupportException
 import org.apache.gluten.execution._
-import org.apache.gluten.extension.CudfNodeValidationRule.{createGPUColumnarExchange, setTagForWholeStageTransformer}
+import org.apache.gluten.extension.CudfNodeValidationRule.{setStageExecutionModeForShuffle, setTagForWholeStageTransformer}
 
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.execution.{ColumnarShuffleExchangeExec, GPUColumnarShuffleExchangeExec, SparkPlan}
+import org.apache.spark.sql.execution._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.util.SparkTestUtil
 
 // Add the node name prefix 'Cudf' to GlutenPlan when can offload to cudf
 case class CudfNodeValidationRule(glutenConf: GlutenConfig) extends Rule[SparkPlan] {
@@ -32,23 +33,19 @@ case class CudfNodeValidationRule(glutenConf: GlutenConfig) extends Rule[SparkPl
     if (!glutenConf.enableColumnarCudf) {
       return plan
     }
-    val transformedPlan = plan.transformUp {
-      case shuffle @ ColumnarShuffleExchangeExec(
-            _,
-            VeloxResizeBatchesExec(w: WholeStageTransformer, _, _, _),
-            _,
-            _,
-            _) =>
-        setTagForWholeStageTransformer(w)
-        createGPUColumnarExchange(shuffle)
-      case shuffle @ ColumnarShuffleExchangeExec(_, w: WholeStageTransformer, _, _, _) =>
-        setTagForWholeStageTransformer(w)
-        createGPUColumnarExchange(shuffle)
+    val taggedPlan = plan.transformUp {
       case transformer: WholeStageTransformer =>
         setTagForWholeStageTransformer(transformer)
         transformer
     }
-    transformedPlan
+
+    if (!SQLConf.get.adaptiveExecutionEnabled) {
+      // Set mapper and reducer stage execution mode for Shuffle.
+      setStageExecutionModeForShuffle(taggedPlan, supportsCudf = false)._1
+    } else {
+      // Will be handled by AdjustStageExecutionMode rule.
+      taggedPlan
+    }
   }
 }
 
@@ -81,17 +78,67 @@ object CudfNodeValidationRule {
     }
   }
 
-  def createGPUColumnarExchange(shuffle: ColumnarShuffleExchangeExec): SparkPlan = {
-    val exec = GPUColumnarShuffleExchangeExec(
-      shuffle.outputPartitioning,
-      shuffle.child,
-      shuffle.shuffleOrigin,
-      shuffle.projectOutputAttributes,
-      shuffle.advisoryPartitionSize)
-    val res = exec.doValidate()
-    if (!res.ok()) {
-      throw new GlutenNotSupportException(res.reason())
+  // supportsCudf is the first parent WholeStageTransformer's `isCudf` for plan.
+  // For WholeStageTransformer, it calls the child's setStageExecutionModeForShuffle with its
+  // `isCudf` for the child shuffle reader,
+  // and returns its `isCudf` for the parent shuffle writer.
+  def setStageExecutionModeForShuffle(
+      plan: SparkPlan,
+      supportsCudf: Boolean): (SparkPlan, Boolean) = {
+    def getStageExecutionMode(supportsCudf: Boolean): StageExecutionMode = {
+      if (supportsCudf) {
+        if (SparkTestUtil.isTesting) {
+          MockGPUStageMode
+        } else {
+          GPUStageMode
+        }
+      } else {
+        CPUStageMode
+      }
+    }
+
+    plan match {
+      case shuffle @ ColumnarShuffleExchangeExec(
+            _,
+            VeloxResizeBatchesExec(child, _),
+            _,
+            _,
+            _,
+            _,
+            _) =>
+        // `supportsCudf` is not decided yet for shuffle writer.
+        val (newChild, mapperStageSupportsCudf) =
+          setStageExecutionModeForShuffle(child, false)
+        val mapperStageMode = getStageExecutionMode(mapperStageSupportsCudf)
+        val reducerStageMode = getStageExecutionMode(supportsCudf)
+        (
+          shuffle.copy(
+            child = VeloxResizeBatchesExec(newChild, Some(mapperStageMode)),
+            mapperStageMode = Some(mapperStageMode),
+            reducerStageMode = Some(reducerStageMode)),
+          supportsCudf)
+      case shuffle: ColumnarShuffleExchangeExec =>
+        val (newChild, mapperStageSupportsCudf) =
+          setStageExecutionModeForShuffle(shuffle.child, false)
+        val mapperStageMode = getStageExecutionMode(mapperStageSupportsCudf)
+        val reducerStageMode = getStageExecutionMode(supportsCudf)
+        (
+          shuffle.copy(
+            child = newChild,
+            mapperStageMode = Some(mapperStageMode),
+            reducerStageMode = Some(reducerStageMode)),
+          supportsCudf)
+      case resizeBatches: VeloxResizeBatchesExec =>
+        val (newChild, _) =
+          setStageExecutionModeForShuffle(resizeBatches.child, supportsCudf)
+        (VeloxResizeBatchesExec(newChild, Some(getStageExecutionMode(supportsCudf))), supportsCudf)
+      case wst: WholeStageTransformer =>
+        val (newChild, _) = setStageExecutionModeForShuffle(wst.child, wst.isCudf)
+        (wst.withNewChildren(Seq(newChild)), wst.isCudf)
+      case other =>
+        val newChildren =
+          other.children.map(child => setStageExecutionModeForShuffle(child, supportsCudf)._1)
+        (other.withNewChildren(newChildren), supportsCudf)
     }
-    exec
   }
 }