fastgrnncuda: precompute sigmoid

Author: MJ10

pytorch/edgeml_pytorch/cuda/fastgrnn_cuda.cpp

@@ -2,35 +2,27 @@
 
 #include <vector>
 
-// CUDA forward declarations
-
 std::vector<torch::Tensor> fastgrnn_cuda_forward(
     torch::Tensor input,
     torch::Tensor w,
     torch::Tensor u,
-    torch::Tensor bias_z,
-    torch::Tensor bias_h_prime,
-    torch::Tensor old_h,
+    torch::Tensor bias_gate,
+    torch::Tensor bias_update,
     torch::Tensor zeta,
-    torch::Tensor nu);
+    torch::Tensor nu,
+    torch::Tensor old_h);
 
 std::vector<torch::Tensor> fastgrnn_cuda_backward(
     torch::Tensor grad_h,
     torch::Tensor input,
     torch::Tensor old_h,
-    torch::Tensor z_t,
-    torch::Tensor h_prime_t,
-    torch::Tensor pre_comp,
+    torch::Tensor zeta,
+    torch::Tensor nu,
     torch::Tensor w,
     torch::Tensor u,
-    torch::Tensor bias_z,
-    torch::Tensor bias_h_prime,
-    torch::Tensor zeta,
-    torch::Tensor nu);
+    torch::Tensor z,
+    torch::Tensor h_prime);
 
-// C++ interface
-
-// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4.
 #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
 #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
 #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
