1 #define TORCH_ASSERT_NO_OPERATORS
2 #define _USE_MATH_DEFINES
3
4 #include <ATen/native/Activation.h>
5
6 #include <cmath>
7
8 #include <thrust/tuple.h>
9
10 #include <ATen/AccumulateType.h>
11 #include <ATen/Dispatch.h>
12 #include <ATen/core/TensorBase.h>
13 #include <c10/core/Scalar.h>
14 #include <c10/cuda/CUDAMathCompat.h>
15 #include <ATen/cuda/ApplyGridUtils.cuh>
16 #include <ATen/cuda/detail/OffsetCalculator.cuh>
17 #include <ATen/native/cuda/Loops.cuh>
18
19 namespace at::native {
20 namespace {
21
leaky_relu_kernel(TensorIteratorBase & iter,const Scalar & negval_)22 void leaky_relu_kernel(TensorIteratorBase& iter, const Scalar& negval_) {
23 AT_DISPATCH_FLOATING_TYPES_AND2(
24 at::ScalarType::Half,
25 at::ScalarType::BFloat16,
26 iter.dtype(),
27 "leaky_relu_cuda",
28 [&]() {
29 using opmath_t = at::opmath_type<scalar_t>;
30 auto negval = negval_.to<opmath_t>();
31 gpu_kernel(iter, [negval] GPU_LAMBDA(scalar_t a) -> scalar_t {
32 opmath_t aop = static_cast<opmath_t>(a);
33 return aop > opmath_t(0) ? aop : aop * negval;
34 });
35 });
36 }
37
leaky_relu_backward_kernel(TensorIteratorBase & iter,const Scalar & negval_)38 void leaky_relu_backward_kernel(
39 TensorIteratorBase& iter,
40 const Scalar& negval_) {
41 AT_DISPATCH_FLOATING_TYPES_AND2(
42 at::ScalarType::Half,
43 at::ScalarType::BFloat16,
44 iter.dtype(),
45 "leaky_relu_backward_cuda",
46 [&]() {
47 using opmath_t = at::opmath_type<scalar_t>;
48 auto negval = negval_.to<opmath_t>();
49 gpu_kernel(
50 iter, [negval] GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
51 opmath_t aop = static_cast<opmath_t>(a);
52 opmath_t bop = static_cast<opmath_t>(b);
53 return aop > opmath_t(0) ? bop : bop * negval;
54 });
55 });
56 }
57 } // namespace
58
59 REGISTER_DISPATCH(leaky_relu_stub, &leaky_relu_kernel);
60 REGISTER_DISPATCH(leaky_relu_backward_stub, &leaky_relu_backward_kernel);
61
62 } // namespace at::native
63