Optimizing GRU CUDA kernel

username: kevin_zhang
cooked: <p>For practice, I implemented the forward() and backward() CUDA kernels for GRU, following this tutorial:<br>
<a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" class="onebox" target="_blank" rel="nofollow noopener">https://pytorch.org/tutorials/advanced/cpp_extension.html</a></p>
<p>However, even though my CUDA implementation is faster than Python code, it still lags behind the PyTorch Library GRU implementation:</p>
<p>Forward pass average time:<br>
Python on GPU:                    323<br>
Fused cuda on GPU:            178<br>
Torch library GRU on GPU:   94</p>
<p>I am trying to see whether there are things I can do to make my implementation match the library performance.</p>
<p>Any advice?</p>
<p>My GRU implementation can be found <a href="https://github.com/kevinghst/GRU-CUDA/tree/master/gru" rel="nofollow noopener">here</a></p>
<p>Below is my code for the forward pass:</p>
<pre><code class="lang-auto">template &lt;typename scalar_t&gt;
__global__ void gru_cuda_forward_kernel(
    const torch::PackedTensorAccessor&lt;scalar_t,3,torch::RestrictPtrTraits,size_t&gt; gate_x,
    const torch::PackedTensorAccessor&lt;scalar_t,3,torch::RestrictPtrTraits,size_t&gt; gate_h,
    const torch::PackedTensorAccessor&lt;scalar_t,2,torch::RestrictPtrTraits,size_t&gt; old_h,
    torch::PackedTensorAccessor&lt;scalar_t,2,torch::RestrictPtrTraits,size_t&gt; resetgate,
    torch::PackedTensorAccessor&lt;scalar_t,2,torch::RestrictPtrTraits,size_t&gt; inputgate,
    torch::PackedTensorAccessor&lt;scalar_t,2,torch::RestrictPtrTraits,size_t&gt; newgate,
    torch::PackedTensorAccessor&lt;scalar_t,2,torch::RestrictPtrTraits,size_t&gt; new_h) {
      //batch index
      const int n = blockIdx.y;
      // column index
      const int c = blockIdx.x * blockDim.x + threadIdx.x;

      if (c &lt; gate_x.size(2) &amp;&amp; c &lt; gate_h.size(2)) {
        resetgate[n][c] = sigmoid(gate_x[n][0][c] + gate_h[n][0][c]);
        inputgate[n][c] = sigmoid(gate_x[n][1][c] + gate_h[n][1][c]);
        newgate[n][c] = tanh(gate_x[n][2][c] + resetgate[n][c] * gate_h[n][2][c]);
        new_h[n][c] = newgate[n][c] + inputgate[n][c] * (old_h[n][c] - newgate[n][c]);
      }
}

std::vector&lt;torch::Tensor&gt; gru_cuda_forward(
    torch::Tensor input,
    torch::Tensor x2h_w,
    torch::Tensor h2h_w,
    torch::Tensor x2h_b,
    torch::Tensor h2h_b,
    torch::Tensor old_h) {
        auto gate_x_weights = torch::addmm(x2h_b, input, x2h_w.transpose(0, 1));
        auto gate_h_weights = torch::addmm(h2h_b, old_h, h2h_w.transpose(0, 1));

        const auto batch_size = old_h.size(0);
        const auto state_size = old_h.size(1);

        auto gate_x = gate_x_weights.reshape({batch_size, 3, state_size});
        auto gate_h = gate_h_weights.reshape({batch_size, 3, state_size});

        auto resetgate = torch::zeros_like(old_h);
        auto inputgate = torch::zeros_like(old_h);
        auto newgate = torch::zeros_like(old_h);
        auto new_h = torch::zeros_like(old_h);

        const int threads = 1024;
        const dim3 blocks((state_size + threads - 1) / threads, batch_size);

        AT_DISPATCH_FLOATING_TYPES(gate_x.type(), "gru_forward_cuda", ([&amp;] {
            gru_cuda_forward_kernel&lt;scalar_t&gt;&lt;&lt;&lt;blocks, threads&gt;&gt;&gt;(
                gate_x.packed_accessor&lt;scalar_t,3,torch::RestrictPtrTraits,size_t&gt;(),
                gate_h.packed_accessor&lt;scalar_t,3,torch::RestrictPtrTraits,size_t&gt;(),
                old_h.packed_accessor&lt;scalar_t,2,torch::RestrictPtrTraits,size_t&gt;(),
                resetgate.packed_accessor&lt;scalar_t,2,torch::RestrictPtrTraits,size_t&gt;(),
                inputgate.packed_accessor&lt;scalar_t,2,torch::RestrictPtrTraits,size_t&gt;(),
                newgate.packed_accessor&lt;scalar_t,2,torch::RestrictPtrTraits,size_t&gt;(),
                new_h.packed_accessor&lt;scalar_t,2,torch::RestrictPtrTraits,size_t&gt;());
          }));
        

        return {
          new_h, resetgate, inputgate, newgate, gate_h
        };
}
</code></pre>
<p>Thanks.</p>