Roofline Bound Check

Problem

Classify a kernel with the roofline model. Arithmetic intensity is flops / bytes_moved and the machine's ridge point is peak_flops_per_s / peak_bandwidth_per_s. If intensity >= ridge the kernel is compute bound with time flops / peak_flops; otherwise it is memory bound with time bytes_moved / peak_bandwidth.

The signature is roofline_bound(flops, bytes_moved, peak_flops_per_s, peak_bandwidth_per_s). Return [bound, time_seconds].

Hint

Transformer decode is memory bound because it moves gigabytes of weights for very few FLOPs per token.

Test cases

case 1

in: [1000000000000,1000000000,100000000000000,1000000000000]

expected: ["compute",0.01]

case 2

in: [10000000000,1000000000,100000000000000,1000000000000]

expected: ["memory",0.001]

case 3

in: [1000000000000,10000000000,100000000000000,1000000000000]

expected: ["compute",0.01]

case 4

in: [1000000000,1000000000000,100000000000000,1000000000000]

expected: ["memory",1]

Starter codePython

def roofline_bound(flops, bytes_moved, peak_flops_per_s, peak_bandwidth_per_s):
    # Your code here
    pass

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