JIT Design and Implementation

This document explains the design and implementation of Julia's JIT, after codegen has finished and unoptimized LLVM IR has been produced. The JIT is responsible for optimizing and compiling this IR to machine code, and for linking it into the current process and making the code available for execution.

Introduction

The JIT is responsible for managing compilation resources, looking up previously compiled code, and compiling new code. It is primarily built on LLVM'sOn-Request-Compilation (ORCv2) technology, which provides support for a number of useful features such as concurrent compilation, lazy compilation, and the ability to compile code in a separate process. Though LLVM provides a basic JIT compiler in the form of LLJIT, Julia uses many ORCv2 APIs directly to create its own custom JIT compiler.

Overview

Codegen produces an LLVM module containing IR for one or more Julia functions from the original Julia SSA IR produced by type inference (labeled as translate on the compiler diagram above). It also produces a mapping of code-instance to LLVM function name. However, though some optimizations have been applied by the Julia-based compiler on Julia IR, the LLVM IR produced by codegen still contains many opportunities for optimization. Thus, the first step the JIT takes is to run a target-independent optimization pipeline^[tdp] on the LLVM module. Then, the JIT runs a target-dependent optimization pipeline, which includes target-specific optimizations and code generation, and outputs an object file. Finally, the JIT links the resulting object file into the current process and makes the code available for execution. All of this is controlled by code insrc/jitlayers.cpp.

Currently, only one thread at a time is permitted to enter the optimize-compile-link pipeline at a time, due to restrictions imposed by one of our linkers (RuntimeDyld). However, the JIT is designed to support concurrent optimization and compilation, and the linker restriction is expected to be lifted in the future when RuntimeDyld has been fully superseded on all platforms.

Optimization Pipeline

The optimization pipeline is based off LLVM's new pass manager, but the pipeline is customized for Julia's needs. The pipeline is defined insrc/pipeline.cpp, and broadly proceeds through a number of stages as detailed below.

1. Early Simplification
   1. These passes are mainly used to simplify the IR and canonicalize patterns so that later passes can identify those patterns more easily. Additionally, various intrinsic calls such as branch prediction hints and annotations are lowered into other metadata or other IR features.SimplifyCFG (simplify control flow graph),DCE (dead code elimination), andSROA (scalar replacement of aggregates) are some of the key players here.
2. Early Optimization
   1. These passes are typically cheap and are primarily focused around reducing the number of instructions in the IR and propagating knowledge to other instructions. For example,EarlyCSE is used to perform common subexpression elimination, andInstCombine andInstSimplify perform a number of small peephole optimizations to make operations less expensive.
3. Loop Optimization
   1. These passes canonicalize and simplify loops. Loops are often hot code, which makes loop optimization extremely important for performance. Key players here includeLoopRotate,LICM, andLoopFullUnroll. Some bounds check elimination also happens here, as a result of theIRCE pass which can prove certain bounds are never exceeded.
4. Scalar Optimization
   1. The scalar optimization pipeline contains a number of more expensive, but more powerful passes such asGVN (global value numbering),SCCP (sparse conditional constant propagation), and another round of bounds check elimination. These passes are expensive, but they can often remove large amounts of code and make vectorization much more successful and effective. Several other simplification and optimization passes intersperse the more expensive ones to reduce the amount of work they have to do.
5. Vectorization
   1. Automatic vectorization is an extremely powerful transformation for CPU-intensive code. Briefly, vectorization allows execution of asingle instruction on multiple data (SIMD), e.g. performing 8 addition operations at the same time. However, proving code to be both capable of vectorization and profitable to vectorize is difficult, and this relies heavily on the prior optimization passes to massage the IR into a state where vectorization is worth it.
6. Intrinsic Lowering
   1. Julia inserts a number of custom intrinsics, for reasons such as object allocation, garbage collection, and exception handling. These intrinsics were originally placed to make optimization opportunities more obvious, but they are now lowered into LLVM IR to enable the IR to be emitted as machine code.
7. Cleanup
   1. These passes are last-chance optimizations, and perform small optimizations such as fused multiply-add propagation and division-remainder simplification. Additionally, targets that do not support half-precision floating point numbers will have their half-precision instructions lowered into single-precision instructions here, and passes are added to provide sanitizer support.

Target-Dependent Optimization and Code Generation

LLVM provides target-dependent optimization and machine code generation in the same pipeline, located in the TargetMachine for a given platform. These passes include instruction selection, instruction scheduling, register allocation, and machine code emission. The LLVM documentation provides a good overview of the process, and the LLVM source code is the best place to look for details on the pipeline and passes.

