[feat] add relu_backward

Author: Yin Hongyun

diopi_test/python/configs/diopi_configs.py

@@ -656,12 +656,11 @@
             args=[
                 {
                     "ins": ['input'],
+                    "requires_grad": [True],
                     "shape": ((), (1024,), (2, 4096), (64, 28, 28),
                               (32, 64, 112, 112), (64, 3, 7, 28, 28),
                               (0,), (256, 0), (8, 0, 128)),
-                    "dtype": [np.float16, np.float32, np.float64,
-                              np.int16, np.int32, np.int64,
-                              np.uint8, np.int8],
+                    "dtype": [np.float16, np.float32, np.float64],
                     "gen_fn": 'Genfunc.randn',
                 },
             ],

diopi_test/python/conformance/diopi_functions.py

@@ -423,6 +423,15 @@ def relu(input, inplace=False) -> Tensor:
     return unary_op(input, inplace, "diopiRelu")
 
 
+def relu_backward(input, grad_outputs, **kwargs) -> Tensor:
+    assert len(grad_outputs) == 1, "only accept 1 gradient to do backward"
+    grad_input = raw_like(input)
+    func = check_function("diopiReluBackward")
+    ret = func(input.context(), grad_input, grad_outputs[0], input)
+    check_returncode(ret)
+    return {"input": grad_input} if grad_input.requires_grad else {}
+
+
 def abs(input, inplace=False) -> Tensor:
     return unary_op(input, inplace, "diopiAbs")
 

impl/torch/functions/functions.cpp

@@ -158,6 +158,18 @@ diopiError_t diopiRelu(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiC
     return diopiSuccess;
 }
 
+diopiError_t diopiReluBackward(diopiContextHandle_t ctx, diopiConstTensorHandle_t grad_in,  diopiTensorHandle_t grad_out, diopiConstTensorHandle_t input){
+    impl::aten::setCurStream(ctx);
+
+    auto atGradOut = impl::aten::buildATen(grad_out);
+    auto atInput = impl::aten::buildATen(input);
+    auto atGradIn = impl::aten::buildATen(grad_in); 
+    auto mask = (atInput > 0).to(atGradOut.dtype());
+    atGradIn.copy_(atGradOut * mask);
+
+    return diopiSuccess;
+}
+
 diopiError_t diopiReluInp(diopiContextHandle_t ctx, diopiTensorHandle_t input) {
     impl::aten::setCurStream(ctx);
     auto atInput = impl::aten::buildATen(input);
@@ -4001,6 +4013,7 @@ diopiError_t diopiLinspace(diopiContextHandle_t ctx, diopiTensorHandle_t out, co
     return diopiSuccess;
 }
 
+
 diopiError_t diopiRoll(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t shifts, diopiSize_t dims) {
     impl::aten::setCurStream(ctx);
     auto atInput = impl::aten::buildATen(input);

proto/include/diopi/functions.h

@@ -236,6 +236,11 @@ DIOPI_API diopiError_t diopiBatchNormBackward(diopiContextHandle_t ctx, diopiTen
  */
 DIOPI_API diopiError_t diopiRelu(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input);
 
+/**
+ * @brief Computes the gradient of the rectified linear unit function.
+ */
+DIOPI_API diopiError_t diopiReluBackward(diopiContextHandle_t ctx, diopiConstTensorHandle_t grad_in, diopiTensorHandle_t grad_out, diopiConstTensorHandle_t input);
+
 /**
  * @brief The in-place version of diopiRelu().
  * @param[in] ctx Context environment.
@@ -701,6 +706,13 @@ DIOPI_API diopiError_t diopiAdaptiveMaxPool2dBackward(diopiContextHandle_t ctx,
  */
 DIOPI_API diopiError_t diopiDropout(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiTensorHandle_t mask, diopiConstTensorHandle_t input, double p,
                                     bool train, diopiGeneratorHandle_t generator);
+
+/**
+ *@brief Randomly zeroes some of the elements of the input tensor with probability p
+ */
+DIOPI_API diopiError_t diopiDropout(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiTensorHandle_t mask, diopiConstTensorHandle_t input, double p,
+                                    bool train, diopiGeneratorHandle_t generator);
+
 /**
  * @brief The in-place version of diopiDropout().
  * @param[in] ctx Context environment.