ocannl

OCANNL is sponsored by Ahrefs! Visit the Ahrefs website.

OCANNL -- OCaml Compiles Algorithms for Neural Networks Learning

• A from-scratch, compiled Deep Learning framework.
• Implements backpropagation (i.e. first-order reverse mode autodiff) and shape inference.
• The long-term goal is to provide several "low-level" backends, aiming to seek inspiration from projects such as tinygrad, TVM, Luminal.
  • OCANNL starts with a high-level representation, but can compile everything down to for loops.
• The library users can compile any amount of code into a routine (i.e. a compilation unit). The user decides explicitly what the scope of a compilation unit is, by putting together the corresponding code. Depending on the use case:
  • the whole training update step can be a single routine,
  • or the step can be composed of a gradient update routine (a forward pass and a backprop pass) and a params update routine (e.g. SGD with momentum, ADAM, etc.),
  • or the user can compile parts of a model separately, manually composing the corresponding forward pass code and the backprop code.
• Tensor axes are split into kinds: batch, input and output. Tensor dimensions have an optional basis.
  • The basis (aka dimension units) ensures a more precise semantics for dimension matching. It's not an axis selection mechanism.
• OCANNL has full support for a significantly extended einsum notation, integrated with shape inference. See comparison with einops for how this relates to the popular einops library. Supports static indexing, with a built-in operation to take a slice of the batch axes, integrated with shape inference. Extensible to more static indexing patterns as needs arise.
  • OCANNL does not have dynamic indexing (using the last axis of one tensor as indices into another tensor). If it's needed, it can be added (we had a prototype once, removed to reduce complexity). Then it would also be integrated with shape inference.
• OCANNL offers two main levels of abstraction.
  • Tensor expressions as differentiable computations, centered around the %op syntax extension.
    • %op stands for "operation", it's meant to express tensors: Tensor.t, and tensor functions.
  • Plain computations, centered around the %cd syntax extension. It integrates the arrayjit backend library with shape inference.
    • %cd stands for "code", it's meant to express assignment computations: Assignments.comp.
• Fully supports mixed-precision computations, with bidirectional precision inference.
  • E.g. higher-precision network components, or gradients at a higher precision than values.
• Should be easily extensible.
• Model surgery should be straightforward (not sure if we are there yet).

Usage

The CUDA backend requires at least CUDA version 12.8. The Metal backend requires at least MSL version 3.1. The HIP backend (AMD GPUs) requires ROCm / the AMD HIP SDK, via the hipjit bindings (opam install hipjit).

API documentation entry point.

A possible route to learning OCANNL:

1. Read the introductory slides.
2. Read: shapes and the generalized einsum beginner-to-advanced slides.
3. Read Tensors and Contexts and, for the runtime API, Context.
4. Read the migration guide.
5. Read the syntax extensions documentation docs/syntax_extensions.md.
6. Read the NN building blocks file lib/nn_blocks.ml and the training recipes lib/train.ml.

• Work through the makemore tutorial — a character-level language-model progression mirroring Andrej Karpathy's Neural Networks: Zero to Hero lectures.

7. Read the introductory part of the shape inference documentation docs/shape_inference.md.
8. For the paper-facing account, read the workshop article docs/ocannl_workshop_article_human.md (an archival artifact: it was written for the OCaml Workshop / FProPer submission and describes the project as of the 0.8 release; no conference submission is currently scheduled), the formal core technical report ocannl-formal-core-technical-report.pdf (LaTeX source in docs/), and the constraint-generation notes docs/shape-constraint-generation.md.
9. Skim the configuration documentation ocannl_config.reference.
10. Improve your understanding by reading or skimming the framework internals: tensor/shape.mli, tensor/tensor.mli, tensor/operation.ml, arrayjit/lib/context.mli.
11. Read the implementation overview:
12. The various tests.
13. Shape inference details docs/shape_inference.md.
14. Backend-independent optimizations docs/lowering_and_inlining.md -- lowering means translating (compiling) from the high-level representation (as assignments) to the low-level representation.
15. Schedules and autotuning docs/schedules_and_autotuning.md -- the loop-nest transform layer (parallelization, tiling, staging, tensor cores) and the empirical search over it.

Using the tracing debugger with CUDA and HIP computations

To use debugging as provided by configuring Utils.settings.debug_log_from_routines <- true with the cuda or hip backend, wrap the code that schedules work and synchronizes the GPU with Utils.capture_stdout_logs. Both GPU APIs expose device-side printf, but not fprintf; the runtime drains the device printing buffer to process stdout around synchronization. Synchronize the context inside the capture window so all device output is available before stdout is restored.

