title: How to generalize fast matrix multiplication on GPU using numba
body: <p>Lately I've been trying to get into programming for GPUs in Python using the Numba library. I have been reading up on it on their website using the tutorial there and currently I'm stuck on their example, which can be found here: <a href="https://numba.pydata.org/numba-doc/latest/cuda/examples.html" rel="nofollow noreferrer">https://numba.pydata.org/numba-doc/latest/cuda/examples.html</a>. I'm attempting to generalize the example for the fast matrix multiplication a bit (which is of the form A*B=C). When testing I noticed that matrices with dimensions that are not perfectly divisible by the number of threads per block (TPB) do not yield a correct answer.</p>
  <p>I copied the code below from the example at <a href="https://numba.pydata.org/numba-doc/latest/cuda/examples.html" rel="nofollow noreferrer">https://numba.pydata.org/numba-doc/latest/cuda/examples.html</a> and created a very small test case with 4 by 4 matrices. If I choose TPB=2 everything is fine, but when I set TPB=3 it goes wrong. I understand that the code goes out of the bounds of the matrices, but I am unable to prevent this from happening (I tried some if statements on <em>ty + i * TPB</em> and <em>tx + i * TPB</em>, but these did not work.</p>
  <pre><code>from numba import cuda, float32
  import numpy as np
  import math
  
  @cuda.jit
  def fast_matmul(A, B, C):
      # Define an array in the shared memory
      # The size and type of the arrays must be known at compile time
      sA = cuda.shared.array(shape=(TPB, TPB), dtype=float32)
      sB = cuda.shared.array(shape=(TPB, TPB), dtype=float32)
  
      x, y = cuda.grid(2)
  
      tx = cuda.threadIdx.x
      ty = cuda.threadIdx.y
      bpg = cuda.gridDim.x    # blocks per grid
  
      if x &gt;= C.shape[0] and y &gt;= C.shape[1]:
          # Quit if (x, y) is outside of valid C boundary
          return
  
      # Each thread computes one element in the result matrix.
      # The dot product is chunked into dot products of TPB-long vectors.
      tmp = 0.
      for i in range(bpg):
          # Preload data into shared memory
          sA[tx, ty] = A[x, ty + i * TPB]
          sB[tx, ty] = B[tx + i * TPB, y]
  
          # Wait until all threads finish preloading
          cuda.syncthreads()
  
          # Computes partial product on the shared memory
          for j in range(TPB):
              tmp += sA[tx, j] * sB[j, ty]
  
          # Wait until all threads finish computing
          cuda.syncthreads()
  
      C[x, y] = tmp
  
  
  
  #%%
  
  x_h = np.arange(16).reshape([4,4])
  y_h = np.ones([4,4])
  z_h = np.zeros([4,4])
  
  x_d = cuda.to_device(x_h)
  y_d = cuda.to_device(y_h)
  z_d = cuda.to_device(z_h)
  
  TPB = 3
  threadsperblock = (TPB, TPB)
  blockspergrid_x = math.ceil(z_h.shape[0] / threadsperblock[0])
  blockspergrid_y = math.ceil(z_h.shape[1] / threadsperblock[1])
  blockspergrid = (blockspergrid_x, blockspergrid_y)
  
  fast_matmul[blockspergrid, threadsperblock](x_d, y_d, z_d)
  z_h = z_d.copy_to_host()
  print(z_h)
  </code></pre>
  <p>I would like to write some code that is not dependent on the matrices A, B, and C having dimensions that are perfectly divisible by the TPB, as these are sometimes out of my control. I understand that GPUs are only faster with matrix multiplication for very large matrices, but I wanted to use small examples to be able to check whether the answer is correct before applying it to actual data.</p>