【Hackathon 7th No.28】为 paddle.clip 进行功能增强

paddle/phi/kernels/clip_tensor_grad_kernel.h

@@ -0,0 +1,30 @@
+// Copyright (c) 2024 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/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ClipTensorGradKernel(const Context& dev_ctx,
+                          const DenseTensor& x,
+                          const DenseTensor& min,
+                          const DenseTensor& max,
+                          const DenseTensor& out_grad,
+                          DenseTensor* x_grad);
+
+}  // namespace phi

paddle/phi/kernels/clip_tensor_kernel.h

@@ -0,0 +1,29 @@
+// Copyright (c) 2024 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/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ClipTensorKernel(const Context& dev_ctx,
+                      const DenseTensor& x,
+                      const DenseTensor& min,
+                      const DenseTensor& max,
+                      DenseTensor* out);
+
+}  // namespace phi

paddle/phi/kernels/cpu/clip_tensor_grad_kernel.cc

@@ -0,0 +1,56 @@
+// 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/clip_tensor_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ClipTensorGradKernel(const Context& dev_ctx,
+                          const DenseTensor& x,
+                          const DenseTensor& min,
+                          const DenseTensor& max,
+                          const DenseTensor& out_grad,
+                          DenseTensor* x_grad) {
+  DenseTensor tem_min = phi::Cast<T, Context>(dev_ctx, min, x.dtype());
+  DenseTensor tem_max = phi::Cast<T, Context>(dev_ctx, max, x.dtype());
+
+  const T* x_data = x.data<T>();
+  const T* min_data = tem_min.data<T>();
+  const T* max_data = tem_max.data<T>();
+  auto numel = x.numel();
+  auto* dout = out_grad.data<T>();
+
+  auto* dx = dev_ctx.template Alloc<T>(x_grad);
+  for (int i = 0; i < numel; i++) {
+    dx[i] = (x_data[i] > min_data[i] && x_data[i] < max_data[i])
+                ? dout[i]
+                : static_cast<T>(0);
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(clip_tensor_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::ClipTensorGradKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}

paddle/phi/kernels/cpu/clip_tensor_kernel.cc

@@ -0,0 +1,56 @@
+// Copyright (c) 2024 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/clip_tensor_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ClipTensorKernel(const Context& dev_ctx,
+                      const DenseTensor& x,
+                      const DenseTensor& min,
+                      const DenseTensor& max,
+                      DenseTensor* out) {
+  DenseTensor tem_min = phi::Cast<T, Context>(dev_ctx, min, x.dtype());
+  DenseTensor tem_max = phi::Cast<T, Context>(dev_ctx, max, x.dtype());
+
+  const T* x_data = x.data<T>();
+  const T* min_data = tem_min.data<T>();
+  const T* max_data = tem_max.data<T>();
+
+  auto x_numel = x.numel();
+
+  T* out_data = dev_ctx.template Alloc<T>(out);
+
+  for (int i = 0; i < x_numel; i++) {
+    out_data[i] = x_data[i] < min_data[i] ? min_data[i] : x_data[i];
+    out_data[i] = out_data[i] > max_data[i] ? max_data[i] : out_data[i];
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(clip_tensor,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::ClipTensorKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}

paddle/phi/kernels/gpu/clip_tensor_grad_kernel.cu

@@ -0,0 +1,76 @@
+// Copyright (c) 2024 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/clip_tensor_grad_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/float16.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/funcs/broadcast_function.h"
+
+namespace phi {
+
+template <typename T>
+__global__ void ClipTensorGradFunctor(const int N,
+                                      const T* out_grad,
+                                      const T* x,
+                                      const T* min,
+                                      const T* max,
+                                      T* x_grad) {
+  int idx = blockDim.x * blockIdx.x + threadIdx.x;
+  for (; idx < N; idx += blockDim.x * gridDim.x) {
+    x_grad[idx] = (x[idx] > min[idx]) && (x[idx] < max[idx])
+                      ? out_grad[idx]
+                      : static_cast<T>(0);
+  }
+}
+
+template <typename T, typename Context>
+void ClipTensorGradKernel(const Context& dev_ctx,
+                          const DenseTensor& x,
+                          const DenseTensor& min,
+                          const DenseTensor& max,
+                          const DenseTensor& out_grad,
+                          DenseTensor* x_grad) {
+  DenseTensor tem_min = phi::Cast<T, Context>(dev_ctx, min, x.dtype());
+  DenseTensor tem_max = phi::Cast<T, Context>(dev_ctx, max, x.dtype());
+
+  const T* x_data = x.data<T>();
+  auto numel = x.numel();
+  const T* min_data = tem_min.data<T>();
+  const T* max_data = tem_max.data<T>();
+  const T* out_grad_data = out_grad.data<T>();
+
+  T* x_grad_data = dev_ctx.template Alloc<T>(x_grad);
+
+  auto stream = dev_ctx.stream();
+  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel);
+  ClipTensorGradFunctor<T>
+      <<<config.block_per_grid.x, config.thread_per_block.x, 0, stream>>>(
+          numel, out_grad_data, x_data, min_data, max_data, x_grad_data);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(clip_tensor_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::ClipTensorGradKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}

paddle/phi/kernels/gpu/clip_tensor_kernel.cu

@@ -0,0 +1,65 @@
+// Copyright (c) 2024 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/clip_tensor_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/common/float16.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/funcs/broadcast_function.h"
+#include "paddle/phi/kernels/funcs/elementwise_functor.h"
+
+namespace phi {
+
+template <typename T>
+struct ClipTensorFunctor {
+  inline HOSTDEVICE T operator()(const T x, const T min_, const T max_) const {
+    T x_ = x < min_ ? min_ : x;
+    T x__ = x_ > max_ ? max_ : x_;
+    return x__;
+  }
+};
+
+template <typename T, typename Context>
+void ClipTensorKernel(const Context& dev_ctx,
+                      const DenseTensor& x,
+                      const DenseTensor& min,
+                      const DenseTensor& max,
+                      DenseTensor* out) {
+  DenseTensor tem_min = phi::Cast<T, Context>(dev_ctx, min, x.dtype());
+  DenseTensor tem_max = phi::Cast<T, Context>(dev_ctx, max, x.dtype());
+
+  std::vector<const DenseTensor*> ins = {&x, &tem_min, &tem_max};
+  std::vector<DenseTensor*> outs = {out};
+  dev_ctx.template Alloc<T>(out);
+
+  ClipTensorFunctor<T> func;
+  funcs::ElementwiseKernel<T, ClipTensorFunctor<T>, 1>(
+      dev_ctx, ins, &outs, func);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(clip_tensor,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::ClipTensorKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}

paddle/phi/kernels/xpu/clip_tensor_grad_kernel.cc

@@ -0,0 +1,77 @@
+// Copyright (c) 2024 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/clip_tensor_grad_kernel.h"
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/infermeta/unary.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+#include "paddle/phi/kernels/compare_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/phi/kernels/logical_kernel.h"
+#include "paddle/phi/kernels/where_kernel.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ClipTensorGradKernel(const Context& dev_ctx,
+                          const DenseTensor& x,
+                          const DenseTensor& min,
+                          const DenseTensor& max,
+                          const DenseTensor& out_grad,
+                          DenseTensor* x_grad) {
+  DenseTensor ex_min = phi::Cast<T, Context>(dev_ctx, min, x.dtype());
+  DenseTensor ex_max = phi::Cast<T, Context>(dev_ctx, max, x.dtype());
+
+  phi::DenseTensor x_ls_min;
+  MetaTensor meta_x_ls_min(&x_ls_min);
+  UnchangedExceptDtypeInferMeta(x, &meta_x_ls_min);
+  meta_x_ls_min.set_dtype(phi::DataType::BOOL);
+  phi::LessThanKernel<T, Context>(dev_ctx, ex_min, x, &x_ls_min);
+
+  phi::DenseTensor x_ls_max;
+  MetaTensor meta_x_ls_max(&x_ls_max);
+  UnchangedExceptDtypeInferMeta(x, &meta_x_ls_max);
+  meta_x_ls_max.set_dtype(phi::DataType::BOOL);
+  phi::LessThanKernel<T, Context>(dev_ctx, x, ex_max, &x_ls_max);
+
+  phi::DenseTensor out;
+  MetaTensor meta_out(&out);
+  UnchangedExceptDtypeInferMeta(x, &meta_out);
+  meta_out.set_dtype(phi::DataType::BOOL);
+  phi::LogicalAndKernel<bool, Context>(dev_ctx, x_ls_min, x_ls_max, &out);
+
+  phi::DenseTensor zero_tensor;
+  MetaTensor meta_zero(&zero_tensor);
+  UnchangedInferMeta(x_grad, &meta_zero);
+  phi::FullKernel<T, Context>(dev_ctx,
+                         common::vectorize(x_grad->dims()),
+                         0.0f,
+                         zero_tensor.dtype(),
+                         &zero_tensor);
+  phi::WhereKernel<T, Context>(dev_ctx, out, out_grad, zero_tensor, x_grad);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(clip_tensor_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::ClipTensorGradKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   int64_t,
+                   int) {}