Introduce numba target extension to RBC

[guilhermeleobas]

This is a follow up PR of #379 as the branch is not updating there anymore.

---

Numba 0.54 introduced a new API for register new synthetic targets. For RBC, this means that one can implement intrinsics and overload functions/methods that is available to a specific target (i.e. omniscidb_cpu or omniscidb_gpu). It also solves a long issue (#343) where Numba would use the CPU bindings in the GPU target.

Summary of changes

• Refactor the target and typing contexts and introduce OmnisciTargetOptions, OmnisciTargetDescr and specific dispatchers.
• Remove the context manager used to replace the internal codegen by the one OmniSciDB uses. With the introduction of specific dispatchers, this is not needed anymore.

Things to discuss

• This is a breaking change and it requires Numba 0.54 or greater to work
  * libdevice support seems to be broken in OmniSciDB 5.9 (or even master)

To-do

• Overloads don't work yet

[pearu]

Marking as onhold as it is BC change for numba versions that rbc still supports. Reconsider when numba 0.57 is released.

[guilhermeleobas]

RBC only support the latest two versions of Numba. So, it should be safe to merge this one once Numba 0.56 is out.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[pearu]

I have a few nits. It also looks like the PR requires refreshing as with the recent heavydb-internal master, I get

rbc/tests/heavydb/test_array.py::test_getitem_float FAILED

================================================================================================ FAILURES =================================================================================================
___________________________________________________________________________________________ test_getitem_float ____________________________________________________________________________________________

heavydb = RemoteHeavyDB(user='admin', password="****************", host='192.168.1.153', port=6274, dbname='heavyai')

    def test_getitem_float(heavydb):
        heavydb.reset()
    
        @heavydb('double(double[], int32)')
        def array_getitem_double(x, i):
            return x[i]
    
        query = f'select f8, array_getitem_double(f8, 2) from {heavydb.table_name}'
        desrc, result = heavydb.sql_execute(query)
        for a, item in result:
            assert a[2] == item
            assert type(a[2]) == type(item)
    
        @heavydb('float(float[], int64)')
        def array_getitem_float(x, i):
            return x[i]
    
        query = f'select f4, array_getitem_float(f4, 2) from {heavydb.table_name}'
>       desrc, result = heavydb.sql_execute(query)

rbc/tests/heavydb/test_array.py:217: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
rbc/heavydb/remoteheavydb.py:852: in sql_execute
    self.register()
rbc/heavydb/remoteheavydb.py:1192: in register
    return self._register()
rbc/heavydb/remoteheavydb.py:1252: in _register
    llvm_module, succesful_fids = compile_to_LLVM(
rbc/irtools.py:310: in compile_to_LLVM
    fname = compile_instance(func, sig, target_info, typing_context,
rbc/irtools.py:184: in compile_instance
    cres = compiler.compile_extra(typingctx=typing_context,
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler.py:693: in compile_extra
    return pipeline.compile_extra(func)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler.py:429: in compile_extra
    return self._compile_bytecode()
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler.py:497: in _compile_bytecode
    return self._compile_core()
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler.py:472: in _compile_core
    raise e
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler.py:463: in _compile_core
    pm.run(self.state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:349: in run
    raise e
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:341: in run
    self._runPass(idx, pass_inst, state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_lock.py:35: in _acquire_compile_lock
    return func(*args, **kwargs)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:296: in _runPass
    mutated |= check(pss.run_pass, internal_state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:269: in check
    mangled = func(compiler_state)
rbc/heavydb/pipeline.py:61: in run_pass
    pm.run(state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:349: in run
    raise e
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:341: in run
    self._runPass(idx, pass_inst, state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_lock.py:35: in _acquire_compile_lock
    return func(*args, **kwargs)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:296: in _runPass
    mutated |= check(pss.run_pass, internal_state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:269: in check
    mangled = func(compiler_state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typed_passes.py:105: in run_pass
    typemap, return_type, calltypes, errs = type_inference_stage(
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typed_passes.py:83: in type_inference_stage
    errs = infer.propagate(raise_errors=raise_errors)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typeinfer.py:1078: in propagate
    errors = self.constraints.propagate(self)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typeinfer.py:177: in propagate
    raise e
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typeinfer.py:155: in propagate
    constraint(typeinfer)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typeinfer.py:704: in __call__
    self.resolve(typeinfer, typeinfer.typevars, fnty=fnty)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typeinfer.py:601: in resolve
    sig = typeinfer.resolve_call(fnty, pos_args, kw_args)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typeinfer.py:1555: in resolve_call
    return self.context.resolve_function_type(fnty, pos_args, kw_args)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typing/context.py:196: in resolve_function_type
    res = self._resolve_user_function_type(func, args, kws)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typing/context.py:248: in _resolve_user_function_type
    return func.get_call_type(self, args, kws)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/types/functions.py:312: in get_call_type
    raise e
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/types/functions.py:308: in get_call_type
    sig = temp.apply(nolitargs, nolitkws)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typing/templates.py:351: in apply
    sig = generic(args, kws)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typing/templates.py:614: in generic
    disp, new_args = self._get_impl(args, kws)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typing/templates.py:713: in _get_impl
    impl, args = self._build_impl(cache_key, args, kws)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typing/templates.py:824: in _build_impl
    disp_type.get_call_type(self.context, args, kws)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/types/functions.py:541: in get_call_type
    self.dispatcher.get_call_template(args, kws)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/dispatcher.py:363: in get_call_template
    self.compile(tuple(args))
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/dispatcher.py:965: in compile
    cres = self._compiler.compile(args, return_type)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/dispatcher.py:125: in compile
    status, retval = self._compile_cached(args, return_type)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/dispatcher.py:139: in _compile_cached
    retval = self._compile_core(args, return_type)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/dispatcher.py:152: in _compile_core
    cres = compiler.compile_extra(self.targetdescr.typing_context,
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler.py:693: in compile_extra
    return pipeline.compile_extra(func)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler.py:429: in compile_extra
    return self._compile_bytecode()
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler.py:497: in _compile_bytecode
    return self._compile_core()
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler.py:472: in _compile_core
    raise e
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler.py:463: in _compile_core
    pm.run(self.state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:349: in run
    raise e
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:341: in run
    self._runPass(idx, pass_inst, state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_lock.py:35: in _acquire_compile_lock
    return func(*args, **kwargs)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:296: in _runPass
    mutated |= check(pss.run_pass, internal_state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/compiler_machinery.py:269: in check
    mangled = func(compiler_state)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/typed_passes.py:394: in run_pass
    lower.lower()
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/lowering.py:196: in lower
    self.lower_normal_function(self.fndesc)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/lowering.py:250: in lower_normal_function
    entry_block_tail = self.lower_function_body()
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/lowering.py:279: in lower_function_body
    self.lower_block(block)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/lowering.py:293: in lower_block
    self.lower_inst(inst)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/lowering.py:438: in lower_inst
    val = self.lower_assign(ty, inst)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/lowering.py:624: in lower_assign
    return self.lower_expr(ty, value)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/lowering.py:1159: in lower_expr
    res = self.lower_call(resty, expr)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/lowering.py:889: in lower_call
    res = self._lower_call_normal(fnty, expr, signature)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/lowering.py:1122: in _lower_call_normal
    impl = self.context.get_function(fnty, signature)
../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/base.py:570: in get_function
    return self.get_function(fn, sig, _firstcall=False)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <rbc.heavydb.heavydb_compiler.JITRemoteTargetContext object at 0x7f6383257040>, fn = Function(<intrinsic heavydb_buffer_ptr_getitem_>)
sig = (ArrayPointer[STRUCT__dPW3WptrLW2WszV5Vbool8W7Wis_nullK]*, int32) -> float64, _firstcall = False

    def get_function(self, fn, sig, _firstcall=True):
        """
        Return the implementation of function *fn* for signature *sig*.
        The return value is a callable with the signature (builder, args).
        """
        assert sig is not None
        sig = sig.as_function()
        if isinstance(fn, types.Callable):
            key = fn.get_impl_key(sig)
            overloads = self._defns[key]
        else:
            key = fn
            overloads = self._defns[key]
    
        try:
            return _wrap_impl(overloads.find(sig.args), self, sig)
        except errors.NumbaNotImplementedError:
            pass
        if isinstance(fn, types.Type):
            # It's a type instance => try to find a definition for the type class
            try:
                return self.get_function(type(fn), sig)
            except NotImplementedError:
                # Raise exception for the type instance, for a better error message
                pass
    
        # Automatically refresh the context to load new registries if we are
        # calling the first time.
        if _firstcall:
            self.refresh()
            return self.get_function(fn, sig, _firstcall=False)
    
>       raise NotImplementedError("No definition for lowering %s%s" % (key, sig))
E       NotImplementedError: No definition for lowering <function heavydb_buffer_ptr_getitem_.<locals>.codegen at 0x7f63833315e0>(ArrayPointer[STRUCT__dPW3WptrLW2WszV5Vbool8W7Wis_nullK]*, int32) -> float64

../../../miniconda3/envs/rbc-dev/lib/python3.9/site-packages/numba/core/base.py:572: NotImplementedError
============================================================================================ warnings summary =============================================================================================
rbc/tests/heavydb/test_array.py::test_ptr[cpu-int8_t i1]
  /home/pearu/git/guilhermeleobas/rbc/rbc/heavydb/remoteheavydb.py:1018: UserWarning: 
    thrift type map: add new member TDatumType.MULTIPOINT to rbc/extension_functions.thrift
    thrift type map: add new member TExtArgumentType.ColumnGeoPoint to rbc/extension_functions.thrift
    thrift type map: add new member TExtArgumentType.DayTimeInterval to rbc/extension_functions.thrift
    thrift type map: add new member TExtArgumentType.GeoMultiPoint to rbc/extension_functions.thrift
    thrift type map: add new member TExtArgumentType.YearMonthTimeInterval to rbc/extension_functions.thrift
    warnings.warn('\n  '.join([''] + messages))

rbc/tests/heavydb/test_array.py::test_len_i32
  /home/pearu/git/guilhermeleobas/rbc/rbc/heavydb/remoteheavydb.py:958: UserWarning: thrift.TExtArgumentType.ColumnGeoPoint(=73) value not in ext_arguments_map
    warnings.warn(f'thrift.TExtArgumentType.{n}(={v}) value not '

rbc/tests/heavydb/test_array.py::test_len_i32
  /home/pearu/git/guilhermeleobas/rbc/rbc/heavydb/remoteheavydb.py:958: UserWarning: thrift.TExtArgumentType.DayTimeInterval(=71) value not in ext_arguments_map
    warnings.warn(f'thrift.TExtArgumentType.{n}(={v}) value not '

rbc/tests/heavydb/test_array.py::test_len_i32
  /home/pearu/git/guilhermeleobas/rbc/rbc/heavydb/remoteheavydb.py:958: UserWarning: thrift.TExtArgumentType.GeoMultiPoint(=70) value not in ext_arguments_map
    warnings.warn(f'thrift.TExtArgumentType.{n}(={v}) value not '

rbc/tests/heavydb/test_array.py::test_len_i32
  /home/pearu/git/guilhermeleobas/rbc/rbc/heavydb/remoteheavydb.py:958: UserWarning: thrift.TExtArgumentType.YearMonthTimeInterval(=72) value not in ext_arguments_map
    warnings.warn(f'thrift.TExtArgumentType.{n}(={v}) value not '

-- Docs: https://docs.pytest.org/en/stable/warnings.html
========================================================================================= short test summary info =========================================================================================
FAILED rbc/tests/heavydb/test_array.py::test_getitem_float - NotImplementedError: No definition for lowering <function heavydb_buffer_ptr_getitem_.<locals>.codegen at 0x7f63833315e0>(ArrayPointer[STRU...

[pearu]

We don't support Python 2, so the code can be adopted to Python 3.

[guilhermeleobas]

This function is not actually used in RBC but it is still required in Numba to generate the LLVM IR.

[pearu]

This may give false-positive result. In fact, locally, I get

$ pytest -sv rbc/tests/test_externals_libdevice.py -x
<snip>
>       assert fname in str(fns[fname])
E       assert '__nv_abs' in "test_externals_libdevice_abs['int32(int32), device=gpu']"
E        +  where "test_externals_libdevice_abs['int32(int32), device=gpu']" = str(Caller(<function define.<locals>.inner.<locals>.fn at 0x7fd3dfbc4310>, Signature('int32(int32), device=gpu')))