Manual vectorize in elemwise ops in numba backend

[aseyboldt]

Currently we use numba.vectorize to implement Elemwise ops. This comes with a few problems howerver:

• It doesn't support multiple output
• It can't easily generalize to Blockwise-like ops
• I think it causes slow compile times (not sure yet...)
• We can't easily extend it to do reductions in the same loops
• We can't take advantage of known broadcasting patterns for performance

I think we could replace this by a custom numba implementation that uses the numba.extending hooks to directly produce llvm code. I've been experimenting a bit, and I think I got to a point where this is starting to seem doable. A basic POC that doesn't deal with broadcasting at all yet:

import numba
import numpy as np
from numba.core import cgutils
from numba.np import arrayobj
from numba import types
from llvmlite import ir

flags = {
    "arcp",  # Allow Reciprocal
    "contract",  # Allow floating-point contraction
    "afn",  # Approximate functions
    "reassoc",
}


@numba.extending.register_jitable(fastmath=flags)
def scalar_func(x, y):
    return np.exp(2 * x + y)


inner_sig = types.float64(types.float64, types.float64)


@numba.extending.intrinsic
def loop_call(typingctx, out, x, y, shape):
    sig = types.void(out, x, y, shape)

    x_type = x
    y_type = y
    out_type = out

    def codegen(context, builder, signature, args):
        [out, x, y, shape] = args

        inner = context.compile_subroutine(builder, scalar_func, inner_sig).fndesc
        shape = [builder.extract_value(shape, idx) for idx in range(shape.type.count)]

        xary = arrayobj.make_array(x_type)(context, builder, x)
        yary = arrayobj.make_array(y_type)(context, builder, y)
        outary = arrayobj.make_array(out_type)(context, builder, out)

        with cgutils.loop_nest(builder, shape, ir.IntType(64)) as idx:
            x_ptr = cgutils.get_item_pointer(context, builder, x_type, xary, idx)
            y_ptr = cgutils.get_item_pointer(context, builder, y_type, yary, idx)
            out_ptr = cgutils.get_item_pointer(context, builder, out_type, outary, idx)

            x_val = builder.load(x_ptr)
            y_val = builder.load(y_ptr)

            out = context.call_internal(builder, inner, inner_sig, [x_val, y_val])
            builder.store(out, out_ptr)
        return
    return sig, codegen


@numba.njit(fastmath=flags)
def vectorized(out, x, y):
    loop_call(out, x, y, x.shape)


rng = np.random.default_rng(42)
x = rng.normal(size=(3, 4))
y = rng.normal(size=(3, 4))
out = np.ones_like(x)

vectorized(out, x, y)
np.allclose(out, np.exp(2 * x + y))

Most of the magic happens in codegen, where we can use normal python code to generate the llvm ir instructions to do the actual computations.

Update

With support for multiple outputs and summation of outputs:

import numba
import numpy as np
from numba.core import cgutils
from numba.np import arrayobj
from numba import types
from llvmlite import ir

import llvmlite.binding as llvm

#llvm.set_option('', '--debug-only=loop-vectorize,iv-descriptors')

# Some fast math flags are necessary for vectorization
flags = {
    "arcp",  # Allow Reciprocal
    "contract",  # Allow floating-point contraction
    "afn",  # Approximate functions
    "reassoc",
}


# The scalar function we want to wrap
@numba.extending.register_jitable(fastmath=flags)
def scalar_func(x, y):
    return (np.exp(2 * x + y), 2 * x)


# And its signature
inner_sig = types.Tuple([types.float64, types.float64])(types.float64, types.float64)

# How should inputs (x and y) be broadcasted?
input_bc_patterns = [
    (False, False),
    (False, True),
]

# Which output dimensions should be summed over?
output_bc_patterns = [
    (True, True),
    (False, True),
]

n_inputs = 2
n_outputs = 2

