Merge branch 'develop' into 104

Author: zeroRains

cmake/cinn.cmake

@@ -164,8 +164,8 @@ cinn_cc_library(
 add_dependencies(cinnapi GEN_LLVM_RUNTIME_IR_HEADER ZLIB::ZLIB)
 add_dependencies(cinnapi GEN_LLVM_RUNTIME_IR_HEADER ${core_deps})
 if(NOT CINN_ONLY)
-  target_link_libraries(cinnapi pd_op_dialect phi)
-  add_dependencies(cinnapi pd_op_dialect phi)
+  target_link_libraries(cinnapi op_dialect_vjp phi)
+  add_dependencies(cinnapi op_dialect_vjp phi)
 endif()
 
 target_link_libraries(cinnapi ${PYTHON_LIBRARIES})
@@ -222,8 +222,8 @@ function(gen_cinncore LINKTYPE)
   add_dependencies(${CINNCORE_TARGET} GEN_LLVM_RUNTIME_IR_HEADER ZLIB::ZLIB)
   add_dependencies(${CINNCORE_TARGET} GEN_LLVM_RUNTIME_IR_HEADER ${core_deps})
   if(NOT CINN_ONLY)
-    target_link_libraries(${CINNCORE_TARGET} pd_op_dialect phi)
-    add_dependencies(${CINNCORE_TARGET} pd_op_dialect phi)
+    target_link_libraries(${CINNCORE_TARGET} op_dialect_vjp phi)
+    add_dependencies(${CINNCORE_TARGET} op_dialect_vjp phi)
   endif()
 
   add_dependencies(${CINNCORE_TARGET} pybind)

cmake/external/xpu.cmake

@@ -26,7 +26,7 @@ set(XPU_XPTI_LIB_NAME "libxpti.so")
 if(NOT DEFINED XPU_BASE_DATE)
   set(XPU_BASE_DATE "20231103")
 endif()
-set(XPU_XCCL_BASE_VERSION "1.0.53.6")
+set(XPU_XCCL_BASE_VERSION "1.1.6.1")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
   set(XPU_XFT_BASE_VERSION "20230602")
 endif()

paddle/cinn/backends/nvrtc/nvrtc_util.cc

@@ -32,6 +32,7 @@
 
 PD_DECLARE_string(cinn_nvcc_cmd_path);
 PD_DECLARE_bool(nvrtc_compile_to_cubin);
+PD_DECLARE_bool(cinn_nvrtc_cubin_with_fmad);
 
 namespace cinn {
 namespace backends {
@@ -106,6 +107,9 @@ std::string Compiler::CompileCudaSource(const std::string& code,
   }
   if (compile_to_cubin_) {
     compile_options.push_back("-arch=sm_" + cc);
+    std::string enable_fmad =
+        FLAGS_cinn_nvrtc_cubin_with_fmad ? "true" : "false";
+    compile_options.push_back("--fmad=" + enable_fmad);
   } else {
     compile_options.push_back("-arch=compute_" + cc);
   }

paddle/cinn/common/context.h

@@ -52,6 +52,22 @@ struct NameGenerator {
   mutable std::mutex mutex_;
 };
 
+struct PrettyNamer {
+  const std::string& GetOrNew(const size_t hash_key,
+                              const std::string& name_hint) {
+    if (pretty_names_.find(hash_key) == pretty_names_.end()) {
+      pretty_names_[hash_key] = name_generator_.New(name_hint);
+    }
+    return pretty_names_.at(hash_key);
+  }
+
+  NameGenerator& GetNameGenerator() { return name_generator_; }
+
+ private:
+  absl::flat_hash_map<size_t, std::string> pretty_names_;
+  NameGenerator name_generator_;
+};
+
 class Context {
  public:
   static Context& Global();
@@ -61,10 +77,15 @@ class Context {
    * @param name_hint The prefix.
    */
   std::string NewName(const std::string& name_hint) {
-    return name_generator_.New(name_hint);
+    return pretty_namer_.GetNameGenerator().New(name_hint);
   }
 
-  void ResetNameId() { name_generator_.ResetID(); }
+  std::string PrettyUniqName(const size_t hash_key,
+                             const std::string& name_hint) {
+    return pretty_namer_.GetOrNew(hash_key, name_hint);
+  }
+
+  void ResetNameId() { pretty_namer_.GetNameGenerator().ResetID(); }
 
   const std::vector<std::string>& runtime_include_dir();
 
@@ -82,7 +103,7 @@ class Context {
  private:
   Context() = default;
 
-  NameGenerator name_generator_;
+  PrettyNamer pretty_namer_;
   std::vector<std::string> runtime_include_dir_;
   mutable std::mutex mutex_;
 

paddle/cinn/hlir/dialect/operator/ir/CMakeLists.txt

@@ -62,7 +62,7 @@ if(NOT CINN_ONLY)
     manual_op.cc
     op_attribute.cc
     DEPS
-    pd_op_dialect)
+    op_dialect_vjp)
 
   target_include_directories(cinn_op_dialect PRIVATE ${CINN_DIALECT_BINARY_DIR})
 endif()

paddle/cinn/hlir/dialect/operator/transforms/CMakeLists.txt

@@ -7,7 +7,7 @@ if(NOT CINN_ONLY)
     cinn_group_lowering_pass.cc
     tensor_node.cc
     DEPS
-    pd_op_dialect
+    op_dialect_vjp
     pir_compiler
     cinn_runtime_dialect)
 
@@ -18,7 +18,7 @@ if(NOT CINN_ONLY)
     DEPS
     drr
     cinn_op_dialect
-    pd_op_dialect)
+    op_dialect_vjp)
 
   cinn_cc_library(
     add_broadcast_to_elementwise_pass
@@ -27,5 +27,5 @@ if(NOT CINN_ONLY)
     DEPS
     pir
     cinn_op_dialect
-    pd_op_dialect)
+    op_dialect_vjp)
 endif()

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