@@ -39,62 +31,44 @@ std::vector<torch::Tensor> fastgrnn_forward(
     torch::Tensor input,
     torch::Tensor w,
     torch::Tensor u,
-    torch::Tensor bias_z,
-    torch::Tensor bias_h_prime,
-    torch::Tensor old_h,
+    torch::Tensor bias_gate,
+    torch::Tensor bias_update,
     torch::Tensor zeta,
-    torch::Tensor nu) {
+    torch::Tensor nu,
+    torch::Tensor old_h) {
   CHECK_INPUT(input);
   CHECK_INPUT(w);
   CHECK_INPUT(u);
-  CHECK_INPUT(bias_z);
-  CHECK_INPUT(bias_h_prime);
-  CHECK_INPUT(old_h);
+  CHECK_INPUT(bias_gate);
+  CHECK_INPUT(bias_update);
   CHECK_INPUT(zeta);
   CHECK_INPUT(nu);
+  CHECK_INPUT(old_h);
 
-  return fastgrnn_cuda_forward(input, w, u, bias_z, bias_h_prime, old_h, zeta, nu);
+  return fastgrnn_cuda_forward(input, w, u, bias_gate, bias_update, zeta, nu, old_h);
 }
 
 std::vector<torch::Tensor> fastgrnn_backward(
     torch::Tensor grad_h,
     torch::Tensor input,
     torch::Tensor old_h,
-    torch::Tensor z_t,
-    torch::Tensor h_prime_t,
-    torch::Tensor pre_comp,
+    torch::Tensor zeta,
+    torch::Tensor nu,
     torch::Tensor w,
     torch::Tensor u,
-    torch::Tensor bias_z,
-    torch::Tensor bias_h_prime,
-    torch::Tensor zeta,
-    torch::Tensor nu) {
+    torch::Tensor z,
+    torch::Tensor h_prime) {
   CHECK_INPUT(grad_h);
   CHECK_INPUT(input);
   CHECK_INPUT(old_h);
-  CHECK_INPUT(z_t);
-  CHECK_INPUT(h_prime_t);
-  CHECK_INPUT(pre_comp);
-  CHECK_INPUT(w);
-  CHECK_INPUT(u);
-  CHECK_INPUT(bias_z);
-  CHECK_INPUT(bias_h_prime);
   CHECK_INPUT(zeta);
   CHECK_INPUT(nu);
+  CHECK_INPUT(z);
+  CHECK_INPUT(h_prime);
+  CHECK_INPUT(w);
+  CHECK_INPUT(u);
 
-  return fastgrnn_cuda_backward(
-    grad_h,
-    input,
-    old_h,
-    z_t,
-    h_prime_t,
-    pre_comp,
-    w,
-    u,
-    bias_z,
-    bias_h_prime,
-    zeta,
-    nu);
+  return fastgrnn_cuda_backward(grad_h, input, old_h, zeta, nu, w, u, z, h_prime);
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {

pytorch/edgeml_pytorch/cuda/fastgrnn_cuda_kernel.cu

@@ -5,76 +5,77 @@
 
 #include <vector>
 
+__forceinline__ torch::Tensor d_sigmoid(torch::Tensor z) {
+  return (1 - z) * z;
+}
+
+__forceinline__ torch::Tensor d_tanh(torch::Tensor z) {
+  return 1 - z.pow(2);
+}
+
+
 namespace {
 template <typename scalar_t>
 __device__ __forceinline__ scalar_t sigmoid(scalar_t z) {
   return 1.0 / (1.0 + exp(-z));
 }
 
 template <typename scalar_t>
-__device__ __forceinline__ scalar_t d_sigmoid(scalar_t z) {
-  const auto s = sigmoid(z);
-  return (1.0 - s) * s;
+__device__ __forceinline__ scalar_t d_sigmoid(scalar_t sig_z) {
+  return (1.0 - sig_z) * sig_z;
 }
 
 template <typename scalar_t>
-__device__ __forceinline__ scalar_t d_tanh(scalar_t z) {
-  const auto t = tanh(z);
-  return 1 - (t * t);
+__device__ __forceinline__ scalar_t d_tanh(scalar_t tan_z) {
+  return 1 - (tan_z * tan_z);
 }
 
 template <typename scalar_t>
 __global__ void fastgrnn_cuda_forward_kernel(
+  torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> new_h,
+  torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> z,
+  torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> h_prime,
     const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> pre_comp,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> old_h,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> new_h,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> z_t,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> h_prime_t,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> bias_z,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> bias_h_prime,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> zeta,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> nu) {
-  //batch index
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> bias_z,
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> bias_h_prime,
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> nu,
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> zeta,
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> old_h) {
   const int n = blockIdx.y;
-  // column index
   const int c = blockIdx.x * blockDim.x + threadIdx.x;
-  if (c < pre_comp.size(1)){
-    z_t[n][c] = sigmoid(pre_comp[n][c] + bias_z[n][c]);
-    h_prime_t[n][c] = tanh(pre_comp[n][c] + bias_h_prime[n][c]);
-    
-    new_h[n][c] = (sigmoid(zeta[0][0]) * (1 - z_t[n][c]) + sigmoid(nu[0][0])) * h_prime_t[n][c] + z_t[n][c] * old_h[n][c];
+  if (c < old_h.size(1)){
+    z[n][c] = sigmoid(pre_comp[n][c] + bias_z[0][c]);
+    h_prime[n][c] = tanh(pre_comp[n][c] + bias_h_prime[0][c]);
+    new_h[n][c] = (zeta[0][0] * (1.0 - z[n][c]) + nu[0][0]) * h_prime[n][c] + old_h[n][c] * z[n][c];
   }
 }
 
+
 template <typename scalar_t>
 __global__ void fastgrnn_cuda_backward_kernel(
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_zeta,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_nu,
     torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_precomp,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_bias_z,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_bias_h_prime_t,
     torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_old_h,
+    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_bias_z,
+    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_bias_h_prime,
+    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_nu,
+    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_zeta,
     const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> grad_h,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> old_h,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> z_t,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> h_prime_t,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> pre_comp,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> bias_z,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> bias_h_prime,
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> z,
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> h_prime,
     const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> zeta,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> nu) {
-  //batch index
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> nu,
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_zeta_sigmoid,
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_nu_sigmoid,
+    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> old_h) {
   const int n = blockIdx.y;
-  // column index
   const int c = blockIdx.x * blockDim.x + threadIdx.x;
-  if (c < d_precomp.size(1)){
-    auto temp_grad = grad_h[n][c] * h_prime_t[n][c];
-    d_zeta[0][0] = temp_grad * (1 - z_t[n][c]) * d_sigmoid(zeta[0][0]);
-    d_nu[0][0] = temp_grad * d_sigmoid(nu[0][0]);
-    d_bias_z[n][c] = grad_h[n][c] * (sigmoid(zeta[0][0]) * -1 * h_prime_t[n][c] + old_h[n][c]) * d_sigmoid(pre_comp[n][c] + bias_z[n][c]);;
-    d_bias_h_prime_t[n][c] = grad_h[n][c] * (sigmoid(zeta[0][0]) * (1 - z_t[n][c]) + sigmoid(nu[0][0])) * d_tanh(pre_comp[n][c] + bias_h_prime[n][c]);
-    d_old_h[n][c] = grad_h[n][c] * z_t[n][c];
-    d_precomp[n][c] = d_bias_z[n][c] + d_bias_h_prime_t[n][c];
+  if (c < old_h.size(1)){ 
+    d_old_h[n][c] = z[n][c] * grad_h[n][c];
+    d_bias_h_prime[n][c] = (zeta[0][0] * (1.0 - z[n][c]) + nu[0][0]) * d_tanh(h_prime[n][c]) * grad_h[n][c];
+    d_bias_z[n][c] = (old_h[n][c] - zeta[0][0] * h_prime[n][c]) * d_sigmoid(z[n][c]) * grad_h[n][c];
+    d_precomp[n][c] = d_bias_z[n][c] + d_bias_h_prime[n][c];
+    d_zeta[n][c] = (1.0 - z[n][c]) * h_prime[n][c]*grad_h[n][c] * d_zeta_sigmoid[0][0];
+    d_nu[n][c] = h_prime[n][c] * grad_h[n][c] * d_nu_sigmoid[0][0];
   }
 }
 } // namespace
@@ -85,88 +86,86 @@ std::vector<torch::Tensor> fastgrnn_cuda_forward(
     torch::Tensor u,
     torch::Tensor bias_z,
     torch::Tensor bias_h_prime,
-    torch::Tensor old_h,
     torch::Tensor zeta,
-    torch::Tensor nu) {
-  auto w_comp = torch::mm(input, w);
-  auto u_comp = torch::mm(old_h, u);
-  auto pre_comp = torch::add(u_comp, w_comp);
-
+    torch::Tensor nu,
+    torch::Tensor old_h) {
+  
+  auto pre_comp = torch::addmm(torch::mm(input, w.transpose(0, 1)), old_h, u.transpose(0, 1));
+  nu = torch::sigmoid(nu);
+  zeta = torch::sigmoid(zeta);
   const auto batch_size = old_h.size(0);
   const auto state_size = old_h.size(1);
-
   auto new_h = torch::zeros_like(old_h);
-  auto z_t = torch::zeros_like(old_h);
-  auto h_prime_t = torch::zeros_like(old_h);
-
+  auto z = torch::zeros_like(old_h);
+  auto h_prime = torch::zeros_like(old_h);
   const int threads = 1024;
   const dim3 blocks((state_size + threads - 1) / threads, batch_size);
-
   AT_DISPATCH_FLOATING_TYPES(pre_comp.type(), "fastgrnn_forward_cuda", ([&] {
     fastgrnn_cuda_forward_kernel<scalar_t><<<blocks, threads>>>(
-        pre_comp.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        old_h.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
         new_h.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        z_t.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        h_prime_t.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        z.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        h_prime.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        pre_comp.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
         bias_z.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
         bias_h_prime.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        nu.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
         zeta.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        nu.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>());
+        old_h.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>());
   }));
-
-  return {new_h, z_t, h_prime_t, pre_comp};
+  return {new_h, z, h_prime};
 }
 
 std::vector<torch::Tensor> fastgrnn_cuda_backward(
     torch::Tensor grad_h,
     torch::Tensor input,
     torch::Tensor old_h,
-    torch::Tensor z_t,
-    torch::Tensor h_prime_t,
-    torch::Tensor pre_comp,
+    torch::Tensor zeta,
+    torch::Tensor nu,
     torch::Tensor w,
     torch::Tensor u,
-    torch::Tensor bias_z,
-    torch::Tensor bias_h_prime,
-    torch::Tensor zeta,
-    torch::Tensor nu) {
-  auto d_precomp = torch::zeros_like(pre_comp);
-  auto d_old_h = torch::zeros_like(old_h);
-  auto d_zeta = torch::zeros_like(zeta);
-  auto d_nu = torch::zeros_like(nu);
-  auto d_bias_z = torch::zeros_like(bias_z);
-  auto d_bias_h_prime = torch::zeros_like(bias_h_prime);
-
-  const auto batch_size = old_h.size(0);
-  const auto state_size = old_h.size(1);
-
-  const int threads = 1024;
-  const dim3 blocks((state_size + threads - 1) / threads, batch_size);
-
-  AT_DISPATCH_FLOATING_TYPES(pre_comp.type(), "fastgrnn_forward_cuda", ([&] {
+    torch::Tensor z,
+    torch::Tensor h_prime) {
+    auto d_precomp = torch::zeros_like(old_h);
+    auto d_bias_z = torch::zeros_like(old_h);
+    auto d_bias_h_prime = torch::zeros_like(old_h);
+    auto d_nu = torch::zeros_like(old_h);
+    auto d_zeta = torch::zeros_like(old_h);
+    auto d_old_h = torch::zeros_like(old_h);
+    zeta = torch::sigmoid(zeta);
+    nu = torch::sigmoid(nu);
+    auto d_nu_sigmoid = d_sigmoid(nu);
+    auto d_zeta_sigmoid = d_sigmoid(zeta);
+    const auto batch_size = old_h.size(0);
+    const auto state_size = old_h.size(1);
+
+    const int threads = 1024;
+    const dim3 blocks((state_size + threads - 1) / threads, batch_size);
+    AT_DISPATCH_FLOATING_TYPES(old_h.type(), "fastgrnn_backward_cuda", ([&] {
     fastgrnn_cuda_backward_kernel<scalar_t><<<blocks, threads>>>(
-        d_zeta.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        d_nu.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
         d_precomp.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        d_old_h.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
         d_bias_z.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
         d_bias_h_prime.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        d_old_h.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        d_nu.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        d_zeta.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
         grad_h.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        old_h.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        z_t.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        h_prime_t.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        pre_comp.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        bias_z.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        bias_h_prime.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        z.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        h_prime.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
         zeta.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        nu.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>());
+        nu.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        d_zeta_sigmoid.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        d_nu_sigmoid.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
+        old_h.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>());
   }));
 
-  d_old_h = torch::add(d_old_h, torch::mm(torch::add(d_bias_h_prime, d_bias_z), u.transpose(0, 1)));
-  auto d_input = torch::mm(d_precomp, w.transpose(0, 1));
-  auto d_w = torch::mm(input.transpose(0, 1), d_precomp);  
-  auto d_u = torch::mm(old_h.transpose(0, 1), d_precomp);
-
-  return {d_old_h, d_input, d_w, d_u, d_bias_z, d_bias_h_prime, d_nu, d_zeta};
+  d_old_h = torch::addmm(d_old_h, d_precomp, u);
+  auto d_input = torch::mm(d_precomp, w);  
+  auto d_w = torch::mm(d_precomp.transpose(0, 1), input);
+  auto d_u = torch::mm(d_precomp.transpose(0, 1), old_h);
+  d_bias_z = d_bias_z.sum(0, true);
+  d_bias_h_prime = d_bias_h_prime.sum(0, true);
+  d_zeta = (d_zeta.sum(0, true)).sum(1, true);
+  d_nu = (d_nu.sum(0, true)).sum(1, true);
+    
+  return {d_input, d_w, d_u, d_bias_z, d_bias_h_prime, d_zeta, d_nu, d_old_h};
 }