jacinle.jit.cython.jit_compile

Classes

CythonCompiledFunction
CythonJITCompiler

Functions

jit_cython([f, name, force_update, ...])

JIT compile a simple Python function with Cython.

Class CythonCompiledFunction

class CythonCompiledFunction

Bases: object

__call__(*args, **kwargs)

Call self as a function.

__init__(name, func_name, build_dir, py_source, extra_args)

__new__(**kwargs)

Class CythonJITCompiler

class CythonJITCompiler

Bases: object

__init__()

__new__(**kwargs)

add_extra_args(func)

build_func(func_name, build_dir)

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)

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)

write_cython(name, func_name, build_dir, source)

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)

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)

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

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

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'

Functions

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