@@ -81,6 +81,19 @@ bool IsSameDim(const phi::DDim& first, const std::vector<int64_t>& second) {
   return false;
 }
 
+std::vector<int64_t> GetBroadcastAxis(const phi::DDim& in_shape,
+                                      const std::vector<int64_t>& out_shape) {
+  std::vector<int64_t> broadcast_axes(in_shape.size(), 0);
+  auto in_shape_size = in_shape.size();
+  if (in_shape_size >= 1) {
+    for (int i = 1; i <= in_shape_size; ++i) {
+      broadcast_axes[in_shape_size - i] = out_shape.size() - i;
+    }
+  }
+
+  return broadcast_axes;
+}
+
 bool ProcessOp(pir::Operation* op, pir::PatternRewriter* rewriter) {
   auto x_dims = op->operand_source(0)
                     .type()
@@ -93,21 +106,21 @@ bool ProcessOp(pir::Operation* op, pir::PatternRewriter* rewriter) {
 
   if (x_dims != y_dims) {
     auto output_shape = GetOutputShape(x_dims, y_dims);
-    std::vector<int64_t> vec_dims;
-    for (int64_t i = 0; i < output_shape.size(); ++i) {
-      vec_dims.push_back(i);
-    }
     if (!IsSameDim(x_dims, output_shape)) {
       // add broadcast to input 0
       auto new_transpose_op = rewriter->Build<cinn::dialect::BroadcastOp>(
-          op->operand_source(0), vec_dims, output_shape);
+          op->operand_source(0),
+          GetBroadcastAxis(x_dims, output_shape),
+          output_shape);
 
       op->operand(0).set_source(new_transpose_op->result(0));
     }
 
     if (!IsSameDim(y_dims, output_shape)) {
       auto new_transpose_op = rewriter->Build<cinn::dialect::BroadcastOp>(
-          op->operand_source(1), vec_dims, output_shape);
+          op->operand_source(1),
+          GetBroadcastAxis(y_dims, output_shape),
+          output_shape);
 
       op->operand(1).set_source(new_transpose_op->result(0));
     }

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

@@ -58,6 +58,9 @@ OpPatternKind GetOpKind(const std::string& op_name) {
 }
 
 phi::DDim GetFirstInputShape(const ::pir::Operation* op) {
+  if (op->num_operands() == 0) {
+    return phi::DDim({});
+  }
   auto in = op->operand_source(0);
 
   return in.type().dyn_cast<paddle::dialect::DenseTensorType>().dims();

paddle/cinn/hlir/framework/CMakeLists.txt

@@ -22,13 +22,12 @@ gather_srcs(
   op_lowering_impl.cc
   accuracy_checker.cc
   visualize_helper.cc
-  compile_error.cc
-  group_scheduler.cc)
+  compile_error.cc)
 
-# TODO(Aurelius84): pir_compiler depends on pd_op_dialect and could
+# TODO(Aurelius84): pir_compiler depends on op_dialect_vjp and could
 # not found under CINN_ONLY mode
 if(NOT CINN_ONLY)
-  cinn_cc_library(pir_compiler SRCS pir_compiler.cc DEPS cinnapi pd_op_dialect)
+  cinn_cc_library(pir_compiler SRCS pir_compiler.cc DEPS cinnapi op_dialect_vjp)
 endif()
 
 if(WITH_CUDA)
@@ -44,8 +43,6 @@ endif()
 if(WITH_CUDA)
   cinn_cc_test(test_hlir_framework_op_lowering SRCS op_lowering_test.cc DEPS
                cinncore decomposer_test_helper)
-  cinn_cc_test(test_group_scheduler SRCS group_scheduler_test.cc DEPS cinncore
-               decomposer_test_helper)
 endif()
 cinn_cc_test(test_hlir_framework_tensor SRCS tensor_test.cc DEPS cinncore)
 cinn_cc_test(test_hlir_framework_scope SRCS scope_test.cc DEPS cinncore)

paddle/cinn/hlir/framework/op_lowering_impl.cc

@@ -17,9 +17,9 @@
 #include "paddle/cinn/ast_gen_ius/tensor_group.h"
 #include "paddle/cinn/hlir/framework/compile_error.h"
 #include "paddle/cinn/hlir/framework/graph_compiler_util.h"
-#include "paddle/cinn/hlir/framework/group_scheduler.h"
 #include "paddle/cinn/hlir/framework/op_lowering_util.h"
 #include "paddle/cinn/hlir/op/external_api_registry.h"
+#include "paddle/cinn/ir/group_schedule/st_shape_group_scheduler.h"
 #include "paddle/cinn/ir/schedule/ir_schedule.h"
 #include "paddle/cinn/optim/transform_gpu_forloop.h"
 #include "paddle/cinn/runtime/flags.h"
@@ -470,8 +470,20 @@ ir::Expr OpLowererImpl::DoGroupSchedule(
     const GroupPtr& group,
     const std::unordered_map<std::string, ir::Tensor>& tensor_map) {
   if (FLAGS_cinn_new_group_scheduler) {
-    GroupScheduler group_scheduler(&ir_sch, group, target_);
-    group_scheduler();
+    std::unordered_set<std::string> output_tensor_names;
+    std::transform(
+        group->output_nodes.begin(),
+        group->output_nodes.end(),
+        std::inserter(output_tensor_names, output_tensor_names.begin()),
+        [](const Node* node) {
+          NodeData* node_data =
+              (*node->outlinks().begin())->sink()->safe_as<NodeData>();
+          CHECK(node_data);
+          return node_data->id();
+        });
+    ir::StaticShapeGroupScheduler group_scheduler(
+        &ir_sch, output_tensor_names, target_);
+    group_scheduler.Schedule();
     return ir_sch.GetModule().GetExprs().at(0);
   }
   // topological order.

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

@@ -810,18 +810,22 @@ void LoopAssignReduceWithoutLast(ir::IRSchedule& ir_sch,  // NOLINT
 }
 
 std::vector<int> GetReducerDimAttr(::pir::Operation* reduce_op) {
-  VLOG(3) << "GetReducerDimAttr from " << reduce_op->name();
-  auto* source_op = reduce_op->operand_source(/*dim_idx=*/1)
-                        .dyn_cast<::pir::OpResult>()
-                        .owner();
-  CHECK(source_op->isa<paddle::dialect::FullIntArrayOp>());
+  int rank = reduce_op->operand_source(0)
+                 .type()
+                 .dyn_cast<::pir::DenseTensorType>()
+                 .dims()
+                 .size();
+
+  auto attr = reduce_op->attributes().at("dim");
+  auto attr_vec = attr.dyn_cast<::pir::ArrayAttribute>().AsVector();
+
   std::vector<int> dim;
-  auto dim_attr = source_op->attributes()
-                      .at("value")
-                      .dyn_cast<::pir::ArrayAttribute>()
-                      .AsVector();
-  for (auto& attr : dim_attr) {
-    dim.push_back(attr.dyn_cast<::pir::Int64Attribute>().data());
+  for (auto vec_element : attr_vec) {
+    auto axis = vec_element.dyn_cast<::pir::Int64Attribute>().data();
+    if (axis < 0) {
+      axis += rank;
+    }
+    dim.push_back(axis);
   }
   return dim;
 }

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

@@ -13,6 +13,8 @@
 // limitations under the License.
 
 #include "paddle/cinn/hlir/framework/pir/op_mapper.h"
+#include "paddle/cinn/hlir/dialect/operator/ir/cinn_op.h"
+#include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 
@@ -25,15 +27,41 @@ namespace {
 
 void AppendAttrForReduceOp(const ::pir::Operation& op,
                            utils::AttributeMap& attrs) {  // NOLINT
-  auto* source_op =
-      op.operand_source(/*dim_idx=*/1).dyn_cast<::pir::OpResult>().owner();
-  CHECK(source_op->isa<paddle::dialect::FullIntArrayOp>());
-  auto dim_val =
-      paddle::dialect::GetInt64Vector(source_op->attributes().at("value"));
-  std::vector<int> dim(dim_val.begin(), dim_val.end());
+  auto attr = op.attributes().at("dim");
+  auto attr_vec = attr.dyn_cast<::pir::ArrayAttribute>().AsVector();
+
+  std::vector<int> dim;
+  for (auto vec_element : attr_vec) {
+    dim.push_back(vec_element.dyn_cast<::pir::Int64Attribute>().data());
+  }
+
   attrs["dim"] = dim;
 }
 
+void AppendAttrForBoadcastToOp(const ::pir::Operation& op,
+                               utils::AttributeMap& attrs) {  // NOLINT
+  auto axes_attr = op.attributes().at("broadcast_axes");
+  auto attr_vec = axes_attr.dyn_cast<::pir::ArrayAttribute>().AsVector();
+
+  std::vector<int> axis;
+  for (auto vec_element : attr_vec) {
+    axis.push_back(vec_element.dyn_cast<::pir::Int64Attribute>().data());
+  }
+
+  attrs["broadcast_axes"] = axis;
+
+  auto out_shape_attr = op.attributes().at("out_shape");
+  auto out_shape_attr_vec =
+      out_shape_attr.dyn_cast<::pir::ArrayAttribute>().AsVector();
+
+  std::vector<int> out_shape;
+  for (auto vec_element : out_shape_attr_vec) {
+    out_shape.push_back(vec_element.dyn_cast<::pir::Int64Attribute>().data());
+  }
+
+  attrs["out_shape"] = out_shape;
+}
+
 }  // namespace
 
 #define REGISTER_OPERAND_RULE(OP, args...)                                    \
@@ -42,18 +70,17 @@ void AppendAttrForReduceOp(const ::pir::Operation& op,
   };
 
 #define REGISTER_ATTR_RULE(OP, func) \
-  attr_funcs_[paddle::dialect::OP::name()] = func;
+  attr_funcs_[cinn::dialect::OP::name()] = func;
 
 void OpMapper::RegisterMapRules() {
   // max(x, dim) -> reduce_max(x)
   REGISTER_OPERAND_RULE(MaxOp, 0);
   REGISTER_OPERAND_RULE(SumOp, 0);
   REGISTER_OPERAND_RULE(MinOp, 0);
   REGISTER_OPERAND_RULE(ProdOp, 0);
-  REGISTER_ATTR_RULE(MaxOp, AppendAttrForReduceOp);
-  REGISTER_ATTR_RULE(SumOp, AppendAttrForReduceOp);
-  REGISTER_ATTR_RULE(MinOp, AppendAttrForReduceOp);
-  REGISTER_ATTR_RULE(ProdOp, AppendAttrForReduceOp);
+  REGISTER_ATTR_RULE(ReduceMaxOp, AppendAttrForReduceOp);
+  REGISTER_ATTR_RULE(ReduceSumOp, AppendAttrForReduceOp);
+  REGISTER_ATTR_RULE(BroadcastOp, AppendAttrForBoadcastToOp);
 }
 
 }  // namespace pir