Skip to content

jit

jit(f, /, ctx=None)

A decorator that marks a function for just-in-time compilation. This allows the function to be compiled and optimized for performance when called.

  • f: The function to be marked for JIT compilation. This function can use basic python control flow and operations (e.g. while, for, if). However, it shouldn't use more complex features like generators, classes, or recursion.
  • ctx: The scheduler to use for computation. Defaults to the result of get_default_scheduler().
  • A transformed function that inserts defer and compute statements to do tracing and optimization.

Example usage: @jit def my_function(A, B, C): D = A @ B while some_condition(D): D = D + C return D

In this example, my_function will be transformed to include defer and compute statements, allowing it to be optimized and executed efficiently when called. def opt_my_function(A, B, C): A, B = defer(A), defer(B) D = A @ B D = compute(D) while some_condition(D): C = defer(C) D = D + C D = compute(D) D = compute(D) return D

Source code in src/finch/finch_fused/jit.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def jit(f, /, ctx=None):
    """
    A decorator that marks a function for just-in-time compilation. This allows the
    function to be compiled and optimized for performance when called.

    Parameters:
    - f: The function to be marked for JIT compilation. This function can use
        basic python control flow and operations (e.g. while, for, if). However,
        it shouldn't use more complex features like generators, classes, or recursion.
    - ctx: The scheduler to use for computation. Defaults to the result of
        `get_default_scheduler()`.

    Returns:
    - A transformed function that inserts defer and compute statements to do tracing
       and optimization.


    Example usage:
    @jit
    def my_function(A, B, C):
        D = A @ B
        while some_condition(D):
            D = D + C
        return D

    In this example, `my_function` will be transformed to include defer and compute
    statements, allowing it to be optimized and executed efficiently when called.
    def opt_my_function(A, B, C):
        A, B = defer(A), defer(B)
        D = A @ B
        D = compute(D)
        while some_condition(D):
            C = defer(C)
            D = D + C
            D = compute(D)
        D = compute(D)
        return D
    """
    if ctx is None:
        ctx = get_default_scheduler()
    fused_fn = parse_fused_function(f)
    transformed_fn = insert_lazy_and_compute(fused_fn)
    return fused_function_to_python_function(transformed_fn)