Jane Street's SPMD Compiler for High-Performance Computing
Jane Street has developed a domain-specific language (DSL) compiler that leverages SPMD programming model to improve the performance and maintainability of SIMD code in OCaml.
Jane Street's summer intern Sai Konkimalla developed a compiler for a domain-specific language (DSL) that makes SIMD programming more ergonomic and widely accessible in OCaml. The firm relies heavily on performance-critical SIMD code for options pricing and low-latency risk checks. However, these programs typically involve complex C code due to OCaml's lack of native SIMD support. Konkimalla's project aimed to simplify this process by creating a DSL that uses the SPMD (single program, multiple data) programming model. This model allows developers to write a kernel once and compile it to multiple architectures without worrying about specific instruction sets (source ID 1). [1]
The DSL leverages OxCaml's recently added runtime metaprogramming features, enabling users to define kernels as regular OxCaml functions and then compile and run them on various architectures. This approach provides a higher level of abstraction, which allows the compiler to make better decisions about how to optimally use the underlying hardware. SPMD is a model popularized by CUDA and ISPC, where the developer writes a program describing what happens at each index of a loop, and the compiler converts it into a parallel computation (source ID 1). [1]
Konkimalla's project involved implementing a front end that translates kernel definitions into a GADT (generalized algebraic data type) representing a directed acyclic graph (DAG). The front end also included support for numeric data types and control flow constructs. To ensure correctness, the project was validated by comparing the JIT implementation against a reference interpreter and by running extensive test suites. Initial benchmarks showed that the JIT compiler was much faster than the interpreter but slower than manual OCaml and C implementations (source ID 1). [1]
To improve performance, an intermediate representation (IR) based on A-normal form (ANF) was implemented. This IR is more amenable to optimization and makes code generation trivial. The IR allows the optimizer to reason about how to re-order or unroll loops, leading to more efficient code. Future work includes adding vectorization, constant hoisting, and scheduling optimizations to further enhance performance (source ID 1). [1]
Advertisement
The project, developed by Sai Konkimalla during his 2026 summer internship at Jane Street, aimed to address the firm’s reliance on low-level C for SIMD (single-instruction, multiple-data) code used in performance-critical applications like options pricing and low-latency risk checks. Jane Street’s existing approach required developers to write complex inner loops in C, which involved manual tuning of unroll factors and vector widths, and was tightly coupled to specific instruction sets. This made porting code to new architectures or processors a laborious process. The project’s goal was to create a domain-specific language (DSL) compiler that would simplify and generalize SIMD programming by leveraging OxCaml’s new runtime metaprogramming features. [1]
The DSL compiler utilizes the SPMD (single program, multiple data) programming model, which abstracts away hardware-specific details and allows the compiler to optimize code for different architectures automatically. SPMD is similar to models used in CUDA for Nvidia GPUs and ISPC for Intel CPUs, where developers define the computation for each element of a data set, and the compiler handles the parallelization. This approach allows for more maintainable and portable code, as the same kernel can be compiled to run on various hardware without modification. The DSL compiler was implemented in OxCaml, with the front end translating kernel definitions into a GADT (generalized algebraic data type) that models a directed acyclic graph (DAG) of operations. [1]
The project leveraged OxCaml’s runtime metaprogramming to enable just-in-time (JIT) compilation, which allows the compiler to generate and execute code during program execution. This approach provides two key benefits: it eliminates the need to precompile for all possible target architectures, and it enables the compiler to use runtime data for optimizations, such as eliminating branch instructions that can slow down execution. Sai’s work was the first real use of OxCaml’s runtime metaprogramming features at Jane Street, and he identified and reported several bugs in its implementation to the OxCaml team. [1]
To validate the correctness and performance of the JIT compiler, Sai compared it against a reference interpreter and benchmarked it against implementations in plain OxCaml, C with AVX2 and AVX512 intrinsics, and SPMD kernels. The JIT was found to be significantly faster than the interpreter but slower than the C implementations, which was expected due to the lack of compiler optimizations. Sai then introduced an intermediate representation (IR) based on A-normal form (ANF), which made it easier to analyze and optimize the code. This IR enabled optimizations such as dead code elimination and loop unrolling, and the project is now working on adding further optimizations like vectorization and constant hoisting to match the performance of C implementations. [1]
Advertisement
The Cali Reporter newsletter
The stories worth your time.
California, the country, and what comes next. Join our free newsletter for reporting that keeps you informed.
More in Technology
- Google Play's Closed Testing Requirements and the Importance of Feedback
Google Play's closed testing rule emphasizes the need for developers to gather meaningful feedback from testers, not just track app installs. This approach ensures that real user experiences are evaluated and addressed…
- World Diabetes Day 2026 Focuses on Early Diagnosis and Awareness
The 2026 World Diabetes Day campaign emphasizes the importance of early diagnosis, equitable access to care, and reducing diabetes stigma, according to the WHO. The theme 'Early action on diabetes' highlights the need…
- Understanding Memory Usage on Linux VPS
This article explains how to accurately interpret memory usage on a Linux VPS, focusing on the importance of the 'available' metric and the potential pitfalls of relying on the 'free' column.