Implement work-around for numba issue causing a segfault on M1 when using literal_unroll() with bools.

[twiecki]

Closes #1023, which is actually several numba issues described in numba/numba#8215.

Implemented with help from @aseyboldt and @stuartarchibald.

Here are a few important guidelines and requirements to check before your PR can be merged:

• There is an informative high-level description of the changes.
• The description and/or commit message(s) references the relevant GitHub issue(s).
• pre-commit is installed and set up.
• The commit messages follow these guidelines.
• The commits correspond to relevant logical changes, and there are no commits that fix changes introduced by other commits in the same branch/BR.
• There are tests covering the changes introduced in the PR.

Don't worry, your PR doesn't need to be in perfect order to submit it. As development progresses and/or reviewers request changes, you can always rewrite the history of your feature/PR branches.

If your PR is an ongoing effort and you would like to involve us in the process, simply make it a draft PR.

[stuartarchibald]

If this is homogenizing the axis type as ((int, int), ^ N ) is it necessary to version the loop body with the unroller? Is there a chance axis can be None perhaps, in which case this seems valid?

[twiecki]

version the loop body with the unroller

I don't understand what you mean here.

[stuartarchibald]

version the loop body with the unroller

I don't understand what you mean here.

Apologies, I'll expand on this.

The need for doing something "special" to handle mixed type containers comes from the difficulties associated with working out the types of variables in the loop body in such a case, for example:

for x in (1, 'a', 2j, 3): # types = (int, char, cmplx, int)
    print(x) # what type is `x`?! It changes on each iteration.

As a result, literal_unroll was created to help:

for x in literal_unroll((1, 'a', 2j, 3)): # types = (int, char, cmplx, int)
    print(x)

which translates to something like (it does all this in LLVM IR and isn't quite like this, but for the purposes of explanation, it's close enough!):

tup = (1, 'a', 2j, 3)
for idx in range(len(tup)):
    if idx in (0, 3):
        print(x) # This is the 'int' print
    elif idx in (1,):
        print(x) # This is the 'char' print
    elif idx in (2,):
        print(x) # This is the 'cmplx' print

i.e. the loop body is "versioned" and this lets Numba "iterate" over a mixed type containers.

If you have a tuple that is already homogeneous in type, then you don't need to do this translation as the loop is type stable.

For example:

for x in (1, 2, 4, 5): # types = (int, int, int, int)
  print(x) # x is always int

Applying this to the present case. Is op_axis now always going to contain some number of tuples which are of type (int, int), if so, standard iteration should "just work".

[twiecki]

I see, thanks for explaining!

[stuartarchibald]

I see, thanks for explaining!

No problem. I wasn't sure in the case above whether you'll sometimes have e.g. a None for axis, or perhaps you need to preserve the literal nature of values captured scope of op_axis, but if you don't, I think you might be able to get away with standard iteration, which will get you a faster loop (as there's no switch table in the body).

[twiecki]

I removed the literal_unroll for now.

[stuartarchibald]

Sounds good. For performance, if you can invent a way to make a loop type stable, that's usually a good approach as it means there's no need for a switch table and duplicated loop bodies. LLVM can often "see" through the "switch table" pattern but it's better, if possible, to avoid creating the risk of it not being able to optimising it away by simply avoiding it in the program source.

Thanks for working through this issue. The Numba folks will take a look at the segfault(s) next week!

[twiecki]

@stuartarchibald Thank you for the super fast help!

[stuartarchibald]

@stuartarchibald Thank you for the super fast help!

No problem! Thanks for using Numba!

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 79.23%. Comparing base (d09e222) to head (b708213).
Report is 527 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1027      +/-   ##
==========================================
- Coverage   79.23%   79.23%   -0.01%     
==========================================
  Files         152      152              
  Lines       47954    47953       -1     
  Branches    10918    10919       +1     
==========================================
- Hits        37996    37995       -1     
  Misses       7449     7449              
  Partials     2509     2509              

Files	Coverage Δ
aesara/link/numba/dispatch/tensor_basic.py	97.93% <100.00%> (-0.03%)	⬇️