Warp Tile Size

Problem

Compute the output elements owned by one warp tile and the accumulator registers per thread, assuming 32 threads per warp: elements = warp_m * warp_n, per_thread = elements // 32.

The signature is warp_tile_size(warp_m, warp_n). Return [elements, per_thread].

Hint

Accumulator registers often dominate register pressure in GEMM kernels.

Test cases

case 1

in: [64,64]

expected: [4096,128]

case 2

in: [32,32]

expected: [1024,32]

case 3

in: [16,8]

expected: [128,4]

case 4

in: [128,64]

expected: [8192,256]

Starter codePython

def warp_tile_size(warp_m, warp_n):
    # Your code here
    pass

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