support leaky_relu half kernel

oneflow/user/kernels/activation_kernels.cpp

@@ -27,6 +27,7 @@ namespace oneflow {
   REGISTER_SILU_KERNEL(DeviceType::kCPU, dtype);        \
   REGISTER_SELU_KERNEL(DeviceType::kCPU, dtype);        \
   REGISTER_SOFTSIGN_KERNEL(DeviceType::kCPU, dtype);    \
+  REGISTER_LEAKYRELU_KERNEL(DeviceType::kCPU, dtype);   \
   REGISTER_RELU_BACKWARD_KERNEL(DeviceType::kCPU, dtype);
 
 REGISTER_ACTIVATION_CPU_KERNEL(float);

oneflow/user/kernels/activation_kernels.cu

@@ -40,6 +40,26 @@ struct EluGradFunctor<half> {
   EluGradFunctor<float> float_functor;
 };
 
+template<>
+struct LeakyReluFunctor<half> {
+  OF_DEVICE_FUNC explicit LeakyReluFunctor(float alpha) : alpha(alpha) {}
+  __device__ half operator()(half x) const {
+    half zero = __float2half(0);
+    return (x > zero) ? x : __float2half(alpha) * x;
+  }
+  const float alpha;
+};
+
+template<>
+struct LeakyReluGradFunctor<half> {
+  OF_DEVICE_FUNC explicit LeakyReluGradFunctor(float alpha) : alpha(alpha) {}
+  __device__ half operator()(half x, half dy) const {
+    half zero = __float2half(0);
+    return (x > zero) ? dy : __float2half(alpha) * dy;
+  }
+  const float alpha;
+};
+
 template<>
 struct CeluFunctor<half> {
   OF_DEVICE_FUNC explicit CeluFunctor(float alpha)
@@ -173,6 +193,7 @@ struct ReluGradFunctor<half> {
   REGISTER_SILU_KERNEL(DeviceType::kCUDA, dtype);        \
   REGISTER_SELU_KERNEL(DeviceType::kCUDA, dtype);        \
   REGISTER_SOFTSIGN_KERNEL(DeviceType::kCUDA, dtype);    \
+  REGISTER_LEAKYRELU_KERNEL(DeviceType::kCUDA, dtype);   \
   REGISTER_RELU_BACKWARD_KERNEL(DeviceType::kCUDA, dtype);
 
 namespace {

oneflow/user/kernels/activation_kernels.h

@@ -19,6 +19,26 @@ limitations under the License.
 
 namespace oneflow {
 
+template<typename T>
+struct LeakyReluFunctor {
+  OF_DEVICE_FUNC explicit LeakyReluFunctor(float alpha) : alpha(alpha) {}
+  OF_DEVICE_FUNC T operator()(T x) const { return (x > 0) ? x : alpha * x; }
+  const T alpha;
+};
+
+template<typename T>
+struct LeakyReluGradFunctor {
+  OF_DEVICE_FUNC explicit LeakyReluGradFunctor(float alpha) : alpha(alpha) {}
+  OF_DEVICE_FUNC T operator()(T x, T dy) const {
+    if (alpha > 0) {
+      return dy > 0 ? dy : dy * alpha;
+    } else {
+      return (x > 0) ? dy : dy * alpha;
+    }
+  }
+  const T alpha;
+};
+
 template<typename T>
 struct EluFunctor {
   OF_DEVICE_FUNC explicit EluFunctor(float alpha) : alpha(alpha) {}
@@ -232,6 +252,20 @@ struct ReluGradFunctor {
       },                                                          \
       "dx", "x", "dy");
 
+#define REGISTER_LEAKYRELU_KERNEL(device, dtype)                            \
+  REGISTER_UNARY_ELEMWISE_USER_KERNEL(                                      \
+      device, "leaky_relu", LeakyReluFunctor, dtype, dtype,                 \
+      [](user_op::KernelComputeContext* ctx) {                              \
+        return LeakyReluFunctor<dtype>(ctx->Attr<float>("alpha"));          \
+      },                                                                    \
+      "y", "x");                                                            \
+  REGISTER_BINARY_ELEMWISE_USER_KERNEL(                                     \
+      device, "leaky_relu_grad", LeakyReluGradFunctor, dtype, dtype, dtype, \
+      [](user_op::KernelComputeContext* ctx) {                              \
+        return LeakyReluGradFunctor<dtype>(ctx->Attr<float>("alpha"));      \
+      },                                                                    \
+      "dx", "x", "dy");
+
 #define REGISTER_CELU_KERNEL(device, dtype)                        \
   REGISTER_UNARY_ELEMWISE_USER_KERNEL(                             \
       device, "celu", CeluFunctor, dtype, dtype,                   \

python/oneflow/test/modules/test_activation.py

@@ -587,6 +587,18 @@ def test_leakyrelu_module_with_random_data(test_case):
         y = m(x)
         return y
 
+    @autotest()
+    @unittest.skipIf(os.getenv("ONEFLOW_TEST_CPU_ONLY"), "only test cpu cases")
+    def test_leakyrelu_module_with_half_random_data(test_case):
+        m = torch.nn.LeakyReLU(negative_slope=random() | nothing())
+        m.train(random())
+        device = random_device()
+        m.to(device)
+        x = random_tensor().to(device)
+        x = x.to(torch.float16)
+        y = m(x)
+        return y
+
     @autotest()
     def test_leakyrelu_module_with_0dim_data(test_case):
         m = torch.nn.LeakyReLU(negative_slope=random() | nothing())