[CINN] Fix arange float16 and bfloat16 support (PaddlePaddle#72669)

Author: Enigmatisms

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

@@ -554,6 +554,12 @@ class ArangeOpPattern
           case phi::DataType::INT32:
             input = phi::Scalar(input.to<int>());
             break;
+          case phi::DataType::FLOAT16:
+            input = phi::Scalar(input.to<float>());
+            break;
+          case phi::DataType::BFLOAT16:
+            input = phi::Scalar(input.to<float>());
+            break;
           default:
             input = phi::Scalar(input.to<int64_t>());
         }

paddle/cinn/hlir/op/elementwise.cc

@@ -1310,8 +1310,20 @@ std::shared_ptr<framework::OpStrategy> StrategyForArangeSymbolic(
     EXPR_FROM_ATTR(int)
   } else if (dtype.is_int(64)) {
     EXPR_FROM_ATTR(int64_t)
+  } else if (dtype.is_bfloat16()) {
+    EXPR_FROM_ATTR(float)
+    start->set_type(cinn::common::BFloat16());
+    step->set_type(cinn::common::BFloat16());
+  } else if (dtype.is_float16()) {
+    EXPR_FROM_ATTR(float)
+    start->set_type(cinn::common::Float16());
+    step->set_type(cinn::common::Float16());
   } else {
-    CINN_NOT_IMPLEMENTED
+    PADDLE_ENFORCE_NOT_NULL(
+        nullptr,
+        ::common::errors::InvalidArgument(
+            "The dtype of arange op should be float32, float64, int32, int64, "
+            "bfloat16 or float16."));
   }
 
 #undef EXPR_FROM_ATTR

paddle/phi/infermeta/nullary.cc

@@ -57,6 +57,10 @@ void ArangeInferMeta(const Scalar& start,
         GET_SIZE_GIVEN_TYPE(double)
       case DataType::INT32:
         GET_SIZE_GIVEN_TYPE(int)
+      case DataType::FLOAT16:
+        GET_SIZE_GIVEN_TYPE(float)
+      case DataType::BFLOAT16:
+        GET_SIZE_GIVEN_TYPE(float)
       default:
         GET_SIZE_GIVEN_TYPE(int64_t)
     }