[GLUTEN-11524][VL] Support CPU/GPU hybrid task scheduling

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

@@ -144,10 +144,10 @@ object VeloxRuleApi {
     injector.injectFinal(c => RemoveGlutenTableCacheColumnarToRow(c.session))
     injector.injectFinal(
       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 => GlutenAutoAdjustStageResourceProfile(new GlutenConfig(c.sqlConf), c.session))
     injector.injectFinal(c => GlutenFallbackReporter(new GlutenConfig(c.sqlConf), c.session))
     injector.injectFinal(_ => RemoveFallbackTagRule())
   }

backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala

@@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreeNodeTag
 import org.apache.spark.sql.execution.AdjustStageExecutionMode.{adjustExecutionMode, unsetTag}
 import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, ColumnarAQEShuffleReadExec, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.joins.BaseJoinExec
 import org.apache.spark.util.SparkTestUtil
 
 // For ShuffleStage, the resource profile is set to ColumnarShuffleExchangeExec.inputColumnarRDD.
@@ -51,30 +52,19 @@ case class AdjustStageExecutionMode(
   }
 
   private def adjustExecutionModeForGPU(plan: SparkPlan): SparkPlan = {
-    val transformers = plan.collect { case t: WholeStageTransformer => t }
-    if (transformers.isEmpty) {
-      return plan
-    }
-    if (transformers.size > 1) {
-      logWarning(s"Not offloading GPU because multiple WholeStageTransformer exist. Remove tags.")
+    if (!AdjustStageExecutionMode.offloadGpu(plan, glutenConf)) {
       unsetTag(plan, CudfTag.CudfTag)
       return plan
     }
 
-    val transformer = transformers.head
-    if (transformer.isCudf) {
-      val gpuStageMode = if (SparkTestUtil.isTesting) {
-        // Only unset for transformer.
-        transformer.unsetTagValue(CudfTag.CudfTag)
-        MockGPUStageMode
-      } else {
-        GPUStageMode
-      }
-
-      adjustExecutionMode(plan, gpuStageMode)
+    val gpuStageMode = if (SparkTestUtil.isTesting) {
+      unsetTag(plan, CudfTag.CudfTag)
+      MockGPUStageMode
     } else {
-      plan
+      GPUStageMode
     }
+
+    adjustExecutionMode(plan, gpuStageMode)
   }
 }
 
@@ -97,9 +87,9 @@ object AdjustStageExecutionMode extends Logging {
         shuffle
           .copy(mapperStageMode = Some(stageExecutionMode))
           .withNewChildren(Seq(adjustExecutionMode(shuffle.child, stageExecutionMode)))
-      case resizeBatches: VeloxResizeBatchesExec =>
+      case VeloxResizeBatchesExec(child @ AQEShuffleReadExec(_, _), _) =>
         VeloxResizeBatchesExec(
-          adjustExecutionMode(resizeBatches.child, stageExecutionMode),
+          adjustExecutionMode(child, stageExecutionMode),
           Some(stageExecutionMode))
       case _ =>
         plan.withNewChildren(plan.children.map(adjustExecutionMode(_, stageExecutionMode)))
@@ -113,4 +103,33 @@ object AdjustStageExecutionMode extends Logging {
       case _ =>
     }
   }
+
+  def offloadGpu(plan: SparkPlan, glutenConf: GlutenConfig): Boolean = {
+    if (glutenConf.gpuOnlyOffloadJoinStage) {
+      if (!plan.exists(_.isInstanceOf[BaseJoinExec])) {
+        logWarning(s"Not offloading GPU because missing offload condition.")
+        return false
+      }
+    }
+
+    val transformers = plan.collect { case t: WholeStageTransformer => t }
+
+    if (transformers.isEmpty) {
+      logWarning(s"Not offloading GPU because no WholeStageTransformer.")
+      return false
+    }
+
+    if (transformers.size > 1) {
+      // Do not offload GPU if the whole stage is broken down into multiple native pipelines.
+      logWarning(s"Not offloading GPU because multiple WholeStageTransformer exist.")
+      return false
+    }
+
+    if (!transformers.head.isCudf) {
+      logWarning(s"Not offloading GPU because WholeStageTransformer is not tagged cudf.")
+      return false
+    }
+
+    true
+  }
 }

backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala

@@ -113,20 +113,17 @@ class StageExecutionModeSuite extends VeloxWholeStageTransformerSuite {
             s"Expected GPU AQE shuffle reader, but got ${reader.executionMode}")
       }
 
-      val shuffleStages = plan.collect {
-        case stage: ShuffleQueryStageExec => stage
+      val shuffleStages = shuffleReaders.map(_.delegate).map {
+        case Left(aqeReader) =>
+          assert(aqeReader.child.isInstanceOf[ShuffleQueryStageExec])
+          aqeReader.child.asInstanceOf[ShuffleQueryStageExec]
+        case Right(shuffleQueryStageExec: ShuffleQueryStageExec) => shuffleQueryStageExec
       }
 
-      val exchanges = shuffleStages.flatMap {
-        _.plan.collect {
-          case exchange: ColumnarShuffleExchangeExec => exchange
-        }
-      }
-
-      assert(exchanges.nonEmpty)
-
-      exchanges.foreach {
-        exchange =>
+      shuffleStages.foreach {
+        shuffleStage =>
+          assert(shuffleStage.shuffle.isInstanceOf[ColumnarShuffleExchangeExec])
+          val exchange = shuffleStage.shuffle.asInstanceOf[ColumnarShuffleExchangeExec]
           assert(
             !exchange.mapperStageMode.contains(MockGPUStageMode),
             s"Expected CPU mapper stage, but got ${exchange.mapperStageMode}")