@numba.extending.intrinsic(cache=False)
def loop_call(typingctx, *args):
    assert len(args) == n_inputs + n_outputs + 1
    out_types = args[:n_outputs]
    in_types = args[n_outputs:n_outputs + n_inputs]
    iter_shape_type = args[-1]

    ndim = iter_shape_type.count
    # TODO I think this is actually fine?
    #if not ndim > 0:
    #    raise ValueError("Vectorize should not be called for scalar operations")

    sig = types.void(types.StarArgTuple([*out_types, *in_types, iter_shape_type]))

    def codegen(context, builder, signature, args):

        safe = (False, False)

        [args] = args

        for i in range(n_outputs):
            builder.function.args[2 + 4 * i].add_attribute("noalias")

        args = cgutils.unpack_tuple(builder, args)

        outputs = args[:n_outputs]
        inputs = args[n_outputs:n_outputs + n_inputs]
        iter_shape = args[-1]

        shape = cgutils.unpack_tuple(builder, iter_shape)

        # Lower the code of the scalar function so that we can use it in the inner loop
        inner = context.compile_subroutine(builder, scalar_func, inner_sig).fndesc

        # Extract shape and stride information from the array.
        # For later use in the loop body to do the indexing
        def extract_array(aryty, ary):
            obj = arrayobj.make_array(aryty)(context, builder, ary)
            shape = cgutils.unpack_tuple(builder, obj.shape)
            strides = cgutils.unpack_tuple(builder, obj.strides)
            data = obj.data
            layout = aryty.layout
            return (data, shape, strides, layout)

        inputs = [
            extract_array(aryty, ary)
            for aryty, ary in zip(in_types, inputs, strict=True)
        ]

        outputs = [
            extract_array(aryty, ary)
            for aryty, ary in zip(out_types, outputs, strict=True)
        ]

        zero = ir.Constant(ir.IntType(64), 0)

        # Setup loops and initialize accumulators for outputs
        loop_stack = []
        loops = []
        output_accumulator = [(None, None)] * n_outputs
        for dim, length in enumerate(shape):
            # Find outputs that only have accumulations left
            for output in range(n_outputs):
                if output_accumulator[output][0] is not None:
                    continue
                if all(output_bc_patterns[output][dim:]):
                    # TODO change type of allocator to type of output
                    accu = cgutils.alloca_once_value(builder, ir.Constant(ir.DoubleType(), 0.))
                    output_accumulator[output] = (accu, dim)

            loop = cgutils.for_range(builder, length)
            loop_stack.append(loop)
            loops.append(loop.__enter__())

        # Code in the inner most loop...
        idxs = [loopval.index for loopval in loops]

        # Load values from input arrays
        input_vals = []
        for array_info, bc in zip(inputs, input_bc_patterns, strict=True):
            idxs_bc = [
                zero if bc else idx for idx, bc in zip(idxs, bc, strict=True)
            ]
            ptr = cgutils.get_item_pointer2(context, builder, *array_info, idxs_bc, *safe)
            val = builder.load(ptr)
            input_vals.append(val)

        # Call scalar function
        output_values = context.call_internal(builder, inner, inner_sig, input_vals)
        output_values = cgutils.unpack_tuple(builder, output_values)

        # Update output value or accumulators respectively
        for i, ((accu, _), value) in enumerate(zip(output_accumulator, output_values, strict=True)):
            if accu is not None:
                new_value = builder.fadd(builder.load(accu), value)
                builder.store(new_value, accu)
            else:
                idxs_bc = [
                    zero if bc else idx for idx, bc in zip(idxs, output_bc_patterns[i], strict=True)
                ]
                ptr = cgutils.get_item_pointer2(context, builder, *outputs[i], idxs_bc)
                builder.store(value, ptr)

        for depth, loop in enumerate(loop_stack[::-1]):
            for output, (accu, accu_depth) in enumerate(output_accumulator):
                if accu_depth == depth:
                    idxs_bc = [
                        zero if bc else idx for idx, bc in zip(idxs, output_bc_patterns[output], strict=True)
                    ]
                    ptr = cgutils.get_item_pointer2(context, builder, *outputs[output], idxs_bc)
                    builder.store(builder.load(accu), ptr)
            loop.__exit__(None, None, None)
        return
    return sig, codegen


@numba.njit(fastmath=flags, no_cpython_wrapper=False, cache=False)
def vectorized(out1, out2, x, y, shape):
    loop_call(*(out1, out2, x, y, shape))


rng = np.random.default_rng(42)
x = rng.normal(size=(500, 300))
y = rng.normal(size=(500, 1))
out1 = np.zeros((1, 1))
out2 = np.zeros((500, 1))

vectorized.compile(
    (
        types.float64[:, ::1],
        types.float64[:, ::1],
        types.float64[:, ::1],
        types.float64[:, ::1],
        types.UniTuple(types.uint64, 2)
    )
)

vectorized(out1, out2, x, y, (500, 300))
assert np.allclose(out1.ravel(), np.exp(2 * x + y).sum().ravel())
assert np.allclose(out2.ravel(), (2 * x).sum(1).ravel())

%timeit vectorized(out1, out2, x, y, (500, 300))
%timeit np.exp(2 * x + y).sum()

# This gives me with svml https://numba.pydata.org/numba-doc/latest/user/performance-tips.html#intel-svml
385 µs ± 11 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
936 µs ± 21.8 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[ricardoV94]

Any reason not to take this approach? If not, I would suggest converting this to an issue soon

[aseyboldt]

I guess the most interesting limitation is that this would not work (or at least I don't know how to make it work) if we don't know broadcasting patterns in advance. So if we go for dynamic broadcasting we couldn't at least use it in those cases, so we'd have to keep a second implementation around.
We can also not use the support for parallel computations that numba.vectorize has (or we'd have to do quite a bit of digging to figure out how to use that on this low level).
And of course we have to deal with reading code that generates llvm ir. :-)

[ricardoV94]

So the alternative is the same code via string composition and parsing we are using elsewhere?