jacinle.jit.cython package

Submodules

jacinle.jit.cython.jit_compile module

class jacinle.jit.cython.jit_compile.CythonCompiledFunction(name, func_name, build_dir, py_source, extra_args)

Bases: object

class jacinle.jit.cython.jit_compile.CythonJITCompiler

Bases: object

add_extra_args(func)

all_funcs = {}

args_re = re.compile("([a-zA-Z_][a-zA-Z_0-9]*)( )*:( )*'?(([a-zA-Z_][a-zA-Z_0-9]*\\.)*([a-zA-Z_][a-zA-Z_0-9]*))(\\[.*?\\])?'?", re.MULTILINE)

build_func(func_name, build_dir)

cdef_function_re = re.compile("^([ \t]*)def (.*?)( )*->( )*'?(([a-zA-Z_][a-zA-Z_0-9]*\\.)*([a-zA-Z_][a-zA-Z_0-9]*))'?(\\[.*?\\])?", re.MULTILINE)

cdef_re = re.compile("^([ \t]+)([a-zA-Z_][a-zA-Z_0-9]*)( )*:( )*'?(([a-zA-Z_][a-zA-Z_0-9]*\\.)*([a-zA-Z_][a-zA-Z_0-9]*))(\\[.*?\\])?'?", re.MULTILINE)

check_module_exist(func_name, build_dir)

check_source_updated(build_dir, source, force_update=False)

cleanup_source(func)

compile(f=None, *, name=None, force_update=False, boundscheck=True, wraparound=True)

decorator_cleanup_re = re.compile('^([ \t]*)@jit_cython.*\n', re.MULTILINE)

dependency_visitor = <jacinle.jit.cython.jit_compile._DependencyNodeVisitor object>

gen_cython_source(func)

gen_pyx_source(func)

gen_pyx_source_cdef(func)

get_build_dir(func)

get_name(func)

get_source(func)

load_func(func_name, build_dir)

optimize_cdef(func)

optimize_cdef_function(func)

resolve_dependencies(func)

setup_template = "\nfrom setuptools import setup, Extension\nimport numpy as np\n\nsetup(\n name='{func_name}',\n ext_modules=[\n Extension('{func_name}', sources=['{func_name}.pyx'], include_dirs=[np.get_include()], annotate=True)\n ]\n)\n"

source_footer = ''

source_header = '\nimport numpy as np\ncimport numpy as np\ncimport cython\n'

write_cython(name, func_name, build_dir, source)

jacinle.jit.cython.jit_compile.jit_cython(f=None, *, name=None, force_update=False, boundscheck=True, wraparound=True)

JIT compile a simple Python function with Cython. Currently, only support functions with built-in Python functions and numpy operations (and arrays). This function uses Python annotations to determine the types of the arguments and return values. Therefore, to maximize performance, you should try to annotate all arguments and variables used. See examples/cython/jit_example.py for more details.

Example

@jit_cython(boundscheck=False, wraparound=False)
def f(x: int):
    return x + 1

@jit_cython
def sum(x: np.ndarray, n: int):
    i: int
    s: float = 0.0
    for i in range(n):
        s += x[i]
    return s