PaddlePaddle
diff --git a/‎paddle/cinn/hlir/dialect/operator/ir/ops.yaml
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/ir/ops.yaml
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_to_pd_util.cc
Copy file name to clipboard
+33Lines changed: 33 additions & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_to_pd_util.cc
Copy file name to clipboard
+33Lines changed: 33 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
Copy file name to clipboard
+63-17Lines changed: 63 additions & 17 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
Copy file name to clipboard
+63-17Lines changed: 63 additions & 17 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/utils.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/hlir/framework/pir/utils.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/cinn/hlir/op/elementwise.cc
Copy file name to clipboard
+48-74Lines changed: 48 additions & 74 deletions b/‎paddle/cinn/hlir/op/elementwise.cc
Copy file name to clipboard
+48-74Lines changed: 48 additions & 74 deletions
diff --git a/‎paddle/cinn/hlir/pe/elementwise.cc
Copy file name to clipboard
+6-9Lines changed: 6 additions & 9 deletions b/‎paddle/cinn/hlir/pe/elementwise.cc
Copy file name to clipboard
+6-9Lines changed: 6 additions & 9 deletions
diff --git a/‎paddle/cinn/hlir/pe/elementwise.h
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions b/‎paddle/cinn/hlir/pe/elementwise.h
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions
@@ -1,3 +1,11 @@
+- op : arange
+  args : (Tensor start_, Tensor end_, Tensor step_, ScalarType start, ScalarType end, ScalarType step, DataType dtype=DataType::INT64)
+  output : Tensor(out)
+  infer_meta :
+    func : ArangeInferMeta
+    param : [start, end, step, dtype]
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
+
 - op : argmax
   args : (Tensor x, int64_t[] axis, bool keepdim=false, DataType dtype=DataType::INT64)
   output : Tensor(out)
 
@@ -22,6 +22,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/utils/general_functions.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/pir/include/core/builtin_dialect.h"
 #include "paddle/pir/include/core/ir_mapping.h"
@@ -154,6 +155,34 @@ ::pir::Operation* ConvertArgMinMaxOp(
   return pd_op;
 }
 
