[CINN] Add the TileTransposeTactic (#70942)

Author: lshpku

paddle/cinn/backends/codegen_gpu_dev.cc

@@ -95,12 +95,23 @@ std::vector<ir::stmt::StmtRef> CodeGenGpuDev::GenerateBufferAliasStmts(
   }
 
   for (auto &t : unique_tensors) {
-    auto data_type = t->type();
-    auto data_ptr_type = data_type;
-    data_ptr_type.set_cpp_handle();
+    auto tensor_type = t->type();
+    auto tensor_ptr_type = tensor_type;
+    tensor_ptr_type.set_cpp_handle();
+
+    auto buffer_type = t->buffer->dtype;
+    auto buffer_ptr_type = buffer_type;
+    buffer_ptr_type.set_cpp_handle();
+
+    Expr t_var = Var(t->name, tensor_ptr_type);
+    Expr buf_var = Var(t->buffer->name, buffer_ptr_type);
+
+    // A tensor and its buffer may have different types when multiple tensors
+    // share the same buffer. In this case, add a Cast before aliasing.
+    if (tensor_type != buffer_type) {
+      buf_var = common::cast(buf_var, tensor_ptr_type);
+    }
 
-    Var t_var(t->name, data_ptr_type);
-    Var buf_var(t->buffer->name, data_ptr_type);
     buffer_alias.push_back(ir::stmt::Let(t_var, buf_var));
   }
 

paddle/cinn/ir/group_schedule/config/group_tile_config.cc

@@ -640,33 +640,20 @@ TileConfigMap BuildStaticReduceConfig(
                                    /* spatial_inner_num = */ 1,
                                    /* vectorize_factor = */ 1,
                                    NoneReduceMethod()};
-    BucketInfo bucket_info__1024_1M{/* sp_lower_bound = */ 1024,
-                                    /* sp_upper_bound = */ 1024 * 1024 - 1,
-                                    /* rb_lower_bound = */ 1,
-                                    /* rb_upper_bound = */ 1,
-                                    /* sp_is_dynamic = */ true,
-                                    /* rb_is_dynamic = */ false};
-    TileConfig tile_config__1024_1M{/* warp_num = */ 32,
-                                    /* tree_reduce_num = */ 1,
-                                    /* grid_reduce_num = */ 1,
-                                    /* spatial_inner_num = */ 4,
-                                    /* vectorize_factor = */ 1,
-                                    NoneReduceMethod()};
-    BucketInfo bucket_info__1M_INF{/* sp_lower_bound = */ 1024 * 1024,
-                                   /* sp_upper_bound = */ kMaxNumel,
-                                   /* rb_lower_bound = */ 1,
-                                   /* rb_upper_bound = */ 1,
-                                   /* sp_is_dynamic = */ true,
-                                   /* rb_is_dynamic = */ false};
-    TileConfig tile_config__1M_INF{/* warp_num = */ 32,
-                                   /* tree_reduce_num = */ 1,
-                                   /* grid_reduce_num = */ 1,
-                                   /* spatial_inner_num = */ 4,
-                                   /* vectorize_factor = */ 1,
-                                   NoneReduceMethod()};
+    BucketInfo bucket_info__1024_INF{/* sp_lower_bound = */ 1024,
+                                     /* sp_upper_bound = */ kMaxNumel,
+                                     /* rb_lower_bound = */ 1,
+                                     /* rb_upper_bound = */ 1,
+                                     /* sp_is_dynamic = */ true,
+                                     /* rb_is_dynamic = */ false};
+    TileConfig tile_config__1024_INF{/* warp_num = */ 32,
+                                     /* tree_reduce_num = */ 1,
+                                     /* grid_reduce_num = */ 1,
+                                     /* spatial_inner_num = */ 4,
+                                     /* vectorize_factor = */ 1,
+                                     NoneReduceMethod()};
     return {{bucket_info__1_1023, tile_config__1_1023},
-            {bucket_info__1024_1M, tile_config__1024_1M},
-            {bucket_info__1M_INF, tile_config__1M_INF}};
+            {bucket_info__1024_INF, tile_config__1024_INF}};
   } else if (base_info->reduce_numel <= 256) {
     BucketInfo bucket_info{/* sp_lower_bound = */ 1,
                            /* sp_upper_bound = */ kMaxNumel,

paddle/cinn/ir/group_schedule/dy_shape_group_scheduler.cc

@@ -21,6 +21,7 @@
 #include "paddle/cinn/ir/group_schedule/tactic/compute_inline_tactic.h"
 #include "paddle/cinn/ir/group_schedule/tactic/tile_broadcast_tactic.h"
 #include "paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.h"
+#include "paddle/cinn/ir/group_schedule/tactic/tile_transpose_tactic.h"
 #include "paddle/cinn/ir/ir_analyzer/ir_analyzer.h"
 #include "paddle/cinn/ir/op/ir_operators.h"
 #include "paddle/common/enforce.h"
@@ -36,6 +37,7 @@ void DynamicShapeGroupScheduler::Init() {
   InitBuckets();
   tactics_.emplace_back(CreateAlignIterSpaceTactic());
   tactics_.emplace_back(CreateTileBroadcastTactic());
+  tactics_.emplace_back(CreateTileTransposeTactic());
   tactics_.emplace_back(CreateTileFirstGeneralTactic());
   tactics_.emplace_back(CreateComputeInlineTactic());
   tactics_.emplace_back(CreateComputeAtReductionTactic());

paddle/cinn/ir/group_schedule/tactic/CMakeLists.txt

@@ -8,4 +8,5 @@ gather_srcs(cinnapi_src SRCS compute_at_reduction_tactic.cc)
 gather_srcs(cinnapi_src SRCS bind_cuda_tactic.cc)
 gather_srcs(cinnapi_src SRCS arrange_storage_tactic.cc)
 gather_srcs(cinnapi_src SRCS tile_broadcast_tactic.cc)
+gather_srcs(cinnapi_src SRCS tile_transpose_tactic.cc)
 gather_srcs(cinnapi_src SRCS tile_first_general_tactic.cc)