PaddlePaddle
diff --git a/‎paddle/fluid/primitive/base/decomp_trans.cc
Copy file name to clipboard
+43-23Lines changed: 43 additions & 23 deletions b/‎paddle/fluid/primitive/base/decomp_trans.cc
Copy file name to clipboard
+43-23Lines changed: 43 additions & 23 deletions
diff --git a/‎paddle/fluid/primitive/base/decomp_trans.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/primitive/base/decomp_trans.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/primitive/base/primitive_ops.h
Copy file name to clipboard
+83-76Lines changed: 83 additions & 76 deletions b/‎paddle/fluid/primitive/base/primitive_ops.h
Copy file name to clipboard
+83-76Lines changed: 83 additions & 76 deletions
@@ -44,17 +44,19 @@ std::unordered_set<std::string> decomp_op_contain_none = {"pd_op.squeeze",
 std::unordered_set<std::string> dynamic_shape_blacklist = {"pd_op.squeeze",
                                                            "pd_op.unsqueeze"};
 
-static bool find_value(const std::vector<int64_t>& vec, int64_t value) {
-  if (std::find(vec.begin(), vec.end(), value) != vec.end()) {
+static bool has_dynamic_shape(const phi::DDim& dims) {
+  std::vector<int64_t> vec = common::vectorize<int64_t>(dims);
+  if (std::find(vec.begin(), vec.end(), -1) != vec.end()) {
     return true;
   } else {
     return false;
   }
 }
 
 static void check_ops(const std::string& op_name) {
-  auto it = primitive_set.find(op_name);
-  if (it == primitive_set.end()) {
+  auto primitives_set = GetPrimitiveOpNames();
+  auto it = primitives_set.find(op_name);
+  if (it == primitives_set.end()) {
     PADDLE_THROW(
         phi::errors::InvalidArgument("[Prim] Currently, decomposed program "
                                      "should not contain none primitive op %s.",
@@ -63,19 +65,36 @@ static void check_ops(const std::string& op_name) {
   return;
 }
 
-static const phi::DDim& GetValueDims(pir::Value value) {
-  if (!value.type()) {
+static const phi::DDim GetValueDims(pir::Value value) {
+  pir::Type origin_type = value.type();
+  if (!origin_type) {
     PADDLE_THROW(phi::errors::InvalidArgument("The type of value is nullptr."));
   }
-  if (value.type().isa<DenseTensorType>()) {
-    return value.type().dyn_cast<DenseTensorType>().dims();
-  } else if (value.type().isa<SelectedRowsType>()) {
-    return value.type().dyn_cast<SelectedRowsType>().dims();
+  auto getdims = [](pir::Type value_type) -> phi::DDim {
+    if (value_type.isa<DenseTensorType>()) {
+      return value_type.dyn_cast<DenseTensorType>().dims();
+    } else if (value_type.isa<SelectedRowsType>()) {
+      return value_type.dyn_cast<SelectedRowsType>().dims();
+    } else {
+      PADDLE_THROW(phi::errors::InvalidArgument(
+          "[Prim] Currently, we can only get shape for dense "
+          "tensor."));
+    }
+  };
+  phi::DDim value_dim;
+  if (origin_type.isa<pir::VectorType>()) {
+    pir::VectorType types = origin_type.dyn_cast<pir::VectorType>();
+    // all tensor dim in VectorType must be the same, expect dynamic shape.
+    for (size_t idx = 0; idx < types.size(); idx++) {
+      value_dim = getdims(types[idx]);
+      if (has_dynamic_shape(value_dim)) {
+        return value_dim;
+      }
+    }
   } else {
-    PADDLE_THROW(phi::errors::InvalidArgument(
-        "[Prim] Currently, we can only get shape for dense "
-        "tensor."));
+    value_dim = getdims(origin_type);
   }
+  return value_dim;
 }
 
 static phi::DataType GetValueDtype(pir::Value value) {
@@ -95,8 +114,7 @@ static phi::DataType GetValueDtype(pir::Value value) {
 static bool check_dynamic_shape(const pir::OpOperand& item,
                                 const pir::Operation& op) {
   auto dims = GetValueDims(item.source());
-  std::vector<int64_t> shape = common::vectorize<int64_t>(dims);
-  if (find_value(shape, -1)) {
+  if (has_dynamic_shape(dims)) {
     VLOG(6) << "[Prim] Decomp op receives dynamic shape [" << dims
             << "] in inputs of op " << op.name();
     return true;
@@ -187,12 +205,11 @@ void DecompProgram::check_decomp_outputs(
               orig_dim,
               decomp_dim));
 
-      std::vector<int64_t> shape = common::vectorize<int64_t>(orig_dim);
-      if (find_value(common::vectorize<int64_t>(orig_dim), -1)) {
+      if (has_dynamic_shape(orig_dim)) {
         VLOG(6) << "[Prim] Decomp op receives dynamic shape [" << orig_dim
                 << "] in " << i << "-index output of origin op " << op_name;
       }
-      if (find_value(common::vectorize<int64_t>(decomp_dim), -1)) {
+      if (has_dynamic_shape(decomp_dim)) {
         VLOG(6) << "[Prim] Decomp op receives dynamic shape [" << decomp_dim
                 << "] in " << i << "-index output of decomp op " << op_name;
       }
@@ -312,9 +329,10 @@ void DecompProgram::decomp_program() {
   }
   std::ostringstream orig_prog_stream;
   program_->Print(orig_prog_stream);
-  // Todo: Use cout instead of VLOG in case of incomplete log.
-  VLOG(4) << "[Prim] Origin program before decomp :\n"
-          << orig_prog_stream.str();
+  if (VLOG_IS_ON(4)) {
+    std::cout << "[Prim] Origin program before decomp :\n"
+              << orig_prog_stream.str() << std::endl;
+  }
 
   if (!paddle::prim::PrimCommonUtils::IsFwdPrimEnabled()) {
     return;
@@ -324,8 +342,10 @@ void DecompProgram::decomp_program() {
   decomp_block(block, orig_vars_dict, tar_vars);
   std::ostringstream decomp_prog_stream;
   program_->Print(decomp_prog_stream);
-  // Todo: Use cout instead of VLOG in case of incomplete log.
-  VLOG(4) << "[Prim] New program after decomp :\n" << decomp_prog_stream.str();
+  if (VLOG_IS_ON(4)) {
+    std::cout << "[Prim] New program after decomp :\n"
+              << decomp_prog_stream.str() << std::endl;
+  }
   if (FLAGS_prim_check_ops) {
     for (auto& op : *block) {
       check_ops(op.name());
 
@@ -72,6 +72,7 @@ class DecompProgram {
   std::vector<pir::Value> dst_vars_;
   std::set<std::string> blacklist_;
   std::set<std::string> whitelist_;
+  std::set<std::string> decomposed_ops_set;
 };
 
 bool has_decomp_rule(const pir::Operation& op);
 
@@ -18,80 +18,87 @@
 #include <string>
 
 namespace paddle {
-static std::set<std::string> primitive_set = {
-    /* primitives*/
-    "pd_op.add",
-    "pd_op.subtract",
-    "pd_op.multiply",
-    "pd_op.divide",
-    "pd_op.less_equal",
-    "pd_op.less_than",
-    "pd_op.equal",
-    "pd_op.not_equal",
-    "pd_op.greater_equal",
-    "pd_op.greater_than",
-    "pd_op.bitwise_and",
-    "pd_op.bitwise_not",
-    "pd_op.bitwise_or",
-    "pd_op.bitwise_xor",
-    "pd_op.exp",
-    "pd_op.scale",
-    "pd_op.matmul",
-    "pd_op.expand",
-    "pd_op.sum",
-    "pd_op.abs",
-    "pd_op.assign",
-    "pd_op.concat",
-    "pd_op.elementwise_pow",
-    "pd_op.floor",
-    "pd_op.gather",
-    "pd_op.gather_nd",
-    "pd_op.log",
-    "pd_op.max",
-    "pd_op.min",
-    "pd_op.maximum",
-    "pd_op.minimum",
-    "pd_op.prod",
-    "pd_op.roll",
-    "pd_op.scatter",
-    "pd_op.scatter_nd_add",
-    "pd_op.tile",
-    "pd_op.transpose",
-    "pd_op.pad",
-    "pd_op.cumsum",
-    "pd_op.put_along_axis",
-    "pd_op.equal",
-    "pd_op.greater_than",
-    "pd_op.less_equal",
-    "pd_op.sin",
-    "pd_op.cos",
-    "pd_op.where",
-    "pd_op.split",
-    "pd_op.reshape",
-    "pd_op.erf",
-    "pd_op.tanh",
-    "pd_op.cast",
-    "pd_op.sign",
-    "pd_op.slice",
-    "pd_op.uniform",
-    "pd_op.shape",
-    "pd_op.full",
-    "pd_op.full_int_array",
-    "pd_op.if",
-    "pd_op.while",
-    /* basic ops by PIR*/
-    "builtin.combine",
-    "builtin.slice",
-    "pd_op.feed",
-    "builtin.set_parameter",
-    "builtin.parameter",
-    "builtin.constant",
-    "pd_op.data",
-    "builtin.shadow_output",
-    /* skip some special ops */
-    "pd_op.squeeze",
-    "pd_op.unsqueeze",
-    "pd_op.top_p_sampling",
-    "pd_op.tril",
-};
+const std::set<std::string>& GetPrimitiveOpNames() {
+  static std::set<std::string> primitive_set = {
+      /* primitives*/
+      "pd_op.arange",
+      "pd_op.add",
+      "pd_op.subtract",
+      "pd_op.multiply",
+      "pd_op.divide",
+      "pd_op.less_equal",
+      "pd_op.less_than",
+      "pd_op.equal",
+      "pd_op.not_equal",
+      "pd_op.greater_equal",
+      "pd_op.greater_than",
+      "pd_op.bitwise_and",
+      "pd_op.bitwise_not",
+      "pd_op.bitwise_or",
+      "pd_op.bitwise_xor",
+      "pd_op.exp",
+      "pd_op.scale",
+      "pd_op.matmul",
+      "pd_op.expand",
+      "pd_op.sum",
+      "pd_op.abs",
+      "pd_op.assign",
+      "pd_op.concat",
+      "pd_op.elementwise_pow",
+      "pd_op.floor",
+      "pd_op.gather",
+      "pd_op.gather_nd",
+      "pd_op.log",
+      "pd_op.max",
+      "pd_op.min",
+      "pd_op.maximum",
+      "pd_op.minimum",
+      "pd_op.prod",
+      "pd_op.roll",
+      "pd_op.scatter",
+      "pd_op.scatter_nd_add",
+      "pd_op.tile",
+      "pd_op.transpose",
+      "pd_op.pad",
+      "pd_op.cumsum",
+      "pd_op.put_along_axis",
+      "pd_op.equal",
+      "pd_op.greater_than",
+      "pd_op.less_equal",
+      "pd_op.sin",
+      "pd_op.cos",
+      "pd_op.where",
+      "pd_op.split",
+      "pd_op.reshape",
+      "pd_op.erf",
+      "pd_op.tanh",
+      "pd_op.cast",
+      "pd_op.sign",
+      "pd_op.slice",
+      "pd_op.uniform",
+      "pd_op.shape",
+      "pd_op.full",
+      "pd_op.full_int_array",
+      "pd_op.if",
+      "pd_op.while",
+      /* basic ops by PIR*/
+      "builtin.combine",
+      "builtin.slice",
+      "builtin.split",
+      "pd_op.feed",
+      "builtin.set_parameter",
+      "builtin.parameter",
+      "builtin.constant",
+      "pd_op.data",
+      "builtin.shadow_output",
+      /* skip some special ops */
+      "pd_op.squeeze",
+      "pd_op.unsqueeze",
+      "pd_op.top_p_sampling",
+      "pd_op.tril",
+      "cf.yield",
+      "pd_op.increment_",
+  };
+  return primitive_set;
+}
 }  // namespace paddle