test_target_extension.py

"""This tests the target extension API to ensure that rudimentary expected
behaviours are present and correct. It uses a piece of fake hardware as a
target, the Dummy Processing Unit (DPU), to do this. The DPU borrows a lot from
the CPU but is part of the GPU class of target. The DPU target has deliberately
strange implementations of fundamental operations so as to make it identifiable
in testing."""

import unittest
from numba.tests.support import TestCase
import contextlib
import ctypes
import operator
from functools import cached_property
import numpy as np
from numba import njit, types
from numba.extending import overload, intrinsic, overload_classmethod
from numba.core.target_extension import (
    JitDecorator,
    target_registry,
    dispatcher_registry,
    jit_registry,
    target_override,
    GPU,
    resolve_dispatcher_from_str,
)
from numba.core import utils, fastmathpass, errors
from numba.core.dispatcher import Dispatcher
from numba.core.descriptors import TargetDescriptor
from numba.core import cpu, typing, cgutils
from numba.core.base import BaseContext
from numba.core.compiler_lock import global_compiler_lock
from numba.core import callconv
from numba.core.codegen import CPUCodegen, JITCodeLibrary
from numba.core.callwrapper import PyCallWrapper
from numba.core.imputils import RegistryLoader, Registry
from numba import _dynfunc
import llvmlite.binding as ll
from llvmlite import ir as llir
from numba.core.runtime import rtsys

from numba.core import compiler
from numba.core.compiler import CompilerBase, DefaultPassBuilder
from numba.core.compiler_machinery import FunctionPass, register_pass
from numba.core.typed_passes import PreLowerStripPhis

# Define a new target, this target extends GPU, this places the DPU in the
# target hierarchy as a type of GPU.


class DPU(GPU):
    ...


# register the dpu target hierarchy token in the target registry, this
# permits lookup and reference in userspace by the string "dpu"
target_registry["dpu"] = DPU

# Create a JIT DPU codegen for the DPU target


class JITDPUCodegen(CPUCodegen):
    # This largely rips off the CPU for ease

    _library_class = JITCodeLibrary

    def _customize_tm_options(self, options):
        # Customize the target machine options.
        options["cpu"] = self._get_host_cpu_name()
        arch = ll.Target.from_default_triple().name
        if arch.startswith("x86"):
            reloc_model = "static"
        elif arch.startswith("ppc"):
            reloc_model = "pic"
        else:
            reloc_model = "default"
        options["reloc"] = reloc_model
        options["codemodel"] = "jitdefault"

        # Set feature attributes (such as ISA extensions)
        # This overrides default feature selection by CPU model above
        options["features"] = self._tm_features

        # Deal with optional argument to ll.Target.create_target_machine
        sig = utils.pysignature(ll.Target.create_target_machine)
        if "jit" in sig.parameters:
            # Mark that this is making a JIT engine
            options["jit"] = True

    def _customize_tm_features(self):
        # For JIT target, we will use LLVM to get the feature map
        return self._get_host_cpu_features()

    def _add_module(self, module):
        self._engine.add_module(module)

    def set_env(self, env_name, env):
        """Set the environment address.

        Update the GlobalVariable named *env_name* to the address of *env*.
        """
        gvaddr = self._engine.get_global_value_address(env_name)
        envptr = (ctypes.c_void_p * 1).from_address(gvaddr)
        envptr[0] = ctypes.c_void_p(id(env))


# This is the function registry for the dpu, it just has one registry, this one!
dpu_function_registry = Registry()

# Implement a new context for the DPU target


