← projects

Bagl compiler

I wrote an OCaml compiler for Bagl, a small typed functional language that checks tensor shapes at compile time and derives the gradient of a function inside the compiler.

OCaml, compilers, type inference

source codetry it in the browserlanguage reference

What it is

Bagl is a small functional language with tensors as a built-in type, where a tensor is an array of floats with one or two dimensions. I wrote its compiler in OCaml, together with a bytecode virtual machine, a REPL (an interactive prompt), a language server for editors, and a browser playground.

  • Tensor shapes are part of the type, so multiplying a 2x3 matrix by a 2x2 matrix is a compile error.
  • grad turns a function into its derivative. For a tensor loss, the compiler writes the gradient code itself, so you can train a model in plain Bagl.

The whole compiler also runs in the browser through js_of_ocaml, as about 160 KB of JavaScript (50 KB compressed). You can try the playground or read the language reference, where you can edit and run the examples.

How it works

A hand-written lexer and a recursive descent parser build a syntax tree. Every node keeps its place in the source, so errors point at the exact line.

Next, the compiler rewrites grad calls into ordinary Bagl code. Scalar functions get symbolic derivatives. Functions of a tensor get reverse mode code, which computes the gradient for all inputs in one backward pass.

Type inference uses Hindley-Milner, a method that works out every type without annotations. Type variables live in a union-find structure, so unifying two types links them in place. Tensor dimensions have their own union-find variables, and that is where shape errors are caught.

The program is then lowered to an intermediate representation made of basic blocks, which are runs of code with no jumps inside. Four optimization passes run on it. Code generation then emits stack bytecode for the virtual machine, and a .baglc file can store it.

Decisions

  • I chose OCaml because algebraic data types and pattern matching keep the definitions of the tree, the intermediate representation and the bytecode short and safe. In Rust, the shared union-find graph would conflict with the borrow checker.
  • grad runs before type inference, so later stages need no special cases, and derivatives are type checked like other code. The cost is a required annotation on tensor parameters. Anything the rewrite can't handle is a compile error, never a wrong gradient.
  • Bagl has no type classes, so + picks int or float from its operands. The checker stays single pass, but fn x -> x + x is typed int -> int.
  • The virtual machine is a stack machine, because code generation for it is close to mechanical and the bytecode is easy to save. A register machine would need a register allocator.

What was hard

The virtual machine has separate int and float instructions, and lowering picks one by running inference again on each subexpression. When a binding was missing from that environment, lowering fell back to int, and correct float programs crashed the virtual machine. I fixed three cases of the bug, once for lambda parameters and twice for letrec. An early grad test found the first one, because the derivative of x * x is x + x, a shape no hand-written test had used.

Other fixes followed.

  • (x + 1.0) + x type checked, but (x + x) + 1.0 failed. Inference now waits to default to int until a binding is generalized.
  • A top level if crashed code generation, because local variables were allocated lazily. The whole test suite had passed, since no test had that shape.
  • 0 - 5 came back from a .baglc file as 4294967291. The reader now extends the sign, and round trip tests cover it.

Results

  • 104 test cases in test/test_bagl.ml cover the lexer, the parser, types, lowering, tensors, gradients, and programs that must be rejected.
  • Tests compare gradients with finite differences, i.e., slopes measured from tiny changes to the input, and they agree within 1e-4.
  • examples/train_xor.bagl runs 200 steps of gradient descent, and a test checks that the weights reach the exact solution [1, 1, -2] within 1e-3.
  • CI runs the tests, the examples, and all 21 examples from the reference page through the browser bundle.

What I would do next

  • Lowering should read types from inference instead of inferring them again. My notes call a typed syntax tree the next refactor.
  • The intermediate representation isn't in SSA form, where each variable is assigned once, so common subexpression elimination stays inside one block. SSA would let it work across blocks.
  • A Num constraint would make the numeric operators fully polymorphic.

← projets

Compilateur Bagl

J'ai écrit un compilateur OCaml pour Bagl, un petit langage fonctionnel typé qui vérifie la forme des tenseurs à la compilation et calcule le gradient d'une fonction dans le compilateur.

OCaml, compilation, inférence de types

code sourcel'essayer dans le navigateurréférence du langage

Ce que c'est

Bagl est un petit langage fonctionnel où les tenseurs sont un type intégré, un tenseur étant un tableau de nombres réels à une ou deux dimensions. J'ai écrit son compilateur en OCaml, avec une machine virtuelle à bytecode, un REPL (une invite interactive), un serveur de langage pour les éditeurs et un terrain de jeu dans le navigateur.

  • La forme des tenseurs fait partie du type, donc multiplier une matrice 2x3 par une matrice 2x2 est une erreur de compilation.
  • grad transforme une fonction en sa dérivée. Pour une perte sur des tenseurs, le compilateur écrit lui-même le code du gradient, donc on peut entraîner un modèle en Bagl pur.

