move trunc to pten (#39543)

* move trunc to pten

* modify according to YuanRisheng's comment

paddle/fluid/operators/trunc_op.cc

@@ -12,7 +12,8 @@ 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/trunc_op.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
 
 namespace paddle {
 namespace operators {
@@ -80,10 +81,3 @@ REGISTER_OPERATOR(trunc, ops::TruncOp, ops::TruncOpMaker,
                   ops::TruncGradOpMaker<paddle::imperative::OpBase>);
 
 REGISTER_OPERATOR(trunc_grad, ops::TruncGradOp);
-
-REGISTER_OP_CPU_KERNEL(trunc, ops::TruncKernel<float>, ops::TruncKernel<double>,
-                       ops::TruncKernel<int>, ops::TruncKernel<int64_t>);
-
-REGISTER_OP_CPU_KERNEL(trunc_grad, ops::TruncGradKernel<float>,
-                       ops::TruncGradKernel<double>, ops::TruncGradKernel<int>,
-                       ops::TruncGradKernel<int64_t>);

paddle/fluid/operators/trunc_op.cu

-/* Copyright (c) 2021 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/fluid/operators/trunc_op.h"
-#include "paddle/fluid/platform/device/gpu/gpu_info.h"
-#include "paddle/fluid/platform/device/gpu/gpu_primitives.h"
-
-namespace paddle {
-namespace operators {
-
-using platform::PADDLE_CUDA_NUM_THREADS;
-
-template <typename T>
-class TruncFunctor {
- public:
-  __device__ TruncFunctor(const T x) : x_(x) {}
-  __device__ T operator()() { return trunc(x_); }
-
- public:
-  const T x_;
-};
-
-template <>
-class TruncFunctor<int> {
- public:
-  __device__ TruncFunctor(const int x) : x_(x) {}
-  __device__ int operator()() { return x_; }
-
- public:
-  const int x_;
-};
-
-template <>
-class TruncFunctor<int64_t> {
- public:
-  __device__ TruncFunctor(const int64_t x) : x_(x) {}
-  __device__ int64_t operator()() { return x_; }
-
- public:
-  const int64_t x_;
-};
-
-template <typename T>
-__global__ void Trunc(const T* x, T* out, int64_t N) {
-  CUDA_KERNEL_LOOP(index, N) {
-    TruncFunctor<T> functor(x[index]);
-    out[index] = functor();
-  }
-}
-
-template <typename T>
-__global__ void TruncGrad(T* dx, int64_t N) {
-  CUDA_KERNEL_LOOP(index, N) { dx[index] = static_cast<T>(0.0); }
-}
-
-template <typename T>
-class TruncCUDAKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext& context) const override {
-    auto* x = context.Input<Tensor>("X");
-    auto* out = context.Output<Tensor>("Out");
-
-    const auto* x_data = x->data<T>();
-    auto* out_data = out->mutable_data<T>(context.GetPlace());
-
-    int64_t numel = x->numel();
-
-    int theads = PADDLE_CUDA_NUM_THREADS;
-    int blocks = (numel + theads - 1) / theads;
-
-    Trunc<<<blocks, theads>>>(x_data, out_data, numel);
-  }
-};
-
-template <typename T>
-class TruncCUDAGradKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext& context) const override {
-    auto* dout = context.Input<Tensor>(framework::GradVarName("Out"));
-    auto* dx = context.Output<Tensor>(framework::GradVarName("X"));
-
-    const auto* dout_data = dout->data<T>();
-    auto* dx_data = dx->mutable_data<T>(context.GetPlace());
-
-    int64_t numel = dout->numel();
-
-    int theads = PADDLE_CUDA_NUM_THREADS;
-    int blocks = (numel + theads - 1) / theads;
-
-    TruncGrad<<<blocks, theads>>>(dx_data, numel);
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-REGISTER_OP_CUDA_KERNEL(trunc, ops::TruncCUDAKernel<float>,
-                        ops::TruncCUDAKernel<double>, ops::TruncCUDAKernel<int>,
-                        ops::TruncCUDAKernel<int64_t>);
-
-REGISTER_OP_CUDA_KERNEL(trunc_grad, ops::TruncCUDAGradKernel<float>,
-                        ops::TruncCUDAGradKernel<double>,
-                        ops::TruncCUDAGradKernel<int>,
-                        ops::TruncCUDAGradKernel<int64_t>);

paddle/fluid/operators/trunc_op.h

-/* Copyright (c) 2021 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. */
-
-#pragma once
-
-#include <math.h>
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/fluid/framework/operator.h"
-
-namespace paddle {
-namespace operators {
-
-using Tensor = framework::Tensor;
-
-template <typename T>
-class TruncKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext& context) const override {
-    const Tensor* x = context.Input<Tensor>("X");
-    Tensor* out = context.Output<Tensor>("Out");
-
-    size_t numel = x->numel();
-    const T* x_data = x->data<T>();
-    T* out_data = out->mutable_data<T>(context.GetPlace());
-
-    for (size_t i = 0; i < numel; i++) {
-      out_data[i] = trunc(x_data[i]);
-    }
-  }
-};
-
-template <typename T>
-class TruncGradKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext& context) const override {
-    auto* dx = context.Output<Tensor>(framework::GradVarName("X"));
-    T* dx_data = dx->mutable_data<T>(context.GetPlace());
-
-    int numel = dx->numel();
-    memset(dx_data, 0.0, numel * sizeof(T));
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle

paddle/pten/kernels/cpu/trunc_grad_kernel.cc

+// 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/pten/kernels/trunc_grad_kernel.h"
+#include "paddle/pten/backends/cpu/cpu_context.h"
+#include "paddle/pten/core/kernel_registry.h"
+
+namespace pten {
+
+template <typename T, typename Context>
+void TruncGradKernel(const Context& dev_ctx,
+                     const DenseTensor& out_grad,
+                     DenseTensor* in_grad) {
+  T* dx_data = dev_ctx.template Alloc<T>(in_grad);
+
+  int numel = in_grad->numel();
+  memset(dx_data, 0.0, numel * sizeof(T));
+}
+
+}  // namespace pten
+
+PT_REGISTER_KERNEL(trunc_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   pten::TruncGradKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}

paddle/pten/kernels/cpu/trunc_kernel.cc

+// 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 <math.h>
+
+#include "paddle/pten/backends/cpu/cpu_context.h"
+#include "paddle/pten/core/kernel_registry.h"
+#include "paddle/pten/kernels/trunc_kernel.h"
+
+namespace pten {
+
+template <typename T, typename Context>
+void TruncKernel(const Context& dev_ctx,
+                 const DenseTensor& x,
+                 DenseTensor* out) {
+  size_t numel = x.numel();
+  const T* x_data = x.data<T>();
+  T* out_data = dev_ctx.template Alloc<T>(out);
+
+  for (size_t i = 0; i < numel; i++) {
+    out_data[i] = trunc(x_data[i]);
+  }
+}
+
+}  // namespace pten
+
+PT_REGISTER_KERNEL(
+    trunc, CPU, ALL_LAYOUT, pten::TruncKernel, float, double, int, int64_t) {}

paddle/pten/kernels/gpu/trunc_grad_kernel.cu

+// 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/fluid/platform/device/gpu/gpu_primitives.h"
+#include "paddle/pten/backends/gpu/gpu_context.h"
+#include "paddle/pten/backends/gpu/gpu_info.h"
+#include "paddle/pten/core/kernel_registry.h"
+#include "paddle/pten/kernels/trunc_grad_kernel.h"
+
+namespace pten {
+
+using paddle::platform::PADDLE_CUDA_NUM_THREADS;
+
+template <typename T>
+__global__ void TruncGrad(T* dx, int64_t N) {
+  CUDA_KERNEL_LOOP(index, N) { dx[index] = static_cast<T>(0.0); }
+}
+
+template <typename T, typename Context>
+void TruncGradKernel(const Context& dev_ctx,
+                     const DenseTensor& out_grad,
+                     DenseTensor* in_grad) {
+  const auto* out_grad_data = out_grad.data<T>();
+  T* in_grad_data = dev_ctx.template Alloc<T>(in_grad);
+
+  int64_t numel = out_grad.numel();
+
+  int theads = PADDLE_CUDA_NUM_THREADS;
+  int blocks = (numel + theads - 1) / theads;
+
+  TruncGrad<<<blocks, theads>>>(in_grad_data, numel);
+}
+
+}  // namespace pten
+
+PT_REGISTER_KERNEL(trunc_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   pten::TruncGradKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}

paddle/pten/kernels/gpu/trunc_kernel.cu

+// 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/fluid/platform/device/gpu/gpu_primitives.h"
+#include "paddle/pten/backends/gpu/gpu_context.h"
+#include "paddle/pten/backends/gpu/gpu_info.h"
+#include "paddle/pten/core/kernel_registry.h"
+#include "paddle/pten/kernels/trunc_kernel.h"
+
+namespace pten {
+
+using paddle::platform::PADDLE_CUDA_NUM_THREADS;
+
+template <typename T>
+class TruncFunctor {
+ public:
+  __device__ TruncFunctor(const T x) : x_(x) {}
+  __device__ T operator()() { return trunc(x_); }
+
+ public:
+  const T x_;
+};
+
+template <>
+class TruncFunctor<int> {
+ public:
+  __device__ TruncFunctor(const int x) : x_(x) {}
+  __device__ int operator()() { return x_; }
+
+ public:
+  const int x_;
+};
+
+template <>
+class TruncFunctor<int64_t> {
+ public:
+  __device__ TruncFunctor(const int64_t x) : x_(x) {}
+  __device__ int64_t operator()() { return x_; }
+
+ public:
+  const int64_t x_;
+};
+
+template <typename T>
+__global__ void Trunc(const T* x, T* out, int64_t N) {
+  CUDA_KERNEL_LOOP(index, N) {
+    TruncFunctor<T> functor(x[index]);
+    out[index] = functor();
+  }
+}
+
+template <typename T, typename Context>
+void TruncKernel(const Context& dev_ctx,
+                 const DenseTensor& x,
+                 DenseTensor* out) {
+  const auto* x_data = x.data<T>();
+  auto* out_data = dev_ctx.template Alloc<T>(out);
+
+  int64_t numel = x.numel();
+
+  int theads = PADDLE_CUDA_NUM_THREADS;
+  int blocks = (numel + theads - 1) / theads;
+
+  Trunc<<<blocks, theads>>>(x_data, out_data, numel);
+}
+
+}  // namespace pten
+
+PT_REGISTER_KERNEL(
+    trunc, GPU, ALL_LAYOUT, pten::TruncKernel, float, double, int, int64_t) {}

paddle/pten/kernels/trunc_grad_kernel.h

+// 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.
+
+#pragma once
+
+#include "paddle/pten/core/dense_tensor.h"
+
+namespace pten {
+
+template <typename T, typename Context>
+void TruncGradKernel(const Context& dev_ctx,
+                     const DenseTensor& out_grad,
+                     DenseTensor* in_grad);
+
+}  // namespace pten

paddle/pten/kernels/trunc_kernel.h

+// 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.
+
+#pragma once
+
+#include "paddle/pten/core/dense_tensor.h"
+
+namespace pten {
+
+template <typename T, typename Context>
+void TruncKernel(const Context& dev_ctx,
+                 const DenseTensor& x,
+                 DenseTensor* out);
+
+}  // namespace pten

paddle/pten/ops/compat/trunc_sig.cc

+// 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/pten/core/compat/op_utils.h"
+
+namespace pten {
+
+KernelSignature TruncOpArgumentMapping(const ArgumentMappingContext& ctx) {
+  return KernelSignature("trunc", {"X"}, {}, {"Out"});
+}
+
+KernelSignature TruncGradOpArgumentMapping(const ArgumentMappingContext& ctx) {
+  return KernelSignature(
+      "trunc_grad", {GradVarName("Out")}, {}, {GradVarName("X")});
+}
+
+}  // namespace pten
+
+PT_REGISTER_ARG_MAPPING_FN(trunc, pten::TruncOpArgumentMapping);
+PT_REGISTER_ARG_MAPPING_FN(trunc_grad, pten::TruncGradOpArgumentMapping);