class DPUContext(BaseContext):
    allow_dynamic_globals = True

    # Overrides
    def create_module(self, name):
        return self._internal_codegen._create_empty_module(name)

    @global_compiler_lock
    def init(self):
        self._internal_codegen = JITDPUCodegen("numba.exec")
        # Initialize NRT runtime
        rtsys.initialize(self)
        self.refresh()

    def refresh(self):
        registry = dpu_function_registry
        try:
            loader = self._registries[registry]
        except KeyError:
            loader = RegistryLoader(registry)
            self._registries[registry] = loader
        self.install_registry(registry)
        # Also refresh typing context, since @overload declarations can
        # affect it.
        self.typing_context.refresh()

    @property
    def target_data(self):
        return self._internal_codegen.target_data

    def codegen(self):
        return self._internal_codegen

    # Borrow the CPU call conv
    @cached_property
    def call_conv(self):
        return callconv.CPUCallConv(self)

    def get_env_body(self, builder, envptr):
        """
        From the given *envptr* (a pointer to a _dynfunc.Environment object),
        get a EnvBody allowing structured access to environment fields.
        """
        body_ptr = cgutils.pointer_add(
            builder, envptr, _dynfunc._impl_info["offsetof_env_body"]
        )
        return cpu.EnvBody(self, builder, ref=body_ptr, cast_ref=True)

    def get_env_manager(self, builder):
        envgv = self.declare_env_global(
            builder.module, self.get_env_name(self.fndesc)
        )
        envarg = builder.load(envgv)
        pyapi = self.get_python_api(builder)
        pyapi.emit_environment_sentry(
            envarg, debug_msg=self.fndesc.env_name,
        )
        env_body = self.get_env_body(builder, envarg)
        return pyapi.get_env_manager(self.environment, env_body, envarg)

    def get_generator_state(self, builder, genptr, return_type):
        """
        From the given *genptr* (a pointer to a _dynfunc.Generator object),
        get a pointer to its state area.
        """
        return cgutils.pointer_add(
            builder,
            genptr,
            _dynfunc._impl_info["offsetof_generator_state"],
            return_type=return_type,
        )

    def post_lowering(self, mod, library):
        if self.fastmath:
            fastmathpass.rewrite_module(mod, self.fastmath)

        library.add_linking_library(rtsys.library)

    def create_cpython_wrapper(
        self, library, fndesc, env, call_helper, release_gil=False
    ):
        wrapper_module = self.create_module("wrapper")
        fnty = self.call_conv.get_function_type(fndesc.restype, fndesc.argtypes)
        wrapper_callee = llir.Function(
            wrapper_module, fnty, fndesc.llvm_func_name
        )
        builder = PyCallWrapper(
            self,
            wrapper_module,
            wrapper_callee,
            fndesc,
            env,
            call_helper=call_helper,
            release_gil=release_gil,
        )
        builder.build()
        library.add_ir_module(wrapper_module)

    def create_cfunc_wrapper(self, library, fndesc, env, call_helper):
        # There's no cfunc wrapper on the dpu
        pass

    def get_executable(self, library, fndesc, env):
        """
        Returns
        -------
        (cfunc, fnptr)

        - cfunc
            callable function (Can be None)
        - fnptr
            callable function address
        - env
            an execution environment (from _dynfunc)
        """
        # Code generation
        fnptr = library.get_pointer_to_function(
            fndesc.llvm_cpython_wrapper_name
        )

        # Note: we avoid reusing the original docstring to avoid encoding
        # issues on Python 2, see issue #1908
        doc = "compiled wrapper for %r" % (fndesc.qualname,)
        cfunc = _dynfunc.make_function(
            fndesc.lookup_module(),
            fndesc.qualname.split(".")[-1],
            doc,
            fnptr,
            env,
            # objects to keepalive with the function
            (library,),
        )
        library.codegen.set_env(self.get_env_name(fndesc), env)
        return cfunc


# Nested contexts to help with isolatings bits of compilations
class _NestedContext(object):
    _typing_context = None
    _target_context = None

    @contextlib.contextmanager
    def nested(self, typing_context, target_context):
        old_nested = self._typing_context, self._target_context
        try:
            self._typing_context = typing_context
            self._target_context = target_context
            yield
        finally:
            self._typing_context, self._target_context = old_nested


