【Hackathon 4 No.20】Add i0 / i0e to paddle

paddle/phi/api/yaml/backward.yaml

@@ -775,6 +775,26 @@
   kernel :
     func : huber_loss_grad
 
+- backward_op : i0_grad
+  forward : i0 (Tensor x) -> Tensor(out)
+  args : (Tensor x, Tensor out_grad)
+  output : Tensor(x_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param : [x]
+  kernel :
+    func : i0_grad
+
+- backward_op : i0e_grad
+  forward : i0e (Tensor x) -> Tensor(out)
+  args : (Tensor x, Tensor out, Tensor out_grad)
+  output : Tensor(x_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param : [x]
+  kernel :
+    func : i0e_grad
+
 - backward_op : imag_grad
   forward : imag (Tensor x) -> Tensor(out)
   args : (Tensor out_grad)

paddle/phi/api/yaml/ops.yaml

@@ -920,6 +920,24 @@
   intermediate : residual
   backward : huber_loss_grad
 
+- op : i0
+  args : (Tensor x)
+  output : Tensor(out)
+  infer_meta :
+    func : UnchangedInferMeta
+  kernel :
+    func : i0
+  backward : i0_grad
+
+- op : i0e
+  args : (Tensor x)
+  output : Tensor(out)
+  infer_meta :
+    func : UnchangedInferMeta
+  kernel :
+    func : i0e
+  backward : i0e_grad
+
 - op : imag
   args : (Tensor x)
   output : Tensor (out)

paddle/phi/kernels/cpu/i0_grad_kernel.cc

@@ -0,0 +1,42 @@
+// Copyright (c) 2023 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/i0_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/for_range.h"
+#include "paddle/phi/kernels/impl/bessel_grad_kernel_impl.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void I0GradKernel(const Context& ctx,
+                  const DenseTensor& x,
+                  const DenseTensor& out_grad,
+                  DenseTensor* x_grad) {
+  auto size = x.numel();
+  auto* x_data = x.data<T>();
+  auto* out_grad_data = out_grad.data<T>();
+  auto* x_grad_data = ctx.template Alloc<T>(x_grad);
+
+  phi::funcs::ForRange<Context> for_range(ctx, size);
+  I0GradFunctor<T> functor(x_data, out_grad_data, x_grad_data, size);
+  for_range(functor);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(i0_grad, CPU, ALL_LAYOUT, phi::I0GradKernel, float, double) {
+}

paddle/phi/kernels/cpu/i0_kernel.cc

@@ -0,0 +1,37 @@
+/* Copyright (c) 2023 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/i0_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/for_range.h"
+#include "paddle/phi/kernels/impl/bessel_kernel_impl.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void I0Kernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
+  const int64_t size = x.numel();
+  const T* x_data = x.data<T>();
+  T* out_data = ctx.template Alloc<T>(out);
+
+  phi::funcs::ForRange<Context> for_range(ctx, size);
+  I0Functor<T> functor(x_data, out_data, size);
+  for_range(functor);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(i0, CPU, ALL_LAYOUT, phi::I0Kernel, float, double) {}

paddle/phi/kernels/cpu/i0e_grad_kernel.cc

@@ -0,0 +1,44 @@
+// Copyright (c) 2023 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/i0e_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/for_range.h"
+#include "paddle/phi/kernels/impl/bessel_grad_kernel_impl.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void I0eGradKernel(const Context& ctx,
+                   const DenseTensor& x,
+                   const DenseTensor& out,
+                   const DenseTensor& out_grad,
+                   DenseTensor* x_grad) {
+  auto size = x.numel();
+  auto* x_data = x.data<T>();
+  auto* out_data = out.data<T>();
+  auto* out_grad_data = out_grad.data<T>();
+  auto* x_gard_data = ctx.template Alloc<T>(x_grad);
+
+  phi::funcs::ForRange<Context> for_range(ctx, size);
+  I0eGradFunctor<T> functor(x_data, out_data, out_grad_data, x_gard_data, size);
+  for_range(functor);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    i0e_grad, CPU, ALL_LAYOUT, phi::I0eGradKernel, float, double) {}

paddle/phi/kernels/cpu/i0e_kernel.cc

@@ -0,0 +1,37 @@
+/* Copyright (c) 2023 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/i0e_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/for_range.h"
+#include "paddle/phi/kernels/impl/bessel_kernel_impl.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void I0eKernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
+  int64_t size = x.numel();
+  const T* x_data = x.data<T>();
+  T* out_data = ctx.template Alloc<T>(out);
+
+  phi::funcs::ForRange<Context> for_range(ctx, size);
+  I0eFunctor<T> functor(x_data, out_data, size);
+  for_range(functor);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(i0e, CPU, ALL_LAYOUT, phi::I0eKernel, float, double) {}

paddle/phi/kernels/gpu/i0_grad_kernel.cu

@@ -0,0 +1,39 @@
+/* Copyright (c) 2023 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/i0_grad_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/impl/bessel_grad_kernel_cuda_impl.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void I0GradKernel(const Context& ctx,
+                  const DenseTensor& x,
+                  const DenseTensor& out_grad,
+                  DenseTensor* x_grad) {
+  ctx.template Alloc<T>(x_grad);
+  std::vector<const DenseTensor*> ins = {&x, &out_grad};
+  std::vector<DenseTensor*> outs = {x_grad};
+  auto functor = CudaI0GradFunctor<T>();
+  phi::funcs::ElementwiseKernel<T>(ctx, ins, &outs, functor);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(i0_grad, GPU, ALL_LAYOUT, phi::I0GradKernel, float, double) {
+}

paddle/phi/kernels/gpu/i0_kernel.cu

@@ -0,0 +1,35 @@
+/* Copyright (c) 2023 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/i0_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/impl/bessel_kernel_cuda_impl.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void I0Kernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
+  ctx.template Alloc<T>(out);
+  std::vector<const DenseTensor*> ins = {&x};
+  std::vector<DenseTensor*> outs = {out};
+  auto functor = CudaI0Functor<T>();
+  phi::funcs::ElementwiseKernel<T>(ctx, ins, &outs, functor);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(i0, GPU, ALL_LAYOUT, phi::I0Kernel, float, double) {}

paddle/phi/kernels/gpu/i0e_grad_kernel.cu

@@ -0,0 +1,40 @@
+/* Copyright (c) 2023 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/i0e_grad_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/impl/bessel_grad_kernel_cuda_impl.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void I0eGradKernel(const Context& ctx,
+                   const DenseTensor& x,
+                   const DenseTensor& out,
+                   const DenseTensor& out_grad,
+                   DenseTensor* x_grad) {
+  ctx.template Alloc<T>(x_grad);
+  std::vector<const DenseTensor*> ins = {&x, &out, &out_grad};
+  std::vector<DenseTensor*> outs = {x_grad};
+  auto functor = CudaI0eGradFunctor<T>();
+  phi::funcs::ElementwiseKernel<T>(ctx, ins, &outs, functor);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    i0e_grad, GPU, ALL_LAYOUT, phi::I0eGradKernel, float, double) {}

paddle/phi/kernels/gpu/i0e_kernel.cu

@@ -0,0 +1,35 @@
+/* Copyright (c) 2023 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/i0e_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/impl/bessel_kernel_cuda_impl.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void I0eKernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
+  ctx.template Alloc<T>(out);
+  std::vector<const DenseTensor*> ins = {&x};
+  std::vector<DenseTensor*> outs = {out};
+  auto functor = CudaI0eFunctor<T>();
+  phi::funcs::ElementwiseKernel<T>(ctx, ins, &outs, functor);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(i0e, GPU, ALL_LAYOUT, phi::I0eKernel, float, double) {}

paddle/phi/kernels/i0_grad_kernel.h

@@ -0,0 +1,28 @@
+// Copyright (c) 2023 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 I0GradKernel(const Context& ctx,
+                  const DenseTensor& x,
+                  const DenseTensor& out_grad,
+                  DenseTensor* x_grad);
+
+}  // namespace phi