Implement most ufunc attributes and ufunc.reduce

[guilhermeleobas]

This PR exposes most ufunc attributes and implements ufunc.reduce for ufuncs created with @vectorize

• ufunc.nin
• ufunc.nout
• ufunc.nargs
• ufunc.ntypes
• ufunc.types
• ufunc.identity
• ufunc.signature (always None)

Implementing ufunc.ntypes or ufunc.types is not possible as Numba needs to call the NumPy ufunc attribute at runtime. However, this is not possible as NumPy doesn't allow creating weak references to an ufunc, which is required by types.BoundFunction.

---

To-do:

• Check how Cython implements ufunc.reduce
• Raise an exception if ufunc doesn't have an identity.

---

Reduce is implemented for the following set of arguments:

def reduce(array, axis=0, dtype=None, initial=None):
    ...

keeping #1269 (comment) as a reference here.

---

Code I used to benchmark:

from numba import vectorize, njit
import numpy as np
import time

@vectorize(identity=0)
def add(a, b):
    return a + b


@njit
def bar(a, axis):
    return add.reduce(a, axis=axis)


def measure_time(fn, *args):
    start = time.time()
    fn(*args)
    end = time.time()
    print(f'{fn.__name__}: {end - start}s')


N = 1_000_000_000
T = (10,) * int(np.log10(N))
a = np.arange(N).reshape(T)
axis = tuple(range(a.ndim))

# jit compile `bar`
bar(a, axis)

# measure time
measure_time(bar, a, axis)
measure_time(bar.py_func, a, axis)

[guilhermeleobas]

@kc611, could you please take a look at this PR once you have a chance? I'd like to discuss possible ways to speedup ufunc.reduce

[mattip]

Cool. It would be nice if numba could do reduce ufuncs.

[mattip]

Does NumPy need to explicitly set the field? From reading the documentation I would think it inherits from the object type.

[guilhermeleobas]

When attempting to create a weak reference to an ufunc, CPython raises an error:

>>> import weakref
>>> import numpy as np
>>> weakref.ref(np.minimum)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: cannot create weak reference to 'numpy.ufunc' object

[kc611]

Question: The compute_flat_idx__ is not being used anywhere. Is that intended?

And is it overall faster or slower than compute_flat_idx?

[guilhermeleobas]

compute_flat_idx__ implements the same code as compute_flat_idx but it is just a bit switch table where each case is computed by hand beforehand. In short, it looks something like this:

def compute_flat_idx(axis, ...):
    if axis == 0:
        # compute flat idx when axis == 0
    else axis == 1:
        # compute flat idx when axis == 1
    ...

And is it overall faster or slower than compute_flat_idx?

Faster

[kc611]

The incapability of accessing types can be opened as a separate issue, to be fixed once there is a solution upstream. Otherwise if we don't expect a solution in future, might be a good idea to remove the commented code.

[kc611]

@guilhermeleobas Thanks for the patch. This PR mostly looks good. Some nitpicks and questions above before we can mark this as RTM.

Edit: This also needs a documentation entry that support for ufunc.reduce on DUFunc objects is now available.

[guilhermeleobas]

CI failure doesn't seem to be related to this PR.

[kc611]

LGTM!

[gmarkall]

gpuci run tests

[gmarkall]

@kc611 Is this good to be RTM (assuming it passes the gpuCI run I just started)?

[kc611]

I think it's good to go. (Unless someone objects)

[gmarkall]

gpuCI failed one CUDA version for some kind of Conda error. I expected if this was going to break CUDA, it'd break for all configurations, so I'm happy to ignore that.

[guilhermeleobas]

Thanks for the review, folks. I'll create the issues to track future work:

• Rewrite ufunc.reduce to not use np.ndenumerate
• Remove experimental feature warning in 2(?) Numba versions?
• Support ufunc.ntypes and ufunc.types attributes

Also, if it is ok, I'd like another commit that updates documentation and adds a test. Can do that in a second PR.