CINN CustomDevice. Bugfix tile_transpose_tactic.cc analysis_predictor.cc (#78512)

Author: YuhanXu

paddle/cinn/ir/group_schedule/tactic/tile_transpose_tactic.cc

@@ -553,15 +553,29 @@ void TileTransposeTactic::TileCacheBlock(ir::IRSchedule* sch,
 
   // Step 3. Do inner-block transpose.
   int offset = high_axis_.size();
-  sch->Split(shared_cache_block_id, offset + 1, {-1, 4, 8});
+#ifdef CINN_WITH_CUSTOM_DEVICE
+  sch->Split(
+      shared_cache_block_id,
+      offset + 1,
+      {-1, static_cast<int>(context_->config.tile_config.warp_size / 8), 8});
   sch->Split(shared_cache_block_id,
              offset,
              {-1, static_cast<int>(context_->config.tile_config.warp_size)});
 
   sch->Split(local_cache_block_id,
              offset + 1,
              {-1, static_cast<int>(context_->config.tile_config.warp_size)});
+  sch->Split(
+      local_cache_block_id,
+      offset,
+      {-1, static_cast<int>(context_->config.tile_config.warp_size / 8), 8});
+#else  // CINN_WITH_CUDA
+  sch->Split(shared_cache_block_id, offset + 1, {-1, 4, 8});
+  sch->Split(shared_cache_block_id, offset, {-1, 32});
+
+  sch->Split(local_cache_block_id, offset + 1, {-1, 32});
   sch->Split(local_cache_block_id, offset, {-1, 4, 8});
+#endif
 
   sch->Reorder(shared_cache_block_id, OffsetVec({0, 2, 3, 4, 1}, offset));
   sch->Reorder(local_cache_block_id, OffsetVec({0, 3, 1, 2, 4}, offset));
@@ -576,10 +590,18 @@ void TileTransposeTactic::TileBlock(ir::IRSchedule* sch,
   CanonicalizeLayout(sch, block_id);
 
   int offset = high_axis_.size();
+#ifdef CINN_WITH_CUSTOM_DEVICE
   sch->Split(block_id,
              offset + 1,
              {-1, static_cast<int>(context_->config.tile_config.warp_size)});
+  sch->Split(
+      block_id,
+      offset,
+      {-1, static_cast<int>(context_->config.tile_config.warp_size / 8), 8});
+#else  // CINN_WITH_CUDA
+  sch->Split(block_id, offset + 1, {-1, 32});
   sch->Split(block_id, offset, {-1, 4, 8});
+#endif
 
   sch->Reorder(block_id, OffsetVec({0, 3, 1, 2, 4}, offset));
 

paddle/fluid/inference/api/analysis_predictor.cc

@@ -924,7 +924,8 @@ void AnalysisPredictor::OptimizeInferencePirProgram() {
       delete_assert_op_pm.Run(pir_program_.get());
     }
 
-    if (config_.use_gpu() && config_.cinn_enabled()) {
+    if ((config_.use_gpu() || config_.use_custom_device()) &&
+        config_.cinn_enabled()) {
       if (!config_.custom_pass_only_) {
         ::pir::PassManager fused_op_pm(::pir::IrContext::Instance(),
                                        config_.pm_opt_level_);