Cacheable Numba Functions

It is useful to write Numba jitted ("Numba-fied") functions that Numba can cache on disk for at least two reasons:

1. Subsequent runs of the same code on the same machine start up quite a bit faster since Numba uses the cached version rather than invoking the LLVM compiler each time.
2. Code being installed into a container can include a "warmup" run during image build time which causes Numba to compile and cache all the jitted functions in the code so the compile time is absorbed into a single container build step rather than paid each time the container image is used.

Sample function which cannot be cached

The following code performs a commonly seen function in LASER which is to count up the number of agents, by node, in parallel. To prevent collisions between threads/cores in the count update, we use thread local counts, i.e., one set of counters for each thread, and then sum them up at the end.

---

The Problem

Allocating tl_counts using nb.get_num_threads() inside the function prevents Numba from caching the function.

---

@nb.njit(parallel=True, cache=True)
def get_counts(num_nodes, num_people, filter_mask, node_ids):
    tl_counts = np.zeros((nb.get_num_threads(), num_nodes), dtype=np.int32)  # Adjust size as needed
    for i in nb.prange(num_people):
        if not filter_mask[i]:
            tl_counts[nb.get_thread_id(), node_ids[i]] += 1  # Local accumulation

    return tl_counts.sum(axis=0)  # Sum across threads to get the final counts

% python3 examples/demo.py
model.py:548: NumbaWarning: Cannot cache compiled function "get_counts" as it uses dynamic globals (such as ctypes pointers and large global arrays)
  @nb.njit(parallel=True, cache=True)

The fix

Option #1 - Wrapper with Same Arguments

The fix is fairly straight forward:

1. We declare a wrapper function with the same arguments as the original Numba-fied function.
2. We hoist the tl_counts allocation out of the Numba function and into the wrapper.
3. We modify the Numba-fied function to accept tl_counts as an external input (and remove the num_nodes argument because it is no longer needed in the Numba-fied function).
4. We call the wrapper function wherever we were calling the Numba-fied version directly.

def get_counts_wrap(num_nodes, num_people, filter_mask, node_ids):
    tl_counts = np.zeros((nb.get_num_threads(), num_nodes), dtype=np.int32)  # Adjust size as needed
    return get_counts(num_people, filter_mask, node_ids, tl_counts)


@nb.njit(parallel=True, cache=True)
def get_counts(num_people, filter_mask, node_ids, tl_counts):
    for i in nb.prange(num_people):
        if not filter_mask[i]:
            tl_counts[nb.get_thread_id(), node_ids[i]] += 1  # Local accumulation

    return tl_counts.sum(axis=0)  # Sum across threads to get the final counts

Option #2 - Pass in n_threads

1. Add n_threads to the argument list.
2. Modify calls to the function to pass n_threads=nb.get_num_threads()

Note, this appears simpler, but now requires all users of the Numba-fied function to think about or be reminded about the Numba usage by having to pass the n_threads argument.

@nb.njit(parallel=True, cache=True)
def get_counts(num_nodes, num_people, filter_mask, node_ids, n_threads):
    tl_counts = np.zeros((n_threads, num_nodes), dtype=np.int32)  # Adjust size as needed
    for i in nb.prange(num_people):
        if not filter_mask[i]:
            tl_counts[nb.get_thread_id(), node_ids[i]] += 1  # Local accumulation

    return tl_counts.sum(axis=0)  # Sum across threads to get the final counts

Option #3 - n_threads Wrapper

Another option would be to leave the tl_counts allocation in the Numba-fied function, but add an n_threads argument which is supplied by the wrapper. This allows external code to call the [wrapper] function without thinking about the Numba issues.

def get_counts_wrap(num_nodes, num_people, filter_mask, node_ids):
    return get_counts(num_nodes, num_people, filter_mask, node_ids, tl_counts, n_threads=nb.get_num_threads())


@nb.njit(parallel=True, cache=True)
def get_counts(num_nodes, num_people, filter_mask, node_ids, tl_counts, n_threads):
    tl_counts = np.zeros((n_threads, num_nodes), dtype=np.int32)  # Adjust size as needed
    for i in nb.prange(num_people):
        if not filter_mask[i]:
            tl_counts[nb.get_thread_id(), node_ids[i]] += 1  # Local accumulation

    return tl_counts.sum(axis=0)  # Sum across threads to get the final counts

Option #4 - Global N_THREADS

Numba is also okay with and will cache the following:

N_THREADS = nb.get_num_threads()

@nb.njit(parallel=True, cache=True)
def get_counts(num_nodes, num_people, filter_mask, node_ids, tl_counts):
    tl_counts = np.zeros((N_THREADS, num_nodes), dtype=np.int32)  # Adjust size as needed
    for i in nb.prange(num_people):
        if not filter_mask[i]:
            tl_counts[nb.get_thread_id(), node_ids[i]] += 1  # Local accumulation

    return tl_counts.sum(axis=0)  # Sum across threads to get the final counts