Implicit GEMM Mapping

Problem

Map a convolution to a GEMM: output spatial P = (H + 2*pad - R) // stride + 1 and Q is computed the same way from W and S; the GEMM is M = N*P*Q, N_dim = K and K_dim = C*R*S.

The signature is implicit_gemm_dimensions(N, H, W, C, K, R, S, pad, stride). Return [M, N_dim, K_dim].

Hint

Implicit GEMM lets cuDNN reuse a single highly tuned matrix-multiply kernel for convolutions.

Test cases

case 1

in: [1,28,28,3,64,3,3,1,1]

expected: [784,64,27]

case 2

in: [2,14,14,16,32,3,3,0,1]

expected: [288,32,144]

case 3

in: [1,8,8,4,8,1,1,0,1]

expected: [64,8,4]

case 4

in: [1,32,32,64,128,3,3,1,2]

expected: [256,128,576]

Starter codePython

def implicit_gemm_dimensions(N, H, W, C, K, R, S, pad, stride):
    # Your code here
    pass

4 test cases · run your code against them in the editor