# Implement a DPU TargetDescriptor, this one borrows bits from the CPU
class DPUTarget(TargetDescriptor):
    options = cpu.CPUTargetOptions
    _nested = _NestedContext()

    @cached_property
    def _toplevel_target_context(self):
        # Lazily-initialized top-level target context, for all threads
        return DPUContext(self.typing_context, self._target_name)

    @cached_property
    def _toplevel_typing_context(self):
        # Lazily-initialized top-level typing context, for all threads
        return typing.Context()

    @property
    def target_context(self):
        """
        The target context for DPU targets.
        """
        nested = self._nested._target_context
        if nested is not None:
            return nested
        else:
            return self._toplevel_target_context

    @property
    def typing_context(self):
        """
        The typing context for CPU targets.
        """
        nested = self._nested._typing_context
        if nested is not None:
            return nested
        else:
            return self._toplevel_typing_context

    def nested_context(self, typing_context, target_context):
        """
        A context manager temporarily replacing the contexts with the
        given ones, for the current thread of execution.
        """
        return self._nested.nested(typing_context, target_context)


# Create a DPU target instance
dpu_target = DPUTarget("dpu")


# Declare a dispatcher for the DPU target
class DPUDispatcher(Dispatcher):
    targetdescr = dpu_target


# Register a dispatcher for the DPU target, a lot of the code uses this
# internally to work out what to do RE compilation
dispatcher_registry[target_registry["dpu"]] = DPUDispatcher

# Implement a dispatcher for the DPU target


class djit(JitDecorator):
    def __init__(self, *args, **kwargs):
        self._args = args
        self._kwargs = kwargs

    def __call__(self, *args):
        assert len(args) < 2
        if args:
            func = args[0]
        else:
            func = self._args[0]
        self.py_func = func
        # wrap in dispatcher
        return self.dispatcher_wrapper()

    def get_dispatcher(self):
        """
        Returns the dispatcher
        """
        return dispatcher_registry[target_registry["dpu"]]

    def dispatcher_wrapper(self):
        disp = self.get_dispatcher()
        # Parse self._kwargs here
        topt = {}
        if "nopython" in self._kwargs:
            topt["nopython"] = True

        # It would be easy to specialise the default compilation pipeline for
        # this target here.
        pipeline_class = compiler.Compiler
        if "pipeline_class" in self._kwargs:
            pipeline_class = self._kwargs["pipeline_class"]
        return disp(
            py_func=self.py_func,
            targetoptions=topt,
            pipeline_class=pipeline_class,
        )


# add it to the decorator registry, this is so e.g. @overload can look up a
# JIT function to do the compilation work.
jit_registry[target_registry["dpu"]] = djit

# The DPU target "knows" nothing, add in some primitives for basic things...

# need to register how to lower dummy for @intrinsic


@dpu_function_registry.lower_constant(types.Dummy)
def constant_dummy(context, builder, ty, pyval):
    return context.get_dummy_value()


# and how to deal with IntegerLiteral to Integer casts
@dpu_function_registry.lower_cast(types.IntegerLiteral, types.Integer)
def literal_int_to_number(context, builder, fromty, toty, val):
    lit = context.get_constant_generic(
        builder, fromty.literal_type, fromty.literal_value,
    )
    return context.cast(builder, lit, fromty.literal_type, toty)


# and how to lower an Int constant
@dpu_function_registry.lower_constant(types.Integer)
def const_int(context, builder, ty, pyval):
    lty = context.get_value_type(ty)
    return lty(pyval)


# and tell the DPU how to lower a float constant
@dpu_function_registry.lower_constant(types.Float)
def const_float(context, builder, ty, pyval):
    lty = context.get_value_type(ty)
    return lty(pyval)


# The DPU actually subtracts when it's asked to 'add'!
@intrinsic(target="dpu")
def intrin_add(tyctx, x, y):
    sig = x(x, y)

    def codegen(cgctx, builder, tyargs, llargs):
        return builder.sub(*llargs)

    return sig, codegen


# Use extending.overload API to register 'add', call the dpu specific intrinsic
@overload(operator.add, target="dpu")
def ol_add(x, y):
    if isinstance(x, types.Integer) and isinstance(y, types.Integer):

        def impl(x, y):
            return intrin_add(x, y)

        return impl


