Merge branch 'develop' into fix_randomness_in_fusion

Author: huangjiyi

.pre-commit-config.yaml

@@ -53,7 +53,7 @@ repos:
     rev: v1.27.3
     hooks:
     -   id: typos
-        args: []
+        args: [--force-exclude]
 # For Python files
 -   repo: https://github.com/psf/black-pre-commit-mirror
     rev: 24.8.0

_typos.toml

@@ -1,3 +1,9 @@
+[files]
+# The following files will be excluded from spell check during commits
+extend-exclude = [
+    "test/dataset/imikolov_test.py"
+]
+
 [default.extend-words]
 # PaddlePaddle specific words
 lod = "lod"
@@ -71,7 +77,6 @@ Prepar = 'Prepar'
 precent = 'precent'
 sheduler = 'sheduler'
 outpus = 'outpus'
-atrribute = 'atrribute'
 normlize = 'normlize'
 Costum = 'Costum'
 differnt = 'differnt'
@@ -154,7 +159,6 @@ CACH = 'CACH'
 endianess = 'endianess'
 VAILD = 'VAILD'
 ues = 'ues'
-aer = 'aer'
 elemenents = 'elemenents'
 CANN = 'CANN'
 pathes = 'pathes'
@@ -293,10 +297,8 @@ wiil = 'wiil'
 configurated = 'configurated'
 perfome = 'perfome'
 consructor = 'consructor'
-attribtue = 'attribtue'
 quitted = 'quitted'
 attribtes = 'attribtes'
-automatical = 'automatical'
 orignal = 'orignal'
 furture = 'furture'
 Indext = 'Indext'
@@ -467,7 +469,6 @@ channnel = 'channnel'
 Suger = 'Suger'
 Actural = 'Actural'
 subsituted = 'subsituted'
-automaticly = 'automaticly'
 Minium = 'Minium'
 sequnece = 'sequnece'
 payed = 'payed'
@@ -615,7 +616,6 @@ theads = 'theads'
 postive = 'postive'
 progrss = 'progrss'
 diffrent = 'diffrent'
-attritube = 'attritube'
 compability = 'compability'
 hge = 'hge'
 Funcion = 'Funcion'
@@ -772,7 +772,6 @@ Localy = 'Localy'
 PARM = 'PARM'
 thi = 'thi'
 Oll = 'Oll'
-Auxillary = 'Auxillary'
 Infor = 'Infor'
 statment = 'statment'
 varn = 'varn'

cmake/external/pybind11.cmake

