Specification

Phase 1, Section 2: Compiler Implementation

Phase 1: Implementation — the shape-typed core compiled end-to-end. Section 2 of 7: Compiler Implementation.

1. What this section commits to

The OCaml compiler — the largest single component in Phase 1. Specifically:

The OCaml project layout: libraries, modules, executable.
The pipeline from source to compiled artifact: parse → elaborate → type-check → lower → emit.
The data structures at each stage: AST, typed IR, lowered representation.
The bidirectional type-checker, including SMT integration for refinement obligations.
The lowering strategy: CPU path through MLIR linalg to LLVM; GPU path through runtime FFI to library kernels.
Error handling: source locations, structured errors, SMT counterexamples surfaced.
Acceptance criteria specific to the compiler.

This section does not address the runtime (Section 3), operation extensions (Section 4), or testing infrastructure (Section 6). Those are distinct concerns and deserve their own treatment.

2. Architecture overview

The compiler is a straightforward batch compiler with five stages:

source.kina          (Kina native syntax; the "input" form for Phase 1)
    ↓ Parse
ast.ml              (untyped AST with source locations)
    ↓ Elaborate
typed_ir.ml         (typed core IR; matches Phase 0 spec)
    ↓ Type-check + SMT discharge
typed_ir.ml         (same, but with all refinement obligations proven)
    ↓ Lower
mlir_text.ml        (string-level MLIR with linalg / arith / func dialects)
    ↓ Emit + invoke external tools
output.so           (compiled shared object via mlir-opt | mlir-translate | llc)

Each stage is a function from the previous stage's data type to the next. The pipeline is implemented as composition; failures at any stage produce structured errors and stop the pipeline. No stage modifies in place; each produces a fresh value. This is standard functional-compiler hygiene and pays off immediately in testing (each stage can be tested independently with synthesized inputs).

The pipeline is synchronous and single-pass. There is no incremental compilation in Phase 1 (deferred to Phase 7+ with the surface language). There are no concurrent passes. There is no caching beyond memoization within the SMT discharge layer. Adding any of those is later work; Phase 1's job is to be correct first.

3. Project structure

The OCaml project lives in compiler/ under the monorepo (per Phase 0 §5.5). Detailed layout:

compiler/
├── dune-project                        — toolchain + dependency manifest
├── compiler.opam                       — for opam install / publishing
├── lib/
│   ├── dune                            — library declaration
│   ├── loc/
│   │   ├── dune
│   │   └── loc.ml                      — source locations (start/end position pairs)
│   ├── error/
│   │   ├── dune
│   │   └── error.ml                    — structured error type, pretty-printing
│   ├── ast/
│   │   ├── dune
│   │   └── ast.ml                      — untyped AST with locations
│   ├── parser/
│   │   ├── dune
│   │   ├── lexer.ml                    — sedlex-based tokenizer
│   │   └── parser.ml                   — Menhir-based parser for Kina surface syntax
│   ├── ir/
│   │   ├── dune
│   │   ├── types.ml                    — type representations
│   │   ├── refinements.ml              — refinement language data types
│   │   └── ir.ml                       — typed IR data types
│   ├── elab/
│   │   ├── dune
│   │   ├── env.ml                      — typing environment
│   │   ├── unify.ml                    — shape unification
│   │   └── elab.ml                     — AST → typed IR + obligation generation
│   ├── smt/
│   │   ├── dune
│   │   ├── presburger.ml               — internal Presburger expression representation
│   │   ├── z3_wrapper.ml               — translation to Z3 + context management
│   │   ├── obligation.ml               — SMT obligation type
│   │   └── cache.ml                    — memoization of discharge results
│   ├── lowering/
│   │   ├── dune
│   │   ├── mlir_emit.ml                — typed IR → MLIR text
│   │   ├── linalg_patterns.ml          — per-operation linalg lowering
│   │   ├── runtime_dispatch.ml         — GPU operation → runtime FFI call
│   │   └── pipeline.ml                 — subprocess invocation of mlir-opt etc.
│   └── driver/
│       ├── dune
│       └── driver.ml                   — top-level: file in, .so out
├── bin/
│   ├── dune
│   └── compiler.ml                     — CLI entry point (uses Cmdliner)
└── test/
    ├── unit/                           — per-module tests (alcotest)
    ├── integration/                    — end-to-end compile-and-run
    └── fixtures/                       — small example programs

