PaddlePaddle
diff --git a/‎paddle/cinn/ir/group_schedule/config/database.cc
Copy file name to clipboard
+7-1Lines changed: 7 additions & 1 deletion b/‎paddle/cinn/ir/group_schedule/config/database.cc
Copy file name to clipboard
+7-1Lines changed: 7 additions & 1 deletion
diff --git a/‎paddle/cinn/ir/group_schedule/config/database.h
Copy file name to clipboard
-2Lines changed: 0 additions & 2 deletions b/‎paddle/cinn/ir/group_schedule/config/database.h
Copy file name to clipboard
-2Lines changed: 0 additions & 2 deletions
diff --git a/‎paddle/cinn/ir/group_schedule/config/filedatabase.cc
Copy file name to clipboard
+34-29Lines changed: 34 additions & 29 deletions b/‎paddle/cinn/ir/group_schedule/config/filedatabase.cc
Copy file name to clipboard
+34-29Lines changed: 34 additions & 29 deletions
diff --git a/‎paddle/cinn/ir/group_schedule/config/filedatabase.h
Copy file name to clipboard
+1-4Lines changed: 1 addition & 4 deletions b/‎paddle/cinn/ir/group_schedule/config/filedatabase.h
Copy file name to clipboard
+1-4Lines changed: 1 addition & 4 deletions
diff --git a/‎paddle/cinn/ir/group_schedule/config/group_tile_config.cc
Copy file name to clipboard
+74-11Lines changed: 74 additions & 11 deletions b/‎paddle/cinn/ir/group_schedule/config/group_tile_config.cc
Copy file name to clipboard
+74-11Lines changed: 74 additions & 11 deletions
@@ -19,10 +19,16 @@ namespace ir {
 
 void NaiveTileConfigDatabase::AddConfig(
     const common::Target& target,
-    const IterSpaceType& iter_space_type,
     const BucketInfo& bucket_info,
     const ScheduleConfig::TileConfig& config,
     int priority) {
+  IterSpaceType iter_space_type = [&] {
+    std::vector<std::pair<std::string, std::string>> res;
+    for (const auto& dim : bucket_info.space) {
+      res.emplace_back(dim.iter_type, (dim.is_dynamic ? "dynamic" : "static"));
+    }
+    return res;
+  }();
   config_map_[iter_space_type][bucket_info] = config;
 }
 
 
@@ -32,7 +32,6 @@ using IterSpaceType = std::vector<std::pair<std::string, std::string>>;
 class TileConfigDatabase {
  public:
   virtual void AddConfig(const common::Target& target,
-                         const IterSpaceType& iter_space_type,
                          const BucketInfo& bucket_info,
                          const ScheduleConfig::TileConfig& config,
                          int priority) = 0;
@@ -45,7 +44,6 @@ class TileConfigDatabase {
 class NaiveTileConfigDatabase final : public TileConfigDatabase {
  public:
   void AddConfig(const common::Target& target,
-                 const IterSpaceType& iter_space_type,
                  const BucketInfo& bucket_info,
                  const ScheduleConfig::TileConfig& config,
                  int priority = 1) override;
 
@@ -39,22 +39,19 @@ namespace ir {
 
 bool TileConfigToProto(group_schedule::config::proto::TileData* tile_data,
                        const TileConfigMap& tile_config_map,
-                       const IterSpaceType& iter_space_type,
                        const int& priority) {
   for (auto& it : tile_config_map) {
-    group_schedule::config::proto::Dimension s_dimension, r_dimension;
-
     // prepare key---convert bucket info to proto::bucket_info
-    s_dimension.set_lower_bound(it.first.sp_lower_bound);
-    s_dimension.set_upper_bound(it.first.sp_upper_bound);
-    s_dimension.set_iter_type(iter_space_type[0].first);
-    s_dimension.set_is_dynamic(iter_space_type[0].second == "dynamic");
-    r_dimension.set_lower_bound(it.first.rb_lower_bound);
-    r_dimension.set_upper_bound(it.first.rb_upper_bound);
-    r_dimension.set_iter_type(iter_space_type[1].first);
-    r_dimension.set_is_dynamic(iter_space_type[1].second == "dynamic");
-    *(tile_data->mutable_bucket_info()->add_dimension()) = s_dimension;
-    *(tile_data->mutable_bucket_info()->add_dimension()) = r_dimension;
+    BucketInfo bucket_info = it.first;
+    int dims = bucket_info.space.size();
+    for (int i = 0; i < dims; i++) {
+      group_schedule::config::proto::Dimension cur_dimension;
+      cur_dimension.set_lower_bound(bucket_info.space[i].lower_bound);
+      cur_dimension.set_upper_bound(bucket_info.space[i].upper_bound);
+      cur_dimension.set_iter_type(bucket_info.space[i].iter_type);
+      cur_dimension.set_is_dynamic(bucket_info.space[i].is_dynamic);
+      *(tile_data->mutable_bucket_info()->add_dimension()) = cur_dimension;
+    }
 
     // prepare value---transfer tile_config to proto::tile_config
     group_schedule::config::proto::TileConfig tc;
@@ -114,18 +111,24 @@ std::string IterSpaceTypeToDir(const common::Target target,
 }
 
 bool FileTileConfigDatabase::Tofile(const common::Target& target,
-                                    const IterSpaceType& iter_space_type,
                                     int priority) {
   // Step1. To proto
   TileConfigMap& tile_config_map = target_config_data_;
   group_schedule::config::proto::TileData tile_data;
-  auto is_success =
-      TileConfigToProto(&tile_data, tile_config_map, iter_space_type, priority);
+  auto is_success = TileConfigToProto(&tile_data, tile_config_map, priority);
   if (is_success == false) {
     PADDLE_THROW(::common::errors::Unavailable(
         "Can't convert tile_config_map to its proto message."));
   }
   // Step2. ToJson
+  IterSpaceType iter_space_type = [&] {
+    std::vector<std::pair<std::string, std::string>> res;
+    auto bucket_info = tile_config_map.begin()->first;
+    for (const auto& dim : bucket_info.space) {
+      res.emplace_back(dim.iter_type, (dim.is_dynamic ? "dynamic" : "static"));
+    }
+    return res;
+  }();
   std::string dump_path = IterSpaceTypeToDir(target, iter_space_type);
   size_t length = tile_config_map.size();
   std::vector<std::string> json_lines(length);
@@ -187,7 +190,7 @@ bool comparepriority(group_schedule::config::proto::TileData tile_data1,
 
 TileConfigMap FileTileConfigDatabase::GetConfigs(
     const common::Target& target, const IterSpaceType& iter_space_type) const {
-  // Step1. ReadFromJsonFile->Message;
+  // Step 1: Read from json file and convert json to proto message
   std::string file_path = IterSpaceTypeToDir(target, iter_space_type);
   auto json_lines = ReadLinesFromFile(file_path);
   size_t line_length = json_lines.size();
@@ -196,39 +199,41 @@ TileConfigMap FileTileConfigDatabase::GetConfigs(
       line_length);
   JsonStringToMessageOfTileConfig(&tile_database, json_lines);
 
-  // Step2. ParseFromProtoMessage();
+  // Step 2: Parse from proto message
   TileConfigMap tile_config_map;
   // order tile_database according to priority
   std::sort(tile_database.begin(), tile_database.end(), comparepriority);
   for (const auto& piece_tileconfig : tile_database) {
     group_schedule::config::proto::BucketInfo its =
         piece_tileconfig.bucket_info();
-    // proto::BucketInfo to  bucketinfo
-    BucketInfo bucket_info;
-    bucket_info.sp_lower_bound = its.dimension(0).lower_bound();
-    bucket_info.sp_upper_bound = its.dimension(0).upper_bound();
-    bucket_info.rb_lower_bound = its.dimension(1).lower_bound();
-    bucket_info.rb_upper_bound = its.dimension(1).upper_bound();
+    //  Step 2.1: Convert proto bucketinfo to source bucketinfo
+    int dims = its.dimension_size();
+    BucketInfo bucket_info(static_cast<size_t>(dims));
+    for (int i = 0; i < dims; i++) {
+      bucket_info.space[i].lower_bound = its.dimension(i).lower_bound();
+      bucket_info.space[i].upper_bound = its.dimension(i).upper_bound();
+      bucket_info.space[i].iter_type = its.dimension(i).iter_type();
+      bucket_info.space[i].is_dynamic = its.dimension(i).is_dynamic();
+    }
+    //  Step 2.2: Convert proto tile_config to source tile_config
     ScheduleConfig::TileConfig tconfig;
     tconfig.tree_reduce_num = piece_tileconfig.tile_config().tree_reduce_num();
     tconfig.spatial_inner_num =
         piece_tileconfig.tile_config().spatial_inner_num();
     tconfig.warp_num = piece_tileconfig.tile_config().warp_num();
     tile_config_map[bucket_info] = tconfig;
-    // Tode[XiaZichao] Add function to cut one lattice into smaller ones.
+    // TODO(XiaZichao): Add function to cut one lattice into smaller ones
   }
-  // ToDo[XiaZichao] update json file using top view of tileconfigMap
+  // TODO(XiaZichao): update json file using top view of tileconfigMap
   return tile_config_map;
 }
 
 void FileTileConfigDatabase::AddConfig(const common::Target& target,
-                                       const IterSpaceType& iter_space_type,
                                        const BucketInfo& bucket_info,
                                        const ScheduleConfig::TileConfig& config,
                                        int priority) {
   target_config_data_[bucket_info] = config;
-  auto status =
-      FileTileConfigDatabase::Tofile(target, iter_space_type, priority);
+  auto status = FileTileConfigDatabase::Tofile(target, priority);
   if (status == true) {
     target_config_data_.clear();
     return;
 
@@ -22,7 +22,6 @@ namespace ir {
 class FileTileConfigDatabase : TileConfigDatabase {
  public:
   void AddConfig(const common::Target& target,
-                 const IterSpaceType& iter_space_type,
                  const BucketInfo& bucket_info,
                  const ScheduleConfig::TileConfig& config,
                  int priority) override;
@@ -31,9 +30,7 @@ class FileTileConfigDatabase : TileConfigDatabase {
 
  private:
   TileConfigMap target_config_data_;
-  bool Tofile(const common::Target& target,
-              const IterSpaceType& iter_space_type,
-              int priority);
+  bool Tofile(const common::Target& target, int priority);
 };
 
 }  // namespace ir
 
@@ -20,6 +20,47 @@ namespace ir {
 
 const int kMaxNumel = INT32_MAX;
 
+BucketInfo::BucketInfo(int sp_lower_bound,
+                       int sp_upper_bound,
+                       int rb_lower_bound,
+                       int rb_upper_bound,
+                       bool sp_is_dynamic = false,
+                       bool rb_is_dynamic = false) {
+  BucketInfo::Dimension sp_dimension(
+      sp_lower_bound, sp_upper_bound, "S", sp_is_dynamic);
+  BucketInfo::Dimension rb_dimension(
+      rb_lower_bound, rb_upper_bound, "R", rb_is_dynamic);
+  this->space.push_back(sp_dimension);
+  this->space.push_back(rb_dimension);
+}
+
+bool BucketInfo::operator==(const BucketInfo& other) const {
+  if (this->space.size() != other.space.size()) {
+    return false;
+  }
+  int length = this->space.size();
+  for (int i = 0; i < length; i++) {
+    if (this->space[i].is_dynamic != other.space[i].is_dynamic ||
+        this->space[i].iter_type != other.space[i].iter_type ||
+        this->space[i].lower_bound != other.space[i].lower_bound ||
+        this->space[i].upper_bound != other.space[i].upper_bound) {
+      return false;
+    }
+  }
+  return true;
+}
+
+std::string BucketInfo::ToString() const {
+  std::stringstream ss;
+  ss << "BucketInfo: [";
+  for (const auto& dim : space) {
+    ss << dim.iter_type << "(" << dim.lower_bound << " - " << dim.upper_bound
+       << "), ";
+  }
+  ss << "]";
+  return ss.str();
+}
+
 int64_t Next2Power(int64_t n) {
   if (n == 1) {
     return 1;
@@ -190,7 +231,9 @@ BuildStaticSpatialConfig(
     BucketInfo bucket_info{/* sp_lower_bound = */ 1,
                            /* sp_upper_bound = */ 1,
                            /* rb_lower_bound = */ 1,
-                           /* rb_upper_bound = */ kMaxNumel};
+                           /* rb_upper_bound = */ kMaxNumel,
+                           /* sp_is_dynamic = */ false,
+                           /* rb_is_dynamic = */ true};
     ScheduleConfig::TileConfig tile_config{
         /* warp_num = */ 8,
         /* tree_reduce_num = */ 256,
@@ -201,7 +244,9 @@ BuildStaticSpatialConfig(
     BucketInfo bucket_info_1_256{/* sp_lower_bound = */ 1,
                                  /* sp_upper_bound = */ kMaxNumel,
                                  /* rb_lower_bound = */ 1,
-                                 /* rb_upper_bound = */ 256};
+                                 /* rb_upper_bound = */ 256,
+                                 /* sp_is_dynamic = */ false,
+                                 /* rb_is_dynamic = */ true};
     ScheduleConfig::TileConfig tile_config_1_256{
         /* warp_num = */ 8,
         /* tree_reduce_num = */ 32,
@@ -211,7 +256,9 @@ BuildStaticSpatialConfig(
     BucketInfo bucket_info_257_2048{/* sp_lower_bound = */ 1,
                                     /* sp_upper_bound = */ kMaxNumel,
                                     /* rb_lower_bound = */ 257,
-                                    /* rb_upper_bound = */ 2048};
+                                    /* rb_upper_bound = */ 2048,
+                                    /* sp_is_dynamic = */ false,
+                                    /* rb_is_dynamic = */ true};
     ScheduleConfig::TileConfig tile_config_257_2048{
         /* warp_num = */ 8,
         /* tree_reduce_num = */ 128,
@@ -221,7 +268,9 @@ BuildStaticSpatialConfig(
     BucketInfo bucket_info_2049_INF{/* sp_lower_bound = */ 1,
                                     /* sp_upper_bound = */ kMaxNumel,
                                     /* rb_lower_bound = */ 2049,
-                                    /* rb_upper_bound = */ kMaxNumel};
+                                    /* rb_upper_bound = */ kMaxNumel,
+                                    /* sp_is_dynamic = */ false,
+                                    /* rb_is_dynamic = */ true};
     ScheduleConfig::TileConfig tile_config_2049_INF{
         /* warp_num = */ 8,
         /* tree_reduce_num = */ 256,
@@ -242,7 +291,9 @@ BuildStaticReduceConfig(
     BucketInfo bucket_info__1_1023{/* sp_lower_bound = */ 1,
                                    /* sp_upper_bound = */ 1023,
                                    /* rb_lower_bound = */ 1,
-                                   /* rb_upper_bound = */ 1};
+                                   /* rb_upper_bound = */ 1,
+                                   /* sp_is_dynamic = */ true,
+                                   /* rb_is_dynamic = */ false};
     ScheduleConfig::TileConfig tile_config__1_1023{
         /* warp_num = */ -1,
         /* tree_reduce_num = */ 1,
@@ -251,7 +302,9 @@ BuildStaticReduceConfig(
     BucketInfo bucket_info__1024_1M{/* sp_lower_bound = */ 1024,
                                     /* sp_upper_bound = */ 1024 * 1024 - 1,
                                     /* rb_lower_bound = */ 1,
-                                    /* rb_upper_bound = */ 1};
+                                    /* rb_upper_bound = */ 1,
+                                    /* sp_is_dynamic = */ true,
+                                    /* rb_is_dynamic = */ false};
     ScheduleConfig::TileConfig tile_config__1024_1M{
         /* warp_num = */ 32,
         /* tree_reduce_num = */ 1,
@@ -260,7 +313,9 @@ BuildStaticReduceConfig(
     BucketInfo bucket_info__1M_INF{/* sp_lower_bound = */ 1024 * 1024,
                                    /* sp_upper_bound = */ kMaxNumel,
                                    /* rb_lower_bound = */ 1,
-                                   /* rb_upper_bound = */ 1};
+                                   /* rb_upper_bound = */ 1,
+                                   /* sp_is_dynamic = */ true,
+                                   /* rb_is_dynamic = */ false};
     ScheduleConfig::TileConfig tile_config__1M_INF{
         /* warp_num = */ 32,
         /* tree_reduce_num = */ 1,
@@ -273,7 +328,9 @@ BuildStaticReduceConfig(
     BucketInfo bucket_info{/* sp_lower_bound = */ 1,
                            /* sp_upper_bound = */ kMaxNumel,
                            /* rb_lower_bound = */ 2,
-                           /* rb_upper_bound = */ 256};
+                           /* rb_upper_bound = */ 256,
+                           /* sp_is_dynamic = */ true,
+                           /* rb_is_dynamic = */ false};
     ScheduleConfig::TileConfig tile_config{
         /* warp_num = */ 8,
         /* tree_reduce_num = */ 32,
@@ -290,7 +347,9 @@ BuildStaticReduceConfig(
     BucketInfo bucket_info{/* sp_lower_bound = */ 1,
                            /* sp_upper_bound = */ kMaxNumel,
                            /* rb_lower_bound = */ 257,
-                           /* rb_upper_bound = */ 2048};
+                           /* rb_upper_bound = */ 2048,
+                           /* sp_is_dynamic = */ true,
+                           /* rb_is_dynamic = */ false};
     ScheduleConfig::TileConfig tile_config{
         /* warp_num = */ warp_num,
         /* tree_reduce_num = */ tree_reduce_num,
@@ -304,7 +363,9 @@ BuildStaticReduceConfig(
     BucketInfo bucket_info{/* sp_lower_bound = */ 1,
                            /* sp_upper_bound = */ kMaxNumel,
                            /* rb_lower_bound = */ 2049,
-                           /* rb_upper_bound = */ kMaxNumel};
+                           /* rb_upper_bound = */ kMaxNumel,
+                           /* sp_is_dynamic = */ true,
+                           /* rb_is_dynamic = */ false};
     ScheduleConfig::TileConfig tile_config{
         /* warp_num = */ warp_num,
         /* tree_reduce_num = */ tree_reduce_num,
@@ -324,7 +385,9 @@ BuildDynamicShapeConfig(
   BucketInfo bucket_info{/* sp_lower_bound = */ 1,
                          /* sp_upper_bound = */ kMaxNumel,
                          /* rb_lower_bound = */ 1,
-                         /* rb_upper_bound = */ kMaxNumel};
+                         /* rb_upper_bound = */ kMaxNumel,
+                         /* sp_is_dynamic = */ true,
+                         /* rb_is_dynamic = */ true};
   ScheduleConfig::TileConfig tile_config{
       /* warp_num = */ warp_num,
       /* tree_reduce_num = */ tree_reduce_num,