Roofline Crossover Batch Size

Problem

For a matmul with inner dimension k, output dimension n and variable batch m, arithmetic intensity is 2 * m * k * n / (bytes_per_value * (m * k + k * n + m * n)). Given peak_flops and memory bandwidth, find the smallest integer batch size whose intensity reaches the machine ridge point peak_flops / bandwidth.

The signature is roofline_crossover_batch(k, n, peak_flops, bandwidth, bytes_per_value=2, max_batch=65536). Return -1 if no batch up to max_batch reaches the ridge point.

Hint

As batch grows the weight term is amortized; if the asymptotic intensity 2kn / (b(k+n)) stays below the ridge point there is no crossover.

Test cases

case 1

in: [4096,4096,312000000000000,2000000000000,2,4096]

expected: 169

case 2

in: [4096,4096,10000000000000000,1000000000000,2,65536]

expected: -1

case 3

in: [1024,2048,100000000000000,1000000000000,2,65536]

expected: 118

case 4

in: [8,8,1000000000000,100000000000000,2,65536]

expected: 1

case 5

in: [4096,4096,312000000000000,2000000000000,2,100]

expected: -1

Starter codePython

import math
def roofline_crossover_batch(k, n, peak_flops, bandwidth, bytes_per_value=2, max_batch=65536):
    # Your code here
    pass

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