class TestTargetHierarchySelection(TestCase):
    """This tests that the target hierarchy is scanned in the right order,
    that appropriate functions are selected based on what's available and that
    the DPU target is distinctly different to the CPU"""

    def test_0_dpu_registry(self):
        """Checks that the DPU registry only contains the things added

        This test must be first to execute among all tests in this file to
        ensure the no lazily loaded entries are added yet.
        """
        self.assertFalse(dpu_function_registry.functions)
        self.assertFalse(dpu_function_registry.getattrs)
        # int literal -> int cast is registered
        self.assertEqual(len(dpu_function_registry.casts), 1)
        # int, float and dummy constants are registered
        self.assertEqual(len(dpu_function_registry.constants), 3)

    def test_specialise_gpu(self):
        def my_func(x):
            pass

        # Can be used by both CPU and DPU
        @overload(my_func, target="generic")
        def ol_my_func1(x):
            def impl(x):
                return 1 + x

            return impl

        # Should be used by the DPU if there's no dpu specific one
        @overload(my_func, target="gpu")
        def ol_my_func2(x):
            def impl(x):
                return 10 + x

            return impl

        @djit()
        def dpu_foo():
            return my_func(7)

        @njit()
        def cpu_foo():
            return my_func(7)

        # DPU chooses the ol_my_func2 as it's most specific, and DPU subtracts
        # for addition, so 10 + x -> 10 - 7 -> 3
        self.assertPreciseEqual(dpu_foo(), 3)
        # CPU uses the generic one function ol_my_func1 and adds
        self.assertPreciseEqual(cpu_foo(), 8)

    def test_specialise_dpu(self):
        def my_func(x):
            pass

        # Can be used by both CPU and DPU
        @overload(my_func, target="generic")
        def ol_my_func1(x):
            def impl(x):
                return 1 + x

            return impl

        # Should be used by the DPU if there's no dpu specific one
        @overload(my_func, target="gpu")
        def ol_my_func2(x):
            def impl(x):
                return 10 + x

            return impl

        # Should be used by the DPU only
        @overload(my_func, target="dpu")
        def ol_my_func3(x):
            def impl(x):
                return 100 + x

            return impl

        @djit()
        def dpu_foo():
            return my_func(7)

        @njit()
        def cpu_foo():
            return my_func(7)

        # DPU chooses the ol_my_func3 as it's most specific, and DPU subtracts
        # for addition, so 100 + x -> 100 - 7 -> 93
        self.assertPreciseEqual(dpu_foo(), 93)
        # CPU uses the generic one function ol_my_func1 and adds
        self.assertPreciseEqual(cpu_foo(), 8)

    def test_no_specialisation_found(self):

        def my_func(x):
            pass

        # only create a cuda specialisation
        @overload(my_func, target='cuda')
        def ol_my_func_cuda(x):
            return lambda x: None

        @djit(nopython=True)
        def dpu_foo():
            my_func(1)

        # new style errors raise UnsupportedError, old style ends up as
        # TypingError
        accept = (errors.UnsupportedError, errors.TypingError)
        with self.assertRaises(accept) as raises:
            dpu_foo()

        msgs = ["Function resolution cannot find any matches for function",
                "test_no_specialisation_found.<locals>.my_func",
                "for the current target:",
                "'numba.tests.test_target_extension.DPU'"]

        for msg in msgs:
            self.assertIn(msg, str(raises.exception))

    def test_invalid_target_jit(self):

        with self.assertRaises(errors.NumbaValueError) as raises:
            @njit(_target='invalid_silicon')
            def foo():
                pass

            foo()

        msg = "No target is registered against 'invalid_silicon'"
        self.assertIn(msg, str(raises.exception))

    def test_invalid_target_overload(self):

        def bar():
            pass

        # This is a typing error at present as it fails during typing when the
        # overloads are walked.
        with self.assertRaises(errors.TypingError) as raises:
            @overload(bar, target='invalid_silicon')
            def ol_bar():
                return lambda : None

            @njit
            def foo():
                bar()

            foo()

        msg = "No target is registered against 'invalid_silicon'"
        self.assertIn(msg, str(raises.exception))

    def test_intrinsic_selection(self):
        """
        Test to make sure that targets can share generic implementations and
        cannot reach implementations that are not in their target hierarchy.
        """

        # NOTE: The actual operation performed by these functions is irrelevant
        @intrinsic(target="generic")
        def intrin_math_generic(tyctx, x, y):
            sig = x(x, y)

            def codegen(cgctx, builder, tyargs, llargs):
                return builder.mul(*llargs)

            return sig, codegen

        @intrinsic(target="dpu")
        def intrin_math_dpu(tyctx, x, y):
            sig = x(x, y)

            def codegen(cgctx, builder, tyargs, llargs):
                return builder.sub(*llargs)

            return sig, codegen

        @intrinsic(target="cpu")
        def intrin_math_cpu(tyctx, x, y):
            sig = x(x, y)

            def codegen(cgctx, builder, tyargs, llargs):
                return builder.add(*llargs)

            return sig, codegen

        # CPU can use the CPU version
        @njit
        def cpu_foo_specific():
            return intrin_math_cpu(3, 4)

        self.assertEqual(cpu_foo_specific(), 7)

        # CPU can use the 'generic' version
        @njit
        def cpu_foo_generic():
            return intrin_math_generic(3, 4)

        self.assertEqual(cpu_foo_generic(), 12)

        # CPU cannot use the 'dpu' version
        @njit
        def cpu_foo_dpu():
            return intrin_math_dpu(3, 4)

        accept = (errors.UnsupportedError, errors.TypingError)
        with self.assertRaises(accept) as raises:
            cpu_foo_dpu()

        msgs = ["Function resolution cannot find any matches for function",
                "intrinsic intrin_math_dpu",
                "for the current target",]
        for msg in msgs:
            self.assertIn(msg, str(raises.exception))

        # DPU can use the DPU version
        @djit(nopython=True)
        def dpu_foo_specific():
            return intrin_math_dpu(3, 4)

        self.assertEqual(dpu_foo_specific(), -1)

        # DPU can use the 'generic' version
        @djit(nopython=True)
        def dpu_foo_generic():
            return intrin_math_generic(3, 4)

        self.assertEqual(dpu_foo_generic(), 12)

        # DPU cannot use the 'cpu' version
        @djit(nopython=True)
        def dpu_foo_cpu():
            return intrin_math_cpu(3, 4)

        accept = (errors.UnsupportedError, errors.TypingError)
        with self.assertRaises(accept) as raises:
            dpu_foo_cpu()

        msgs = ["Function resolution cannot find any matches for function",
                "intrinsic intrin_math_cpu",
                "for the current target",]
        for msg in msgs:
            self.assertIn(msg, str(raises.exception))

    def test_overload_allocation(self):
        def cast_integer(context, builder, val, fromty, toty):
            # XXX Shouldn't require this.
            if toty.bitwidth == fromty.bitwidth:
                # Just a change of signedness
                return val
            elif toty.bitwidth < fromty.bitwidth:
                # Downcast
                return builder.trunc(val, context.get_value_type(toty))
            elif fromty.signed:
                # Signed upcast
                return builder.sext(val, context.get_value_type(toty))
            else:
                # Unsigned upcast
                return builder.zext(val, context.get_value_type(toty))

        @intrinsic(target='dpu')
        def intrin_alloc(typingctx, allocsize, align):
            """Intrinsic to call into the allocator for Array
            """
            def codegen(context, builder, signature, args):
                [allocsize, align] = args

                # XXX: error are being eaten.
                #      example: replace the next line with `align_u32 = align`
                align_u32 = cast_integer(context, builder, align,
                                         signature.args[1], types.uint32)
                meminfo = context.nrt.meminfo_alloc_aligned(builder, allocsize,
                                                            align_u32)
                return meminfo

            from numba.core.typing import signature
            mip = types.MemInfoPointer(types.voidptr)  # return untyped pointer
            sig = signature(mip, allocsize, align)
            return sig, codegen

        @overload_classmethod(types.Array, '_allocate', target='dpu',
                              jit_options={'nopython':True})
        def _ol_arr_allocate_dpu(cls, allocsize, align):
            def impl(cls, allocsize, align):
                return intrin_alloc(allocsize, align)
            return impl

        @overload(np.empty, target='dpu', jit_options={'nopython':True})
        def ol_empty_impl(n):
            def impl(n):
                return types.Array._allocate(n, 7)
            return impl

        def buffer_func():
            pass

        @overload(buffer_func, target='dpu', jit_options={'nopython':True})
        def ol_buffer_func_impl():
            def impl():
                return np.empty(10)
            return impl

        from numba.core.target_extension import target_override

        # XXX: this should probably go inside the dispatcher
        with target_override('dpu'):
            @djit(nopython=True)
            def foo():
                return buffer_func()
            r = foo()
        from numba.core.runtime import nrt
        self.assertIsInstance(r, nrt.MemInfo)


