Consider deprecating numba.generated_jit

[stuartarchibald]

The decorator @numba.generated_jit provides similar functionality to the much more powerful @numba.extending.overload decorator. This issue discusses issues with the @generated_jit functionality and is to track decisions, if any, over deprecation in favour of just using @overload.

The @generated_jit decorator essentially provides functionality to allow type-based selection of an implementation at compile time. Example:

from numba import njit, generated_jit, types

@generated_jit
def select(x):
    if isinstance(x, types.Float):
        def impl(x):
            return x + 1
        return impl
    elif isinstance(x, types.UnicodeType):
        def impl(x):
            return x + " the number one"
        return impl
    else:
        raise TypeError("Unsupported Type")

@njit
def foo(x):
    return select(x)

print(foo(1.))
print(foo("a string"))

A few issues exist with the concept/implementation:

1. generated_jit use tends to break "JIT transparency", i.e. if the @jit family decorators are removed or the environment variable NUMBA_DISABLE_JIT is set, then the source code does not execute in the same way as it would were the JIT compiler present (xref: NUMBA_DISABLE_JIT breaks @generated_jit #4759).
2. Various flags etc are not supported by @generated_jit, for example inlining (xref: generated_jit doesn't get passed correct arguments with inline='always' argument #4692).
3. I've not proved this as yet, but would guess based on the implementation, that the target specific overload functionality would probably not work as expected.

Conceptually, @generated_jit is like @overload but the overloaded function is the decorated function, as a result, there already exists a replacement for @generated_jit, it's @overload. Taking the example above and turning it into an @overload:

from numba import njit, types
from numba.extending import overload

def select(x):
    if isinstance(x, float):
        return x + 1
    elif isinstance(x, str):
        return x + " the number one"
    else:
        raise TypeError("Unsupported Type")

@overload(select)
def ol_select(x):
    if isinstance(x, types.Float):
        def impl(x):
            return x + 1
        return impl
    elif isinstance(x, types.UnicodeType):
        def impl(x):
            return x + " the number one"
        return impl
    else:
        raise TypeError("Unsupported Type")

@njit
def foo(x):
    return select(x)

print(foo(1.))
print(foo("a string"))

Further, since preliminary support for isinstance has been added, along with advances in branch pruning, the necessity for @generated_jit in basic cases like the above has decreased, for example, this is a shorter version of the above.

from numba import njit, types
from numba.extending import overload

def select(x):
    if isinstance(x, float):
        return x + 1
    elif isinstance(x, str):
        return x + " the number one"
    else:
        raise TypeError("Unsupported Type")

overload(select)(lambda x: select)

@njit
def foo(x):
    return select(x)

print(foo(1.))
print(foo("a string"))

Final points...

• The CUDA target (and any other target implemented through the target specific functionality) now supports @overload and @intrinsic and will inherit all the effort that has been put into those decorators, trying to make @generated_jit work seems unnecessary given the replacement already exists.
• From a maintenance perspective, reducing the public facing API reduces maintenance burden, however, @generated_jit is not particularly involved in part because it is not as heavily developed as @overload.

Considerations:

1. Is there any functionality that @generated_jit provides that @overload does not?
2. Does @overload perhaps need an update to assist with 1. ?
3. How much of an impact would deprecation have against existing code bases?
   1. Need to find prevalence of use.
   2. Need to work out how hard the translation would be/should a wrapper through @overload be provided?
4. If deprecated, how long a cycle is required?

[gmarkall]

#8467 removes internal uses of generated_jit.

[stuartarchibald]

From discussion at the Numba public meeting today.

1. It was agreed that @numba.extending.overload does indeed provide the same functionality but is much better supported on a number of fronts.
2. An attempt is going to be made to try and implement @generated_jit using @overload internally so as to try and make it more stable.
3. Invariant of 2. the generated_jit function is going to be deprecated, probably with a long-ish cycle as there seems to be quite a bit of use but this is quite hard to quantify exactly.

[stuartarchibald]

With reference to #8466 (comment).

Items 1. and 3. are implemented in #8499.

It turns out that 2. can be done with:

def generated_jit(function=None, cache=False,
                  pipeline_class=None, **options):
    """
    This decorator allows flexible type-based compilation
    of a jitted function.  It works as `@jit`, except that the decorated
    function is called at compile-time with the *types* of the arguments
    and should return an implementation function for those types.
    """
    from numba.extending import overload
    jit_options = dict()
    if pipeline_class is not None:
        jit_options['pipeline_class'] = pipeline_class
    jit_options['cache'] = cache
    jit_options |= options

    if function is not None:
        overload(function, jit_options=jit_options,
                    strict=False)(function)
        return function
    else:
        def wrapper(func):
            overload(func, jit_options=jit_options,
                        strict=False)(func)
            return func
        return wrapper

However, this leads to issues as it slightly changes the behaviour of generated_jit:

a. In the case of a call like generated_jit(options)(function) the returned value is a closure, this will not serialize, it therefore causes issues in caching.

b. The use of @overload enforces more strictness (and correctness!) with respect to signatures. The functions returned by the @generated_jit decorated function have to comply with the exact-match signature semantics of @overload, i.e. the following is not accepted via a @generated_jit implementation using @overload (as above) but is in the original @generated_jit implementation:

@generated_jit
def foo(a, b=1):
  if isinstance(a, types.Float):
    def impl(a, b): # Note that b is not declared as a kwarg
      return a + b
    return impl

[sklam]

a. In the case of a call like generated_jit(options)(function) the returned value is a closure, this will not serialize, it therefore causes issues in caching.

Do you have an example of the caching issue? I tried a few things and the caching is working.

[stuartarchibald]

a. In the case of a call like generated_jit(options)(function) the returned value is a closure, this will not serialize, it therefore causes issues in caching.

Do you have an example of the caching issue? I tried a few things and the caching is working.

I think it was test:

numba.tests.test_serialize.TestDispatcherPickling.test_call_generated

[sklam]

notes from out-of-band discussion for #8466 (comment):

The failure of the test actually came from the replacement implementation of @generated_jit (#8466 (comment)) lacking the ability to be called standalone as a dispatcher.

[stuartarchibald]

Some notes from another OOB discussion...

If the @generated_jit decorator is implemented as in #8466 (comment) but also returns a numba.core.decorators.jit wrapped object such that it can be called as a "standalone" dispatcher (noted here: #8466 (comment)), the effect is that the jit object will be a Dispatcher instance and Numba "knows" how to handle these. It then attempts to "augment" a CPUDispatcher with a Function(CPUDispatcher), which fails. The conclusion from this observation is that it is unlikely it can ever work like this and also maintain exactly the same semantics as the existing decorator.

Conclusions were:

1. If a function needs to be used as a pure Python and a JIT compiled function, then @register_jitable provides this behaviour.
2. If a function needs type based dispatch:
   1. If the dispatch can be described via use of isinstance, then use that with @register_jitable or @jit as desired.
   2. If the dispatch cannot be described via the use of isinstance, then use @overload with a separate pure Python and Numba-@overload implementation.
3. It seems like a partial application based pattern using @overload can be used to replicate the current @register_jitable behaviour and will be documented.