Eight library subdirectories: loc, error, ast, parser, ir, elab, smt, lowering, driver. Each compiles independently (Dune's library granularity); each has its own unit tests; each has clear interface boundaries.

The driver library is what bin/compiler.ml uses. CLI argument parsing (input file, output path, flags) is handled via Cmdliner.

4. Parser and AST

Kina surface syntax is parsed using an ocamllex-based tokenizer and a Menhir parser generator.

4.1 Lexer

Tokens:

type token =
  | LPAREN | RPAREN | LBRACK | RBRACK
  | IDENT of string
  | INT of int64
  | FLOAT of float
  | STRING of string
  | EOF

Single-character tokens for parens and brackets; identifiers (with restricted character set), integer and float literals, string literals (for symbolic operation names), EOF. Comments use ; (Lisp-style line comments).

The lexer is ~100 lines of sedlex. Source locations are tracked per-token using Lexing.position.

4.2 AST

The untyped AST mirrors the Kina surface syntax:

type expr = {
  desc: expr_desc;
  loc: Loc.t;
}

and expr_desc =
  | EInt of int64
  | EFloat of float
  | EVar of string
  | ELet of binding list * expr
  | ECall of string * expr list
  | EFn of fn_decl
  | EOp of string * expr list      (* primitive operations like matmul, broadcast *)
  | EAnnot of expr * type_expr     (* expression with explicit type annotation *)

and binding = {
  name: string;
  typ_annot: type_expr option;
  value: expr;
}

and fn_decl = {
  name: string;
  shape_vars: string list;          (* forall (B M K) ... *)
  requires: type_expr list;         (* requires clauses *)
  params: param list;
  ret_type: type_expr;
  body: expr;
}

and param = {
  name: string;
  typ: type_expr;
}

and type_expr =
  | TTensor of dtype * dim list * pred list * placement
  | TScalar of dtype
  | TUnit
  | TFn of type_expr * type_expr

and dim = ...
and pred = ...
and placement = ...

The AST carries locations on every expression. Type expressions appear as syntactic forms; they are not yet resolved or checked. Elaboration does both.

5. Elaboration and the typed IR

Elaboration is where most of the type-checker's work happens. Five jobs:

Resolve names: variables, function references, operation names. Lookup failures produce structured errors with location and "did you mean?" suggestions.
Translate type expressions: convert AST type_expr to typed-IR Type.t, validating well-formedness (per Phase 0 §1).
Synthesize / check: bidirectional type-checking per Phase 0 §3.7.
Generate SMT obligations: record refinement obligations for batch discharge.
Build the typed IR: produce a typed AST that the lowering pipeline consumes.

5.1 The typed IR

Plain ADTs, not GADTs. Decision rationale: GADT-typed IRs are beautiful but make refactoring painful, and Phase 1 will refactor extensively. Plain ADTs with extensive testing is the pragmatic choice.

module Ir = struct
  type t = {
    desc: t_desc;
    typ: Type.t;
    loc: Loc.t;
  }

  and t_desc =
    | Var of var_id
    | Const of const_value
    | Let of var_id * t * t
    | App of fn_id * t list * subst    (* substitution from polymorphic instantiation *)
    | Op of op * t list                (* primitive operation *)

  and var_id = int                     (* de Bruijn or unique-id; pick one *)
  and fn_id = string
  and op = ...
end

Well-formedness is an invariant maintained by the type-checker. A t value with a wrong typ field is a bug, not a type error in the OCaml type system. The discipline pays off in flexibility: passes can construct intermediate IRs that are temporarily ill-formed during transformation, as long as they restore well-formedness by end-of-pass.

5.2 Bidirectional type-checking

Two mutually-recursive functions:

val synthesize : Env.t -> Ast.expr -> Ir.t * Type.t * Subst.t * Obligation.t list

val check : Env.t -> Type.t -> Ast.expr -> Ir.t * Subst.t * Obligation.t list

synthesize infers a type for an expression, returning the typed IR fragment, the type, any shape-variable substitutions discovered during inference, and a list of SMT obligations. check verifies an expression against an expected type, similarly returning IR + substitutions + obligations.

The standard pattern from bidirectional typing (Pierce & Turner 2000, Dunfield & Krishnaswami 2013) applies. Synthesis at variable references, literals, and operator applications. Checking at function-call sites, return statements, and let-bindings with annotations.

The interface returns four things rather than just typed IR because:

The substitution is needed by the caller to instantiate types in subsequent expressions.
Obligations are batched and discharged in one pass at the top level — discharging per-expression is unnecessarily slow.

5.3 Shape unification

Phase 0 §2.2 noted that shape equality discharges via SMT in general but is often syntactic. The unifier exploits this:

val unify_shapes :
  Env.t -> Shape.t -> Shape.t -> [`Same | `Sub of Subst.t | `Smt of Obligation.t]

Three outcomes:

Same: the shapes are syntactically identical. No work.
Sub: the shapes can be made equal by substituting unconstrained shape variables. Returns the substitution.
Smt: equality requires arithmetic reasoning; returns an obligation for batch discharge.

Most shape unifications hit Same or Sub; only the genuinely arithmetic cases (e.g., B * H = X) become SMT obligations. This keeps SMT calls focused on cases that actually need them.

6. SMT integration

The Z3 wrapper is a thin OCaml module over ocaml-z3. Its job is translating internal Presburger expressions to Z3 ASTs, managing contexts, and discharging obligations efficiently.

6.1 Discharge interface

type result =
  | Discharged
  | Refuted of counterexample
  | Timeout

val discharge :
  context:Context.t ->
  assumptions:Predicate.t list ->
  goal:Predicate.t ->
  result

Three outcomes — proven, refuted with counterexample, or timeout — match the type-checker's three-way decision logic. The counterexample carries an assignment Map<ShapeVar, int64> showing the failure case, which is what makes refinement-typed errors usable (per Phase 0 §3.8).

6.2 Implementation details

Context lifecycle: one Z3 context per elaboration session. Reused across obligations to amortize Z3's setup cost.
Caching: results memoized by (canonical_hash(assumptions), canonical_hash(goal)). Same obligation discharged twice (which happens often in real programs) is a hashmap lookup.
Timeout enforcement: 2-second per-obligation budget (Phase 0 §2.14). On timeout, the result is Timeout and the type-checker rejects the operation.
Batch discharge: obligations collected during elaboration are discharged together at the end of type-checking the function or top-level binding. This lets the cache hit on related obligations (e.g., two operations both requiring D mod 8 = 0).

6.3 The Presburger encoding

Internal Presburger expressions are translated to Z3 as Int sorts. Mod operations encode via Z3's mod primitive (Z3 supports congruence reasoning on integers). Division by constants is supported; division by variables triggers a workaround (introduce a fresh variable, assert the multiplication relationship, leave Z3 to figure it out).

Variable products outside Presburger (Phase 0 §1.3) are handled by the controlled escape hatch: the user introduces a fresh variable representing the product and asserts the relationship axiomatically. The compiler does not verify the assertion; the runtime catches mistakes via shape checks at the actual reshape opcode.

7. Lowering: CPU path

The CPU lowering is the more complex of the two paths. It produces MLIR text that then flows through standard MLIR/LLVM tools.

7.1 Strategy

Each Phase 0 operation has a per-operation lowering pattern in lowering/linalg_patterns.ml. The pattern is a function:

val lower : Ir.t -> MlirText.t

producing a chunk of MLIR text in the linalg-on-tensors form. The patterns are written by hand, one per operation, against the MLIR linalg dialect.

7.2 The dialects

func for function definitions and entry points.
tensor for tensor types and basic shape ops.
linalg for the bulk of the work: matmul, generic elementwise, reduce.
arith for scalar arithmetic.
scf for any control flow that emerges from operation lowerings (Phase 1 operations don't have user-visible control flow, but linalg sometimes lowers to scf).

The compiler emits linalg-on-tensors form. Subsequent passes (-linalg-bufferize, -convert-linalg-to-loops, -convert-scf-to-cf, -convert-arith-to-llvm, etc.) are run by mlir-opt per the standard pass pipeline. The compiler does not need to know about post-linalg MLIR.

7.3 The pipeline invocation

lowering/pipeline.ml invokes external tools as subprocesses:

compiler emits → input.mlir
mlir-opt --linalg-bufferize --convert-linalg-to-loops ... < input.mlir > buffered.mlir
mlir-translate --mlir-to-llvmir < buffered.mlir > output.ll
llc -filetype=obj output.ll -o output.o
clang -shared output.o -o output.so

Each step is a subprocess. Errors at any step are caught (stderr captured, return code checked) and surfaced as structured errors with the offending IR included for debugging.

The pipeline is configurable: a debug build keeps intermediate .mlir and .ll files for inspection; a release build deletes them.

7.4 Symbolic shapes in MLIR

Phase 0 admits symbolic shapes ([B, S, D] with bound shape variables). MLIR's linalg supports symbolic shapes via ? markers: a tensor declared tensor<?x?x768xf32> is symbolic in the first two dimensions, fixed in the third.

The compiler translates each Phase 0 shape to MLIR's symbolic form. Constants stay constant; symbolic variables become ?. Where a concrete value is known at compile time (e.g., a power-of-2 constant tile size), it stays explicit.

Refinement obligations that reduced to compile-time-known facts during type-checking become MLIR attributes for downstream optimization. Refinements that remained symbolic after type-checking become runtime checks at function entry — small assertion blocks that fail loudly if violated.

8. Lowering: GPU path

The GPU path in Phase 1 is dramatically simpler than the CPU path because we don't write custom kernels. We dispatch to library kernels (cuBLAS, cuDNN) via runtime FFI calls.

8.1 Strategy

For each Phase 0 operation when both inputs are placed on cuda:global(d):

Don't emit linalg operations.
Emit a function call to a runtime entry point: e.g., runtime_gpu_matmul(a_ptr, b_ptr, c_ptr, m, n, k).
The runtime (Section 3) implements runtime_gpu_matmul as a cuBLAS GEMM dispatch.

The compiler's job for GPU operations is much smaller than for CPU operations: emit the right runtime call with the right arguments. No kernel generation, no scheduling, no register allocation.

8.2 The placement-driven dispatch

The compiler walks the typed IR. For each operation:

if any input has placement cuda:global(d):
    require all inputs have the same device
    emit runtime FFI call
else:
    emit linalg lowering for CPU

Operations between mismatched placements (host → cuda:global) either type-check via copy_to_device / copy_to_host already in the IR (the user inserted them, or elaboration inserted them) or fail earlier with a placement-mismatch error.

8.3 The runtime FFI surface

A small set of extern "C" functions in the runtime:

void runtime_gpu_matmul(
    void* out, const void* a, const void* b,
    int64_t batch, int64_t m, int64_t k, int64_t n,
    int dtype_code);

void runtime_gpu_elementwise_unary(
    void* out, const void* in, int64_t n_elements,
    int dtype_code, int op_code);

void runtime_gpu_elementwise_binary(
    void* out, const void* a, const void* b, int64_t n_elements,
    int dtype_code, int op_code);

void runtime_gpu_reduce(...);
void runtime_gpu_softmax(...);
void runtime_gpu_gather(...);

void runtime_gpu_copy_to_device(...);
void runtime_gpu_copy_to_host(...);

About 15–20 entry points, each parameterized by dtype and operation code. The compiler emits calls to these; the runtime implements them.

8.4 Why this is enough for Phase 1

Phase 1's success criterion is "transformer block forward pass on GPU within 3-5x of PyTorch performance." That's achievable with library-kernel dispatch. cuBLAS for matmul, cuDNN for elementwise / reductions where they have optimized implementations, naive in-house kernels for the rest. Within 3-5x of PyTorch is the realistic benchmark for this strategy.

The full kernel-authoring story (custom tile-IR kernels, capability typing, FlashAttention-class fused kernels) is Phase 5. Phase 1 doesn't try; that's the deliberate scope limit.

9. Compiler driver

The driver (lib/driver/driver.ml) ties the pipeline together. Pseudo-OCaml:

let compile (input : string) (output : string) : (unit, Error.t list) result =
  let ( let* ) = Result.bind in
  let* tokens = Lexer.lex input in
  let* ast = Parser.parse tokens in
  let* ir, obligations = Elab.elaborate ast in
  let* () = Smt.discharge_all obligations in
  let mlir_text = Lowering.MlirEmit.emit ir in
  let* () = Lowering.Pipeline.run mlir_text ~output in
  Ok ()

Five stages, monadic Result plumbing, structured errors throughout. The driver is the entry point invoked by bin/compiler.ml.

The CLI accepts:

compiler [--debug] [--target=cpu|gpu|auto] -o output.so input.kina

--debug keeps intermediate files. --target overrides automatic placement inference (defaulting to auto based on input tensor placements). Default output path is the input with .so extension.

10. Error handling and source locations

The single most important developer-experience concern. Phase 0 §3.8 noted that refinement-type errors are usable only when the counterexample is surfaced. This section operationalizes that.

10.1 The error type

module Error = struct
  type t = {
    severity: severity;                  (* Error | Warning *)
    location: Loc.t;
    message: string;
    suggestion: string option;
    counterexample: counterexample option;
    related_locations: (Loc.t * string) list;
  }
end

Every error has a primary location, a human-readable message, optionally a suggestion ("did you mean?"), optionally an SMT counterexample, and optionally related locations (e.g., the function definition that was being called).

10.2 Error categories

Category	Examples
Parse	Unmatched paren; unexpected token; invalid identifier
Name resolution	Unknown variable; unknown function; unknown operation
Type structure	Rank mismatch; dtype mismatch; arity mismatch
Refinement	SMT obligation refuted (with counterexample); SMT timeout
Placement	Operation between mismatched device placements
Internal	Bug; please report

Each category has its own error-construction helper to ensure consistent formatting.

10.3 Output format

Errors print to stderr in a Rust-compiler-inspired format:

error: refinement obligation refuted
  --> example.kina:5:7
   |
 5 |   (binop add x y)
   |       ^^^ shape mismatch at position 2
   |
   = note: left dimension is `D`, right dimension is `D + 1`
   = counterexample: D = 1 → left = 1, right = 2
   = help: insert an explicit broadcast or reshape to align shapes

ASCII-only by default; ANSI colors when stdout is a TTY (toggled by --color=auto|always|never).

This formatting investment is real engineering work — pretty-printing source spans with carets and context is fiddly — but the payoff is dramatic. Compiler ergonomics live or die on error message quality.

11. Implementation milestones (compiler-specific)

Section 1's three milestones, refined for the compiler:

11.1 Milestone 1 (months 1–2): CPU compilation

Parser + AST: complete, with comprehensive parse-error testing.
Elaboration: complete, handling all Phase 0 operations.
Type-checker: complete, with bidirectional discipline and SMT integration.
CPU lowering through MLIR: complete, working for all Phase 0 operations.
Driver: working end-to-end on at least one test program (e.g., MLP forward pass).
All on CPU; no GPU work yet.

11.2 Milestone 2 (months 3–4): GPU dispatch

Placement axis fully integrated through elaboration.
GPU lowering: emits runtime FFI calls for cuda:global operations.
copy_to_device / copy_to_host lowered correctly.
DLPack interop demo (in conjunction with Section 3 runtime work).
The same MLP from Milestone 1 now runs on GPU with cuBLAS for matmul.

11.3 Milestone 3 (months 5–6): Operation extensions + polish

Variadic-batch matmul (extending the matmul typing rule and lowering).
Gather (new operation: typing rule + CPU and GPU lowerings).
Softmax-as-primitive (lowering decomposes into the six-op composite internally).
repeat_interleave (typing rule + lowerings).
Error message quality polished: every error has a counterexample where applicable, every error has a help line.
Performance: within 3–5x of PyTorch on the transformer block benchmark.

12. Acceptance criteria specific to the compiler

Phase 1 cannot ship until all of the following pass:

All Phase 0 operations compile correctly for both CPU and GPU placements.
The four operation extensions (variadic matmul, gather, softmax, repeat_interleave) work end-to-end.
Type errors include source locations and (where relevant) SMT counterexamples.
The compiler produces valid LLVM IR that links correctly into a shared object loadable by the runtime.
The compiler is under 15K lines of OCaml (rough budget; if it grows beyond, refactor before continuing).
Compile time for the transformer block test program is under 2 seconds (sanity check that SMT and MLIR aren't blowing up).
All compiler tests pass in CI; no --ignore-test flags allowed in main.

13. Open issues introduced by this section

GADT vs. plain ADT for the typed IR. Committed to plain ADT. Worth re-examining if the bug rate from manually-maintained well-formedness becomes burdensome. GADT migration is mechanical but tedious; doable if needed.
Source maps through MLIR. Compiler errors at the elaboration stage have good source locations. Errors that surface from mlir-opt (wrong IR shape, type mismatch in linalg) lose source connection — they reference the MLIR text, not the original .sexp. Solution: emit MLIR loc(...) attributes preserving original locations, parse them back from MLIR diagnostics. Real engineering work, deferred to Milestone 3 polish.
Symbolic shape lowering granularity. When a refinement is satisfied at type-check time, it can become an MLIR attribute (which optimizations can use). When it's symbolic, it becomes a runtime check (which has overhead). The threshold for converting symbolic → attribute is implementation-dependent and worth tuning. Heuristic; not load-bearing for correctness.
Library kernel coverage. cuBLAS covers matmul; cuDNN covers convolution and some normalization but not everything we want. For operations without a library implementation, in-house naive CUDA kernels are the fallback. The list of "operations needing in-house kernels" is the boundary of Phase 1 vs. Phase 5 work; needs explicit enumeration.
Cmdliner argument design. The CLI needs to be future-proof for Phase 2+ flags (e.g., --enable-effects, --shard-config). Worth designing the flag space deliberately rather than accumulating ad-hoc.
OCaml 5 effect handlers in the compiler implementation itself. The compiler doesn't need OCaml 5 effects in Phase 1 (Phase 2 will use them for the AD machinery in the language being compiled). But OCaml 5 is committed (Phase 0 §5.2) and code organization that accidentally precludes effect-handler use later would be a constraint. Worth flagging without committing to specific patterns.

14. Section 3 preview

Section 3 covers the runtime in Rust:

The MemoryResource trait and four implementations.
The Tensor struct with DLPack-compatible #[repr(C)] prefix.
The kernel-dispatch FFI surface (the entry points the compiler emits calls to).
CUDA toolchain integration: cuBLAS, cuDNN linkage; stream management; device allocation.
The DLPack interop demo (export to PyTorch, run an op there, import back).
Acceptance criteria for the runtime.

Sections 2 and 3 together produce the Phase 1 working compiler + runtime. Section 4 adds the operation extensions. The remaining sections (5, 6, 7) handle verification, testing, and exit criteria.

End of Phase 1, Section 2.