Add cuda implementation for `prelu` backward pass (#18633)

* Add GPU implementation for `prelu` backward pass

test=develop

* Fix logic error in `prelu` GPU backward and simplify a bit

test=develop

* Fix `prelu` backward CUDA implementation

test=develop

CPU version was not used actually, so test passed

paddle/fluid/operators/prelu_op.cc

@@ -113,7 +113,7 @@ class PReluGradOp : public framework::OperatorWithKernel {
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext &ctx) const override {
     return framework::OpKernelType(ctx.Input<Tensor>("X")->type(),
-                                   platform::CPUPlace());
+                                   ctx.device_context());
   }
 };
 

paddle/fluid/operators/prelu_op.cu

@@ -14,11 +14,15 @@ limitations under the License. */
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/operators/math/prelu.h"
 #include "paddle/fluid/operators/prelu_op.h"
+#include "paddle/fluid/operators/reduce_ops/cub_reduce.h"
 #include "paddle/fluid/platform/cuda_primitives.h"
 
 namespace paddle {
 namespace operators {
 
+static const int CUDA_NUM_THREADS = 1024;
+static const int CUDA_MAX_NUM_BLOCKS = 65535;
+
 using Tensor = framework::Tensor;
 
 template <typename DeviceContext, typename T>
@@ -55,6 +59,145 @@ class CUDAPReluKernel : public framework::OpKernel<T> {
   }
 };
 
+namespace prelu {
+struct ElementWiseMode {};
+struct ChannelMode {};
+struct ScalarMode {};
+} /* namespace prelu */
+
+template <typename T, typename M>
+struct AlphaFunctor {
+  HOSTDEVICE inline T operator()(const T* alpha, size_t channel,
+                                 size_t spatial_size, size_t idx) const {}
+};
+
+template <typename T>
+struct AlphaFunctor<T, prelu::ElementWiseMode> {
+  HOSTDEVICE inline T operator()(const T* alpha, size_t channel,
+                                 size_t spatial_size, size_t idx) const {
+    return alpha[blockIdx.x * spatial_size + idx];
+  }
+};
+
+template <typename T>
+struct AlphaFunctor<T, prelu::ChannelMode> {
+  HOSTDEVICE inline T operator()(const T* alpha, size_t channel,
+                                 size_t spatial_size, size_t idx) const {
+    return alpha[blockIdx.x % channel];
+  }
+};
+
+template <typename T>
+struct AlphaFunctor<T, prelu::ScalarMode> {
+  HOSTDEVICE inline T operator()(const T* alpha, size_t channel,
+                                 size_t spatial_size, size_t idx) const {
+    return alpha[0];
+  }
+};
+
+template <typename T, typename M>
+__global__ void PReluGradElementWiseKernel(const T* x_ptr, const T* y_ptr,
+                                           const T* alpha_ptr, const T* dy_ptr,
+                                           T* dx_ptr, T* dalpha_ptr,
+                                           size_t channel,
+                                           size_t spatial_size) {
+  size_t offset = blockIdx.x * spatial_size;
+  AlphaFunctor<T, M> alpha_func;
+
+  for (size_t i = threadIdx.x; i < spatial_size; i += blockDim.x) {
+    T y = y_ptr[offset + i];
+    T x = x_ptr[offset + i];
+    T dy = dy_ptr[offset + i];
+    T alpha = alpha_func(alpha_ptr, channel, spatial_size, i);
+    if (dx_ptr != nullptr) dx_ptr[offset + i] = (y > 0) ? dy : alpha * dy;
+    if (dalpha_ptr != nullptr) dalpha_ptr[offset + i] = (x > 0) ? 0 : x * dy;
+  }
+}
+
+template <typename T, typename M>
+class PreluGradElementwiseFunctor {
+ public:
+  void operator()(cudaStream_t stream, const T* x, const T* y, const T* alpha,
+                  const T* dy, T* dx, T* dalpha, std::vector<int> input_shape) {
+    size_t unroll = input_shape[0] * input_shape[1];
+    size_t spatial_size = input_shape[2] * input_shape[3];
+    CHECK_LT(unroll, CUDA_MAX_NUM_BLOCKS);
+    PReluGradElementWiseKernel<T, M><<<unroll, CUDA_NUM_THREADS, 0, stream>>>(
+        x, y, alpha, dy, dx, dalpha, input_shape[1], spatial_size);
+  }
+};
+
+template <typename T>
+struct IdentityFunctor {
+  HOSTDEVICE inline T operator()(const T& x) const { return x; }
+};
+
+template <typename DeviceContext, typename T>
+class CUDAPReluGradKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& context) const override {
+    auto* x = context.Input<Tensor>("X");
+    auto* y = context.Input<Tensor>("Out");
+    auto* alpha = context.Input<Tensor>("Alpha");
+    auto* dx = context.Output<Tensor>(framework::GradVarName("X"));
+    auto* dy = context.Input<Tensor>(framework::GradVarName("Out"));
+    auto* dalpha = context.Output<Tensor>(framework::GradVarName("Alpha"));
+
+    const T* x_ptr = x->data<T>();
+    const T* y_ptr = y->data<T>();
+    const T* alpha_ptr = alpha->data<T>();
+    const T* dy_ptr = dy->data<T>();
+    T* dx_ptr = dx ? dx->mutable_data<T>(context.GetPlace()) : nullptr;
+    T* dalpha_ptr =
+        dalpha ? dalpha->mutable_data<T>(context.GetPlace()) : nullptr;
+
+    if (!dx && !dalpha) return;
+
+    auto& mode = context.Attr<std::string>("mode");
+
+    int numel = x->numel();
+    auto dim = x->dims();
+    std::vector<int> input_shape = framework::vectorize2int(dim);
+    auto stream = context.cuda_device_context().stream();
+
+    T* dalpha_tmp_ptr;
+    Tensor dalpha_tmp;
+    if (mode == "element" || dalpha_ptr == nullptr) {
+      dalpha_tmp_ptr = dalpha_ptr;
+    } else {
+      auto& dev_ctx = context.template device_context<DeviceContext>();
+      dalpha_tmp = context.AllocateTmpTensor<T, DeviceContext>(dim, dev_ctx);
+      dalpha_tmp_ptr = dalpha_tmp.mutable_data<T>(context.GetPlace());
+    }
+
+    if (mode == "element") {
+      PreluGradElementwiseFunctor<T, prelu::ElementWiseMode> prelu_grad;
+      prelu_grad(stream, x_ptr, y_ptr, alpha_ptr, dy_ptr, dx_ptr,
+                 dalpha_tmp_ptr, input_shape);
+    } else if (mode == "channel") {
+      PreluGradElementwiseFunctor<T, prelu::ChannelMode> prelu_grad;
+      prelu_grad(stream, x_ptr, y_ptr, alpha_ptr, dy_ptr, dx_ptr,
+                 dalpha_tmp_ptr, input_shape);
+    } else {
+      PreluGradElementwiseFunctor<T, prelu::ScalarMode> prelu_grad;
+      prelu_grad(stream, x_ptr, y_ptr, alpha_ptr, dy_ptr, dx_ptr,
+                 dalpha_tmp_ptr, input_shape);
+    }
+
+    if (mode == "element" || dalpha_tmp_ptr == nullptr) return;
+
+    std::vector<int> reduce_dims;
+    for (size_t i = 0; i < input_shape.size(); i++) {
+      if (mode == "channel" && i == 1) continue;
+      reduce_dims.push_back(i);
+    }
+
+    TensorReduce<T, T, cub::Sum, IdentityFunctor<T>>(
+        dalpha_tmp, dalpha, reduce_dims, static_cast<T>(0), cub::Sum(),
+        IdentityFunctor<T>(), stream);
+  }
+};
+
 }  // namespace operators
 }  // namespace paddle
 
@@ -62,3 +205,7 @@ namespace ops = paddle::operators;
 REGISTER_OP_CUDA_KERNEL(
     prelu, ops::CUDAPReluKernel<paddle::platform::CUDADeviceContext, float>,
     ops::CUDAPReluKernel<paddle::platform::CUDADeviceContext, double>);
+REGISTER_OP_CUDA_KERNEL(
+    prelu_grad,
+    ops::CUDAPReluGradKernel<paddle::platform::CUDADeviceContext, float>,
+    ops::CUDAPReluGradKernel<paddle::platform::CUDADeviceContext, double>);