Profiling with a decorator and @njit

Numba Community Support
28

Hi @nelson2005

This should do the job:

from numba import njit, objmode, extending, types
from numba.core.cgutils import get_or_insert_function
import numpy as np
import ctypes
import time
from llvmlite import ir

CLOCK_MONOTONIC = 0x1
clock_gettime_proto = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int,
                                       ctypes.POINTER(ctypes.c_long))
pybind = ctypes.CDLL(None)

clock_gettime_addr = pybind.clock_gettime
clock_gettime_fn_ptr = clock_gettime_proto(clock_gettime_addr)

@extending.intrinsic
def clock_gettime(typingctx, clockid_t , timespec):
    def codegen(context, builder, sig, args):
        fnty = ir.FunctionType(
            ir.IntType(64), 
            (
                context.get_value_type(sig.args[0]), 
                context.get_value_type(sig.args[1]),
            )
        )
        fn = get_or_insert_function(builder.module, fnty, "clock_gettime")
        return builder.call(fn, args)

    sig = types.int64(clockid_t , timespec)
    return sig, codegen

@njit(cache=True)
def timenow():
    timespec = np.zeros(2, dtype=np.int64)
    clock_gettime(CLOCK_MONOTONIC, timespec.ctypes)
    ts = timespec[0]
    tns = timespec[1]
    return np.float64(ts) + 1e-9 * np.float64(tns)

@njit(cache=True)
def pointless_delay(seconds):
    with objmode():
        s = time.time()
        e = 0
        while (e < seconds):
            e = time.time() - s

@njit(cache=True)
def do_stuff(n):
    t0 = timenow()
    pointless_delay(n)
    print("Elapsed", timenow() - t0)

do_stuff(1)
do_stuff(2)
do_stuff(3.21)
1 Like