<jupyter_start><jupyter_text>The life of a Numba kernel - from Python source to kernel launch**Author:** <NAME> ([@gmarkall](https://twitter.com/gmarkall))This notebook gives an overview of the compiler and runtime pipeline that Numba uses to take Python source code, turn it into a CUDA kernel, and launch it. It pulls in a variety of Numba internals to illustrate how the different parts of the pipeline work. OverviewThe Numba CUDA target (or CUDA Python, as it's sometimes referred to) has three components:1. A NumPy-like array library backed by CUDA2. A Python-to-PTX compiler that uses NVVM3. A Python interface to the CUDA driver APIThe first part of this notebook is about the Python-to-PTX compiler that uses NVVM, and this follows on to an overview of the Python interface to the CUDA driver API. We will follow through the process that starts with a Python function, and ends with the execution of a kernel on the GPU. A simple kernelThroughout the notebook we'll use a simple L1 BLAS-like kernel for the example. Let's import Numba and NumPy, and define the kernel:<jupyter_code>from numba import cuda
import numpy as np
@cuda.jit
def axpy(r, a, x, y): 
    i = cuda.grid(1)

    if i < len(r):
        r[i] = a * x[i] + y[i]<jupyter_output><empty_output><jupyter_text>Now let's create some data:<jupyter_code>N = 10000
a = 5.0 
x = np.arange(N).astype(np.int32)
y = np.random.randint(100, size=N).astype(np.int32)
r = np.zeros_like(x)<jupyter_output><empty_output><jupyter_text>and launch the kernel using that data:<jupyter_code>threads = 256 
blocks = (N // threads) + 1 
axpy[blocks, threads](r, a, x, y)<jupyter_output><empty_output><jupyter_text>A quick sanity check to make sure it's done as we expected:<jupyter_code>assert(np.all(r == a * x + y))<jupyter_output><empty_output><jupyter_text>What just happened?When we configured and launched the kernel, Numba jumped in and compiled the kernel, transferred data to the GPU, launched the kernel on the data, and transferred the result back. How did Numba compile and load the kernel to the GPU? The process is outlined in this flowchart:<jupyter_code>from IPython.core.display import SVG
SVG(filename='numba-pipeline.svg')<jupyter_output><empty_output><jupyter_text>We'll follow through each of the steps in the notebook, mostly in code. Some of the following code uses various bits of Numba internals to get things done - this is to illustrate a little more clearly what Numba is doing, but be aware that in general it's not possible to drive Numba's pipeline from an external program. Bytecode compilationNumba's input language is not Python source code / ASTs, but instead Python bytecode. The CPython interpreter compiles the source to Python bytecode - we can see what this looks like for our `axpy` function:<jupyter_code>from dis import dis
dis(axpy.py_func)<jupyter_output>3           0 LOAD_GLOBAL              0 (cuda)
              2 LOAD_METHOD              1 (grid)
              4 LOAD_CONST               1 (1)
              6 CALL_METHOD              1
              8 STORE_FAST               4 (i)

  5          10 LOAD_FAST                4 (i)
             12 LOAD_GLOBAL              2 (len)
             14 LOAD_FAST                0 (r)
             16 CALL_FUNCTION            1
             18 COMPARE_OP               0 (<)
             20 POP_JUMP_IF_FALSE       46

  6          22 LOAD_FAST                1 (a)
             24 LOAD_FAST                2 (x)
             26 LOAD_FAST                4 (i)
             28 BINARY_SUBSCR
             30 BINARY_MULTIPLY
             32 LOAD_FAST                3 (y)
             34 LOAD_FAST                4 (i)
             36 BINARY_SUBSCR
             38 BINARY_ADD
             40 LOAD_FAST                0 (r)
             42 LOAD_FAST                4 (i)
             44 STORE_SUBSCR
        [...]<jupyter_text>The CPython bytecode is a sequence of instructions for a simple stack machine. The line numbers corresponding to each sequence of instructions is given in the first column (3, 5, and 6). The next column is the offset of each instruction in the bytecode - e.g. at location 20 we have `POP_JUMP_IF_FALSE 46`, which branches to location 46. The next column is the opcode - all opcodes are listed in [this section of the Python dis module documentation](https://docs.python.org/3/library/dis.htmlpython-bytecode-instructions). The next column is an operand, whose meaning varies between the opcodes - for example the operand to `POP_JUMP_IF_FALSE` is the location to jump to, whereas the operand for the `COMPARE_OP` is the type of comparison to perform - in this case, `0` indicates `<` - note that many instructions also pop operands from the stack - for example, the `BINARY_SUBSCR` uses the three items loaded to the top of the stack by the `LOAD_FAST` instructions at locations 22, 24, and 26.Further reading: * [Numba manual: Bytecode Analysis](http://numba.pydata.org/numba-doc/dev/developer/architecture.htmlstage-1-analyze-bytecode) - provides a slightly larger bytecode example.* [Stack Machines on Wikipedia](https://en.wikipedia.org/wiki/Stack_machine) - not the greatest introduction, but contains a lot of information. Bytecode AnalysisNumba's first step is to translate the Python Bytecode into Untyped Numba IR - this is a representation closer to the LLVM IR that Numba will eventually hand to NVVM. Its design provides* An abstract machine with infinite registers* A representation that type inference can be executed on* Support for rewriting the IR - transformations that make it easier to implement Python semantics in a lower-level representation.Numba does two forms of analysis on the Python bytecode, which it then uses to construct the Numba IR:* Control Flow Analysis* Data Flow AnalysisTo construct these, we use Numba's *Bytecode Interpreter*:<jupyter_code>from numba.core.bytecode import ByteCode, FunctionIdentity
from numba.core.interpreter import Interpreter

fi = FunctionIdentity.from_function(axpy.py_func)
interp = Interpreter(fi)
bc = ByteCode(fi)
ir = interp.interpret(bc)<jupyter_output><empty_output><jupyter_text>Let's first examine the Control Flow Analysis:<jupyter_code>interp.cfa.dump()<jupyter_output>CFG adjacency lists:
{0: [24, 48], 24: [48], 48: []}
CFG dominators:
{0: {0}, 24: {24, 0}, 48: {48, 0}}
CFG post-dominators:
{0: {0, 48}, 24: {24, 48}, 48: {48}}
CFG back edges: []
CFG loops:
{}
CFG node-to-loops:
{0: [], 24: [], 48: []}
CFG backbone:
{0, 48}<jupyter_text>CFA provides the control flow graph in terms of the basic blocks of the function - in this function there are three basic blocks, numbered 0, 24, and 48. The adjacency list gives the control flow between the basic blocks - the existence of an edge implies control can flow from the source block to the destination block.Knowing which nodes dominate which other nodes is also useful for various analyses - a node A dominates the node B if all control flow paths to B must pass through A.Post-dominators are a bit like the reverse of a dominator - a node C post-dominates a node D if all paths starting at D must pass through C to exit the graph.The text-based representation contains lots of information but is not that easy to interpret - Numba also provides a visual rendering of the CFG to help with interpreting it:<jupyter_code>dg = interp.cfa.graph.render_dot()
dg<jupyter_output><empty_output><jupyter_text>The graph in this case is rather simple - in block 0, we call `cuda.grid(1)` - if `i < len(r)`, control flow goes to block 24 then 48, otherwise it goes straight to 48 (the exit block).The data flow analysis can also be viewed, but there is not such a convenient way to view it:<jupyter_code>from pprint import pprint

for block, abi in interp.dfa.infos.items():
    print(f"Block {block}:\n")
    print("Instructions:")
    pprint(abi.insts)
    print()<jupyter_output>Block 0:

Instructions:
((0, {}),
 (2, {'res': '$2load_global.0'}),
 (4, {'item': '$2load_global.0', 'res': '$4load_method.1'}),
 (6, {'res': '$const6.2'}),
 (8,
  {'args': ['$const6.2'], 'func': '$4load_method.1', 'res': '$8call_method.3'}),
 (10, {'value': '$8call_method.3'}),
 (12, {'res': '$i12.4'}),
 (14, {'res': '$14load_global.5'}),
 (16, {'res': '$r16.6'}),
 (18,
  {'args': ['$r16.6'],
   'func': '$14load_global.5',
   'res': '$18call_function.7'}),
 (20, {'lhs': '$i12.4', 'res': '$20compare_op.8', 'rhs': '$18call_function.7'}),
 (22, {'pred': '$20compare_op.8'}))

Block 24:

Instructions:
((24, {'res': '$a24.0'}),
 (26, {'res': '$x26.1'}),
 (28, {'res': '$i28.2'}),
 (30, {'index': '$i28.2', 'res': '$30binary_subscr.3', 'target': '$x26.1'}),
 (32,
  {'lhs': '$a24.0',
   'res': '$32binary_multiply.4',
   'rhs': '$30binary_subscr.3'}),
 (34, {'res': '$y34.5'}),
 (36, {'res': '$i36.6'}),
 (38, {'index': '$i36.6', 'res': '$38binary_subscr.7', 'target': '$y34.5'}),
 (40,
  {'lhs': '[...]<jupyter_text>The Numba IRWe can view the Numba IR too:<jupyter_code>ir.dump()<jupyter_output>label 0:
    r = arg(0, name=r)                       ['r']
    a = arg(1, name=a)                       ['a']
    x = arg(2, name=x)                       ['x']
    y = arg(3, name=y)                       ['y']
    $2load_global.0 = global(cuda: <module 'numba.cuda' from '/home/gmarkall/numbadev/numba/numba/cuda/__init__.py'>) ['$2load_global.0']
    $4load_method.1 = getattr(value=$2load_global.0, attr=grid) ['$2load_global.0', '$4load_method.1']
    $const6.2 = const(int, 1)                ['$const6.2']
    $8call_method.3 = call $4load_method.1($const6.2, func=$4load_method.1, args=[Var($const6.2, <ipython-input-2-2505fed93738>:3)], kws=(), vararg=None) ['$4load_method.1', '$8call_method.3', '$const6.2']
    i = $8call_method.3                      ['$8call_method.3', 'i']
    $14load_global.5 = global(len: <built-in function len>) ['$14load_global.5']
    $18call_function.7 = call $14load_global.5(r, func=$14load_global.5, args=[Var(r, <ipython-input-2-2505fed93738>:3)], kws=(), [...]<jupyter_text>Examining the IR shows some correspondences between it and:* The Python bytecode - we can see some similar semantics, for example the loading of globals and methods, and the calling of methods* The names of the variables in the IR and the names in the Data Flow Analysis* The control flow analysis - the structure of control flow of the Numba IR matches that in the bytecode, and produced by the CFA.However, it is also a little different from the previous items:* The Numba IR abstract machine is not a stack machine, but instead has an infinite number of variables (`$2load_global0`, etc).* It doesn't explicitly hold the results of analysis, like the dominator and post-dominator analysis.We can also view the Numba IR as a graph:<jupyter_code>ir_dg = ir.render_dot()
ir_dg<jupyter_output><empty_output><jupyter_text>Rewrite passesSome rewrites on the untyped IR occur at this point - e.g.* Some constant propagation, e.g. where one side of a binary operation has a constant value* Rewrite getitems where the index is known to be a constant* Rewrite raise expressions with constant arguments so they can be used in nopython mode.There is no easy way to look at these in the notebook, but most of them can be found in the `numba.core.rewrites` module, decorated with `@register_rewrite('before-inference')`. Type inferenceNext is type inference. This starts with the types of the arguments and propagates type information through the function until everything has been typed. It uses:* Internally stored mappings of input types to output types - e.g. "an add of a float32 and a float64 will have a return type of float64"* The data flow graphWhen there are two incoming types for a variable, e.g. when two different values can be assigned to it because of branching, Numba constructs a set of possible types for the variable, and attempts to "unify" them - that is, choose a type that all of the types in the set can be safely coerced to.For example:```pythondef select(a, b, threshold, value):    if threshold < value:        r = a    r: float32    else:        r = b    r: int32    return r     r: {float32, int32} unifies to float64a = np.float32(1)b = np.int32(2)select(a, b, 10, 11)  Call with (float32, int32, int64, int64)```In the above example `r` is typed as `float32` in the taken branch, and `int32` in the not-taken branch. On the line with the return, the control flow from the if-else re-joins, so the types of `r` from both sides of it are used to construct the set of types. Since the set of types contains different types (`int32` and `float32`), Numba uses a built-in rule for unification - in this case, to ensure the best representation of values from all branches, it chooses `float64`.Further reading: [Numba Enhancement Proposal 1: Changes in Integer Typing](http://numba.pydata.org/numba-doc/latest/proposals/integer-typing.html) explains some of the rationale behind some of the decisions made by Numba's type system. Inspecting the result of type inferenceWe can use the `inspect_types()` function of a jitted kernel to see Numba's typing of all values in the function:<jupyter_code>axpy.inspect_types()<jupyter_output>_ZN6cudapy8__main__8axpy$241E5ArrayIiLi1E1C7mutable7alignedEd5ArrayIiLi1E1C7mutable7alignedE5ArrayIiLi1E1C7mutable7alignedE (array(int32, 1d, C), float64, array(int32, 1d, C), array(int32, 1d, C))
--------------------------------------------------------------------------------
# File: <ipython-input-2-2505fed93738>
# --- LINE 1 --- 

@cuda.jit

# --- LINE 2 --- 

def axpy(r, a, x, y):

    # --- LINE 3 --- 
    # label 0
    #   r = arg(0, name=r)  :: array(int32, 1d, C)
    #   a = arg(1, name=a)  :: float64
    #   x = arg(2, name=x)  :: array(int32, 1d, C)
    #   y = arg(3, name=y)  :: array(int32, 1d, C)
    #   $2load_global.0 = global(cuda: <module 'numba.cuda' from '/home/gmarkall/numbadev/numba/numba/cuda/__init__.py'>)  :: Module(<module 'numba.cuda' from '/home/gmarkall/numbadev/numba/numba/cuda/__init__.py'>)
    #   $4load_method.1 = getattr(value=$2load_global.0, attr=grid)  :: Macro(<class 'numba.cuda.cudadecl.Cuda_grid'>)
    #   del $4load_method.1
    #   del $2load_g[...]<jupyter_text>We only see one typing because we have called the function with one set of argument types.<jupyter_code>axpy.definitions<jupyter_output><empty_output><jupyter_text>Definitions are indexed by a tuple of `(compute capability, argument types)`. Let's call it with a different set, and see how this influences the set of definitions:<jupyter_code>x = x.astype(np.float32)
y = y.astype(np.float32)
r = r.astype(np.float32)
axpy[blocks, threads](r, a, x, y)

axpy.definitions<jupyter_output><empty_output><jupyter_text>We now see that there are two different definitions, which will have two different typings. Because the typings are different, the generated code for the definitions also differ from each other - we now have two different kernel objects.If the kernel is called again for a compute capability and set of argument types that is already in the definitions, the pre-compiled kernel is retrieved from the cache and called directly. If the compute capability or argument types are new and have not been seen before, the compiler pipeline is run again to compile a new specialisation. Generating LLVM IRThe next stage is lowering, which makes the LLVM IR.LLVM optimizes the IR. When we inspect the IR we see the optimized - unoptimized is very verbose (I would usually only look at it for debugging Numba's code generation).The process of translating the typed Numba IR into LLVM IR is relatively mechanical - there is a large number of functions registered inside Numba for "lowering" the Typed Numba IR into LLVM IR, which mostly proceed through translation one IR instruction at a time, generating the equivalent LLVM IR.The LLVM IR can be thought of as a machine-independent assembly language, and it has a relatively small instruction set of fairly low-level / primitive instructions (add, multiply, load, etc.). LLVM IR can be translated into many different target assembly languages - e.g. x86, ARM, PowerPC, PTX, etc - however, note that LLVM IR is not completely portable, and there are slight differences between LLVM IR for each machine. There are additional differences between the IR that is generated for PTX vs. other ISAs, in order to efficiently implement the mapping of work items to threads.<jupyter_code>for ((cc, types), llir) in axpy.inspect_llvm().items():
    print(f'Compute capability {cc} / argtypes {types}:\n')
    print(llir)
    print()<jupyter_output>Compute capability (7, 5) / argtypes (array(int32, 1d, C), float64, array(int32, 1d, C), array(int32, 1d, C)):

source_filename = "<string>"
target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
target triple = "nvptx64-nvidia-cuda"

@"_ZN6cudapy8__main__8axpy$241E5ArrayIiLi1E1C7mutable7alignedEd5ArrayIiLi1E1C7mutable7alignedE5ArrayIiLi1E1C7mutable7alignedE__errcode__" = local_unnamed_addr global i32 0
@"_ZN6cudapy8__main__8axpy$241E5ArrayIiLi1E1C7mutable7alignedEd5ArrayIiLi1E1C7mutable7alignedE5ArrayIiLi1E1C7mutable7alignedE__tidx__" = local_unnamed_addr global i32 0
@"_ZN6cudapy8__main__8axpy$241E5ArrayIiLi1E1C7mutable7alignedEd5ArrayIiLi1E1C7mutable7alignedE5ArrayIiLi1E1C7mutable7alignedE__ctaidx__" = local_unnamed_addr global i32 0
@"_ZN6cudapy8__main__8axpy$241E5ArrayIiLi1E1C7mutable7alignedEd5ArrayIiLi1E1C7mutable7alignedE5ArrayIiLi1E1C7mutable7alignedE__tidy__" = local_unnamed_addr [...]<jupyter_text>LLVM to PTXLLVM has a built-in JIT Compiler that is used for the CPU target. For PTX, we use NVVM because it has much better support for NVIDIA devices than the in-tree PTX backend, in terms of performance (and possibly feature support) but it does lag behind the LLVM release used by Numba (Numba uses LLVM 8, NVVM is somewhere between LLVM 3.4 and 5.0)We will pull in various internals of Numba to assemble a function that can be used to drive NVVM and get PTX out:<jupyter_code>from numba.core.compiler_lock import global_compiler_lock
from numba.cuda.cudadrv import nvvm
from numba.cuda.compiler import compile_cuda
from numba import float32, int32, void


with global_compiler_lock:
    argtys = (float32[:], int32, float32[:], float32[:])
    returnty = void
    cres = compile_cuda(axpy.py_func, void, argtys, debug=False, inline=False)
    fname = cres.fndesc.llvm_func_name
    lib, kernel = cres.target_context.prepare_cuda_kernel(cres.library, fname,
                                                          cres.signature.args,
                                                          debug=False)
    llvm_module = lib._final_module

    cc = (5, 2)
    arch = nvvm.get_arch_option(*cc)
    llvmir = str(llvm_module)
    ptx = nvvm.llvm_to_ptx(llvmir, opt=3, arch=arch)

print(ptx.decode('utf-8'))<jupyter_output>//
// Generated by NVIDIA NVVM Compiler
//
// Compiler Build ID: CL-24817639
// Cuda compilation tools, release 10.0, V10.0.130
// Based on LLVM 3.4svn
//

.version 6.3
.target sm_52
.address_size 64

	// .globl	_ZN6cudapy8__main__8axpy$244E5ArrayIfLi1E1A7mutable7alignedEi5ArrayIfLi1E1A7mutable7alignedE5ArrayIfLi1E1A7mutable7alignedE
.visible .global .align 4 .u32 _ZN6cudapy8__main__8axpy$244E5ArrayIfLi1E1A7mutable7alignedEi5ArrayIfLi1E1A7mutable7alignedE5ArrayIfLi1E1A7mutable7alignedE__errcode__;
.visible .global .align 4 .u32 _ZN6cudapy8__main__8axpy$244E5ArrayIfLi1E1A7mutable7alignedEi5ArrayIfLi1E1A7mutable7alignedE5ArrayIfLi1E1A7mutable7alignedE__tidx__;
.visible .global .align 4 .u32 _ZN6cudapy8__main__8axpy$244E5ArrayIfLi1E1A7mutable7alignedEi5ArrayIfLi1E1A7mutable7alignedE5ArrayIfLi1E1A7mutable7alignedE__ctaidx__;
.visible .global .align 4 .u32 _ZN6cudapy8__main__8axpy$244E5ArrayIfLi1E1A7mutable7alignedEi5ArrayIfLi1E1A7mutable7alignedE5ArrayIfLi1E1A7mutable7alignedE__tidy__;
.vi[...]<jupyter_text>If we're not interested in the way the internals are drive, we can also use `inspect_asm()` to get it for a function and argument types - here we just iterate over them all:<jupyter_code>for ((cc, types), asm) in axpy.inspect_asm().items():
    print(f'Compute capability {cc} / argtypes {types}:\n')
    print(asm)
    print()<jupyter_output>Compute capability (7, 5) / argtypes (array(int32, 1d, C), float64, array(int32, 1d, C), array(int32, 1d, C)):

//
// Generated by NVIDIA NVVM Compiler
//
// Compiler Build ID: CL-24817639
// Cuda compilation tools, release 10.0, V10.0.130
// Based on LLVM 3.4svn
//

.version 6.3
.target sm_70
.address_size 64

	// .globl	_ZN6cudapy8__main__8axpy$241E5ArrayIiLi1E1C7mutable7alignedEd5ArrayIiLi1E1C7mutable7alignedE5ArrayIiLi1E1C7mutable7alignedE
.visible .global .align 4 .u32 _ZN6cudapy8__main__8axpy$241E5ArrayIiLi1E1C7mutable7alignedEd5ArrayIiLi1E1C7mutable7alignedE5ArrayIiLi1E1C7mutable7alignedE__errcode__;
.visible .global .align 4 .u32 _ZN6cudapy8__main__8axpy$241E5ArrayIiLi1E1C7mutable7alignedEd5ArrayIiLi1E1C7mutable7alignedE5ArrayIiLi1E1C7mutable7alignedE__tidx__;
.visible .global .align 4 .u32 _ZN6cudapy8__main__8axpy$241E5ArrayIiLi1E1C7mutable7alignedEd5ArrayIiLi1E1C7mutable7alignedE5ArrayIiLi1E1C7mutable7alignedE__ctaidx__;
.visible .global .align 4 .u32 _ZN6cudapy8__main__8axpy[...]<jupyter_text>Creating a moduleNumba's compilation pipeline has finished with the generation of PTX. The next step towards execution is to use the driver API wrappers to create and load a module. All subsequent operations are completed with the driver API and Numba's wrappers around it.To create a module, we use a `Linker` object, that provides a Pythonic interface to the module management functions in the driver API. We add our PTX to the linker, and then inform the linker that we're finished to get back a cubin and some other information:<jupyter_code>from numba.cuda.cudadrv.driver import Linker

linker = Linker()
linker.add_ptx(ptx)
cubin, size = linker.complete()

compile_info = linker.info_log

print(f'Module size: {size}')
print(compile_info)<jupyter_output>Module size: 5992
info    : 40 bytes gmem
info    : Function properties for '_ZN6cudapy8__main__8axpy$244E5ArrayIfLi1E1A7mutable7alignedEi5ArrayIfLi1E1A7mutable7alignedE5ArrayIfLi1E1A7mutable7alignedE':
info    : used 18 registers, 0 stack, 0 bytes smem, 528 bytes cmem[0], 0 bytes lmem<jupyter_text>Further reading: [The Module Management section](https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html) of the CUDA Driver API lists the functions that are involved in creating and linking a module.  SASSThe module we have just created contains SASS code - the assembly that runs on the GPU. This was JIT-compiled by the driver during the link.There's no easy way to dump the assembly from a cubin from within Python or the Driver API, so to inspect it, we are forced to write to to a file and run `cuobjdump` on it:<jupyter_code>import ctypes

cubinarr = np.ctypeslib.as_array(ctypes.cast(cubin,ctypes.POINTER(ctypes.c_char)), shape=(size,))

with open('axpy.cubin', 'wb') as f:
    f.write(cubinarr)
!file axpy.cubin
!cuobjdump -sass axpy.cubin<jupyter_output>code for sm_75
		Function : _ZN6cudapy8__main__8axpy$244E5ArrayIfLi1E1A7mutable7alignedEi5ArrayIfLi1E1A7mutable7alignedE5ArrayIfLi1E1A7mutable7alignedE
	.headerflags    @"EF_CUDA_SM75 EF_CUDA_PTX_SM(EF_CUDA_SM75)"
        /*0000*/                   IMAD.MOV.U32 R1, RZ, RZ, c[0x0][0x28] ;               /* 0x00000a00ff017624 */
                                                                                         /* 0x000fd000078e00ff */
        /*0010*/                   S2R R0, SR_TID.X ;                                    /* 0x0000000000007919 */
                                                                                         /* 0x000e280000002100 */
        /*0020*/                   S2R R3, SR_CTAID.X ;                                  /* 0x0000000000037919 */
                                                                                         /* 0x000e240000002500 */
        /*0030*/                   IMAD R0, R3, c[0x0][0x0], R0 ;                        /* 0x000000[...]<jupyter_text>There is no PTX in the module (note the empty output):<jupyter_code>!cuobjdump -ptx axpy.cubin<jupyter_output><empty_output><jupyter_text>Loading the moduleThe module has been generated, but the code needs to be loaded to the GPU via the driver API. Here we need to get a context to load it in.Numba mangles the names of Python functions when it JIT compiles them to distinguish between compiled implementations of the same function for different types, functions of the same name in different modules, etc. - this is a similar idea to C++ name mangling. The Itanium name-mangling ABI is followed by Numba, as it is commonly used on many platforms.Numba also prepends the namespace `cudapy` to names, so we must do the same to be able to retrieve a reference to the compiled function.<jupyter_code>from numba.cuda.cudadrv.driver import load_module_image
from numba.core import itanium_mangler

ctx = cuda.get_current_device().get_primary_context()
module = load_module_image(ctx, cubin)

mangled_name = itanium_mangler.prepend_namespace(fname, ns='cudapy')
mangled_name<jupyter_output><empty_output><jupyter_text>Now we can get a reference to the compiled function, which will come back in a Python wrapper:<jupyter_code>cufunc = module.get_function(mangled_name)

type(cufunc)<jupyter_output><empty_output><jupyter_text>This reference can be used to launch a kernel on the GPU, but we must first prepare its arguments. Preparing kernel argumentsFirst we need to do a little bit of work to move data to the GPU, that Numba would usually do for us:<jupyter_code># Copy our arrays to the device
d_r = cuda.to_device(r)
d_x = cuda.to_device(x)
d_y = cuda.to_device(y)<jupyter_output><empty_output><jupyter_text>Next we need to prepare the list of arguments - passing a device array to a kernel is done by passing individual members of Numba's array object struct as parameters. These are:- `meminfo`: Not used on CUDA - on the CPU target it is used as part of Numba's reference counting implementation for managing allocations- `parent`: Not used in CUDA - on the CPU target, it is used for a pointer to the pyobject from which the array was created, in case object mode fallback is needed.- `nitems`: The number of items in the array (regardless of shape)- `itemsize`: The size of one item in the array- `data`: A pointer to the data- Shape and strides: The size of each dimension of the shape and passed first, then the stride of each dimension.This differs from how an array argument would usually be passed to a CUDA C kernel, where only a pointer to the data would be passed in, and the programmer would have a free choice about how to pass in any necessary information for iterating over the pointed-to data. Creating array argumentsWe first need to define a helper function for generating the list of parameters for a given device array. This is based on code in `numba.cuda.compiler.CUDAKernel._prepare_args`:<jupyter_code># A couple of helpers from Numba's CUDA driver implementation
from numba.cuda.cudadrv.driver import (device_pointer, is_device_memory,
                                       device_ctypes_pointer)     # noqa
from ctypes import addressof, c_void_p                            # noqa


# Helper function to create all the args for an array structure
def make_array_args(arr):
    args = []
    c_intp = ctypes.c_ssize_t

    meminfo = ctypes.c_void_p(0)
    parent = ctypes.c_void_p(0)
    nitems = c_intp(arr.size)
    itemsize = c_intp(arr.dtype.itemsize)
    data = ctypes.c_void_p(device_pointer(arr))

    args.append(meminfo)
    args.append(parent)
    args.append(nitems)
    args.append(itemsize)
    args.append(data)

    for ax in range(arr.ndim):
        args.append(c_intp(arr.shape[ax]))
    for ax in range(arr.ndim):
        args.append(c_intp(arr.strides[ax]))

    return args<jupyter_output><empty_output><jupyter_text>Creating the kernel argument listNext we create the list of arguments using the helper:<jupyter_code># Create the list of arguments - we compiled for float32[:], int32, float32[:],
# float32[:], so our second argument will be of int type.
args = []
args += make_array_args(d_r)
args += [ctypes.c_int(13)]  
args += make_array_args(d_x)
args += make_array_args(d_y)

# Make a list of pointers to the arguments
param_vals = []
for arg in args:
    if is_device_memory(arg):
        param_vals.append(addressof(device_ctypes_pointer(arg)))
    else:
        param_vals.append(addressof(arg))

params = (c_void_p * len(param_vals))(*param_vals)<jupyter_output><empty_output><jupyter_text>Launching the kernelAt this point we have everything ready to call `cuLaunchKernel` from the driver API. We will make use of the `driver` instance from `numba.cuda.cudadrv.driver`, as it handles the use of this function with `ctypes`, as well as wrapping it in a function that converts non-zero return codes into Python exceptions (which saves us having to check the return code).This code is inspired by that in `numba.cuda.cudadrv.driver.launch_kernel` (usually called from `Function.__call__`, which is called by `numba.cuda.compiler.CUDAKernel._kernel_call`).<jupyter_code># (see cudadrv.launch_kernel)

# CUresult cuLaunchKernel(CUfunction f,
#                        unsigned int gridDimX,
#                        unsigned int gridDimY,
#                        unsigned int gridDimZ,
#                        unsigned int blockDimX,
#                        unsigned int blockDimY,
#                        unsigned int blockDimZ,
#                        unsigned int sharedMemBytes,
#                        CUstream hStream, void **kernelParams,
#                        void ** extra)

from numba.cuda.cudadrv.driver import driver

driver.cuLaunchKernel(cufunc.handle,
                      blocks, 1, 1,
                      threads, 1, 1,
                      0,
                      0,
                      params,
                      None)<jupyter_output><empty_output><jupyter_text>Next we should try copying the result `d_r` back to the host and check it is as expected:<jupyter_code>h_r = d_r.copy_to_host()
reference = x * 13 + y 
print(h_r)
print(reference)

# Sanity check
np.testing.assert_allclose(reference, h_r)<jupyter_output>[4.50000e+01 8.20000e+01 7.20000e+01 ... 1.30016e+05 1.30073e+05
 1.30001e+05]
[4.50000e+01 8.20000e+01 7.20000e+01 ... 1.30016e+05 1.30073e+05
 1.30001e+05]
{'fraction_non_alphanumeric': 0.09127882187018269, 'fraction_numerical': 0.04234314014890616, 'mean_word_length': 3.7096076614148905, 'pattern_counts': {'":': 0, '<': 117, '<?xml version=': 0, '>': 112, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '28', 'dir': 'jupyter-structured-clean-dedup', 'id': '563231', 'license': 'MIT', 'n_tokens_mistral': 10413, 'n_tokens_neox': 9751, 'n_words': 3368, 'path': 'numba/Life_of_a_Numba_Kernel.ipynb', 'repo_name': 'avillasenorh/my_notebooks', 'size': '30497'}