Update ops for Sigmoid and Tanh

Author: apaszke

Sigmoid.cu

@@ -4,22 +4,27 @@
 #include <THC/THCApply.cuh>
 
 template <typename T>
-struct sigmoidupdateOutput_functor
-{
- __device__ void operator()(T *output, const T *input) const
- {
- *output = ScalarConvert<double, T>::to(1./(1.+ exp(-*input)));
+struct SigmoidGradInputOp {
+ __device__ __forceinline__ void operator()(T* gradInput, const T *output, const T *gradOutput) const {
+ *gradInput = *gradOutput * (1.f - *output) * (*output);
  }
 };
 
-template <typename T>
-struct sigmoidupdateGradInput_functor
-{
- __device__ void operator()(T *gradInput, const T *output, const T *gradOutput) const
- {
- *gradInput = ScalarConvert<double, T>::to(*gradOutput * (1.-*output) * (*output));
+#ifdef CUDA_HALF_TENSOR
+template <>
+struct SigmoidGradInputOp<half> {
+ __device__ __forceinline__ void operator()(half* gradInput, const half *output, const half *gradOutput) const {
+#ifdef CUDA_HALF_INSTRUCTIONS
+ half one = __float2half(1.f);
+ *gradInput = __hmul(*gradOutput, __hmul(__hadd(one, __hneg(*output)), *output));
+#else
+ float out = __half2float(*output);
+ float go = __half2float(*gradOutput);
+ *gradInput = __float2half(go * (1.f - out) * out);
+#endif
  }
 };
+#endif
 
 #include "generic/Sigmoid.cu"
 #include "THCGenerateFloatTypes.h"

Tanh.cu

@@ -4,22 +4,32 @@
 #include <THC/THCApply.cuh>
 
 template <typename T>
-struct tanhupdateOutput_functor
+struct TanhGradInputOp
 {
- __device__ void operator()(T *output, const T *input) const
- {
- *output = tanh(*input);
+ __device__ __forceinline__ void operator()(T *gradInput,
+  const T *output, const T *gradOutput) const {
+ *gradInput = *gradOutput * (1.f - *output * *output);
  }
 };
 
-template <typename T>
-struct tanhupdateGradInput_functor
+#ifdef CUDA_HALF_TENSOR
+template <>
+struct TanhGradInputOp<half>
 {
- __device__ void operator()(T *gradInput, const T *output, const T *gradOutput) const
- {
- *gradInput = *gradOutput * (1 - *output * *output);
+ __device__ __forceinline__ void operator()(half *gradInput,
+ const half *output, const half *gradOutput) const {
+#ifdef CUDA_HALF_INSTRUCTIONS
+ const half one = __float2half(1.f);
+ const half out_square = __hmul(*output, *output);
+ *gradInput = __hmul(*gradOutput, __hadd(one, __hneg(out_square)));
+#else
+ float out = __half2float(*output);
+ float go = __half2float(*gradOutput);
+ *gradInput = __float2half(go * (1.f - out * out));
+#endif
  }
 };
+#endif
 
 #include "generic/Tanh.cu"
 #include "THCGenerateFloatTypes.h"

generic/Sigmoid.cu

@@ -10,8 +10,7 @@ void THNN_(Sigmoid_updateOutput)(
  THCTensor *output)
 {
  THCUNN_assertSameGPU(state, 2, input, output);
- THCTensor_(resizeAs)(state, output, input);
- THC_pointwiseApply2(state, output, input, sigmoidupdateOutput_functor<real>());
+ THCTensor_(sigmoid)(state, output, input);
 }
 
 void THNN_(Sigmoid_updateGradInput)(
@@ -24,7 +23,7 @@ void THNN_(Sigmoid_updateGradInput)(
  THCUNN_check_nElement(state, input, gradOutput);
  THCUNN_assertSameGPU(state, 3, output, gradOutput, gradInput);
  THCTensor_(resizeAs)(state, gradInput, output);
- THC_pointwiseApply3(state, gradInput, output, gradOutput, sigmoidupdateGradInput_functor<real>());
+ THC_pointwiseApply3(state, gradInput, output, gradOutput, SigmoidGradInputOp<real>());
 }
 
 #endif

generic/Tanh.cu

@@ -11,7 +11,7 @@ void THNN_(Tanh_updateOutput)(
 {
  THCUNN_assertSameGPU(state, 2, input, output);
  THCTensor_(resizeAs)(state, output, input);
- THC_pointwiseApply2(state, output, input, tanhupdateOutput_functor<real>());
+ THCTensor_(tanh)(state, output, input);
 }
 
 void THNN_(Tanh_updateGradInput)(
@@ -24,7 +24,7 @@ void THNN_(Tanh_updateGradInput)(
  THCUNN_check_shape(state, output, gradOutput);
  THCUNN_assertSameGPU(state, 3, output, gradOutput, gradInput);
  THCTensor_(resizeAs)(state, gradInput, output);
- THC_pointwiseApply3(state, gradInput, output, gradOutput, tanhupdateGradInput_functor<real>());
+ THC_pointwiseApply3(state, gradInput, output, gradOutput, TanhGradInputOp<real>());
 }
 
 #endif