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 and @generated_jit
decorators. 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.FunctionPassfor describing a pass that operates on a function-at-once level and may mutate the IR state.numba.compiler_machinery.AnalysisPassfor describing a pass that performs analysis only.numba.compiler_machinery.LoweringPassfor 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.