numba_vulkan.autotune#

Choosing the fastest of several ways to launch a computation.

Which kernel variant or workgroup size is fastest depends on the device and the size of the problem. autotune turns a generator of candidate launches into a function that times every candidate on the first call for a key, and runs only the fastest one afterwards. The choices are kept per device and driver in autotune.json in the cache directory (~/.cache/numba-vulkan), unless NUMBA_VULKAN_CACHE=0, so later processes skip the timing.

Functions#

autotune([function, key, repeat])

Run the fastest of several candidate launches.

Module Contents#

numba_vulkan.autotune.autotune(function=None, *, key=None, repeat=3)#

Run the fastest of several candidate launches.

Parameters:

function (generator function) – Called with the arguments of a call, it yields (name, launch) pairs: a name for each candidate, and a function of no arguments that runs it. Every candidate must leave the same results.
key (callable, optional) – Called with the arguments, returns what the choice depends on, such as the shapes of the arrays (hashable, and the same in every process). By default the shapes and types of the array arguments.
repeat (int) – Timed runs of each candidate, after one untimed run; the median counts.

Returns:

The tuned function. Its choices attribute maps the keys seen so far to the names chosen, and timings holds the seconds each candidate took when it last tuned.

Return type:

callable

Notes

While tuning, every candidate runs repeat + 1 times, so arrays the candidates write to are overwritten several times: candidates that update an array in place (a += b) would accumulate. Pass copies, or tune on a representative call first.

Examples

>>> @nv.autotune(key=lambda a, b, c: (a.shape, b.shape))
... def product(a, b, c):
...     m, n = c.shape
...     yield "tiled", lambda: matmul[(n // 16, m // 16), (16, 16)](a, b, c)
...     yield "blocked", lambda: blocked[(n // 64, m // 64), (16, 16)](a, b, c)
>>> product(a, b, c)    # times both; later calls run the faster one