Array shape effects numba compilation

[Durman]

It seems that shape of an array forces numba to recompile. In my case it makes arrays useless to use with numba.
Without using newaxis it works as expected. Is it a bug or limitation?

import numpy as np
import awkward as ak
import numba as nb
from timeit import timeit


@nb.njit
def test(arr):
    return


if __name__ == "__main__":
    a1 = ak.Array([1, 2])
    a2 = ak.Array([1, 2, 3])
    a3 = ak.Array([1, 2, 3, 4])

    t = timeit("test(a1[np.newaxis])", number=1, globals=globals())
    print(f"{t*1000:.2f} ms")
    t = timeit("test(a1[np.newaxis])", number=1, globals=globals())
    print(f"{t*1000:.2f} ms")
    t = timeit("test(a2[np.newaxis])", number=1, globals=globals())
    print(f"{t*1000:.2f} ms")
    t = timeit("test(a2[np.newaxis])", number=1, globals=globals())
    print(f"{t*1000:.2f} ms")
    t = timeit("test(a3[np.newaxis])", number=1, globals=globals())
    print(f"{t*1000:.2f} ms")
    t = timeit("test(a3[np.newaxis])", number=1, globals=globals())
    print(f"{t*1000:.2f} ms")

156.01 ms
0.34 ms
35.95 ms
0.34 ms
36.53 ms
0.36 ms

[agoose77]

Regular dimensions in Awkward (if your type has ... * N * ..., you have a regular dimension) include the shape information in their type. I believe that Numba does this differently for native NumPy arrays; the shape metadata is runtime-only. We have discussed this here before, and the decision at the time was "this is a compromise": #1343

In this case, your np.newaxis inserts a new regular axis at axis=0, meaning that your innermost shape changes with the list length. If you make that extra dimension ragged, e.g. with ak.from_regular(..., axis=...), then you will end up with an array whose Numba type (you can see this via array.numba_type` ) does not change with its length.

In general, this should only matter if the regular sizes of your array are changing between compiles. Usually that doesn't make sense; it's often a user error, but if it is needed in this case, then from_regular might be of use.

[jpivarski]

Different systems take different opinions on whether array lengths and shapes are compile-time or runtime data, and that determines what kinds of changes imply recompilations.

• NumPy takes all shapes, even the number of dimensions, to be runtime.
• Numba takes the number of dimensions to be compile-time, but the length of each to be runtime.
• Awkward and JAX take the lengths of regular-length dimensions to be compile-time (but Awkward also has variable-length lists, so @agoose77's suggestion of using ak.from_regular would work).

Before getting into the below, note that @agoose77's suggestion of using ak.from_regular is probably the solution to your problem, unless you can arrange your workflow into a smaller number of bigger runs.

---

Frequent recompilation is a performance hit, as you've shown, but so are runtime lengths. Knowing a loop length at compile time allows the compiler to perform more optimizations. (We can use performance as a motivation in this case because you're using Numba; of course you want it to go faster.) Which performance hit matters more depends on how many Numba runs you do and how big each array that you pass to it are, since time-to-compile scales with the number of compilations and time-to-run-one-array scales with the size of the array.

I wanted to see this first-hand, so I made an example in the Godbolt explorer. Here's a short function that depends on a runtime length:

void runtime_length(int* data, int num_steps) {
    for (int i = 0;  i < num_steps;  i++) {
        data[i]++;
    }
}

It compiles to this x86 assembly:

        test    esi, esi
        jle     .L1
        lea     eax, [rsi-1]
        cmp     eax, 2
        jbe     .L6
        mov     edx, esi
        movdqa  xmm1, XMMWORD PTR .LC0[rip]
        mov     rax, rdi
        shr     edx, 2
        sal     rdx, 4
        add     rdx, rdi
.L4:
        movdqu  xmm0, XMMWORD PTR [rax]
        add     rax, 16
        paddd   xmm0, xmm1
        movups  XMMWORD PTR [rax-16], xmm0
        cmp     rax, rdx
        jne     .L4
        mov     edx, esi
        and     edx, -4
        test    sil, 3
        je      .L9
.L3:
        movsx   rax, edx
        lea     ecx, [rdx+1]
        sal     rax, 2
        add     DWORD PTR [rdi+rax], 1
        cmp     esi, ecx
        jle     .L1
        add     edx, 2
        add     DWORD PTR [rdi+4+rax], 1
        cmp     esi, edx
        jle     .L1
        add     DWORD PTR [rdi+8+rax], 1
.L1:
        ret
.L9:
        ret
.L6:
        xor     edx, edx
        jmp     .L3

And here's the same function, but with a known num_steps:

void compiletime_length(int* data) {
    for (int i = 0;  i < 1000;  i++) {
        data[i]++;
    }
}

Its x86 assembly is:

        movdqa  xmm1, XMMWORD PTR .LC0[rip]
        lea     rax, [rdi+4000]
.L11:
        movdqu  xmm0, XMMWORD PTR [rdi]
        add     rdi, 16
        paddd   xmm0, xmm1
        movups  XMMWORD PTR [rdi-16], xmm0
        cmp     rax, rdi
        jne     .L11
        ret
.LC0:
        .long   1
        .long   1
        .long   1
        .long   1

It doesn't have to do nearly as much bookkeeping to know if it has reached the end of its loops. (Empirically, it runs faster in microbenchmarks, too.)

So there's an advantage to having the ability to set compile-time lengths, and Awkward's regular-list type does have the list size as part of the type description. But if it's a disadvantage in your case because you have to frequently change that size and that triggers recompilations because it's a new type, you can use ak.from_regular to say that it's a variable-length list type.

[Durman]

Thanks. Now it's clear. I will try from_regular in my case.