@@ -27,7 +27,7 @@ set(PYBIND_PATCH_COMMAND "")
 if(LINUX
    AND (CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
    AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS 9)
-  set(PYBIND_TAG v2.12.0)
+  set(PYBIND_TAG v2.13.6)
   file(TO_NATIVE_PATH
        ${PADDLE_SOURCE_DIR}/patches/pybind/detail/internals.h.patch native_dst)
   # Note: [Why calling some `git` commands before `patch`?]

paddle/cinn/hlir/dialect/operator/ir/generate_shape_util.cc

@@ -539,7 +539,7 @@ void CollectSymbolNames(const std::vector<symbol::DimExpr>& dim_exprs,
 
 template <typename SymbolBindingsT>
 void AppendSymbolBindings(const std::vector<symbol::DimExpr>& dim_exprs,
-                          const std::set<std::string>& symbol_names,
+                          std::set<std::string>* remain_symbol_names_to_bind,
                           int in_tensor_idx,
                           SymbolBindings* symbol_bindings) {
   for (int in_tensor_dim_idx = 0; in_tensor_dim_idx < dim_exprs.size();
@@ -550,7 +550,10 @@ void AppendSymbolBindings(const std::vector<symbol::DimExpr>& dim_exprs,
                        "The type of dim_expr is not atomic"));
     if (!dim_expr.isa<std::string>()) continue;
     const auto& sym_name = dim_expr.dyn_cast<std::string>();
-    if (symbol_names.find(sym_name) == symbol_names.end()) continue;
+    if (remain_symbol_names_to_bind->find(sym_name) ==
+        remain_symbol_names_to_bind->end())
+      continue;
+    remain_symbol_names_to_bind->erase(sym_name);
     symbol_bindings->emplace_back(SymbolBindingsT{
         /*.symbol_name=*/sym_name,
         /*.input_tensor_idx=*/in_tensor_idx,
@@ -564,14 +567,17 @@ void GenerateSymbolBindings(
     const std::vector<pir::Value>& input_tensors,
     const std::set<std::string>& symbol_names,
     SymbolBindings* symbol_bindings) {
+  std::set<std::string> remain_symbol_names_to_bind = symbol_names;
   for (int i = 0; i < input_tensors.size(); ++i) {
     const auto& input_tensor = input_tensors.at(i);
     const auto& dim_exprs = ShapeOrDataDimExprs4Value(input_tensor);
     AppendSymbolBindings<ShapeSymbolBinding>(
-        dim_exprs.shape(), symbol_names, i, symbol_bindings);
+        dim_exprs.shape(), &remain_symbol_names_to_bind, i, symbol_bindings);
     if (dim_exprs.data().has_value()) {
-      AppendSymbolBindings<DataSymbolBinding>(
-          dim_exprs.data().value(), symbol_names, i, symbol_bindings);
+      AppendSymbolBindings<DataSymbolBinding>(dim_exprs.data().value(),
+                                              &remain_symbol_names_to_bind,
+                                              i,
+                                              symbol_bindings);
     }
   }
 }

paddle/cinn/hlir/dialect/operator/transforms/fuse_parallel_matmul_pass.cc

@@ -44,6 +44,7 @@ class MergeParallelMatmulPattern
       if (!op->dyn_cast<paddle::dialect::MatmulOp>()) {
         return false;
       }
+
       bool trans_x =
           op->attribute("transpose_x").dyn_cast<pir::BoolAttribute>().data();
       bool trans_y =
@@ -54,6 +55,10 @@ class MergeParallelMatmulPattern
       return false;
     }
 
+    auto IsFirstInput = [&](pir::Operation* op, pir::Value in_x) -> bool {
+      return in_x == op->operand_source(0);
+    };
+
     auto VectorPrefixEqual = [](const std::vector<std::int64_t>& a,
                                 const std::vector<std::int64_t>& b) {
       return std::vector<std::int64_t>(a.begin(), a.end() - 1) ==
@@ -74,6 +79,10 @@ class MergeParallelMatmulPattern
         if (!ValidMatmulTranspose(it->owner())) {
           continue;
         }
+
+        if (!IsFirstInput(it->owner(), input_x)) {
+          continue;
+        }
         if (!pre_dim.has_value()) {
           pre_dim = ::common::vectorize(
               it->owner()

paddle/cinn/hlir/framework/pir/trivial_op_util.cc

@@ -14,6 +14,7 @@
 
 #include "paddle/cinn/hlir/framework/pir/trivial_op_util.h"
 
+#include "paddle/cinn/common/dim_expr_converter.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
 #include "paddle/cinn/hlir/framework/compile_error.h"
 #include "paddle/cinn/hlir/framework/pir/op_lowering_util.h"
@@ -547,9 +548,6 @@ ExprTransformer RemoveVarInScheduleBlockRealize(const ir::Var& target_vars,
    * remove it in axes.bind()
    */
   const auto& f = [=](const ir::Expr& e) -> ir::Expr {
-    VLOG(4) << "Start RemoveVarInScheduleBlockRealize(" << target_vars << ", "
-            << replaced_expr << ")";
-    VLOG(4) << "      Input is " << e;
     PADDLE_ENFORCE_NE(
         e.As<ir::ScheduleBlockRealize>(),
         nullptr,
@@ -562,22 +560,11 @@ ExprTransformer RemoveVarInScheduleBlockRealize(const ir::Var& target_vars,
     auto block_bound_vars = copied_ir.As<ir::ScheduleBlockRealize>()
                                 ->schedule_block.As<ir::ScheduleBlock>()
                                 ->iter_vars;
-    for (const auto& i_var : schedule_block_iter_vars) {
-      PADDLE_ENFORCE_EQ(
-          i_var.is_var(),
-          true,
-          ::common::errors::InvalidArgument("RemoveVarInScheduleBlockRealize: "
-                                            "axes.bind rhs is is not a Var."));
-    }
     // find replace idx
     int target_idx = -1;
     for (int i = 0; i < schedule_block_iter_vars.size(); ++i) {
-      VLOG(4) << "RemoveVarInScheduleBlockRealize: compare with "
-              << schedule_block_iter_vars[i] << " vs " << target_vars
-              << ", and equality is: "
-              << (schedule_block_iter_vars[i].as_var()->name ==
-                  target_vars->name);
-      if (schedule_block_iter_vars[i].as_var()->name == target_vars->name) {
+      if (schedule_block_iter_vars[i].is_var() &&
+          schedule_block_iter_vars[i].as_var()->name == target_vars->name) {
         target_idx = i;
       }
     }
@@ -688,8 +675,6 @@ ExprTransformer RemoveOneTransformer(int one) {
                                      .GetSingle(copied);
     const ir::Expr& target_block =
         ExprSetFinderUtils::DirectlyFather(copied).GetSingle(target_for);
-    VLOG(4) << "RemoveOneTransformer: directly target_block of for is "
-            << target_block;
     if (target_block.As<ir::ScheduleBlockRealize>() != nullptr) {
       VLOG(4) << "RemoveOneTransformer: father block is root realize";
       ir::Expr shedule_block =
@@ -708,7 +693,6 @@ ExprTransformer RemoveOneTransformer(int one) {
         shedule_block.As<ir::ScheduleBlock>()->body = for_body;
       }
     } else if (target_block.As<ir::Block>() != nullptr) {
-      VLOG(4) << "RemoveOneTransformer: father block is Block";
       std::vector<ir::Expr> new_bodies;
       for (const auto& expr : target_block.As<ir::Block>()->stmts) {
         if (expr != target_for) {
@@ -728,7 +712,6 @@ ExprTransformer RemoveOneTransformer(int one) {
           "RemoveOneTransformer: target for father should be a ir::Block or "
           "ir::ScheduleBlockRealize."));
     }
-    VLOG(4) << "Remove Var to 0 in ScheduleBlockRealizer: " << copied;
     // Remove var to 0 in ScheduleBlockRealizer
     InplaceMutateSingleExpr(
         &copied,
@@ -949,6 +932,10 @@ std::vector<ir::Var> GetAllLoopVars(const ir::Expr& root) {
 
 ir::Expr GetBodyBlock(const ir::Expr& root) {
   const auto& iters = GetNonReduceLoopVars(root);
+  if (iters.empty()) {
+    return ir::Block::Make(
+        {ExprSetFinderUtils::ChildScheduleBlockRealizes.GetSingle(root)});
+  }
   const size_t reduce_size =
       std::count_if(iters.begin(), iters.end(), [](const ir::Var& v) {
         return v->is_reduce_axis;
@@ -965,6 +952,74 @@ ir::Expr GetBodyBlock(const ir::Expr& root) {
       ->body;
 }
 
+ir::Expr ReshapeLoop(const ir::Expr& root,
+                     const std::vector<symbol::DimExpr>& in_shape,
+                     const std::vector<symbol::DimExpr>& out_shape) {
+  auto copied = ir::ir_utils::IRCopy(root);
+
+  ir::ModuleExpr mod_expr({copied});
+  ir::IRSchedule ir_sch(
+      mod_expr, -1, false, cinn::utils::ErrorMessageLevel::kGeneral, true);
+
+  const auto block_realize =
+      (ExprSetFinderUtils::ChildScheduleBlockRealizes).GetSingle(copied);
+  const auto block_name = block_realize.As<ir::ScheduleBlockRealize>()
+                              ->schedule_block.As<ir::ScheduleBlock>()
+                              ->name;
+  const auto shape_partion = fusion::PartionReshapeAxes(in_shape, out_shape);
+
+  for (int idx = shape_partion.size() - 1; idx > 0; --idx) {
+    const auto& in_s = shape_partion[idx - 1].first;
+    const auto& in_e = shape_partion[idx].first;
+    const auto& out_s = shape_partion[idx - 1].second;
+    const auto& out_e = shape_partion[idx].second;
+
+    std::vector<int> fuse_indices;
+    for (int i = in_e - 1; i >= in_s; --i) {
+      if (in_shape[i] != symbol::DimExpr(1)) {
+        fuse_indices.insert(fuse_indices.begin(), i);
+      } else {
+        VLOG(4) << "Remove index[" << i << "]: " << in_shape[i]
+                << " for expr: \n"
+                << copied;
+        copied = ExprTransformerUtils::RemoveOneTransformer(i)(copied);
+        ir_sch.SetExprs({copied});
+        for (auto& index : fuse_indices) {
+          index--;
+        }
+      }
+    }
+    if (fuse_indices.size() > 1) {
+      VLOG(4) << "fuse_indices: " << cinn::utils::Join(fuse_indices, ",");
+      ir_sch.Fuse(block_name, fuse_indices);
+    }
+
+    std::vector<ir::Expr> split_shapes;
+    for (int i = out_s; i < out_e; ++i) {
+      if (out_shape[i] != symbol::DimExpr(1)) {
+        split_shapes.push_back(
+            cinn::common::DimExprConverter().ConvertToIrExpr(out_shape[i]));
+      }
+    }
+    if (split_shapes.size() > 1) {
+      ir_sch.Split(ir_sch.GetLoops(block_name)[in_s], split_shapes)[0];
+    }
+  }
+
+  std::vector<int> insert_axis;
+  std::vector<ir::Var> ones_var;
+  for (int i = 0; i < out_shape.size(); ++i) {
+    if (out_shape[i] == symbol::DimExpr(1)) {
+      insert_axis.push_back(i);
+      ones_var.push_back(ir::Var(1, "one_" + std::to_string(ones_var.size())));
+    }
+  }
+  copied = ExprTransformerUtils::InsertForsTransformer(insert_axis,
+                                                       ones_var)(copied);
+
+  return copied;
+}
+
 }  // namespace trivial_fusion_detail
 }  // namespace pir
 }  // namespace framework

paddle/cinn/hlir/framework/pir/trivial_op_util.h

@@ -297,6 +297,10 @@ std::vector<ir::Var> GetAllLoopVars(const ir::Expr& root);
 
 ir::Expr GetBodyBlock(const ir::Expr& root);
 
+ir::Expr ReshapeLoop(const ir::Expr& root,
+                     const std::vector<symbol::DimExpr>& in_shape,
+                     const std::vector<symbol::DimExpr>& out_shape);
+
 }  // namespace trivial_fusion_detail
 }  // namespace pir
 }  // namespace framework

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

@@ -63,8 +63,8 @@ bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,
     auto* block = expr_block.As<ir::ScheduleBlockRealize>();
     auto& iter_vars = block->schedule_block.As<ir::ScheduleBlock>()->iter_vars;
     for (int i = 0; i < iter_vars.size(); i++) {
-      ir::Var loop_var = block->iter_values[i].as_var_ref();
-      if (reduce_loop_vars.count(loop_var->name) > 0) {
+      if (block->iter_values[i].is_var() &&
+          reduce_loop_vars.count(block->iter_values[i].as_var()->name) > 0) {
         reduce_iter_vars.insert(iter_vars[i]->name);
       }
     }

paddle/cinn/operator_fusion/fusion_tracker/expr_utils.cc

@@ -31,7 +31,7 @@ ir::Expr ApplyItersTransform::operator()(const TransposeItersTransform& trans) {
 
 ir::Expr ApplyItersTransform::operator()(const RemoveOnesTransform& trans) {
   VLOG(4) << "[ItersTransform] Before RemoveOnesTransform("
-          << utils::Join(trans.ones_, ",") << "'): " << expr_;
+          << utils::Join(trans.ones_, ",") << "): " << expr_;
   auto result = RemoveOnesTransformer(trans.ones_)(expr_);
   VLOG(4) << "[ItersTransform] After  RemoveOnesTransform: " << result;
   return result;

paddle/cinn/operator_fusion/fusion_tracker/interpreter.cc

@@ -159,6 +159,20 @@ void RunItersTransformInstr(const std::shared_ptr<ItersTransformInstr>& instr,
   interpreter->scope[instr->target_] = new_pattern;
 }
 
+void RunReshapeAlignInstr(const std::shared_ptr<ReshapeAlignInstr>& instr,
+                          FusionInterpreter* interpreter) {
+  const auto expr = std::visit(
+      FusibleOp2Expr(), interpreter->scope[instr->input_]->fusion_ops[0])[0];
+  VLOG(4) << "Before RunReshapeAlignInstr: \n" << expr;
+  auto result = cinn::hlir::framework::pir::trivial_fusion_detail::ReshapeLoop(
+      expr, instr->in_shape_, instr->out_shape_);
+
+  auto new_pattern = std::make_shared<ScopeElement>();
+  new_pattern->fusion_ops.emplace_back(TrivialOp(result));
+  interpreter->scope[instr->result_] = new_pattern;
+  VLOG(4) << "After ReshapeAlignInstr: \n" << result;
+}
+
 void RunPaddingInstr(const std::shared_ptr<PaddingInstr>& instr,
                      FusionInterpreter* interpreter) {
   ScopeElementPtr new_pattern = std::make_shared<ScopeElement>();
@@ -229,6 +243,10 @@ std::vector<ir::Expr> FusionInterpreter::Run() {
         RunItersTransformInstr(
             dynamic_cast_instr_with_err<ItersTransformInstr>(instr), this);
         break;
+      case T_ReshapeAlign:
+        RunReshapeAlignInstr(
+            dynamic_cast_instr_with_err<ReshapeAlignInstr>(instr), this);
+        break;
       default:
         PADDLE_THROW(
             ::common::errors::Unavailable("Unsupported Fusion Instrution"));