Register Pressure Estimate

Problem

Estimate registers per thread for a GEMM warp tile: accumulator registers = (warp_m * warp_n) // threads, plus fragment_a + fragment_b + addressing registers.

The signature is register_pressure_estimate(warp_m, warp_n, threads, fragment_a, fragment_b, addressing=8). Return an integer.

Hint

Spilling accumulator registers to local memory can erase the benefit of a bigger tile.

Test cases

case 1

in: [64,64,32,4,2,8]

expected: 142

case 2

in: [32,32,32,4,2,8]

expected: 46

case 3

in: [128,64,32,8,4,16]

expected: 284

case 4

in: [16,16,32,2,2,0]

expected: 12

Starter codePython

def register_pressure_estimate(warp_m, warp_n, threads, fragment_a, fragment_b, addressing=8):
    # Your code here
    pass

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