Docs
Jit Compilation

JIT Compilation (experimental)

Endia provides Just-In-Time (JIT) compilation capabilities to optimize and cache function execution for improved performance. The JIT compiler traces function calls, optimizes subgraphs, and caches the results for future use.

By default the JIT compiler uses a set of home made optimizations to improve performance. However, you can also leverage Modular's MAX Engine 🔥. In certain scenarios, using MAX can provide additional performance benefits.

jit

def jit(f: Variant[Callable, def (List[Array]) -> Array], compile_with_MAX: Bool = False) -> Callable
Applies Just-In-Time compilation to the given function or Callable.
Parameters
  • f (Variant[Callable, def (List[Array]) -> Array]): The function or Callable to be JIT-compiled.
  • compile_with_MAX (bool, optional): Whether to compile subgraphs with MAX. Default is False.
Returns
  • Callable: A JIT-compiled version of the input function.
Example
import endia as nd
 
def my_function(args: List[Array]):
    return nd.sum(args[0] * args[1])
 
x = nd.array('[1.0, 2.0, 3.0]')
y = nd.array('[4.0, 5.0, 6.0]')
my_function_jit = jit(my_function)
result = my_function_jit(List(x, y))

JIT Compilation Process

Learn about the JIT compilation process

  1. Tracing: The function is traced, capturing all operations performed on the input arrays.

  2. Branch Handling: The JIT compiler can handle functions with conditional statements. It compares operations during execution with previously captured traces, branching when necessary and storing new paths.

  3. Graph Building: The traced operations form a computation graph, with each node representing an operation or an array.

  4. Subgraph Optimization: The graph is divided into subgraphs marked by breakpoints. These subgraphs are optimized by fusing elementwise operations and applying other optimizations.

  5. Caching: Optimized subgraphs are cached for future use. When the same breakpoint is encountered in subsequent executions, the cached, optimized subgraph is used instead of recomputing.

Performance Considerations

Important performance considerations for JIT compilation
  • JIT compilation can significantly improve performance for functions that are called multiple times with similar input shapes.
  • The first call to a JIT-compiled function may be slower due to the tracing and optimization process.
  • For functions with highly dynamic control flow that changes frequently based on input data, JIT compilation may not provide significant benefits.
Custom OpsOptimization