class TestTargetOffload(TestCase):
    """In this use case the CPU compilation pipeline is extended with a new
     compilation pass that runs just prior to lowering. The pass looks for
     function calls and when it finds one it sees if there's a DPU function
     available that is a valid overload for the function call. If there is one
     then it swaps the CPU implementation out for a DPU implementation. This
     producing an "offload" effect.
    """

    def test_basic_offload(self):

        _DEBUG = False

        # This is the DPU function for sin, it'll return a pi-like constant
        @overload(np.sin, target="dpu")
        def ol_np_sin_DPU(x):
            def dpu_sin_impl(x):
                return 314159.0

            return dpu_sin_impl

        # Check the DPU reports the correct overload value
        @djit(nopython=True)
        def foo(x):
            return np.sin(x)

        self.assertPreciseEqual(foo(5), 314159.0)

        # Check the CPU call is correct

        @njit
        def foo(x):
            return np.sin(x)

        self.assertPreciseEqual(foo(5), np.sin(5))

        @register_pass(mutates_CFG=False, analysis_only=False)
        class DispatcherSwitcher(FunctionPass):
            _name = "DispatcherSwitcher"

            def __init__(self):
                FunctionPass.__init__(self)

            def run_pass(self, state):
                func_ir = state.func_ir
                mutated = False
                for blk in func_ir.blocks.values():
                    # find the assignment nodes in the block and walk them, if
                    # there's a DPU version then swap out for a call to that
                    for call in blk.find_exprs("call"):
                        function = state.typemap[call.func.name]
                        tname = "dpu"

                        # Note: `target_override` context driven compilation can
                        # be done here, the DPU target is in use.
                        with target_override(tname):
                            try:
                                sig = function.get_call_type(
                                    state.typingctx,
                                    state.calltypes[call].args,
                                    {},
                                )
                                disp = resolve_dispatcher_from_str(tname)
                                # force compile check
                                hw_ctx = disp.targetdescr.target_context
                                hw_ctx.get_function(function, sig)
                            except Exception as e:
                                if _DEBUG:
                                    msg = (
                                        f"Failed to find and compile an "
                                        f"overload for {function} for {tname} "
                                        f"due to {e}"
                                    )
                                    print(msg)
                                continue

                            # This is a necessary hack at present so as to
                            # generate code into the same library. I.e. the DPU
                            # target is going to do code gen into the CPUs lib.
                            hw_ctx._codelib_stack = (
                                state.targetctx._codelib_stack
                            )

                            # All is good, so switch IR node for one targeting
                            # this target. Should generate this, but for now
                            # just mutate as:
                            # ir.Expr.call(call.func, call.args, call.kws,
                            #              call.loc, target='dpu')
                            call.target = tname
                            mutated = True
                # return True if the IR was mutated, False if not.
                return mutated

        # DPU compiler pipeline, compiles with offload to the DPU target
        class DPUOffloadCompiler(CompilerBase):
            def define_pipelines(self):
                pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
                pm.add_pass_after(DispatcherSwitcher, PreLowerStripPhis)
                pm.finalize()
                return [pm]

        # Now compile for CPU, but with the DispatcherSwitcher pass in place
        # that switches CPU calls for DPU calls
        @njit(pipeline_class=DPUOffloadCompiler)
        def foo(x):
            return np.sin(x), np.cos(x)  # np.sin is DPU, np.cos is CPU

        self.assertPreciseEqual(foo(5), (314159.0, np.cos(5)))


if __name__ == "__main__":
    unittest.main()