Customizing the Compiler

Warning

The custom pipeline feature is for expert use only. Modifying the compiler behavior can invalidate internal assumptions in the numba source code.

For library developers looking for a way to extend or modify the compiler behavior, you can do so by defining a custom compiler by inheriting from numba.compiler.CompilerBase. The default Numba compiler is defined as numba.compiler.Compiler, implementing the .define_pipelines() method, which adds the nopython-mode, object-mode and interpreted-mode pipelines. For convenience these three pipelines are defined in numba.compiler.DefaultPassBuilder by the methods:

  • .define_nopython_pipeline()

  • .define_objectmode_pipeline()

  • .define_interpreted_pipeline()

respectively.

To use a custom subclass of CompilerBase, supply it as the pipeline_class keyword argument to the @jit decorator. By doing so, the effect of the custom pipeline is limited to the function being decorated.

Implementing a compiler pass

Numba makes it possible to implement a new compiler pass and does so through the use of an API similar to that of LLVM. The following demonstrates the basic process involved.

Compiler pass classes

All passes must inherit from numba.compiler_machinery.CompilerPass, commonly used subclasses are:

  • numba.compiler_machinery.FunctionPass for describing a pass that operates on a function-at-once level and may mutate the IR state.

  • numba.compiler_machinery.AnalysisPass for describing a pass that performs analysis only.

  • numba.compiler_machinery.LoweringPass for describing a pass that performs lowering only.

In this example a new compiler pass will be implemented that will rewrite all ir.Const(x) nodes, where x is a subclass of numbers.Number, such that the value of x is incremented by one. There is no use for this pass other than to serve as a pedagogical vehicle!

The numba.compiler_machinery.FunctionPass is appropriate for the suggested pass behavior and so is the base class of the new pass. Further, a run_pass method is defined to do the work (this method is abstract, all compiler passes must implement it).

First the new class:

from numba import njit
from numba.core import ir
from numba.core.compiler import CompilerBase, DefaultPassBuilder
from numba.core.compiler_machinery import FunctionPass, register_pass
from numba.core.untyped_passes import IRProcessing
from numbers import Number

# Register this pass with the compiler framework, declare that it will not
# mutate the control flow graph and that it is not an analysis_only pass (it
# potentially mutates the IR).
@register_pass(mutates_CFG=False, analysis_only=False)
class ConstsAddOne(FunctionPass):
    _name = "consts_add_one" # the common name for the pass

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

    # implement method to do the work, "state" is the internal compiler
    # state from the CompilerBase instance.
    def run_pass(self, state):
        func_ir = state.func_ir # get the FunctionIR object
        mutated = False # used to record whether this pass mutates the IR
        # walk the blocks
        for blk in func_ir.blocks.values():
            # find the assignment nodes in the block and walk them
            for assgn in blk.find_insts(ir.Assign):
                # if an assignment value is a ir.Consts
                if isinstance(assgn.value, ir.Const):
                    const_val = assgn.value
                    # if the value of the ir.Const is a Number
                    if isinstance(const_val.value, Number):
                        # then add one!
                        const_val.value += 1
                        mutated |= True
        return mutated # return True if the IR was mutated, False if not.

Note also that the class must be registered with Numba’s compiler machinery using @register_pass. This in part is to allow the declaration of whether the pass mutates the control flow graph and whether it is an analysis only pass.

Next, define a new compiler based on the existing numba.compiler.CompilerBase. The compiler pipeline is defined through the use of an existing pipeline and the new pass declared above is added to be run after the IRProcessing pass.

class MyCompiler(CompilerBase): # custom compiler extends from CompilerBase

    def define_pipelines(self):
        # define a new set of pipelines (just one in this case) and for ease
        # base it on an existing pipeline from the DefaultPassBuilder,
        # namely the "nopython" pipeline
        pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
        # Add the new pass to run after IRProcessing
        pm.add_pass_after(ConstsAddOne, IRProcessing)
        # finalize
        pm.finalize()
        # return as an iterable, any number of pipelines may be defined!
        return [pm]

Finally update the @njit decorator at the call site to make use of the newly defined compilation pipeline.

@njit(pipeline_class=MyCompiler) # JIT compile using the custom compiler
def foo(x):
    a = 10
    b = 20.2
    c = x + a + b
    return c

print(foo(100)) # 100 + 10 + 20.2 (+ 1 + 1), extra + 1 + 1 from the rewrite!

Debugging compiler passes

Observing IR Changes

It is often useful to be able to see the changes a pass makes to the IR. Numba conveniently permits this through the use of the environment variable NUMBA_DEBUG_PRINT_AFTER. In the case of the above pass, running the example code with NUMBA_DEBUG_PRINT_AFTER="ir_processing,consts_add_one" gives:

----------------------------nopython: ir_processing-----------------------------
label 0:
    x = arg(0, name=x)                       ['x']
    $const0.1 = const(int, 10)               ['$const0.1']
    a = $const0.1                            ['$const0.1', 'a']
    del $const0.1                            []
    $const0.2 = const(float, 20.2)           ['$const0.2']
    b = $const0.2                            ['$const0.2', 'b']
    del $const0.2                            []
    $0.5 = x + a                             ['$0.5', 'a', 'x']
    del x                                    []
    del a                                    []
    $0.7 = $0.5 + b                          ['$0.5', '$0.7', 'b']
    del b                                    []
    del $0.5                                 []
    c = $0.7                                 ['$0.7', 'c']
    del $0.7                                 []
    $0.9 = cast(value=c)                     ['$0.9', 'c']
    del c                                    []
    return $0.9                              ['$0.9']
----------------------------nopython: consts_add_one----------------------------
label 0:
    x = arg(0, name=x)                       ['x']
    $const0.1 = const(int, 11)               ['$const0.1']
    a = $const0.1                            ['$const0.1', 'a']
    del $const0.1                            []
    $const0.2 = const(float, 21.2)           ['$const0.2']
    b = $const0.2                            ['$const0.2', 'b']
    del $const0.2                            []
    $0.5 = x + a                             ['$0.5', 'a', 'x']
    del x                                    []
    del a                                    []
    $0.7 = $0.5 + b                          ['$0.5', '$0.7', 'b']
    del b                                    []
    del $0.5                                 []
    c = $0.7                                 ['$0.7', 'c']
    del $0.7                                 []
    $0.9 = cast(value=c)                     ['$0.9', 'c']
    del c                                    []
    return $0.9                              ['$0.9']

Note the change in the values in the const nodes.

Pass execution times

Numba has built-in support for timing all compiler passes, the execution times are stored in the metadata associated with a compilation result. This demonstrates one way of accessing this information based on the previously defined function, foo:

compile_result = foo.overloads[foo.signatures[0]]
nopython_times = compile_result.metadata['pipeline_times']['nopython']
for k in nopython_times.keys():
    if ConstsAddOne._name in k:
        print(nopython_times[k])

the output of which is, for example:

pass_timings(init=1.914000677061267e-06, run=4.308700044930447e-05, finalize=1.7400006981915794e-06)

this displaying the pass initialization, run and finalization times in seconds.