[CINN] Fix compute inline with constant index (#69060)

huangjiyi · web-flow · commit 858e877c0495 · 2024-10-31T14:19:12.000+08:00
diff --git a/paddle/cinn/ir/schedule/ir_schedule.cc b/paddle/cinn/ir/schedule/ir_schedule.cc
@@ -93,47 +93,31 @@ void BaseInliner::Visit(const ir::Block* expr, Expr* op) {
   IRMutator::Visit(expr, op);
 }
 
-bool BaseInliner::UpdateAndCheckIndexVars(const std::vector<Expr>& indices,
-                                          int expected_ndim) {
-  int n = indices.size();
-  if (n != expected_ndim) {
-    return false;
-  }
-  std::vector<Var> result;
-  result.reserve(n);
-  for (auto& i : indices) {
-    if (i.as_var()) {
-      result.push_back(i.as_var_ref());
-    } else {
-      return false;
-    }
-  }
-  int n_distinct = std::set<Var, CompVar>(result.begin(), result.end()).size();
-  if (n != n_distinct) {
-    return false;
-  }
-  if (idx_vars_.empty()) {
-    idx_vars_ = std::move(result);
+bool BaseInliner::UpdateAndCheckIndexVars(const std::vector<Expr>& indices) {
+  if (idx_expr_.empty()) {
+    idx_expr_ = std::move(indices);
   } else {
-    if (idx_vars_.size() != result.size()) return false;
-    for (int i = 0; i < result.size(); ++i) {
-      if (Expr(idx_vars_[i]) != Expr(result[i])) return false;
+    if (idx_expr_.size() != indices.size()) return false;
+    for (int i = 0; i < indices.size(); ++i) {
+      if (Expr(idx_expr_[i]) != Expr(indices[i])) return false;
     }
   }
   return true;
 }
 
 void BaseInliner::SetIndexSubstitution(const std::vector<Expr>& indices) {
   PADDLE_ENFORCE_EQ(indices.size(),
-                    idx_vars_.size(),
+                    idx_expr_.size(),
                     ::common::errors::InvalidArgument(
-                        "The size of indices should be equal to idx_vars_"));
-  int n = idx_vars_.size();
-  idx_sub_var_.reserve(n);
-  idx_sub_expr_.reserve(n);
+                        "The size of indices should be equal to idx_expr_"));
+  int n = indices.size();
+  idx_sub_var_.clear();
+  idx_sub_expr_.clear();
   for (int i = 0; i < n; ++i) {
-    idx_sub_var_.push_back(idx_vars_[i]);
-    idx_sub_expr_.push_back(indices[i]);
+    if (idx_expr_[i].is_var()) {
+      idx_sub_var_.push_back(idx_expr_[i].as_var_ref());
+      idx_sub_expr_.push_back(indices[i]);
+    }
   }
 }
 
@@ -145,18 +129,7 @@ bool ComputeInliner::BodyPatternAllowInline() {
       inlined_store_.As<Store>(),
       ::common::errors::NotFound(
           "Param inlined store should be Store node! Please check."));
-  auto find_vars = ir::ir_utils::CollectIRNodesWithoutTensor(
-      inlined_store_, [&](const Expr* x) { return x->as_var(); });
-  std::set<Var, CompVar> vars_set;
-  for (auto& i : find_vars) {
-    if (i.as_var_ref()->name[0] == 'S') continue;
-    // if (i.as_var_ref()->name == "S0" || i.as_var_ref()->name == "S1")
-    // continue;
-    vars_set.insert(i.as_var_ref());
-  }
-
-  int n_vars = vars_set.size();
-  if (!UpdateAndCheckIndexVars(inlined_store_.As<Store>()->indices, n_vars)) {
+  if (!UpdateAndCheckIndexVars(inlined_store_.As<Store>()->indices)) {
     return false;
   }
   return true;
@@ -222,7 +195,7 @@ bool ReverseComputeInliner::BodyPatternAllowInline() {
   std::set<Var, CompVar> vars_set;
   for (auto& i : find_vars) vars_set.insert(i.as_var_ref());
   int n_vars = vars_set.size();
-  if (!UpdateAndCheckIndexVars(inlined_store_.As<Store>()->indices, n_vars)) {
+  if (!UpdateAndCheckIndexVars(inlined_store_.As<Store>()->indices)) {
     return false;
   }
   return true;
@@ -258,15 +231,17 @@ Expr ReverseComputeInliner::ReplaceInlinedTensor(Expr* load) {
 Expr ReverseComputeInliner::ReplaceTargetTensor(Expr* store) {
   auto indices = inlined_load_.As<ir::Load>()->indices;
   PADDLE_ENFORCE_EQ(indices.size(),
-                    idx_vars_.size(),
+                    idx_expr_.size(),
                     ::common::errors::InvalidArgument(
-                        "The size of indices should be equal to idx_vars_"));
-  size_t n = idx_vars_.size();
-  idx_sub_var_.reserve(n);
-  idx_sub_expr_.reserve(n);
+                        "The size of indices should be equal to idx_expr_"));
+  size_t n = idx_expr_.size();
+  idx_sub_var_.clear();
+  idx_sub_expr_.clear();
   for (int i = 0; i < n; ++i) {
-    idx_sub_var_.emplace_back(indices[i].as_var_ref());
-    idx_sub_expr_.emplace_back(idx_vars_[i]);
+    if (indices[i].is_var()) {
+      idx_sub_var_.emplace_back(indices[i].as_var_ref());
+      idx_sub_expr_.emplace_back(idx_expr_[i]);
+    }
   }
 
   Expr value_copy = ir::ir_utils::IRCopy(target_store_);
diff --git a/paddle/cinn/ir/schedule/ir_schedule.h b/paddle/cinn/ir/schedule/ir_schedule.h
@@ -501,9 +501,8 @@ class BaseInliner : public ir::IRMutator<> {
   void Visit(const ir::Block* expr, Expr* op) override;
 
  protected:
-  //! Check if indices are validate. If so, set idx_vars_ properly.
-  bool UpdateAndCheckIndexVars(const std::vector<Expr>& indices,
-                               int expected_ndim);
+  //! Check if indices are validate. If so, set idx_expr_ properly.
+  bool UpdateAndCheckIndexVars(const std::vector<Expr>& indices);
 
   void SetIndexSubstitution(const std::vector<Expr>& indices);
 
@@ -513,7 +512,7 @@ class BaseInliner : public ir::IRMutator<> {
   //! The body of the block to be inlined
   Expr inlined_store_{nullptr};
   //! The indices used for indexing the buffer to be inlined
-  std::vector<Var> idx_vars_;
+  std::vector<Expr> idx_expr_;
   //! Replacing vars(idx_sub_var_) in indices to corresponding
   //! expr(idx_sub_expr_)
   std::vector<Var> idx_sub_var_;
