get strides of loop variables (#66742)

lshpku · web-flow · commit e9794e746384 · 2024-08-06T11:12:26.000+08:00
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
@@ -95,7 +95,7 @@ std::shared_ptr<GroupInfo> OpLowererImpl::GetGroupInfo(
     const std::unordered_map<::pir::Value, ir::Tensor>& tensor_map) {
   std::shared_ptr<GroupInfo> group_info = std::make_shared<GroupInfo>();
   group_info->data_space = fusion_group_info.loop_ranges;
-  group_info->loop_transform_map = fusion_group_info.loop_transform_map;
+  group_info->loop_strides = fusion_group_info.loop_strides;
   group_info->reduce_axis = fusion_group_info.reduce_axis;
   group_info->reduce_var_names =
       std::set<std::string>(fusion_group_info.reduce_var_name.begin(),
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.h b/paddle/cinn/hlir/framework/pir/op_lowering_impl.h
@@ -50,7 +50,7 @@ typedef bool (OpLowererImpl::*ScheduleDetermineFunction)(::pir::Operation*);
 struct GroupInfo {
   std::vector<int64_t> data_space;
   std::vector<int64_t> reduce_axis;
-  std::vector<int64_t> loop_transform_map;
+  std::vector<int64_t> loop_strides;
   std::set<std::string> reduce_var_names;
   std::set<std::string> shared_var_names;
   std::set<std::string> direct_output_var_names;
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc b/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
@@ -611,53 +611,114 @@ std::vector<ir::Var> GetAllForIters(const ir::Expr& expr) {
 
 }  // namespace trivial_fusion_detail
 
-void InitLoopTransformMap(FusionGroupInfo* group_info, const ir::Expr& body) {
-  using trivial_fusion_detail::ExprSetFinderUtils::ChildScheduleBlockRealizes;
-  using trivial_fusion_detail::ExprSetFinderUtils::ChildTensorLoads;
-  using trivial_fusion_detail::ExprSetFinderUtils::
-      ScheduleBlockRealizeIsSplitTransform;
+struct VarReplacer : public ir::IRMutator<ir::Expr*> {
+  std::unordered_set<ir::Var> iter_vars;
+  ir::Var inspecting_var;
+
+  explicit VarReplacer(const std::vector<ir::Var>& _iter_vars)
+      : iter_vars(_iter_vars.begin(), _iter_vars.end()) {}
+
+  virtual void Visit(const ir::_Var_* op, ir::Expr* expr) {
+    ir::Var var = op->Copy().as_var_ref();
+    if (inspecting_var.defined() && var == inspecting_var) {
+      *expr = ir::Expr(1);
+    } else if (iter_vars.find(var) != iter_vars.end()) {
+      *expr = ir::Expr(0);
+    } else {
+      // We can replace shape variables (e.g. S0) with any constant, and here
+      // we just choose to replace them with 32.
+      *expr = ir::Expr(32);
+    }
+  }
+};
+
+std::vector<int64_t> GetVarStrides(ir::Expr load_offset,
+                                   const std::vector<ir::Var>& iter_vars) {
+  VarReplacer replacer(iter_vars);
+
+  const auto Evaluate = [&](const ir::Var var) {
+    ir::Expr expr = ir::ir_utils::IRCopy(load_offset);
+    replacer.inspecting_var = var;
+    replacer.IRMutator::Visit(&expr, &expr);
+    ir::Expr res = common::AutoSimplify(expr);
+    if (res.is_constant()) {
+      return res.as_int64();
+    }
+    return int64_t(0);
+  };
+
+  const int64_t base = Evaluate(ir::Var());
 
-  std::vector<ir::Expr> split_transform_block =
-      (ChildScheduleBlockRealizes * ScheduleBlockRealizeIsSplitTransform)(body);
-  if (split_transform_block.empty()) {
-    return;
+  std::vector<int64_t> strides;
+  for (const auto& var : iter_vars) {
+    int64_t stride = Evaluate(var) - base;
+    strides.push_back(stride);
   }
-  CHECK_EQ(split_transform_block.size(), 1);
-  VLOG(4) << "found split_transform block: " << body;
+  return strides;
+}
+
+ir::Expr GetLargestLoad(const std::vector<ir::Expr>& exprs) {
+  common::cas_intervals_t var_intervals =
+      common::CollectVarIntervalsOfExprs(exprs);
+  common::SymbolicExprAnalyzer symbolic_expr_analyzer(var_intervals);
 
-  auto* block = split_transform_block[0].As<ir::ScheduleBlockRealize>();
+  const auto GetLoadSize = [](const ir::Expr& expr) {
+    auto* load = expr.As<ir::Load>();
+    auto* tensor = load->tensor.As<ir::_Tensor_>();
+    if (tensor->shape.size() == 0) {
+      return ir::Expr(1);
+    }
+    ir::Expr size = tensor->shape[0];
+    for (size_t i = 1; i < tensor->shape.size(); i++) {
+      size = size * tensor->shape[i];
+    }
+    return common::AutoSimplify(size);
+  };
+
+  ir::Expr res = exprs[0];
+  ir::Expr res_size = GetLoadSize(res);
+  for (size_t i = 1; i < exprs.size(); i++) {
+    ir::Expr cur_size = GetLoadSize(exprs[i]);
+    std::optional<bool> gt = symbolic_expr_analyzer.ProveGT(cur_size, res_size);
+    if (gt.has_value() && gt.value()) {
+      res = exprs[i];
+      res_size = cur_size;
+    }
+  }
+  return res;
+}
+
+std::vector<int64_t> GetLoopStrides(const ir::Expr& body,
+                                    const ir::Expr& expr_block) {
+  using trivial_fusion_detail::ExprSetFinderUtils::ChildTensorLoads;
+
+  auto* block = expr_block.As<ir::ScheduleBlockRealize>();
   auto& iter_values = block->iter_values;
   auto& iter_vars = block->schedule_block.As<ir::ScheduleBlock>()->iter_vars;
   CHECK_EQ(iter_values.size(), iter_vars.size());
+  const std::vector<ir::Var> for_iters =
+      trivial_fusion_detail::GetAllForIters(body);
 
-  // Map iter_vars (e.g. inner_block_0) to loops.
-  std::unordered_map<ir::Var, size_t> iter_var_map;
-  std::vector<ir::Var> for_iters = trivial_fusion_detail::GetAllForIters(body);
-  for (size_t i = 0; i < iter_vars.size(); i++) {
-    auto iter = std::find(
-        for_iters.begin(), for_iters.end(), iter_values[i].as_var_ref());
+  const auto GetLoopIndex = [&](size_t var_index) {
+    auto iter = std::find(for_iters.begin(),
+                          for_iters.end(),
+                          iter_values[var_index].as_var_ref());
     CHECK(iter != for_iters.end());
-    iter_var_map[iter_vars[i]] = std::distance(for_iters.begin(), iter);
-  }
+    return std::distance(for_iters.begin(), iter);
+  };
 
-  // Map loops to the load indices of the first input.
-  // TODO(liangshuhao): Choose the most appropriate input when there are
-  // multiple inputs.
-  const auto& all_loads = ChildTensorLoads(split_transform_block[0]);
-  group_info->loop_transform_map.assign(for_iters.size(), -1);
-  for (const auto& load : all_loads) {
-    const auto& indices = load.As<ir::Load>()->indices;
-    VLOG(4) << "load indices: " << utils::Join(indices, ", ");
-    for (size_t i = 0; i < indices.size(); i++) {
-      if (indices[i].as_var()) {
-        size_t loop_idx = iter_var_map[indices[i].as_var_ref()];
-        group_info->loop_transform_map[loop_idx] = i;
-      }
-    }
-    VLOG(4) << "loop_transform_map: "
-            << utils::Join(group_info->loop_transform_map, ", ");
-    break;
+  const auto& all_loads = ChildTensorLoads(expr_block);
+  std::vector<int64_t> loop_strides(for_iters.size());
+  if (all_loads.empty()) {
+    return loop_strides;
   }
+  const ir::Expr largest_load = GetLargestLoad(all_loads);
+  ir::Expr load_offset = largest_load.As<ir::Load>()->index();
+  std::vector<int64_t> var_strides = GetVarStrides(load_offset, iter_vars);
+  for (size_t i = 0; i < iter_vars.size(); i++) {
+    loop_strides[GetLoopIndex(i)] = var_strides[i];
+  }
+  return loop_strides;
 }
 
 FusionGroupInfo GetFusionGroupInfo(
@@ -667,11 +728,24 @@ FusionGroupInfo GetFusionGroupInfo(
   using trivial_fusion_detail::IsReduceBody;
   using trivial_fusion_detail::ReduceOp;
   using trivial_fusion_detail::ComposeUtils::ConcatVector;
+  using trivial_fusion_detail::ExprSetFinderUtils::ChildScheduleBlockRealizes;
+  using trivial_fusion_detail::ExprSetFinderUtils::ScheduleBlockRealizeIsInit;
+  using trivial_fusion_detail::ExprSetFinderUtils::
+      ScheduleBlockRealizeIsSplitTransform;
 
   FusionGroupInfo group_info = FusionGroupInfo();
 
+  const auto GetSplitTransformBlock = [](const ir::Expr& expr_body) {
+    return (ChildScheduleBlockRealizes *
+            ScheduleBlockRealizeIsSplitTransform)(expr_body);
+  };
+
   for (const auto& body : op_compute_bodies) {
-    InitLoopTransformMap(&group_info, body);
+    std::vector<ir::Expr> split_transform_block = GetSplitTransformBlock(body);
+    if (!split_transform_block.empty()) {
+      CHECK_EQ(split_transform_block.size(), 1);
+      group_info.loop_strides = GetLoopStrides(body, split_transform_block[0]);
+    }
 
     if (IsReduceBody(body)) {
       ReduceOp op = ReduceOp(body);
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_impl.h b/paddle/cinn/hlir/framework/pir/trivial_op_impl.h
@@ -165,12 +165,13 @@ std::vector<ir::Var> GetAllForIters(const ir::Expr& expr);
 
 struct FusionGroupInfo {
   std::vector<int64_t> loop_ranges;
-  std::vector<int64_t> loop_transform_map;
+  std::vector<int64_t> loop_strides;
   std::vector<int64_t> reduce_axis;
   std::vector<std::string> reduce_var_name;
 
   std::string DebugPrint() {
     return "GroupInfo\nloop_ranges: " + cinn::utils::Join(loop_ranges, " ") +
+           "\nloop_strides: " + cinn::utils::Join(loop_strides, ", ") +
            "\nreduce_axis: " + cinn::utils::Join(reduce_axis, " ") +
            "\nreduce_var_name: " + cinn::utils::Join(reduce_var_name, " ");
   }
diff --git a/paddle/cinn/ir/group_schedule/config/group_tile_config.cc b/paddle/cinn/ir/group_schedule/config/group_tile_config.cc
@@ -97,7 +97,7 @@ std::shared_ptr<ScheduleConfig::BaseInfo> InitBasicInfo(
   base_info->shared_var_names = group_info->shared_var_names;
   base_info->direct_output_var_names = group_info->direct_output_var_names;
   base_info->data_rank = group_info->data_space.size();
-  base_info->loop_transform_map = group_info->loop_transform_map;
+  base_info->loop_strides = group_info->loop_strides;
 
   std::set<int64_t> reduce_dim_loc;
   for (int64_t dim : group_info->reduce_axis) {
diff --git a/paddle/cinn/ir/group_schedule/config/group_tile_config.h b/paddle/cinn/ir/group_schedule/config/group_tile_config.h
@@ -31,7 +31,7 @@ using IterSpaceType = std::vector<std::pair<std::string, std::string>>;
 struct ScheduleConfig {
   struct BaseInfo {
     std::vector<int64_t> reduce_axis;
-    std::vector<int64_t> loop_transform_map;
+    std::vector<int64_t> loop_strides;
     int64_t data_rank;
     int64_t reduce_numel;
     int64_t spatial_numel;
diff --git a/paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.cc b/paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.cc
@@ -46,14 +46,25 @@ bool IsWarpReduce(const ScheduleConfig& config) {
 }
 
 bool UseContinuousDataTile(const ScheduleConfig& config) {
+  if (config.base_info->reduce_axis.empty()) {
+    return true;
+  }
+  int64_t min_stride = INT_MAX;
+  int64_t min_reduce_stride = INT_MAX;
   int64_t last_axis = 0;
   int64_t last_reduce_axis = 0;
-  for (auto axis : config.base_info->loop_transform_map) {
-    last_axis = std::max(axis, last_axis);
+  for (size_t i = 0; i < config.base_info->loop_strides.size(); i++) {
+    if (config.base_info->loop_strides[i] < min_stride &&
+        config.base_info->loop_strides[i] != 0) {
+      min_stride = config.base_info->loop_strides[i];
+      last_axis = i;
+    }
   }
-  for (auto axis : config.base_info->reduce_axis) {
-    last_reduce_axis =
-        std::max(config.base_info->loop_transform_map[axis], last_reduce_axis);
+  for (int64_t axis : config.base_info->reduce_axis) {
+    if (config.base_info->loop_strides[axis] < min_reduce_stride) {
+      min_reduce_stride = config.base_info->loop_strides[axis];
+      last_reduce_axis = axis;
+    }
   }
   return last_axis == last_reduce_axis;
 }
@@ -113,11 +124,12 @@ void TileFirstGeneralTactic::Init(ScheduleContext* context) {
   vec_spatial_axis_last_.clear();
 
   if (!context_->config.base_info->reduce_axis.empty()) {
-    int64_t first_reduce_axis = context_->config.base_info->data_rank - 1;
+    int64_t first_reduce_axis = context_->config.base_info->reduce_axis.front();
     for (auto axis : context_->config.base_info->reduce_axis) {
-      first_reduce_axis =
-          std::min(context_->config.base_info->loop_transform_map[axis],
-                   first_reduce_axis);
+      if (context->config.base_info->loop_strides[axis] >
+          context->config.base_info->loop_strides[first_reduce_axis]) {
+        first_reduce_axis = axis;
+      }
     }
     for (int32_t i = 0; i < reduce_start_idx; ++i) {
       if (i < first_reduce_axis) {
@@ -281,8 +293,8 @@ void TileFirstGeneralTactic::ApplyContinuousDataTile(
 
 void TileFirstGeneralTactic::AlignToReduceInput(ir::IRSchedule* sch,
                                                 const std::string& block_id) {
-  auto& loop_transform_map = context_->config.base_info->loop_transform_map;
-  if (loop_transform_map.empty()) {
+  const auto& loop_strides = context_->config.base_info->loop_strides;
+  if (loop_strides.empty()) {
     return;
   }
 
@@ -298,7 +310,7 @@ void TileFirstGeneralTactic::AlignToReduceInput(ir::IRSchedule* sch,
 
   std::sort(loop_perm.begin(), loop_perm.end(), [&](int64_t a, int64_t b) {
     if (IsReduce(a) == IsReduce(b)) {
-      return loop_transform_map[a] < loop_transform_map[b];
+      return loop_strides[a] > loop_strides[b];
     }
     return IsReduce(b);
   });
@@ -309,7 +321,7 @@ void TileFirstGeneralTactic::AlignToReduceInput(ir::IRSchedule* sch,
   for (auto i : loop_perm) {
     if (IsReduce(i)) {
       rd_loops.push_back(loops[i]);
-    } else if (loop_transform_map[i] != -1) {
+    } else if (loop_strides[i] != 0) {
       sp_loops.push_back(loops[i]);
     }
   }