Linking

Currently, Julia is transitioning between two linkers: the older RuntimeDyld linker, and the newerJITLink linker. JITLink contains a number of features that RuntimeDyld does not have, such as concurrent and reentrant linking, but currently lacks good support for profiling integrations and does not yet support all of the platforms that RuntimeDyld supports. Over time, JITLink is expected to replace RuntimeDyld entirely. Further details on JITLink can be found in the LLVM documentation.

Execution

Once the code has been linked into the current process, it is available for execution. This fact is made known to the generating codeinst by updating theinvoke,specsigflags, andspecptr fields appropriately. Codeinsts support upgradinginvoke,specsigflags, andspecptr fields, so long as every combination of these fields that exists at any given point in time is valid to be called. This allows the JIT to update these fields without invalidating existing codeinsts, supporting a potential future concurrent JIT. Specifically, the following states may be valid:

1. invoke is NULL,specsigflags is 0b00,specptr is NULL
   1. This is the initial state of a codeinst, and indicates that the codeinst has not yet been compiled.
2. invoke is non-null,specsigflags is 0b00,specptr is NULL
   1. This indicates that the codeinst was not compiled with any specialization, and that the codeinst should be invoked directly. Note that in this instance,invoke does not read either thespecsigflags orspecptr fields, and therefore they may be modified without invalidating theinvoke pointer.
3. invoke is non-null,specsigflags is 0b10,specptr is non-null
   1. This indicates that the codeinst was compiled, but a specialized function signature was deemed unnecessary by codegen.
4. invoke is non-null,specsigflags is 0b11,specptr is non-null
   1. This indicates that the codeinst was compiled, and a specialized function signature was deemed necessary by codegen. Thespecptr field contains a pointer to the specialized function signature. Theinvoke pointer is permitted to read bothspecsigflags andspecptr fields.

In addition, there are a number of different transitional states that occur during the update process. To account for these potential situations, the following write and read patterns should be used when dealing with these codeinst fields.

1. When writinginvoke,specsigflags, andspecptr:
   1. Perform an atomic compare-exchange operation of specptr assuming the old value was NULL. This compare-exchange operation should have at least acquire-release ordering, to provide ordering guarantees of the remaining memory operations in the write.
   2. Ifspecptr was non-null, cease the write operation and wait for bit 0b10 ofspecsigflags to be written.
   3. Write the new low bit ofspecsigflags to its final value. This may be a relaxed write.
   4. Write the newinvoke pointer to its final value. This must have at least a release memory ordering to synchronize with reads ofinvoke.
   5. Set the second bit ofspecsigflags to 1. This must be at least a release memory ordering to synchronize with reads ofspecsigflags. This step completes the write operation and announces to all other threads that all fields have been set.
2. When reading all ofinvoke,specsigflags, andspecptr:
   1. Read theinvoke field with at least an acquire memory ordering. This load will be referred to asinitial_invoke.
   2. Ifinitial_invoke is NULL, the codeinst is not yet executable.invoke is NULL,specsigflags may be treated as 0b00,specptr may be treated as NULL.
   3. Read thespecptr field with at least an acquire memory ordering.
   4. Ifspecptr is NULL, then theinitial_invoke pointer must not be relying onspecptr to guarantee correct execution. Therefore,invoke is non-null,specsigflags may be treated as 0b00,specptr may be treated as NULL.
   5. Ifspecptr is non-null, theninitial_invoke might not be the finalinvoke field that usesspecptr. This can occur ifspecptr has been written, butinvoke has not yet been written. Therefore, spin on the second bit ofspecsigflags until it is set to 1 with at least acquire memory ordering.
   6. Re-read theinvoke field with at least an acquire memory ordering. This load will be referred to asfinal_invoke.
   7. Read thespecsigflags field with any memory ordering.
   8. invoke isfinal_invoke,specsigflags is the value read in step 7,specptr is the value read in step 3.
3. When updating aspecptr to a different but equivalent function pointer:
   1. Perform a release store of the new function pointer tospecptr. Races here must be benign, as the old function pointer is required to still be valid, and any new ones are also required to be valid as well. Once a pointer has been written tospecptr, it must always be callable whether or not it is later overwritten.

Although these write, read, and update steps are complicated, they ensure that the JIT can update codeinsts without invalidating existing codeinsts, and that the JIT can update codeinsts without invalidating existinginvoke pointers. This allows the JIT to potentially reoptimize functions at higher optimization levels in the future, and also will allow the JIT to support concurrent compilation of functions in the future.