NOTE: debug logging from CUDA or HIP in complex settings is a bit tricky, as it involves another thread (domain) intercepting and filtering stdout. If facing issues, try the setting never_capture_stdout=true (see ocannl_config.reference).

Milestones

See ROADMAP.md for the detailed schedule. GitHub issue assignments are the source of truth for release scope. v1.0 was released on August 13, 2026; the next target is v1.1, soft-dated August 24, 2026. Release dates are now project-internal and aspirational — through v1.0 they were pinned to conference deadlines.

Note (July 2026): v0.7 shipped on July 3, 2026 as the consolidated paper-ready release. v0.6.4 was skipped as a release — its work (concatenation, RoPE, transformer toy) shipped inside v0.7 — and v0.7.2 was consolidated into v0.7. v0.7.1 was dissolved: its AMD HIP backend (#411) shipped in v0.8; completed examples and tokenizer work landed subsequently, while remaining work follows the current GitHub milestone assignments.

Update (August 2026): v0.9 shipped on August 3, 2026 with its milestone fully closed, so v1.0 becomes the next paper-deadline release and v1.1 follows it. CUDA/HIP graph capture (#488) moved from v0.9 to v1.0, and the training/deployment utilities (#96, #97, #122, #465, #467) plus the lib/ design study (#435) moved from v1.0 to v1.1.

Venue history (August 2026): the OCaml Workshop submission was not accepted — the article reads as a research report rather than an introductory demonstration, which put it outside that audience's scope. IFL 2026 was then considered and decided against as a poor fit. No conference submission is currently scheduled; the workshop article stays in the repository unchanged, as a historical artifact of the project at v0.8, and the formal core technical report continues as live work.

Update (August 13, 2026): v1.0 shipped, with its milestone fully closed (49 issues). Release dates are no longer pinned to paper deadlines: v1.1's soft target is August 24, 2026 (the OCaml Workshop date, used as an anchor rather than as a submission). v1.1 and v1.2 were rebalanced along a different seam than the original split — v1.1 is the compiler work plus the training-loop mechanics it needs, v1.2 the consumers and explorations, including the training experience a user sees. The sequence is 0.9 → 1.0 → 1.1 → 1.2.

• 0.7 (Jul 3, 2026, released): Frontend finalization + compiler optimizations. The consolidated paper-ready release for workshop submissions (OCaml Workshop, FProPer). Absorbs the former v0.6.4/v0.6.5/v0.7.0 frontend work and the former v0.7.2 optimization work.
  • Migrate from the "hosted tensor" idea to always requiring a context when accessing tensors and dealing with devices directly; remove the array field of Tnode.t and the hosted memory mode (#333).
  • Tensor saving, loading, and restoring (#373).
  • Axis concatenation in the einsum syntax (a^b), generalizing tensor stacking; shifting (1^i=>i) and padding (i=>1^i) as fixed-index special cases (#49).
  • RoPE and other non-learned position embeddings (#398); decoder-only autoregressive transformer toy (#57).
  • Ternary einsum notation (#305); loop-invariant hoisting (#350) and common subexpression elimination (#351).
  • Universal Pool Allocator across backends (#344): per-context-delta working pools, per-device constant pools, reserved merge pool, and pooled Metal bindings.
  • Sharding primitives, data-parallel training driver, and zero-copy leading-axis slice views (#293).
  • Workshop article, formal core technical report, and shape-constraint-generation notes.
• 0.8 (Jul 13, 2026, released): Parallel schedules, autotuning, tensor cores, and AMD HIP.
  • Schedule transforms and generated CPU/GPU kernels harvested from the Böhm CPU/CUDA matmul articles and llm.c (#412).
  • Kernel fission, hardware-mapped loop axes, shared staging, packed/register-tiled Tile_mma, and explicit SIMD codegen.
  • Measured schedule autotuning with caches, sketch seeds, and per-segment candidates.
  • CUDA WMMA/inline-PTX, Metal simdgroup-matrix, and HIP rocWMMA tensor-core paths.
  • HIP backend for AMD hardware via the independent hipjit bindings (#411).
  • Native Windows support via mingw-w64; an additional MSVC toolchain was evaluated and closed as not planned (#313).
• 0.9 (Aug 3, 2026, released): Schedule quality, deterministic parallelism, and convolution performance.
  • Constraint-based schedule legality (#494), an analytic default-schedule cost model (#491), and the cross-machine benchmark/tuning sweep (#476, re-measured under #538).
  • Deterministic split reductions (#484, #537) and a mixed-precision training recipe: precision policy, master weights, loss scaling, and forward-only reduced precision (#492).
  • Convolution schedule families and boundary handling: implicit-GEMM sketches, blocked tiles, epilogue twins, compact strided staging, and clamped windows (#493, #500, #501, #502, #504).
  • Correct overlapping-window tropical/einmax1 gradients (#512), with the non-overlapping fast path restored (#527).
  • Tensor-core hardening delivered tf32 policy, CUDA 13 support, pad-to-tile scheduling, static partitioning, and packed-uniform retirement (#478, #482, #485, #508, #509).
  • Search survivability: typed candidate-failure containment (#536), HIP scratch pre-validation (#533), no unparallelized GPU dispatches (#532), a complete decline census (#541, #543), and reachable GPU mma candidates (#521).
  • CNN classifiers (#54), GPT-2 inference (#377), and the TVM, Tiramisu, superoptimizer, and Lean Attention studies (#242, #267, #261, #263).
• 1.0 (Aug 13, 2026, released): Advanced compiler tiers and schedule-quality follow-through.
  • Branch-and-bound schedule inference: a refinement tree over partial schedules, legality verdicts with witnesses as fathoming, and admissible cost floors (#514).
  • Inlining as a first-class searchable decision (#555), on top of retiring the concrete-index tracer for affine access relations (#554, #560, #561).
  • CUDA tensor-core completeness, software pipelining, CUDA/HIP graph capture, and budget-driven rematerialization (#481, #487, #488, #498).
  • The gpt2_mini arc: attribute the step (#531), then judge companion coverage at the site's arity (#569) — tuned step 107.4 → 52.4 ms on CUDA, 45.6 → 25.4 ms on HIP; batched/rank-3 sites seeded (#528).
  • CPU reduced precision: 16-bit storage with f32 compute and native fp16 arithmetic (#516, #517); cc worker-pool uniformity on hybrid CPUs (#530).
  • Frontend, configuration and diagnostics: use-site row resolution and stretch (#544), config profiles (#559), the tuner's honest reference point (#552), routine-name collision policy (#513).
• 1.1 (Aug 24, 2026, soft target): Performance carry-overs, algebraic rewrites, and training-loop utilities.
  • Training-loop mechanics: LR schedules, global-norm clipping, gradient accumulation, and mmap-backed checkpoint loading (#465, #467).
  • Fused attention via online softmax and the remaining convolution tiers (#483, #503, #505).
  • Search follow-ups from the v1.0 evaluation: builder preconditions as tree verdicts, a fittable memory leg, profitability in enablement promotion, sketch-family extraction (#577, #578, #579, #580).
  • Narrow-operand tensor-core tiling (#575) and the gpt2_mini residue (#573, #574).
  • Robustness and test seams: digest completeness, fault-injection inventory, executed parity for analysis-level passes (#562, #571, #572, #582).
• 1.2 (undated): Consumers and explorations.
  • Training experience: resumable checkpoints, experiment tracking, and plot polish (#96, #122, #103).
  • Models, reproductions and demos: model surgery, LSTM, Bonsai RNN, digit addition, BERT/ModernBERT, DisTrO, and a Gemma 3 demo target (#33, #60, #182, #427, #297, #278, #570).
  • Frontend design, library and deployment: shape schemes, the Simply/NanoDO study for lib/, PoPE, and inference binaries (#404, #435, #444, #97).
  • Integrations and external-framework study: Polars, krnl/autograph (#219, #277).
  • Performance gated on hardware or on evidence not yet in hand: pressure-aware device memory management, async-copy staging refinements, HIP CDNA tensor cores, and CUDA pinned/constant host memory (#565, #576, #477, #170, #195).

Releases

For more details, see CHANGES.

• 1.0: Advanced compiler tiers and schedule-quality follow-through.
  • Schedule inference as branch-and-bound: a refinement tree over partial schedules, legality verdicts carrying their witnesses, and admissible cost floors, so a subtree is refuted or priced before any of its members is built.
  • Inlining joined scheduling as a searchable decision surface, once the concrete-index tracer was retired in favor of the affine access relations.
  • CUDA/HIP graph capture of the fissioned step, software-pipelined double-buffered staging (cp.async on CUDA), budget-driven rematerialization, and the CUDA tensor-core profile's remaining shapes (fp8, ldmatrix over swizzled staging).
  • Reduced precision on CPU: 16-bit storage with f32 compute, and native fp16 arithmetic where the hardware has it (~2x on an M-series).
  • The gpt2_mini step roughly halved on both GPU backends by judging companion coverage at the site's arity, after a time-attribution profile located 70% of the step in five declined kernels.
  • Config profiles (reproducible / performance), cc worker-pool uniformity on hybrid CPUs, and use-site row resolution narrowed to the leaf-tensor rule it always was.
• 0.9: Program search and optimization.
  • Native affine program analysis and a Schedule.op_legality oracle; an analytic roofline cost model that picks untuned defaults and pre-filters the autotune beam (advisory throughout).
  • Deterministic split reductions, pad-to-tile scheduling (PADTO), static partitioning, epilogue fusion, and launch-time symbolic extents.
  • Convolution schedule families: implicit GEMM via packing Stage, blocked tile flavors, epilogue twins, compacting strided-row staging, and clamped-window pooling.
  • Mixed precision: a precision-assignment policy, master weights with cast twins, dynamic loss scaling with a fused on-device gate, tf32 matmuls, and forward-only reduced precision by load-time conversion.
  • Liveness-based buffer aliasing; a packed uniform that is total over shapes and now backs default parameter initialization.
  • Search survivability: typed candidate-failure containment across Metal/CUDA/HIP, HIP scratch pre-validation, no unparallelized GPU dispatches, and a decline census that accounts for every refusal.
  • A cross-machine benchmark sweep on Metal, CUDA and HIP, with the reports checked in under benchmarks/.
• 0.8: Parallel schedules, autotuning, tensor cores, and the AMD HIP backend.
  • Automatic GPU schedules: hardware axis types render to grid/block/thread loops with launch dimensions, barriers, and shared-memory tiles, validated against per-backend hardware limits.
  • Kernel fission with aligned cross-nest parallelism; CPU kernel-level parallelism through a thread pool; backends renamed to cc / multidev_cc.
  • Register-tiled Tile_mma microkernels, SIMD vector-extension codegen, and CUDA WMMA / Metal simdgroup-matrix / HIP rocWMMA tensor-core paths.
  • Measured schedule search (Autotune.tune) with a digest-guarded cache, per-segment schedules, sketch seeding, and placement A/B tuning.
  • The AMD HIP backend via the independent hipjit bindings, and a cross-framework benchmark suite gated on loss parity.
• 0.7: Frontend finalization, compiler optimizations, and paper-ready formal docs.
  • Removed hosted tensors in favor of explicit context-mediated access.
  • Added axis concatenation/block tensors, RoPE, the decoder-only transformer toy, ternary einsum, sharding primitives, and zero-copy leading-axis slice views.
  • Added loop hoisting, CSE, broader virtual-node inlining, and the universal pool allocator across backends.
  • Added the workshop article, formal core technical report, and shape-constraint-generation notes.
• 0.6.3: Padding inference for convolutions.
  • Padding inference during shape inference.
  • Toy CNN example: circle counting.
• 0.6.2: "you forgot to specify a hidden dimension".
  • Menhir einsum parser.
  • Detection of user errors where there is missing information about a hidden dimension: disables guessing "no axes" or "dimension 1" for shapes of parameters.
• 0.6.1: Syntax extension improvements, transformer building blocks.
  • Heterogeneous precision operations.
  • Counter-based randomness via threefry, second pass (pointwise and weak-but-efficient variants); normal distribution operation.
  • New syntax for inline parameter definitions; record-based syntax instead of string-based.
  • Add transformer and convnet building blocks.
  • Better shape error messages.
• 0.6: more precisions, initialization, counter-based randomness, strided iteration.
  • BF16, FP8.
  • Extended expressivity of projections and the generalized einsum notation to cover strided iteration and convolution.
  • Parameter initialization on devices.
  • Counter-based randomness via threefry, first pass (vectorized and cryptographic strength).
  • Better precision inference, including top-down propagation.
• 0.5.3: Apple Metal backend.
  • Also, CUDA backend works on native Windows.
• 0.5.2: More primitive operations.
  • Supports a lot of primitive operations (including ternary ops), and ternary tensor operations.
  • %cd and %op support both curried and uncurried operator application syntax.
  • More flexible gradient construction via the %cd syntax (better projections inference).
  • Works on Native Windows with the C compiler backend (but CUDA backend blocked by cudajit still).
• 0.5.1: Automatic synchronization and transfers between host and devices.
• 0.5.0: Stream-to-stream synchronization at the buffer level.
  • Support for CUDA events, and Condition-based events for CPU backends.
  • Overhaul of the backend interfaces, both user-facing but especially internal: full code sharing.
  • Automatic stream-to-stream synchronization on a per-tensor-node basis.
• 0.4.1 Half precision, mixed precision, CUDA virtual devices (virtual devices renamed to streams in 0.5.0)
  • Half precision. Maybe improvements for mixed-precision computations.
  • Resolve remaining issues with the new scheduler.
  • Initial version of lib/nn_blocks.ml.
• v0.4 Merge buffers, C-syntax backend builder: a significant refactoring of the API.
• v0.3 Shape inference, jitted routines: a major rewrite of the whole project.
  • v0.3.3: continuous integration and opam release.
  • v0.3.2: new shape inference feature: tracking leftmost axes -- complete inference for splicing, ellipsis-in-the-middle allowed in einsum notation.
  • v0.3.1: sanitizing code inclusion (rootness checks).
  • v0.3.0: declarative shape inference; replaced the session interface with a "jitted code routines" API. Cuda defunct.
• v0.2 Inching toward GPU:
  • v0.2.1 naive-cuda: a Cuda backend where blocks and threads are exposed via dedicated axis types.
  • v0.2.0 stack-as-device: treating the C function stack as the "device memory".
• v0.1 GCCJIT backend:
  • v0.1.2: multicore computations using a thread-local "task id" index.
  • v0.1.1: inlining scalar constants, improved inlining for virtual nodes.
  • v0.1.0: a Gccjit backend, single and double precision floats, code compiled as a monolithic update step function.
• v0.0 Untagged: basic design around shape inference, high-level and low-level code representation. Now-abandoned Meta-OCaml and OCaml backends.

Why not just use OWL?

OCANNL follows different design choices than OWL. For example:

• OCANNL is not functorized, except that it uses first-class modules for backends.
• OCANNL has fewer abstraction layers.
• OCANNL has a more powerful shape inference.
• OCANNL only supports backpropagation, while OWL supports full forward and backward auto-diff.
• Some aspects are more centralized in OCANNL than in OWL and form the "infrastructure":
  • Tensor indexing mechanisms are not extensible, other than changing OCANNL code.
  • Shape inference is fully handled by OCANNL and not extensible, other than changing OCANNL code.
  • Tensor implements "putting pieces together".
  • Train has the optimization "frontend" and utilities.
  • arrayjit, which may one day become a standalone library: generates the code, performs backend-agnostic optimizations (virtual nodes whose computation is inlined), implements the backends.
• Some aspects that are more core to OWL are less encapsulated in OCANNL, so it should be more natural to extend them.
  • Specifically, Operation and Train are just collections of functions.
• OCANNL provides lower-level compilation backends than OWL, it is more self-contained in this sense.

Installation

Although the project is called ocannl, the main package is called neural_nets_lib, to avoid the (opam linter's) complaint that the name can be confused with other packages. This also clarifies that ocannl is composed of arrayjit and neural_nets_lib.

The dependency on cudajit is optional so you have to install it first to enable the CUDA backend. The dependency on metal is MacOS-specific but automatic.

Code Organization

The codebase is organized to separate user-facing recipes from framework internals:

• lib/: User-facing recipes and utilities

  • train.ml - Training utilities and optimizers
  • nn_blocks.ml - Neural network building blocks (transformers, attention, convolution, etc.)
  • ocannl.ml - Re-exports for backward compatibility

• tensor/: Framework internals (separate library ocannl_tensor)

  • tensor.ml/mli - Core tensor type and operations
  • shape.ml/mli - Shape inference system
  • operation.ml - Tensor operations and DSL modules
  • ppx_*.ml - Syntax extensions implementation

• arrayjit/: Low-level optimizing compiler with multiple backends

Development

NOTE TO POTENTIAL CONTRIBUTORS: while I am might be slowly starting to work with PRs in separate branches rather than just a stream of commits on the main branch, design migrations will be broken into small PRs to avoid main (master) branch staleness; and many changes will still be commits on the main branch. We allow for failing tests on the main branch, although going forward this would hopefully be happening less. Tagged i.e. released versions of the code are guaranteed to work as well as the given stage of the project permitted, the policy is that all tests must pass for releases with the backend cc and must have the behavior expected of a backend with all other backends. We try to minimize discrepancy across backends but prefer more stringent tests even if some backends only pass them "in spirit" rather than with exact expectations of the cc backend.

OCANNL uses ppx_minidebug for debugging. Currently, we migrated to a per-file opt-in scheme for enabling ppx_minidebug at compile time (via environment variables, see the top of .ml files in question), and then a unified log level configuration (ocannl_log_level) for tuning logging at runtime. Due to the compile-time nature of the per-file settings, run dune clean after setting/exporting one of these environment variables.