Le compilateur entier tourne aussi dans le navigateur grâce à js_of_ocaml, en environ 160 Ko de JavaScript (50 Ko compressés). On peut essayer le terrain de jeu ou lire la référence du langage, où les exemples se modifient et s'exécutent.

Comment ça marche

Un analyseur lexical écrit à la main et un analyseur syntaxique récursif descendant construisent un arbre de syntaxe. Chaque nœud garde sa place dans le source, donc les erreurs pointent la ligne exacte.

Ensuite, le compilateur réécrit les appels à grad en code Bagl ordinaire. Les fonctions scalaires ont des dérivées symboliques. Les fonctions d'un tenseur ont du code en mode inverse, qui calcule le gradient pour toutes les entrées en une seule passe arrière.

L'inférence de types utilise Hindley-Milner, une méthode qui trouve tous les types sans annotations. Les variables de type vivent dans une structure union-find, donc unifier deux types les relie sur place. Les dimensions des tenseurs ont leurs propres variables union-find, et c'est là que les erreurs de forme sont détectées.

Le programme est ensuite traduit vers une représentation intermédiaire faite de blocs de base, des suites de code sans saut à l'intérieur. Quatre passes d'optimisation s'y appliquent. La génération de code produit ensuite du bytecode à pile pour la machine virtuelle, et un fichier .baglc peut le stocker.

Décisions

  • J'ai choisi OCaml parce que les types algébriques et le filtrage par motif gardent courtes et sûres les définitions de l'arbre, de la représentation intermédiaire et du bytecode. En Rust, le graphe union-find partagé entrerait en conflit avec le vérificateur d'emprunts.
  • grad s'exécute avant l'inférence de types, donc les étapes suivantes n'ont aucun cas particulier, et les dérivées sont vérifiées comme le reste du code. Le prix est une annotation obligatoire sur les paramètres de type tenseur. Tout ce que la réécriture ne sait pas traiter est une erreur de compilation, jamais un gradient faux.
  • Bagl n'a pas de classes de types, donc + choisit entier ou réel selon ses opérandes. Le vérificateur reste en une seule passe, mais fn x -> x + x a le type int -> int.
  • La machine virtuelle est une machine à pile, parce que la génération de code y est presque mécanique et que le bytecode est facile à sauvegarder. Une machine à registres aurait besoin d'un allocateur de registres.

Ce qui était difficile

La machine virtuelle a des instructions séparées pour les entiers et les réels, et la traduction en choisit une en relançant l'inférence sur chaque sous-expression. Quand une liaison manquait dans cet environnement, la traduction prenait l'entier par défaut, et des programmes réels corrects faisaient planter la machine virtuelle. J'ai corrigé trois cas de ce bogue, une fois pour les paramètres de lambda et deux fois pour letrec. Un des premiers tests de grad a trouvé le premier cas, parce que la dérivée de x * x est x + x, une forme qu'aucun test écrit à la main n'utilisait.

D'autres corrections ont suivi.

  • (x + 1.0) + x passait la vérification de types, mais (x + x) + 1.0 échouait. L'inférence attend maintenant qu'une liaison soit généralisée avant de choisir l'entier par défaut.
  • Un if au niveau supérieur faisait planter la génération de code, parce que les variables locales étaient allouées tardivement. Toute la suite de tests passait, car aucun test n'avait cette forme.
  • 0 - 5 revenait d'un fichier .baglc sous la forme 4294967291. Le lecteur étend maintenant le signe, et des tests d'aller-retour le couvrent.

Résultats

  • 104 cas de test dans test/test_bagl.ml couvrent l'analyseur lexical, l'analyseur syntaxique, les types, la traduction, les tenseurs, les gradients et les programmes qui doivent être rejetés.
  • Les tests comparent les gradients avec des différences finies, c'est-à-dire des pentes mesurées sur de très petits changements de l'entrée, et l'écart reste sous 1e-4.
  • examples/train_xor.bagl fait 200 pas de descente de gradient, et un test vérifie que les poids atteignent la solution exacte [1, 1, -2] à 1e-3 près.
  • L'intégration continue exécute les tests, les exemples et les 21 exemples de la page de référence avec la version pour navigateur.

La suite

  • La traduction devrait lire les types de l'inférence au lieu de les inférer à nouveau. Mes notes désignent un arbre de syntaxe typé comme la prochaine refonte.
  • La représentation intermédiaire n'est pas en forme SSA, où chaque variable est assignée une seule fois, donc l'élimination des sous-expressions communes reste dans un seul bloc. La forme SSA lui permettrait de traverser les blocs.
  • Une contrainte Num rendrait les opérateurs numériques entièrement polymorphes.