+::pir::Operation* ConvertArangeOp(::pir::Operation* op,
+                                  ::pir::IrMapping& ir_mapping,        // NOLINT
+                                  ::pir::PatternRewriter& rewriter) {  // NOLINT
+  using paddle::dialect::FullOp;
+  VLOG(6) << "transform " << op->name() << " from cinn_op to pd_op";
+
+  std::array<double, 3> input_list = {0, 0, 1};
+  for (int i = 0; i < 3; i++) {
+    const FullOp full_op = ::pir::CastDefinedTo<FullOp>(op, i);
+    // pd_op.paddle has fixed args type, so we need to convert it to double
+    input_list[i] = full_op.attribute("value")
+                        .dyn_cast<paddle::dialect::ScalarAttribute>()
+                        .data()
+                        .to<double>();
+  }
+
+  ::phi::DataType dtype = op->attributes()
+                              .at("dtype")
+                              .dyn_cast<paddle::dialect::DataTypeAttribute>()
+                              .data();
+  auto pd_op = rewriter.Build<paddle::dialect::ArangeOp>(
+      input_list[0], input_list[1], input_list[2], dtype);
+  for (uint32_t i = 0; i < op->num_results(); ++i) {
+    ir_mapping.Add(op->result(i), pd_op->result(i));
+  }
+  return pd_op;
+}
+
 ::pir::Operation* ConvertSliceOp(::pir::Operation* op,
                                  ::pir::IrMapping& ir_mapping,        // NOLINT
                                  ::pir::PatternRewriter& rewriter) {  // NOLINT
@@ -445,6 +474,10 @@ REGISTER_TRANSFORM_RULES(
     cinn::dialect::ArgmaxOp::name(),
     cinn::dialect::details::ConvertArgMinMaxOp<paddle::dialect::ArgmaxOp>);
 
+REGISTER_TRANSFORM_RULES(arange_op,
+                         cinn::dialect::ArangeOp::name(),
+                         cinn::dialect::details::ConvertArangeOp);
+
 REGISTER_TRANSFORM_RULES(slice_op,
                          cinn::dialect::SliceOp::name(),
                          cinn::dialect::details::ConvertSliceOp);
 
@@ -38,26 +38,11 @@ namespace ir {
 using CompatibleInfo = cinn::hlir::framework::pir::CompatibleInfo;
 using paddle::dialect::FullIntArrayOp;
 using paddle::dialect::FullOp;
+using ::pir::CastDefinedTo;
+using ::pir::IsDefinedBy;
 
 namespace {
 
-template <typename TargetOpT, typename SourceOpT>
-bool IsDefinedBy(const SourceOpT &op, const size_t idx) {
-  const pir::Operation *defined_op = op->operand_source(idx).defining_op();
-  return defined_op && defined_op->isa<TargetOpT>();
-}
-
-template <typename TargetOpT, typename SourceOpT>
-TargetOpT CastDefinedTo(const SourceOpT &op, const size_t idx) {
-  PADDLE_ENFORCE_EQ(IsDefinedBy<TargetOpT>(op, idx),
-                    true,
-                    ::common::errors::PreconditionNotMet(
-                        "Required defined op shall not be nullptr and can cast "
-                        "to target type."));
-  pir::Operation *defined_op = op->operand_source(idx).defining_op();
-  return defined_op->dyn_cast<TargetOpT>();
-}
-
 template <typename T = int>
 std::vector<T> GetVectorFromIntArrayAttribute(
     const pir::ArrayAttribute &array_attr) {
@@ -529,6 +514,66 @@ class SliceOpPattern : public pir::OpRewritePattern<paddle::dialect::SliceOp> {
   }
 };
 
+class ArangeOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::ArangeOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::ArangeOp>::OpRewritePattern;
+
+  bool Match(paddle::dialect::ArangeOp op) const override {
+    // ArangeOp for CINN must have static start, end, step to calculate
+    // the shape of output tensor. Otherwise, it will be denied
+    // due to CauseNewSymbolicShape returning false
+    bool is_denied = CompatibleInfo::IsDeniedForCinn(*op.operation());
+    return !is_denied && IsDefinedBy<FullOp>(op, 0) &&
+           IsDefinedBy<FullOp>(op, 1) && IsDefinedBy<FullOp>(op, 2);
+  }
+
+  void Rewrite(paddle::dialect::ArangeOp op,
+               pir::PatternRewriter &rewriter) const override {
+    const auto &dtype = op.attributes()
+                            .at("dtype")
+                            .dyn_cast<paddle::dialect::DataTypeAttribute>()
+                            .data();
+
+    std::array<phi::Scalar, 3> input_list;
+    for (int i = 0; i < 3; i++) {
+      const FullOp full_op = CastDefinedTo<FullOp>(op, i);
+      phi::Scalar input = full_op.attribute("value")
+                              .dyn_cast<paddle::dialect::ScalarAttribute>()
+                              .data();
+      if (input.dtype() != dtype) {
+        // FullOp creates a tensor (scalar) with fp64 type by default
+        // therefore, we might need to perform type casting
+        switch (dtype) {
+          case phi::DataType::FLOAT32:
+            input = phi::Scalar(input.to<float>());
+            break;
+          case phi::DataType::FLOAT64:
+            input = phi::Scalar(input.to<double>());
+            break;
+          case phi::DataType::INT32:
+            input = phi::Scalar(input.to<int>());
+            break;
+          default:
+            input = phi::Scalar(input.to<int64_t>());
+        }
+      }
+      input_list[i] = input;
+    }
+    auto cinn_arange =
+        rewriter.Build<cinn::dialect::ArangeOp>(op->operand_source(0),
+                                                op->operand_source(1),
+                                                op->operand_source(2),
+                                                input_list[0],
+                                                input_list[1],
+                                                input_list[2],
+                                                dtype);
+    cinn_arange.result(0).set_type(op.result(0).type());
+    rewriter.ReplaceAllUsesWith(op.result(0), cinn_arange.result(0));
+    rewriter.EraseOp(op);
+  }
+};
+
 class ConcatOpPattern
     : public pir::OpRewritePattern<paddle::dialect::ConcatOp> {
  public:
@@ -1385,6 +1430,7 @@ pir::RewritePatternSet PdOpToCinnOpPass::InitializePatterns(
   ps.Add<
       ArgMinMaxOpPattern<paddle::dialect::ArgmaxOp, cinn::dialect::ArgmaxOp>>(
       context);
+  ps.Add<ArangeOpPattern>(context);
   ps.Add<ProdOpPattern>(context);
   ps.Add<ReshapeOpPattern>(context);
   ps.Add<PowOpPattern>(context);
 
@@ -145,7 +145,6 @@ class OpTransInfo {
                                                     "matmul_grad",
                                                     "embedding_grad",
                                                     "embedding",
-                                                    "arange",
                                                     "argsort",
                                                     "assign_value",
                                                     "one_hot",
@@ -552,6 +551,7 @@ const std::unordered_set<std::string> TOCINN_OPS = {
     PD_OP_NAME(ScaleOp),
     PD_OP_NAME(Pool2dOp),
     PD_OP_NAME(IscloseOp),
+    PD_OP_NAME(ArangeOp),
     // PD_OP_NAME(SliceOp),
     PD_OP_NAME(ConcatOp),
     PD_OP_NAME(SplitOp),
 
@@ -1248,67 +1248,6 @@ std::shared_ptr<framework::OpStrategy> StrategyForGenerateShapeSymbolic(
   return strategy;
 }
 
-std::shared_ptr<framework::OpStrategy> StrategyForArange(
-    const framework::NodeAttr &attrs,
-    const std::vector<ir::Tensor> &inputs,
-    const std::vector<Type> &out_type,
-    const std::vector<std::vector<int>> &output_shapes,
-    const Target &target) {
-  auto attr_store = attrs.attr_store;
-  PADDLE_ENFORCE_EQ(
-      attr_store.count("start"),
-      true,
-      ::common::errors::InvalidArgument(
-          "No start attribute in attrs.attr_store! Please check."));
-  PADDLE_ENFORCE_EQ(
-      attr_store.count("stop"),
-      true,
-      ::common::errors::InvalidArgument(
-          "No stop attribute in attrs.attr_store! Please check."));
-  PADDLE_ENFORCE_EQ(
-      attr_store.count("step"),
-      true,
-      ::common::errors::InvalidArgument(
-          "No step attribute in attrs.attr_store! Please check."));
-  PADDLE_ENFORCE_EQ(
-      attr_store.count("dtype"),
-      true,
-      ::common::errors::InvalidArgument(
-          "No dtype attribute in attrs.attr_store! Please check."));
-
-  auto start = std::get<float>(attr_store.at("start"));
-  auto stop = std::get<float>(attr_store.at("stop"));
-  auto step = std::get<float>(attr_store.at("step"));
-  auto dtype =
-      cinn::common::Str2Type(std::get<std::string>(attr_store.at("dtype")));
-
-  framework::CINNCompute arange_compute(
-      [=](lang::Args args, lang::RetValue *ret) {
-        PADDLE_ENFORCE(!args.empty(),
-                       ::common::errors::InvalidArgument(
-                           "The input argument of arange compute is empty! "
-                           "Please check."));
-        CINNValuePack pack_args = args[0];
-
-        PADDLE_ENFORCE_EQ(
-            pack_args.size(),
-            1U,
-            ::common::errors::InvalidArgument("the size of pack_args should be "
-                                              "equal to 1, but got %d.",
-                                              pack_args.size()));
-        std::string tensor_name = pack_args[0].operator std::string();
-
-        auto out = pe::Arange(start, stop, step, dtype, tensor_name);
-        std::vector<cinn::common::CINNValue> res;
-        res.push_back(cinn::common::CINNValue(out));
-        *ret = CINNValuePack{res};
-      });
-
-  auto strategy = std::make_shared<framework::OpStrategy>();
-  strategy->AddImpl(arange_compute, "strategy.reshape.x86", 1);
-  return strategy;
-}
-
 std::shared_ptr<framework::OpStrategy> StrategyForArangeSymbolic(
     const framework::NodeAttr &attrs,
     const std::vector<ir::Tensor> &inputs,
@@ -1320,7 +1259,7 @@ std::shared_ptr<framework::OpStrategy> StrategyForArangeSymbolic(
                     0U,
                     ::common::errors::InvalidArgument(
                         "No start attribute in arange Op! Please check."));
-  PADDLE_ENFORCE_GT(attr_store.count("stop"),
+  PADDLE_ENFORCE_GT(attr_store.count("end"),
                     0U,
                     ::common::errors::InvalidArgument(
                         "No stop attribute in arange Op! Please check."));
@@ -1333,12 +1272,50 @@ std::shared_ptr<framework::OpStrategy> StrategyForArangeSymbolic(
                     ::common::errors::InvalidArgument(
                         "No dtype attribute in arange Op! Please check."));
 
-  auto start = std::get<float>(attr_store.at("start"));
-  auto stop = std::get<float>(attr_store.at("stop"));
-  auto step = std::get<float>(attr_store.at("step"));
+  auto GetArangeSize = [](auto start, auto end, auto step) -> int64_t {
+    using ElementType = std::decay_t<decltype(start)>;
+    PADDLE_ENFORCE_NE(step,
+                      0,
+                      ::common::errors::InvalidArgument(
+                          "The step of range op should not be 0."));
+
+    if ((start < end && step < 0) || (start > end && step > 0)) {
+      return 0;
+    } else {
+      return std::is_integral_v<ElementType>
+                 ? ((std::abs(end - start) + std::abs(step) - 1) /
+                    std::abs(step))
+                 : std::ceil(std::abs((end - start) / step));
+    }
+  };
+
   auto dtype =
       cinn::common::Str2Type(std::get<std::string>(attr_store.at("dtype")));
 
+#define EXPR_FROM_ATTR(type)                            \
+  type start_ = std::get<type>(attr_store.at("start")); \
+  type end_ = std::get<type>(attr_store.at("end"));     \
+  type step_ = std::get<type>(attr_store.at("step"));   \
+  arange_size = GetArangeSize(start_, end_, step_);     \
+  start = Expr(start_);                                 \
+  step = Expr(step_);
+
+  Expr start, step;
+  int64_t arange_size = 0;
+  if (dtype.is_float(32)) {
+    EXPR_FROM_ATTR(float)
+  } else if (dtype.is_float(64)) {
+    EXPR_FROM_ATTR(double)
+  } else if (dtype.is_int(32)) {
+    EXPR_FROM_ATTR(int)
+  } else if (dtype.is_int(64)) {
+    EXPR_FROM_ATTR(int64_t)
+  } else {
+    CINN_NOT_IMPLEMENTED
+  }
+
+#undef EXPR_FROM_ATTR
+
   framework::CINNCompute arange_compute([=](lang::Args args,
                                             lang::RetValue *ret) {
     PADDLE_ENFORCE_EQ(
@@ -1349,20 +1326,19 @@ std::shared_ptr<framework::OpStrategy> StrategyForArangeSymbolic(
     CINNValuePack pack_args = args[0];
 
     PADDLE_ENFORCE_EQ(pack_args.size(),
-                      1U,
+                      4U,
                       ::common::errors::InvalidArgument(
-                          "The number of input argument of arange should be at "
-                          "last 1. Please check."));
-    std::string tensor_name = pack_args[0].operator std::string();
-
-    auto out = pe::Arange(start, stop, step, dtype, tensor_name);
+                          "The number of input argument of arange should be 4"
+                          "(start, end, step, result). Please check."));
+    std::string tensor_name = pack_args[3].operator std::string();
+    auto out = pe::Arange(start, step, dtype, arange_size, tensor_name);
     std::vector<cinn::common::CINNValue> res;
     res.push_back(cinn::common::CINNValue(out));
     *ret = CINNValuePack{res};
   });
 
   auto strategy = std::make_shared<framework::OpStrategy>();
-  strategy->AddImpl(arange_compute, "strategy.reshape.x86", 1);
+  strategy->AddImpl(arange_compute, "strategy.arange.x86", 1);
   return strategy;
 }
 
@@ -1821,8 +1797,6 @@ CINN_REGISTER_HELPER(elementwise_ops) {
       .describe("Returns evenly spaced values within a given interval.")
       .set_num_inputs(0)
       .set_num_outputs(1)
-      .set_attr<cinn::hlir::framework::StrategyFunction>(
-          "CINNStrategy", cinn::hlir::op::StrategyForArange)
       .set_attr<cinn::hlir::framework::StrategyFunctionSymbolic>(
           "CINNStrategySymbolic", cinn::hlir::op::StrategyForArangeSymbolic)
       .set_attr<cinn::hlir::framework::OpPatternKind>(
 
@@ -307,19 +307,16 @@ ir::Tensor Store(const ir::Tensor& A, const std::string& name) {
   return res;
 }
 
-ir::Tensor Arange(const float start,
-                  const float stop,
-                  const float step,
+ir::Tensor Arange(Expr start,
+                  Expr step,
                   const Type& dtype,
+                  const int64_t size,
                   const std::string& output_name) {
-  int num = static_cast<int>(std::ceil((stop - start) / step));
   ir::Tensor res = lang::Compute(
-      {Expr(num)},
+      {Expr(size)},
       [=](const std::vector<ir::Expr>& indices) {
-        return ir::Cast::Make(
-            dtype,
-            Expr(start) +
-                Expr(step) * ir::Cast::Make(cinn::common::F32(), indices[0]));
+        return ir::Cast::Make(dtype,
+                              start + step * ir::Cast::Make(dtype, indices[0]));
       },
       output_name);
   return res;
 
@@ -149,10 +149,10 @@ ir::Tensor Store(const ir::Tensor& A,
                  const std::string& name = UniqName("T_Elementwise_Store_out"));
 
 ir::Tensor Arange(
-    const float start,
-    const float stop,
-    const float step,
+    Expr start,
+    Expr step,
     const Type& dtype,
+    const int64_t arange_size,
     const std::string& name = UniqName("T_Elementwise_Arange_out"));
 
 ir::Tensor Tril(const ir::Tensor& A,