[Phi] Migrate linspace op to phi (#40124)

* [Phi] Migrate linspace op * [Phi] Migrate linspace op * [Phi] Fix linspace op * [PHI] rename data_tranform to data_type_transform * [PHI] Fix DECLARE and PD

Author: chenenquan

paddle/fluid/operators/linspace_op.cc

@@ -12,9 +12,13 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
-#include "paddle/fluid/operators/linspace_op.h"
 #include <string>
+
+#include "paddle/fluid/framework/infershape_utils.h"
+#include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/framework/op_version_registry.h"
+#include "paddle/phi/core/infermeta_utils.h"
+#include "paddle/phi/infermeta/ternary.h"
 
 namespace paddle {
 namespace operators {
@@ -23,33 +27,6 @@ class LinspaceOp : public framework::OperatorWithKernel {
  public:
  using framework::OperatorWithKernel::OperatorWithKernel;
 
- void InferShape(framework::InferShapeContext *ctx) const override {
- OP_INOUT_CHECK(ctx->HasInput("Start"), "Input", "Start", "linspace");
- OP_INOUT_CHECK(ctx->HasInput("Stop"), "Input", "Stop", "linspace");
- OP_INOUT_CHECK(ctx->HasInput("Num"), "Input", "Num", "linspace");
- OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "linspace");
-
- auto s_dims = ctx->GetInputDim("Start");
- PADDLE_ENFORCE_EQ((s_dims.size() == 1) && (s_dims[0] == 1), true,
- platform::errors::InvalidArgument(
- "The shape of Input(Start) must be [1],"
- "but received input shape is [%s].",
- s_dims));
- auto e_dims = ctx->GetInputDim("Stop");
- PADDLE_ENFORCE_EQ((e_dims.size() == 1) && (e_dims[0] == 1), true,
- platform::errors::InvalidArgument(
- "The shape of Input(Stop) must be [1],"
- "but received input shape is [%s].",
- e_dims));
- auto step_dims = ctx->GetInputDim("Num");
- PADDLE_ENFORCE_EQ(
- (step_dims.size() == 1) && (step_dims[0] == 1), true,
- platform::errors::InvalidArgument("The shape of Input(Num) must be [1],"
- "but received input shape is [%s].",
- step_dims));
- ctx->SetOutputDim("Out", {-1});
- }
-
  protected:
  framework::OpKernelType GetExpectedKernelType(
  const framework::ExecutionContext &ctx) const override {
@@ -88,11 +65,13 @@ class LinspaceOpMaker : public framework::OpProtoAndCheckerMaker {
 } // namespace paddle
 
 namespace ops = paddle::operators;
-REGISTER_OP_WITHOUT_GRADIENT(linspace, ops::LinspaceOp, ops::LinspaceOpMaker);
-REGISTER_OP_CPU_KERNEL(linspace, ops::CPULinspaceKernel<float>,
- ops::CPULinspaceKernel<int32_t>,
- ops::CPULinspaceKernel<int64_t>,
- ops::CPULinspaceKernel<double>);
+DECLARE_INFER_SHAPE_FUNCTOR(linspace, LinspaceInferShapeFunctor,
+ PD_INFER_META(phi::LinspaceInferMeta));
+REGISTER_OPERATOR(
+ linspace, ops::LinspaceOp, ops::LinspaceOpMaker,
+ paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
+ paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>,
+ LinspaceInferShapeFunctor);
 
 REGISTER_OP_VERSION(linspace)
  .AddCheckpoint(

paddle/fluid/operators/linspace_op.cu

@@ -209,4 +209,33 @@ void LerpInferMeta(const MetaTensor& x,
  out->share_lod(x);
 }
 
+void LinspaceInferMeta(const MetaTensor& start,
+ const MetaTensor& stop,
+ const MetaTensor& number,
+ MetaTensor* out) {
+ auto s_dims = start.dims();
+ PADDLE_ENFORCE_EQ(
+ (s_dims.size() == 1) && (s_dims[0] == 1),
+ true,
+ phi::errors::InvalidArgument("The shape of Input(Start) must be [1],"
+ "but received input shape is [%s].",
+ s_dims));
+ auto e_dims = stop.dims();
+ PADDLE_ENFORCE_EQ(
+ (e_dims.size() == 1) && (e_dims[0] == 1),
+ true,
+ phi::errors::InvalidArgument("The shape of Input(Stop) must be [1],"
+ "but received input shape is [%s].",
+ e_dims));
+ auto step_dims = number.dims();
+ PADDLE_ENFORCE_EQ(
+ (step_dims.size() == 1) && (step_dims[0] == 1),
+ true,
+ phi::errors::InvalidArgument("The shape of Input(Num) must be [1],"
+ "but received input shape is [%s].",
+ step_dims));
+ out->set_dims(phi::make_ddim({-1}));
+ out->set_dtype(start.dtype());
+}
+
 } // namespace phi

paddle/fluid/operators/linspace_op.h

@@ -58,4 +58,9 @@ void LerpInferMeta(const MetaTensor& x,
  const MetaTensor& weight,
  MetaTensor* out);
 
+void LinspaceInferMeta(const MetaTensor& start,
+ const MetaTensor& stop,
+ const MetaTensor& number,
+ MetaTensor* out);
+
 } // namespace phi

paddle/phi/infermeta/ternary.cc

@@ -0,0 +1,71 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/linspace_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/data_type_transform.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void LinspaceKernel(const Context& ctx,
+ const DenseTensor& start,
+ const DenseTensor& stop,
+ const DenseTensor& number,
+ DataType dtype,
+ DenseTensor* out) {
+ int32_t num = number.data<int32_t>()[0];
+ auto start_t = phi::funcs::TransDataType(ctx, start, dtype);
+ auto stop_t = phi::funcs::TransDataType(ctx, stop, dtype);
+
+ T start_data = start_t.template data<T>()[0];
+ T stop_data = stop_t.template data<T>()[0];
+ PADDLE_ENFORCE_GT(
+ num,
+ 0,
+ phi::errors::InvalidArgument("The num of linspace op should be larger "
+ "than 0, but received num is %d",
+ num));
+
+ out->Resize(phi::make_ddim({num}));
+ T* out_data = ctx.template Alloc<T>(out);
+
+ if (num > 1) {
+ // step should be of double type for all types
+ double step = (static_cast<double>(stop_data - start_data)) / (num - 1);
+ int half_num = num / 2;
+ for (int i = 0; i < num; ++i) {
+ if (i < half_num) {
+ out_data[i] = static_cast<T>(start_data + step * i);
+ } else {
+ out_data[i] = static_cast<T>(stop_data - step * (num - i - 1));
+ }
+ }
+ } else {
+ out_data[0] = static_cast<T>(start_data);
+ }
+}
+
+} // namespace phi
+
+PD_REGISTER_KERNEL(linspace,
+ CPU,
+ ALL_LAYOUT,
+ phi::LinspaceKernel,
+ float,
+ int32_t,
+ int64_t,
+ double) {}