optimize prelu alpha grad

oneflow/core/functional/impl/activation_functor.cpp

@@ -91,7 +91,13 @@ class PReluGradFunctor {
   }
   Maybe<TensorTuple> operator()(const std::shared_ptr<Tensor>& dy, const std::shared_ptr<Tensor>& x,
                                 const std::shared_ptr<Tensor>& alpha) const {
-    return OpInterpUtil::Dispatch<one::TensorTuple>(*op_, {dy, x, alpha});
+    MutableAttrMap attrs;
+    if(alpha->requires_grad()){
+      JUST(attrs.SetAttr<bool>("alpha_requires_grad", true));
+    }else{
+      JUST(attrs.SetAttr<bool>("alpha_requires_grad", false));
+    }
+    return OpInterpUtil::Dispatch<one::TensorTuple>(*op_, {dy, x, alpha}, attrs);
   }
 
  private:

oneflow/ir/include/OneFlow/OneFlowUserOps.td

@@ -580,6 +580,9 @@ def OneFlow_PreluOp : OneFlow_BaseOp<"prelu", [NoSideEffect, DeclareOpInterfaceM
   let output = (outs
     OneFlow_Tensor:$y
   );
+  let attrs = (ins
+    DefaultValuedAttr<BoolAttr, "false">:$alpha_requires_grad
+  );
   let has_logical_tensor_desc_infer_fn = 1;
   let has_physical_tensor_desc_infer_fn = 1;
   let has_get_sbp_fn = 1;

oneflow/user/kernels/prelu_kernel.cu

@@ -380,6 +380,7 @@ class GpuPReluGradKernel final : public user_op::OpKernel {
     user_op::Tensor* dx = ctx->Tensor4ArgNameAndIndex("dx", 0);
     user_op::Tensor* alpha_diff = ctx->Tensor4ArgNameAndIndex("alpha_diff", 0);
     user_op::Tensor* tmp_buffer = ctx->Tensor4ArgNameAndIndex("tmp_buffer", 0);
+    const bool alpha_requires_grad = ctx->Attr<bool>("alpha_requires_grad");
     const int32_t elem_cnt = x->shape().elem_cnt();
     T* broadcasted_alpha_diff = tmp_buffer->mut_dptr<T>();
     T* reduce_sum_tmp_buf = reinterpret_cast<T*>(tmp_buffer->mut_dptr<char>()
@@ -400,10 +401,12 @@ class GpuPReluGradKernel final : public user_op::OpKernel {
                                     alpha->dptr<T>(), dy->dptr<T>(), dx->mut_dptr<T>(),
                                     broadcasted_alpha_diff);
     }
-    NdarrayUtil<DeviceType::kCUDA, T>::ReduceSum(
+    if(alpha_requires_grad){
+      NdarrayUtil<DeviceType::kCUDA, T>::ReduceSum(
         ctx->stream(), XpuVarNdarray<T>(left_extended_shape, alpha_diff->mut_dptr<T>()),
         XpuVarNdarray<const T>(x->shape(), broadcasted_alpha_diff),
         XpuVarNdarray<T>(x->shape(), reduce_sum_tmp_buf));
+    }
   }
   bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
 };

oneflow/user/ops/prelu_op.cpp

@@ -118,6 +118,7 @@ REGISTER_USER_OP_GRAD("prelu").SetGenBackwardOpConfFn([](const user_op::UserOpWr
                                              .Input("x", op.input("x", 0))
                                              .Input("dy", op.GetGradTensorWithOpOutput("y", 0))
                                              .Input("alpha", op.input("alpha", 0))
+                                             .Attr("alpha_requires_grad", op.attr<bool>("alpha_requires_grad"))
                                              .Output("dx")
                                              .Output("alpha_diff")
                                              .Build();