Language specification

The specification gives the full rules of Bagl, chapter by chapter. The language reference covers the same ground with runnable examples, and it is the better place to start.

1. Introduction

1.1 Overview

Bagl (Basic Array/Graph Language) is a small functional programming language with static types and built-in tensors. A tensor is an array with one or more dimensions, e.g., a vector or a matrix. The type of a tensor includes its shape, so the compiler checks that tensor shapes fit together before the program runs. Bagl is meant for numerical code that works on arrays.

1.2 Key features

• Bagl uses Hindley-Milner type inference with let-polymorphism. Hindley-Milner is the standard method for working out types without annotations, and let-polymorphism means a function bound with let can be used at more than one type.
• Tensors are built into the language, and the compiler checks their shapes at compile time.
• Functions are values, so you can pass them as arguments and return them from other functions. A function can also capture variables from the place where it was defined, and the function together with the captured variables is called a closure.
• Shapes can contain dimension variables, e.g., 'n, so one function can work on tensors of many sizes.
• grad turns a function into a new function that computes its derivative, which is called automatic differentiation.
• The compiler reports a shape mismatch at compile time, before the program runs.

1.3 Design goals

1. Safety means that the compiler catches dimension mismatches and type errors at compile time.
2. Expressiveness means that the language supports the common patterns of functional programming.
3. Simplicity means that the core of the language is small, but complete.
4. Predictability means that every construct has clear and deterministic behavior.

1.4 Implementation

The Bagl compiler is written in OCaml, and its source in src/ is about 5,200 lines. It runs a program through the following stages:

Source Code → Lexer → Parser → AST → grad Expansion → Type Inference → IR → Optimization → Bytecode → VM

The implementation has the following parts:

• A lexer written by hand, which records the source location of every token.
• A recursive descent parser, which has one function for each grammar rule.
• A grad expansion pass, which rewrites each grad (fn x -> ...) into an ordinary function that computes the derivative. The pass runs before type inference, so the derivative is type checked like any other code.
• Hindley-Milner type inference, extended with dimension variables.
• An intermediate representation (IR) built from basic blocks. A basic block is a run of instructions with no jumps into or out of its middle.
• Four optimization passes, which are constant folding, dead code elimination (DCE), copy propagation, and common subexpression elimination (CSE).
• A virtual machine (VM) that keeps its working values on a stack and supports closures.

1.5 Example program

// Matrix multiplication with compile-time shape checking
let a: tensor<float>[2,3] = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] in
let b: tensor<float>[3,2] = [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]] in
let c = dot(a, b) in  // Result has shape [2,2]
transpose(c)          // Result has shape [2,2]

// Higher-order functions
let compose = fn f -> fn g -> fn x -> f (g x) in
let double = fn x -> x * 2 in
let inc = fn x -> x + 1 in
compose double inc 5  // Result: 12

// Recursive functions
letrec fact = fn n ->
  if n == 0 then 1
  else n * fact (n - 1)
in fact 5  // Result: 120

1.6 Document structure

The specification has nine chapters, listed in the table below.

1.7 Notation conventions

• Text in code font is literal syntax or an identifier.
• Text in italics is a metavariable, which is a placeholder for any piece of syntax of the named kind.
• [...] marks an optional element.
• {...} marks an element that can repeat zero or more times.
• | separates alternatives.

2. Lexical structure

Chapter 2 defines the tokens of Bagl. A token is the smallest unit the lexer reads from the source text, e.g., a keyword, a number, or an operator.

2.1 Source encoding

Bagl source files are UTF-8 text, and their file extension is .bagl.

2.2 Whitespace and comments

Whitespace

The whitespace characters are space, tab, newline, and carriage return. Whitespace separates tokens, and the lexer otherwise ignores it.

Comments

Bagl has two kinds of comments.

A line comment starts with // and runs to the end of the line.

let x = 5 in x  // This is a comment

A block comment starts with /* and ends with */.

/* This is a
   multi-line comment */
let x = 5 in x

Block comments do not nest, so the first */ ends the comment.

2.3 Keywords

The identifiers in the table below are reserved keywords.

Type keywords

Tensor operation keywords

The lexer also reserves exp, log, sqrt, relu, and step as keywords. Each of them takes a float or a tensor, and on a tensor it works on each element.

2.4 Literals

Integer literals

An integer literal is a sequence of decimal digits.

INT_LIT ::= [0-9]+

A literal never includes a minus sign. A leading minus sign is a separate MINUS token, which the parser reads as negation. For example, -17 is the negation of 17.

Examples: 0, 42, -17, 1000

Float literals

A float literal has a decimal point, an exponent, or both.

FLOAT_LIT ::= [0-9]+ ('.' [0-9]+)? ([eE] [+-]? [0-9]+)?

A number with neither a decimal point nor an exponent is an integer literal instead. As with integers, a leading minus sign is a separate token.

Examples: 3.14, -0.5, 1.0e10, 2.5e-3

Boolean literals

BOOL_LIT ::= 'true' | 'false'

String literals

A string literal is text between double quotes. It can contain escape sequences, which are a backslash followed by one character.

STRING_LIT ::= '"' (char | escape)* '"'
escape     ::= '\\' ('n' | 't' | 'r' | '\\' | '"')

A string literal can't contain a newline.

Examples: "hello", "line\nbreak", "say \"hi\""

2.5 Identifiers

An identifier starts with a letter or an underscore. The rest of it can be letters, digits, and underscores.

IDENT ::= [a-zA-Z_] [a-zA-Z0-9_]*

A keyword can't be used as an identifier. The name grad is not a keyword, but the compiler gives special meaning to grad applied to a function, as section 5.12 describes.

Examples: x, myVar, matrix_2d, _helper

2.6 Operators

Arithmetic operators

Comparison operators

Logical operators

Other operators

2.7 Delimiters

2.8 Operator precedence

The table lists the operators from the lowest precedence to the highest. An operator with higher precedence binds more tightly, so 1 + 2 * 3 means 1 + (2 * 3).

2.9 Token examples

let sum: int = 1 + 2 in sum

The lexer turns the line above into the following tokens.

LET IDENT("sum") COLON INT_TYPE EQ INT_LIT(1) PLUS INT_LIT(2) IN IDENT("sum") EOF

3. Formal grammar

Chapter 3 gives the complete grammar of Bagl in Extended Backus-Naur Form (EBNF). EBNF is a standard notation for grammars, and each rule in it names a piece of syntax and lists what that piece can contain.

3.1 Notation

::=     Definition
|       Alternative
[ ]     Optional (0 or 1)
{ }     Repetition (0 or more)
( )     Grouping
'x'     Terminal symbol

3.2 Complete grammar

Program structure

program     ::= { decl [';'] }

decl        ::= expr

Expressions

expr        ::= let_expr
              | letrec_expr
              | fn_expr
              | if_expr
              | binary_expr

let_expr    ::= 'let' IDENT [':' type_annot] '=' expr 'in' expr

letrec_expr ::= 'letrec' IDENT [':' type_annot] '=' expr 'in' expr

fn_expr     ::= 'fn' IDENT [':' simple_type] '->' expr

if_expr     ::= 'if' expr 'then' expr 'else' expr

binary_expr ::= unary_expr { binop unary_expr }

unary_expr  ::= '-' unary_expr
              | '!' unary_expr
              | application

application ::= primary { argument }

argument    ::= INT_LIT
              | FLOAT_LIT
              | BOOL_LIT
              | STRING_LIT
              | IDENT
              | '(' expr ')'
              | tensor_lit

primary     ::= INT_LIT
              | FLOAT_LIT
              | BOOL_LIT
              | STRING_LIT
              | IDENT
              | '(' expr ')'
              | tensor_lit
              | tensor_op
              | math_op

grad_expr   ::= 'grad' '(' fn_expr ')'

The annotation on a fn parameter is a simple_type, so an arrow type there needs parentheses, e.g., fn f: (int -> int) -> f 5. Without the parentheses, the parser would read the arrow of the type as the arrow of the function.

An argument is a smaller set than primary. To pass the result of dot, transpose, reshape, or a math function to a function, put it in parentheses, e.g., f (dot(a, b)).

The parser has no separate rule for grad_expr. It reads grad as an ordinary IDENT and grad (fn x -> e) as an application. The grad pass then requires the argument to be a function literal in parentheses, as grad_expr shows, and it reports a compile error otherwise. Section 5.12 describes grad.

Binary operators

binop       ::= '||' | '&&'
              | '==' | '!='
              | '<' | '>' | '<=' | '>='
              | '+' | '-'
              | '*' | '/'

Tensor literals and operations

tensor_lit  ::= '[' [tensor_contents] ']' [':' shape]

tensor_contents ::= tensor_row { ',' tensor_row }
                  | expr { ',' expr }

tensor_row  ::= '[' expr { ',' expr } ']'

tensor_op   ::= 'dot' '(' expr ',' expr ')'
              | 'transpose' '(' expr ')'
              | 'reshape' '(' expr ',' shape ')'

math_op     ::= ('exp' | 'log' | 'sqrt' | 'relu' | 'step') '(' expr ')'

Type annotations

type_annot  ::= arrow_type

arrow_type  ::= simple_type [ '->' arrow_type ]

simple_type ::= 'int'
              | 'float'
              | 'bool'
              | 'string'
              | tensor_type
              | type_var
              | '(' type_annot ')'

tensor_type ::= 'tensor' '<' type_annot '>' shape

type_var    ::= "'" IDENT

shape       ::= '[' [ dim { ',' dim } ] ']'

dim         ::= INT_LIT
              | "'" IDENT

Lexical elements

INT_LIT     ::= digit { digit }

FLOAT_LIT   ::= digit { digit } '.' digit { digit } [ exponent ]
              | digit { digit } exponent

exponent    ::= ('e' | 'E') ['+' | '-'] digit { digit }

BOOL_LIT    ::= 'true' | 'false'

STRING_LIT  ::= '"' { string_char } '"'

string_char ::= <any character except '"', '\', or a newline>
              | '\' escape_char

escape_char ::= 'n' | 't' | 'r' | '\\' | '"'

IDENT       ::= letter { letter | digit | '_' }

letter      ::= 'a'..'z' | 'A'..'Z' | '_'

digit       ::= '0'..'9'

3.3 Grammar properties

Precedence climbing

The parser reads binary operators with precedence climbing, which gives each operator the right precedence and associativity. Precedence climbing is a parsing method that uses the precedence of each operator to decide how to group the operands. The levels are listed below.

Precedence 1: ||          (left-associative)
Precedence 2: &&          (left-associative)
Precedence 3: == !=       (left-associative)
Precedence 4: < > <= >=   (left-associative)
Precedence 5: + -         (left-associative)
Precedence 6: * /         (left-associative)

Function application

Function application has the highest precedence, and it groups to the left. The example below shows the grouping.

f x y z   // Parsed as ((f x) y) z

Arrow types

The arrow in a function type groups to the right. The example below shows the grouping.

int -> int -> int   // Parsed as int -> (int -> int)

3.4 Example parses

Let expression

let x: int = 5 in x + 1

The parser builds the following tree.

ELet
├── name: "x"
├── annot: TAInt
├── value: EInt(5)
└── body: EBinop(Add)
    ├── EVar("x")
    └── EInt(1)

Function definition

fn x -> fn y -> x + y

The parser builds the following tree.

EFn
├── param: "x"
└── body: EFn
    ├── param: "y"
    └── body: EBinop(Add)
        ├── EVar("x")
        └── EVar("y")

Tensor operation

dot(a, b)

The parser builds the following tree.

ETensorOp(TensorDot)
├── EVar("a")
└── EVar("b")

3.5 Syntactic sugar

Bagl has very little syntactic sugar, which is syntax that the compiler rewrites into other syntax. Each construct maps directly to one node of the abstract syntax tree (AST), as the table shows.

4. Type system

Bagl has a static type system based on Hindley-Milner, with two extensions. The first extension adds tensor types. The second adds dimension polymorphism, which lets one function accept tensors of many sizes.

4.1 Base types

Primitive types

Type syntax

int                    // Integer type
float                  // Float type
bool                   // Boolean type
string                 // String type

4.2 Function types

A function type is written with an arrow, ->.

T1 -> T2

In the type above, T1 is the parameter type, and T2 is the return type.

Properties

• The arrow groups to the right, so A -> B -> C means A -> (B -> C).
• Functions are curried. A function of several arguments is a function of one argument that returns another function, which takes the next argument.

Examples

int -> int                      // int to int function
int -> int -> int               // Curried two-argument function
(int -> int) -> int             // Function taking a function
tensor<float>[2,3] -> float     // Tensor to scalar

4.3 Tensor types

A tensor is an array of numbers, and the compiler knows its shape at compile time. A tensor type has the following form.

tensor<E>[D1, D2, ..., Dn]

In the form above, E is the element type, and each Di gives the size of one dimension. The element type must be float, because the type checker requires float elements in every tensor value. The VM also stores every tensor element as a float.

Dimension specifications

A dimension can take one of two forms:

1. A concrete size, which is an integer constant, e.g., 2, 3, or 100.
2. A dimension variable, e.g., 'n or 'm, which stands for a size that the compiler works out from how the tensor is used.

Examples

tensor<float>[3]           // 1D tensor with 3 elements
tensor<float>[2, 3]        // 2x3 matrix
tensor<float>[4, 4, 4]     // 3D tensor
tensor<float>['n, 'm]      // Matrix with variable dimensions
tensor<float>['n, 'n]      // Square matrix

Shape rank

The rank of a tensor is the number of its dimensions.

• A tensor of rank 0 is a scalar, and its shape is empty, [].
• A tensor of rank 1 is a vector, with shape [n].
• A tensor of rank 2 is a matrix, with shape [m, n].
• A tensor of rank 3 or more has no special name.

4.4 Type variables

A type variable stands for any type. Type variables make polymorphism possible, which means one function can work on values of different types.

'a                    // Type variable
'a -> 'a              // Identity function type
'a -> 'b -> 'a        // Constant function type

Type variable naming

• A type variable starts with a quote, e.g., 'a, 'b, or 'foo.
• By convention, the name after the quote is a single letter, e.g., 'a, 'b, or 'c.

4.5 Type schemes (polymorphism)

A type scheme is a type with a list of type variables that can each be replaced by any type. The forall in front lists the variables.

forall 'a 'b 'n. T

The compiler creates a type scheme when it checks a let binding. The process is called let-polymorphism, and the example below shows it.

let id = fn x -> x in      // id : forall 'a. 'a -> 'a
let _ = id 5 in            // Instantiated to int -> int
let _ = id true in         // Instantiated to bool -> bool
id

Generalization rules

Generalization is the step that turns a type into a type scheme. A type variable is generalized when both of the following hold:

1. The variable is not free in the surrounding environment, which means no variable already in scope has it in its type.
2. The variable was created at a deeper level than the current binding.

4.6 Type annotations

A type annotation limits the type that the compiler infers. If the inferred type doesn't match the annotation, the compiler reports an error.

Binding annotations

let x: int = 5 in
let f: int -> int = fn x -> x + 1 in
f x  // Result: 6

Parameter annotations

let inc = fn x: int -> x + 1 in
let apply5 = fn f: (int -> int) -> f 5 in
apply5 inc  // Result: 6

Tensor annotations

let v: tensor<float>[3] = [1.0, 2.0, 3.0] in
let m: tensor<float>[2, 3] = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] in
dot(m, v)  // Result: [14, 32], with type tensor<float>[2]

4.7 Type compatibility

Subtyping

Bagl does not have subtyping, so two types must match exactly after unification. Unification is the step that makes two types equal by filling in their type variables.

Numeric coercion

Bagl never converts between int and float on its own. For scalars, both operands of an arithmetic operator must have the same type. The first example adds two ints and then two floats, and the second example mixes an int with a float.

let a = 1 + 2 in      // int + int = int
let b = 1.0 + 2.0 in  // float + float = float
b                     // Result: 3.0

1 + 2.0  // Error: Cannot unify int with float

Tensor element types

Tensor operations keep the element type of their operands.

let a: tensor<float>[2,3] = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] in
let b: tensor<float>[3,2] = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]] in
dot(a, b)      // Result: tensor<float>[2,2]

4.8 Type equality

Two types are equal when they have the same structure, as the table shows.

4.9 Type errors

The sections below show the most common type errors.

Type mismatch

let x: int = true in x  // Error: Cannot unify int with bool

Dimension mismatch

let a: tensor<float>[2,3] = [[1.0, 2.0], [3.0, 4.0]] in
a  // Error: Dimension mismatch: 3 vs 2

Unbound variable

x + 1    // Error: Unbound variable: x

Shape rank mismatch

let v: tensor<float>[3] = [1.0, 2.0, 3.0] in
transpose(v)    // Error: transpose requires 2D tensor

4.10 Type representation (internal)

The compiler stores types with the following OCaml definitions.

type ty =
  | TInt
  | TFloat
  | TBool
  | TString
  | TUnit
  | TTensor of ty * shape
  | TArrow of ty * ty
  | TVar of tvar ref        (* Unification variable *)

type shape_dim =
  | SDimConst of int        (* Concrete dimension *)
  | SDimVar of dim_var ref  (* Dimension variable *)

Each type variable is a mutable reference. When unification solves a variable, the compiler points the reference at the solution. Path compression then shortens chains of such links, so later lookups are fast.

5. Expressions

Every Bagl program is an expression, and running the program evaluates the expression to a value. Chapter 5 describes each kind of expression.

5.1 Literals

Each line in the blocks of section 5.1 is a separate literal, and each literal on its own is a complete program.

Integer literals

42
-17
0

The type is int. The -17 line is the negation of the literal 17, as section 2.4 explains.

Float literals

3.14
-0.5
1.0e10
2.5e-3

The type is float.

Boolean literals

true
false

The type is bool.

String literals

"hello"
"line\nbreak"
""

The type is string.

Tensor literals

Tensor elements must be floats, and the type checker rejects a tensor literal with int elements.

A literal with one level of brackets is a 1D tensor, which is a vector.

[1.0, 2.0, 3.0]  // Type: tensor<float>[3]

A literal with two levels of brackets is a 2D tensor, which is a matrix.

[[1.0, 2.0], [3.0, 4.0]]  // Type: tensor<float>[2, 2]

A literal can also carry an explicit shape annotation.

[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] : [2, 3]

A literal with int elements fails to type check.

[1, 2, 3]  // Error: Cannot unify int with float

5.2 Variables

A variable refers to a value that has been bound to a name. The names below are examples of variables.

x
myVariable
_helper

A variable must be bound before it is used. A let binding, a letrec binding, or a function parameter can bind it.

5.3 Let expressions

A let expression binds a value to a name. It has one of the two forms below.

let name = value in body
let name: type = value in body

Semantics

1. The program evaluates value.
2. It binds the result to name.
3. It evaluates body, with name in scope.
4. The result of body is the result of the whole expression.

Examples

let x = 5 in x + 1  // Result: 6

let x = 5 in
let y = 10 in
x + y  // Result: 15

let double = fn x -> x * 2 in
double 21  // Result: 42

Let-polymorphism

A value bound with let can be used at more than one type, as the example shows.

let id = fn x -> x in
let a = id 5 in     // id used at int -> int
let b = id true in  // id used at bool -> bool
a                   // Result: 5

5.4 Recursive let (letrec)

A letrec expression binds a recursive function. It has one of the two forms below.

letrec name = value in body
letrec name: type = value in body

The value must be a function, written with fn, and the function can call itself by name.

Examples

letrec fact = fn n ->
  if n == 0 then 1
  else n * fact (n - 1)
in fact 5  // Result: 120

letrec fib = fn n ->
  if n <= 1 then n
  else fib (n - 1) + fib (n - 2)
in fib 10  // Result: 55

5.5 Function expressions

A function expression defines a function that has no name. It has one of the two forms below.

fn param -> body
fn param: type -> body

Semantics

A function expression creates a closure. The closure keeps the values of the variables that the function body uses from the surrounding scope.

Examples

Each example binds the function to a name and returns it, so the result is the function and its type.

let inc = fn x -> x + 1 in inc  // Type: int -> int

let add = fn x -> fn y -> x + y in add  // Type: int -> int -> int

let twice = fn f -> fn x -> f (f x) in twice  // Type: ('a -> 'a) -> 'a -> 'a

The compiler may print a different letter for the type variable in the last example.

Currying

A function of several arguments is written as nested functions of one argument each. Applying such a function to fewer arguments than it takes is called partial application, and the result is a function that takes the remaining arguments.

let add = fn x -> fn y -> x + y in
add 3 5  // Result: 8

let add = fn x -> fn y -> x + y in
let add3 = add 3 in  // Partial application
add3 5               // Result: 8

5.6 Function application

To apply a function, write the function followed by its argument. Application groups to the left.

f x
f x y z  // Same as ((f x) y) z

Semantics

1. The program evaluates the function expression.
2. It evaluates the argument expression.
3. It binds the argument to the parameter.
4. It evaluates the function body.
5. The result of the body is the result of the application.

Examples

let double = fn x -> x * 2 in
double 21  // Result: 42

let double = fn x -> x * 2 in
let compose = fn f -> fn g -> fn x -> f (g x) in
let inc = fn x -> x + 1 in
compose double inc 5  // Result: 12

5.7 Conditional expressions

A conditional expression has the form below.

if condition then expr1 else expr2

Semantics

1. The program evaluates condition, which must have type bool.
2. If the condition is true, the program evaluates expr1 and returns its value.
3. If the condition is false, the program evaluates expr2 and returns its value.

Both branches must have the same type.

Examples

if true then 1 else 2  // Result: 1

let max = fn x -> fn y ->
  if x > y then x else y
in max 5 3  // Result: 5

let abs = fn x ->
  if x < 0 then -x else x
in abs (-5)  // Result: 5

5.8 Binary operators

A binary operator goes between its two operands.

expr1 op expr2

Arithmetic operators

Division by zero raises a runtime error. The arithmetic operators also work element by element on two tensors of the same shape, and on a tensor and a float scalar, as chapter 7 describes.

1 / 0  // Runtime error: Division by zero

Comparison operators

The table shows the int form of <, >, <=, and >=, and the same operators also accept two floats. The operators == and != don't accept tensors or functions.

[1.0] == [1.0]  // Error: Equality is not defined for tensors

Logical operators

In the current implementation, the logical operators always evaluate both operands. They don't stop early when the left operand already decides the result.

Examples

1 + 2 * 3  // Result: 7, because * binds more tightly than +

(1 + 2) * 3  // Result: 9

5 == 5  // Result: true

3 < 5 && 5 < 10  // Result: true

5.9 Unary operators

A unary operator goes before its operand.

-expr
!expr

Examples

-5  // Result: -5

-3.14  // Result: -3.14

!true  // Result: false

!(5 > 3)  // Result: false

5.10 Tensor operations

Dot product

The dot product has the form below.

dot(a, b)

dot multiplies matrices and vectors. The compiler works out the result shape from the operand shapes, as the table shows.

Transpose

The transpose has the form below.

transpose(a)

transpose swaps the two dimensions of a 2D tensor, so shape [m, n] becomes [n, m].

Reshape

A reshape has the form below, where the second argument is a shape.

reshape(a, [d1, d2])

reshape gives a tensor new dimensions. The new shape must have the same total number of elements as the old one.

Examples

let a = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] in       // [2,3]
let b = [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]] in  // [3,2]
dot(a, b)  // Result: [[58.0, 64.0], [139.0, 154.0]], a [2,2] matrix

let m = [[1.0, 2.0], [3.0, 4.0]] in
transpose(m)  // Result: [[1.0, 3.0], [2.0, 4.0]]

let v = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] in
reshape(v, [2, 3])  // Result: [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], a 2x3 matrix

5.11 Parentheses

Parentheses override the usual precedence, as the examples show.

(1 + 2) * 3  // Result: 9, not 7

(fn x -> x + 1) 5  // Applies a function that has no name. Result: 6

5.12 Automatic differentiation with grad

grad does automatic differentiation, which means that it computes the derivative of a function. The expression grad (fn x -> e) is a new function, and that function returns the derivative of e with respect to x. The compiler implements grad as a source-to-source transform. A pass rewrites each grad expression into ordinary Bagl code before type inference, and the rest of the compiler then checks and runs the derivative like any other function.

grad must be applied directly to a function literal in parentheses. The name grad is not a value, so grad on its own is an unbound variable. If you write grad f for a function f bound elsewhere, the compiler reports "grad must be applied directly to a function literal".

Modes

The pass chooses a mode from the annotation on the parameter.

• An unannotated parameter, or a float parameter whose body uses no tensors, gets scalar mode. Scalar mode is symbolic forward mode. The pass builds an expression for the derivative with the sum, product, quotient, and chain rules, and the result has type float -> float.
• A parameter annotated tensor<float>[...] gets reverse mode. The body must compute a scalar, e.g., a loss. The result is a function that returns the gradient, which is a tensor with the same shape as the parameter.
• A float parameter whose body uses tensors gets pathwise mode. The pass applies the reverse mode rules to the scalar parameter, so it can differentiate a scalar through tensor code.

An unannotated parameter whose body uses tensors is a compile error, because the pass needs the annotation to choose a mode.

Supported expressions

In scalar mode, the body can use literals, the parameter, and other variables, which count as constants. It can also use +, -, *, /, unary minus, if, let, and the math functions exp, log, sqrt, relu, and step. The pass differentiates each branch of an if, and it leaves the condition as it is.

Reverse mode and pathwise mode support the same expressions. They also support dot, transpose, element-wise arithmetic, and arithmetic between a tensor and a scalar.

Unsupported forms

The pass reports a compile error for each form it can't differentiate correctly, so it doesn't produce a wrong derivative for those forms. The unsupported forms in every mode include the following:

• A call to another function inside the body
• A letrec or a nested fn inside the body
• A boolean or comparison expression where the pass expects a number

Scalar mode also rejects any tensor operation. Reverse mode and pathwise mode also reject the following forms:

• reshape
• A tensor literal whose elements mention the parameter
• An if whose condition mentions the parameter
• The matrix side of a matrix-vector or vector-matrix dot, because its gradient needs an outer product
• A scalar built from the parameter and broadcast against a matrix
• A body that doesn't compute a scalar

Examples

In scalar mode, the derivative of x * x * x is 3 * x * x, which is 12.0 at 2.0.

let f = grad (fn x -> x * x * x) in
f 2.0  // Result: 12.0

In reverse mode, the gradient of the squared dot product dot(w, x) * dot(w, x) is 2 * dot(w, x) * x. At w0, dot(w0, x) is 4.5, so the gradient is 9.0 * x.

let x = [0.5, -1.0, 2.0] in
let w0 = [1.0, 2.0, 3.0] in
let g = grad (fn w: tensor<float>[3] -> dot(w, x) * dot(w, x)) in
g w0  // Result: [4.5, -9.0, 18.0]

In pathwise mode, dot(x * s, x) is s times dot(x, x). Its derivative with respect to s is therefore dot(x, x), which is 14.0.

let x = [1.0, 2.0, 3.0] in
let g = grad (fn s: float -> dot(x * s, x)) in
g 2.0  // Result: 14.0

6. Type inference

Bagl uses Hindley-Milner type inference, extended with dimension variables for tensor shapes. Chapter 6 describes the algorithm.

6.1 Overview

Type inference works out the type of every expression, so most programs need no type annotations. Some programs still need one. For example, a function parameter that is used as a tensor gets a 2D tensor type unless an annotation says otherwise, so a parameter that should be a vector needs an annotation such as tensor<float>['n].

The algorithm takes the following steps.

1. It gives each unknown type a fresh type variable.
2. It collects constraints from the structure of each expression. A constraint is a rule that two types must be equal.
3. It solves the constraints with unification.
4. It generalizes types at let bindings.

6.2 Type environments

A type environment maps each variable name in scope to a type scheme.

Γ = { x₁ : σ₁, x₂ : σ₂, ... }

In the definition above, each σ is a type scheme of the form ∀ α₁...αₙ. τ.

Operations

• Lookup, written Γ(x), returns the scheme bound to x.
• Extension, written Γ, x : σ, adds a binding.
• FV(Γ) returns all the free type variables in the environment. A free type variable is one that no ∀ binds.

6.3 Inference rules

Each rule below means that if the facts above the line hold, then the fact below the line holds. The notation Γ ⊢ e : τ means that in the environment Γ, the expression e has type τ.

Literals

─────────────────
Γ ⊢ n : int        (integer literal)

─────────────────
Γ ⊢ f : float      (float literal)

─────────────────
Γ ⊢ b : bool       (boolean literal)

Variables

x : σ ∈ Γ    τ = instantiate(σ)
────────────────────────────────
        Γ ⊢ x : τ

Let binding

Γ ⊢ e₁ : τ₁    Γ, x : generalize(Γ, τ₁) ⊢ e₂ : τ₂
─────────────────────────────────────────────────────
            Γ ⊢ let x = e₁ in e₂ : τ₂

Recursive let

Γ, x : τ ⊢ e₁ : τ₁    unify(τ, τ₁)    Γ, x : generalize(Γ, τ) ⊢ e₂ : τ₂
──────────────────────────────────────────────────────────────────────────
                  Γ ⊢ letrec x = e₁ in e₂ : τ₂

Function

Γ, x : τ₁ ⊢ e : τ₂
─────────────────────────
Γ ⊢ fn x -> e : τ₁ -> τ₂

Application

Γ ⊢ e₁ : τ₁    Γ ⊢ e₂ : τ₂    unify(τ₁, τ₂ -> α)
──────────────────────────────────────────────────
              Γ ⊢ e₁ e₂ : α

Conditional

Γ ⊢ e₁ : τ₁    Γ ⊢ e₂ : τ₂    Γ ⊢ e₃ : τ₃
unify(τ₁, bool)    unify(τ₂, τ₃)
────────────────────────────────────────────
      Γ ⊢ if e₁ then e₂ else e₃ : τ₂

6.4 Unification

Unification finds a substitution that makes two types equal. A substitution is a choice of a type for each type variable.

Algorithm

unify(τ₁, τ₂):
  τ₁' = find(τ₁)  // Follow links
  τ₂' = find(τ₂)  // Follow links

  match (τ₁', τ₂'):
    | (int, int) -> ok
    | (float, float) -> ok
    | (bool, bool) -> ok
    | (α, τ) or (τ, α) where α is unbound ->
        if occurs(α, τ) then error "recursive type"
        link(α, τ)
    | (τ₁ -> τ₂, τ₃ -> τ₄) ->
        unify(τ₁, τ₃)
        unify(τ₂, τ₄)
    | (tensor<e₁>[s₁], tensor<e₂>[s₂]) ->
        unify(e₁, e₂)
        unify_shape(s₁, s₂)
    | _ -> error "cannot unify"

Path compression

The compiler stores type variables in a union-find structure, which tracks which variables have been made equal. It also uses path compression. When the compiler follows a chain of links to find what a variable stands for, it points the variable straight at the end of the chain, so the next lookup is faster.

let rec find_ty = function
  | TVar ({ contents = Link t } as r) ->
      let t' = find_ty t in
      r := Link t';  (* Path compression *)
      t'
  | t -> t

6.5 Dimension unification

The compiler unifies dimensions in the same way as types.

unify_dim(d₁, d₂):
  d₁' = find_dim(d₁)
  d₂' = find_dim(d₂)

  match (d₁', d₂'):
    | (n, m) where n = m -> ok
    | (n, m) where n ≠ m -> error "dimension mismatch"
    | (α, d) or (d, α) where α is unbound ->
        link(α, d)

Shape unification

Two shapes unify when they have the same rank and each pair of dimensions unifies.

unify_shape(s₁, s₂):
  if length(s₁) ≠ length(s₂) then
    error "shape rank mismatch"
  for each (d₁, d₂) in zip(s₁, s₂):
    unify_dim(d₁, d₂)

6.6 Generalization

Generalization turns a type into a type scheme. It does so by adding a ∀ for each type variable that is free in the type and not free in the environment.

generalize(Γ, τ):
  free_vars = FV(τ) - FV(Γ)
  return ∀ free_vars. τ

Bagl has no type classes, so it needs a rule for arithmetic on operands of unknown type. Before it generalizes a binding, the compiler gives the type int to any such operand that is still unknown. For example, let f = fn x -> x + x in ... gives f the type int -> int. A float constraint anywhere in the same binding takes priority, so let g = fn x -> (x + x) + 1.0 in ... gives g the type float -> float.

Level-based generalization

Bagl uses levels to find the variables to generalize without scanning the whole environment. A level is a counter of how deeply the current let binding is nested.

1. Each type variable records the level at which it was created.
2. enter_level() adds one to the current level before the compiler checks the value of a binding.
3. leave_level() subtracts one from the level after the compiler checks the value.
4. A variable whose level is greater than the current level can be generalized.

6.7 Instantiation

Instantiation is the reverse of generalization. It replaces each variable bound by ∀ with a fresh type variable.

instantiate(∀ α₁...αₙ. τ):
  substitution = { αᵢ → fresh_var() for i in 1..n }
  return apply(substitution, τ)

6.8 Tensor operation types

Dot product

infer_dot(shape₁, shape₂):
  match (shape₁, shape₂):
    | ([m, k₁], [k₂, n]) ->
        unify_dim(k₁, k₂)
        return [m, n]
    | ([k₁], [k₂, n]) ->
        unify_dim(k₁, k₂)
        return [n]
    | ([m, k₁], [k₂]) ->
        unify_dim(k₁, k₂)
        return [m]
    | ([k₁], [k₂]) ->
        unify_dim(k₁, k₂)
        return []  // scalar

Transpose

infer_transpose(shape):
  match shape:
    | [m, n] -> [n, m]
    | _ -> error "transpose requires 2D tensor"

6.9 Example inference

Example 1: Identity function

let id = fn x -> x in id 5

The compiler takes the following steps.

1. For fn x -> x, it assigns x : α. The body also has type α, so the function has type α -> α.
2. It generalizes the type to ∀α. α -> α.
3. For id 5, it instantiates the scheme to β -> β and unifies that with int -> γ.
4. Unification gives β = int and γ = int, so the final type is int.

Example 2: Tensor dot

let a: tensor<float>[2,3] = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] in
let b: tensor<float>[3,2] = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]] in
dot(a, b)

The compiler takes the following steps.

1. It assigns a : tensor<float>[2,3].
2. It assigns b : tensor<float>[3,2].
3. For dot(a, b), it checks the shapes [2,3] and [3,2].
4. It unifies the inner dimensions, and 3 = 3 holds.
5. The result shape is [2,2].
6. The final type is tensor<float>[2,2].

Example 3: Polymorphic function

let compose = fn f -> fn g -> fn x -> f (g x) in
compose (fn x -> x + 1) (fn x -> x * 2) 5

The compiler takes the following steps.

1. For fn f -> fn g -> fn x -> f (g x), it finds the following types:
   • f : α -> β, g : γ -> α, and x : γ.
   • The whole function has type (α -> β) -> (γ -> α) -> γ -> β.
2. It generalizes the type to ∀αβγ. (α -> β) -> (γ -> α) -> γ -> β.
3. At the application, it instantiates the scheme and unifies it with the arguments:
   • f = fn x -> x + 1 : int -> int
   • g = fn x -> x * 2 : int -> int
   • x = 5 : int
4. The result type is int.

6.10 Error messages

Each type error includes the source location of the problem. For example, the compiler reports an error for the program below, saved as example.bagl.

let y = 1 in
let x: int = true in
x + y

The first line of the report gives the message, and the second line gives the file, line, and column. The compiler then prints the line with the error between its neighbors, and it marks the part that caused the error with ^.

error: Cannot unify int with bool
  --> example.bagl:2:14
  |
1 | let y = 1 in
2 | let x: int = true in
  |              ^^^^
3 | x + y
  |

7. Tensor semantics

Chapter 7 defines what tensor types and tensor operations mean in Bagl.

7.1 Tensor representation

Logical model

A tensor is an array of elements with one or more dimensions. Its type has the following form.

tensor<E>[d₁, d₂, ..., dₙ]

• The element type E is the type of each element. It must be float.
• The shape [d₁, d₂, ..., dₙ] lists the size of each dimension.
• The rank is the number of dimensions, which is the length of the shape.
• The size is the total number of elements, which is the product of the dimensions.

Physical storage

The VM stores each tensor as one flat array in row-major order, which is the order C uses. In row-major order, the elements of the first row come first, then the elements of the second row, and so on.

[[a, b, c],     stored as    [a, b, c, d, e, f]
 [d, e, f]]                   indices: 0  1  2  3  4  5

Strides

A stride is the number of positions in the flat array that you move when one index goes up by one. The strides turn an index with several parts into a position in the flat array.

For a shape [d₁, d₂, ..., dₙ], the strides are as follows.

stride[i] = d[i+1] × d[i+2] × ... × d[n]
stride[n] = 1

For example, the shape [2, 3] has the following strides.

• stride[0] = 3, because each row holds 3 elements.
• stride[1] = 1, because the elements of a row are next to each other.

The index [i, j] maps to the flat position i × 3 + j × 1.

7.2 Tensor literals

Tensor elements must be floats. Section 7.6 shows what happens with int elements.

1D tensors (vectors)

[1.0, 2.0, 3.0]  // Type: tensor<float>[3]

2D tensors (matrices)

[[1.0, 2.0],  // Type: tensor<float>[2, 2]
 [3.0, 4.0]]

Type inference for literals

The compiler checks a tensor literal in the following steps.

1. It infers the element type from the first element, and it requires that type to be float.
2. It checks that every element has the same type.
3. It checks that every row has the same length.
4. It infers the shape from the brackets. One level of brackets gives a vector, and two levels give a matrix, even when the matrix has only one row.

Explicit shape annotation

[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] : [2, 3]

The annotation must match the shape that the compiler infers.

7.3 Shape checking

Compile-time verification

Bagl checks tensor shapes at compile time. In the first example, the annotation matches the literal.

let a: tensor<float>[2,3] = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] in
a  // OK

In the second example, the annotation says the matrix has 2 columns, but the literal has 3.

let b: tensor<float>[2,2] = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] in
b  // Error: Dimension mismatch: 2 vs 3

Dimension variables

Dimension variables make shape polymorphism possible, which means one function can accept tensors of different sizes.

// square works on a square matrix of any size
let square: tensor<float>['n, 'n] -> tensor<float>['n, 'n] = fn m -> dot(m, m) in
square [[1.0, 2.0], [3.0, 4.0]]  // Result: [[7.0, 10.0], [15.0, 22.0]]

The same function rejects a matrix that isn't square, because 'n must stand for one size.

let square: tensor<float>['n, 'n] -> tensor<float>['n, 'n] = fn m -> dot(m, m) in
square [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]  // Error: Dimension mismatch: 2 vs 3

Shape constraints

Each tensor operation adds a constraint on the shapes of its operands.

• dot(A, B) requires the inner dimensions to match.
• transpose(A) requires a 2D tensor.
• reshape(A, s) requires the same total number of elements. The compiler checks the count only when every dimension is a known number, and the VM checks it again at run time.
• Element-wise arithmetic on two tensors requires the two shapes to be equal.

7.4 Tensor operations

7.4.1 Dot product

The syntax is dot(a, b). The table gives the shape rules.

Each element of the result is a sum of products. For a matrix times a matrix ([m,k] × [k,n] → [m,n]), the formula is as follows.

C[i,j] = Σ(k) A[i,k] × B[k,j]

For two vectors ([k] · [k] → scalar), the formula is as follows.

result = Σ(i) A[i] × B[i]

Example:

let a = [[1.0, 2.0], [3.0, 4.0]] in       // [2,2]
let b = [[5.0, 6.0], [7.0, 8.0]] in       // [2,2]
dot(a, b)  // Result: [[19.0, 22.0], [43.0, 50.0]]
// C[0,0] = 1×5 + 2×7 = 19
// C[0,1] = 1×6 + 2×8 = 22
// C[1,0] = 3×5 + 4×7 = 43
// C[1,1] = 3×6 + 4×8 = 50

7.4.2 Transpose

The syntax is transpose(a), and the shape rule is [m, n] → [n, m]. transpose works only on 2D tensors. Element [i, j] of the result is element [j, i] of the input, as the formula shows.

B[i,j] = A[j,i]

Example:

let a = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] in  // [2,3]
transpose(a)  // Result: [[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]], with shape [3,2]

7.4.3 Reshape

The syntax is reshape(a, [d₁, d₂, ..., dₙ]). The new shape must have the same total number of elements as the old one.

d₁ × d₂ × ... × dₙ = original size

reshape keeps the flat data as it is, and it reads that data with the new shape.

Example:

let a = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] in  // [6]
reshape(a, [2, 3])  // Result: [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], with shape [2,3]

let a = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] in  // [6]
reshape(a, [3, 2])  // Result: [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], with shape [3,2]

7.5 Shape inference examples

Example 1: Matrix chain

let a: tensor<float>[2, 3] = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]] in
let b: tensor<float>[3, 4] = [[1.0, 1.0, 1.0, 1.0], [2.0, 2.0, 2.0, 2.0], [3.0, 3.0, 3.0, 3.0]] in
let c: tensor<float>[4, 2] = [[1.0, 0.0], [0.0, 1.0], [1.0, 0.0], [0.0, 1.0]] in
let ab = dot(a, b) in    // [2,3] × [3,4] = [2,4]
dot(ab, c)               // [2,4] × [4,2] = [2,2]

Example 2: Dimension polymorphism

// Generic matrix multiply function
let mm = fn a -> fn b -> dot(a, b) in
let x: tensor<float>[2,3] = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] in
let y: tensor<float>[3,5] = [[1.0, 0.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0, 1.0]] in
mm x y    // Type: tensor<float>[2,5]

Example 3: Shape error

let a: tensor<float>[2, 3] = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] in
let b: tensor<float>[4, 2] = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]] in
dot(a, b)    // Error: Dimension mismatch: 3 vs 4

7.6 Element type handling

Storage

The VM stores every tensor element as a float. The type checker requires float elements, so a program can't build an int tensor. If hand-written bytecode passes an int to TENSOR_CREATE, the VM converts it to a float.

Type inference

The compiler infers the element type of a literal from the first element, and it requires that type to be float. A literal of floats has a float element type.

[1.0, 2.0, 3.0]  // Type: tensor<float>[3]

A literal of ints is a type error.

[1, 2, 3]  // Error: Cannot unify int with float

A literal that mixes ints and floats also fails, because its first element is an int.

[1, 2.0, 3]  // Error: Cannot unify int with float

7.7 Tensor representation (internal)

The VM stores a tensor with the following OCaml record.

type tensor = {
  data: float array;    (* Flat storage *)
  shape: int list;      (* Dimensions *)
  strides: int list;    (* Access strides *)
}

Creating tensors

let create_tensor shape init_val =
  let size = List.fold_left ( * ) 1 shape in
  {
    data = Array.make size init_val;
    shape;
    strides = compute_strides shape;
  }

Indexing

let flat_index indices strides =
  List.fold_left2 (fun acc i s -> acc + i * s) 0 indices strides

let tensor_get t indices =
  t.data.(flat_index indices t.strides)

7.8 Future extensions

The sections below list possible additions to tensor semantics. Two of them already exist in part, as each section says.

Element-wise operations

Element-wise arithmetic is already implemented. The operators +, -, *, and / work element by element on two tensors of the same shape.

let a = [1.0, 2.0, 3.0] in
let b = [4.0, 5.0, 6.0] in
a * b  // Result: [4.0, 10.0, 18.0]

Reduction operations

Bagl doesn't have reduction operations yet. The block below shows a possible syntax, and the compiler doesn't accept it.

sum(a)     // Sum all elements
mean(a)    // Average of elements
max(a)     // Maximum element

Slicing

Bagl doesn't have slicing yet. The block below shows a possible syntax, and the compiler doesn't accept it.

a[0:2, 1:3]    // Submatrix extraction

Broadcasting

Bagl already broadcasts a float scalar against a tensor, so a * 2.0 multiplies each element of a by 2.0.

let a = [1.0, 2.0, 3.0] in
a * 2.0  // Result: [2.0, 4.0, 6.0]

Broadcasting between two tensors of different shapes isn't implemented. A broadcast of the [3, 1] tensor and the [1, 4] tensor below would give a [3, 4] tensor. The compiler rejects the program instead, because the shapes differ.

let a: tensor<float>[3, 1] = [[1.0], [2.0], [3.0]] in
let b: tensor<float>[1, 4] = [[1.0, 2.0, 3.0, 4.0]] in
a + b  // Error: Dimension mismatch: 3 vs 1

8. Runtime system

Chapter 8 describes the Bagl virtual machine (VM) and how it runs a program.

8.1 Virtual machine architecture

Bagl runs programs on a stack-based virtual machine. A stack-based VM takes the operands of each instruction from the top of a stack, and it pushes the result back onto the stack.

Components

┌─────────────────────────────────────────┐
│              VM State                    │
├─────────────────────────────────────────┤
│  Program    │ Bytecode chunks (one per  │
│             │ function)                  │
├─────────────┼───────────────────────────┤
│  Stack      │ Operand stack for         │
│             │ computation                │
├─────────────┼───────────────────────────┤
│  Locals     │ Local variable storage    │
│             │ for current frame          │
├─────────────┼───────────────────────────┤
│  Frames     │ Call stack for function   │
│             │ returns                    │
├─────────────┼───────────────────────────┤
│  PC         │ Program counter           │
├─────────────┼───────────────────────────┤
│  Chunk ID   │ Current function index    │
└─────────────┴───────────────────────────┘

Initial state

The VM starts with the following state. The locals array has 256 slots, or more if the entry function needs more.

{
  program;                    (* Loaded bytecode *)
  stack = Array.make 1024 VUnit;
  sp = 0;                     (* Stack pointer *)
  frames = [];                (* Call stack *)
  pc = 0;                     (* Program counter *)
  chunk_id = program.entry;   (* Start at entry point *)
  locals = Array.make 256 VUnit;
}

8.2 Runtime values

Value types

type value =
  | VInt of int           (* Integer *)
  | VFloat of float       (* Float *)
  | VBool of bool         (* Boolean *)
  | VString of string     (* String *)
  | VUnit                 (* Unit value *)
  | VTensor of tensor     (* Tensor *)
  | VClosure of closure   (* Function closure *)

Tensor values

type tensor = {
  data: float array;      (* Flat element storage *)
  shape: int list;        (* Dimension sizes *)
  strides: int list;      (* Access strides *)
}

Closure values

type closure = {
  func_id: int;           (* Function chunk index *)
  captures: value array;  (* Captured variables *)
}

8.3 Stack operations

Push

push(value):
  stack[sp] = value
  sp = sp + 1

Pop

pop():
  sp = sp - 1
  return stack[sp]

Stack discipline

The stack holds the following kinds of values:

• The intermediate results of a computation
• Function arguments
• Temporary values

8.4 Call frames

A call frame records what the VM needs in order to return from a function call.

Frame structure

type frame = {
  return_addr: int;         (* Where to resume *)
  return_chunk: int;        (* Which function *)
  base_ptr: int;            (* Not currently used *)
  closure: closure option;  (* For captured vars *)
  saved_locals: value array; (* Caller's locals *)
}

Function call

A function call takes the following steps.

1. The VM pops the closure from the stack.
2. It pops the arguments from the stack.
3. It saves the current state in a new frame.
4. It pushes the frame onto the call stack.
5. It switches to the chunk of the called function. A chunk is the bytecode of one function.
6. It creates a new locals array and copies the arguments into it.

Function return

A function return takes the following steps.

1. The VM pops the return value from the stack.
2. It pops the frame from the call stack.
3. It restores the state of the caller.
4. It pushes the return value onto the stack.

8.5 Closures

Closure creation

A closure is created in the following steps.

1. The compiler finds the free variables of the function, which are the variables it uses from the enclosing scope.
2. At run time, the VM copies the values of the free variables from the current environment.
3. The VM creates a closure that holds the function ID and the copied values, which are called the captures.

Closure invocation

When the VM calls a closure, it stores the closure in the new frame. The function body then reads its captures with the LOAD_CAPTURE instruction.

Recursive closures

A recursive closure needs to capture itself, so the VM builds it in the following steps.

1. The bytecode stores the index of the capture that will hold the closure itself.
2. The VM creates the closure with an empty slot at that index.
3. The VM then puts the closure into its own empty slot.

8.6 Execution model

Fetch-execute cycle

The VM runs a loop that fetches the next instruction and then executes it.

while running:
  instr = chunk.code[pc]
  pc = pc + 1
  execute(instr)

Instruction execution

Each instruction takes the following steps.

1. It reads its operands from the stack or the locals.
2. It does its computation.
3. It pushes the result onto the stack.
4. It tells the loop whether to continue.

Termination

Execution ends in any of the following cases:

• The VM executes a HALT instruction.
• The VM executes RETURN when the call stack is empty.
• A runtime error occurs.

8.7 Memory management

Stack allocation

• The stack has a fixed size of 1024 values.
• A stack overflow raises a runtime error.

Local variables

• Each frame gets its own locals array, with 256 slots or more if the function needs more.
• During a call, the frame keeps the locals of the caller, and the VM restores them on return.

Garbage collection

The VM has no garbage collector of its own. It is written in OCaml, and its values are ordinary OCaml values. The OCaml garbage collector frees a value once nothing refers to it.

8.8 Tensor operations at runtime

Tensor creation

The VM creates a tensor in the following steps.

1. It pops the elements from the stack.
2. It puts them in a flat array.
3. It computes the strides.
4. It pushes the tensor value onto the stack.

Dot product

The code below is a simplified form of the matrix case. The VM also handles a vector times a matrix, a matrix times a vector, and two vectors. For two vectors, it pushes the result as a float.

let tensor_dot a b =
  match a.shape, b.shape with
  | [m; k1], [k2; n] when k1 = k2 ->
      let result = create_tensor [m; n] 0.0 in
      for i = 0 to m - 1 do
        for j = 0 to n - 1 do
          let sum = ref 0.0 in
          for k = 0 to k1 - 1 do
            sum := !sum +. a[i,k] *. b[k,j]
          done;
          result[i,j] <- !sum
        done
      done;
      result

Transpose

let tensor_transpose t =
  let [m; n] = t.shape in
  let result = create_tensor [n; m] 0.0 in
  for i = 0 to m - 1 do
    for j = 0 to n - 1 do
      result[j,i] <- t[i,j]
    done
  done;
  result

8.9 Error handling

Runtime errors

exception Runtime_error of string

The runtime errors include the following:

• Stack overflow or stack underflow
• Division by zero
• A value of the wrong type at run time
• Tensor shapes that don't fit the operation
• A program counter outside the code of the current function
• log of a number that isn't positive, or sqrt of a negative number

Error messages

Runtime error: Stack overflow
Runtime error: Division by zero
Runtime error: Expected int, got bool
Runtime error: Invalid shapes for dot: 2x3 and 4x2

8.10 Debugging support

Instruction tracing

The VM has no instruction tracer. The listing below is an example of what a trace would show for let x = 5 in x + 1, with one line for each instruction the VM runs.

0000: PUSH_INT 5
0001: STORE_LOCAL 0
0002: LOAD_LOCAL 0
0003: PUSH_INT 1
0004: IADD
0005: RETURN

Stack inspection

The VM has no command that prints the stack or the locals. The listing below is an example of what such a printout would show.

Stack: [VInt 5, VInt 1]
Locals: [VInt 5, VUnit, ...]

Bytecode disassembly

The --dump-bytecode option prints the bytecode of each function in the following format. The listing below is simplified. The real listing for a program is often longer, because the compiler stores each intermediate result in a local.

=== Function 0 (locals=2, params=0, captures=0) ===
   0: PUSH_INT 5
   1: STORE_LOCAL 0
   2: LOAD_LOCAL 0
   3: PUSH_INT 1
   4: IADD
   5: RETURN

8.11 Performance characteristics

Time complexity

Space complexity

9. Bytecode specification

Chapter 9 describes the Bagl bytecode instruction set. Bytecode is the list of simple instructions that the compiler produces and the VM runs.

The chapter doesn't yet cover every opcode in src/bytecode.ml. It leaves out the following opcodes:

• The float comparisons FLT, FGT, FLE, and FGE
• The element-wise tensor operations TADD, TSUB, TMUL, and TDIV
• The math functions MEXP, MLOG, MSQRT, MRELU, and MSTEP

9.1 Bytecode format

Program structure

A program is a list of chunks, with one chunk for each function, and the index of the function that runs first.

Program = {
  chunks: Chunk[]    // One chunk per function
  entry: int         // Index of main function
}

Chunk structure

A chunk holds the instructions of one function and three counts.

Chunk = {
  code: Opcode[]     // Instruction sequence
  num_locals: int    // Local variable count
  num_params: int    // Parameter count
  num_captures: int  // Captured variable count
}

9.2 Instruction categories

Stack operations

• PUSH_* pushes a value onto the stack.
• POP removes the top of the stack.
• DUP duplicates the top of the stack.

Variable access

• LOAD_LOCAL and STORE_LOCAL read and write local variables.
• LOAD_GLOBAL and STORE_GLOBAL read and write global variables.
• LOAD_CAPTURE reads a value that a closure captured.

Arithmetic

• The integer instructions are IADD, ISUB, IMUL, IDIV, and INEG.
• The float instructions are FADD, FSUB, FMUL, FDIV, and FNEG.

Comparison

• IEQ and INEQ test equality.
• ILT, IGT, ILE, and IGE compare order.

Logical

• The logical instructions are AND, OR, and NOT.

Control flow

• JUMP and JUMP_IF_FALSE move to another instruction.
• CALL and RETURN enter and leave a function.

Closures

• MAKE_CLOSURE and MAKE_REC_CLOSURE create closures.

Tensors

• The tensor instructions are TENSOR_CREATE, TENSOR_DOT, TENSOR_TRANSPOSE, and TENSOR_RESHAPE.

9.3 Instruction reference

Each entry shows the stack before and after the instruction, with the top of the stack on the right.

Stack operations

PUSH_INT n

Pushes the integer n onto the stack.

Stack: ... → ..., VInt(n)

PUSH_FLOAT f

Pushes the float f onto the stack.

Stack: ... → ..., VFloat(f)

PUSH_BOOL b

Pushes the boolean b onto the stack.

Stack: ... → ..., VBool(b)

PUSH_STRING s

Pushes the string s onto the stack.

Stack: ... → ..., VString(s)

PUSH_UNIT

Pushes the unit value onto the stack.

Stack: ... → ..., VUnit

POP

Removes the top value from the stack.

Stack: ..., v → ...

DUP

Duplicates the top value of the stack.

Stack: ..., v → ..., v, v

Local variable operations

LOAD_LOCAL i

Pushes the local variable at index i.

Stack: ... → ..., locals[i]

STORE_LOCAL i

Pops a value and stores it in local i.

Stack: ..., v → ...
Effect: locals[i] = v

Global variable operations

LOAD_GLOBAL i

Pushes the global variable at index i.

Stack: ... → ..., globals[i]

The VM doesn't implement LOAD_GLOBAL yet, and it raises a runtime error when it runs one.

STORE_GLOBAL i

Pops a value and stores it in global i.

Stack: ..., v → ...
Effect: globals[i] = v

The VM doesn't implement STORE_GLOBAL yet, and it raises a runtime error when it runs one.

Closure operations

MAKE_CLOSURE func_id num_captures

Creates a closure.

Stack: ..., capture₀, ..., captureₙ₋₁ → ..., VClosure

The instruction pops num_captures values, and it creates a closure that refers to function func_id and holds the popped values.

MAKE_REC_CLOSURE func_id num_captures self_idx

Creates a recursive closure.

Stack: ..., capture₀, ..., captureₙ₋₁ → ..., VClosure

The instruction works like MAKE_CLOSURE, and it also puts the closure itself into capture slot self_idx.

LOAD_CAPTURE i

Pushes the captured value at index i of the current closure.

Stack: ... → ..., current_closure.captures[i]

Integer arithmetic

IADD

Adds two integers.

Stack: ..., a, b → ..., (a + b)
Types: VInt, VInt → VInt

ISUB

Subtracts two integers.

Stack: ..., a, b → ..., (a - b)
Types: VInt, VInt → VInt

IMUL

Multiplies two integers.

Stack: ..., a, b → ..., (a * b)
Types: VInt, VInt → VInt

IDIV

Divides two integers.

Stack: ..., a, b → ..., (a / b)
Types: VInt, VInt → VInt
Error: Division by zero if b = 0

INEG

Negates an integer.

Stack: ..., a → ..., (-a)
Types: VInt → VInt

Float arithmetic

FADD

Adds two floats.

Stack: ..., a, b → ..., (a +. b)
Types: VFloat, VFloat → VFloat

FSUB

Subtracts two floats.

Stack: ..., a, b → ..., (a -. b)
Types: VFloat, VFloat → VFloat

FMUL

Multiplies two floats.

Stack: ..., a, b → ..., (a *. b)
Types: VFloat, VFloat → VFloat

FDIV

Divides two floats.

Stack: ..., a, b → ..., (a /. b)
Types: VFloat, VFloat → VFloat
Error: Division by zero if b = 0.0

FNEG

Negates a float.

Stack: ..., a → ..., (-.a)
Types: VFloat → VFloat

Comparison operations

IEQ

Tests whether two values are equal.

Stack: ..., a, b → ..., (a = b)
Types: 'a, 'a → VBool

IEQ works on ints, floats, bools, strings, and the unit value. Two values of different kinds are never equal.

INEQ

Tests whether two values are not equal.

Stack: ..., a, b → ..., (a ≠ b)
Types: 'a, 'a → VBool

ILT

Tests whether a is less than b.

Stack: ..., a, b → ..., (a < b)
Types: VInt, VInt → VBool

IGT

Tests whether a is greater than b.

Stack: ..., a, b → ..., (a > b)
Types: VInt, VInt → VBool

ILE

Tests whether a is less than or equal to b.

Stack: ..., a, b → ..., (a ≤ b)
Types: VInt, VInt → VBool

IGE

Tests whether a is greater than or equal to b.

Stack: ..., a, b → ..., (a ≥ b)
Types: VInt, VInt → VBool

Logical operations

AND

Computes the logical and of two booleans.

Stack: ..., a, b → ..., (a && b)
Types: VBool, VBool → VBool

OR

Computes the logical or of two booleans.

Stack: ..., a, b → ..., (a || b)
Types: VBool, VBool → VBool

NOT

Computes the logical not of a boolean.

Stack: ..., a → ..., (!a)
Types: VBool → VBool

Control flow

JUMP addr

Jumps to addr every time.

Effect: pc = addr

JUMP_IF_FALSE addr

Pops a boolean, and jumps to addr if the boolean is false.

Stack: ..., cond → ...
Effect: if !cond then pc = addr
Types: VBool → (consumed)

CALL n_args

Calls a function.

Stack: ..., arg₀, ..., argₙ₋₁, closure → ...
Effect:
  1. Save current frame
  2. Switch to closure's function
  3. Copy args to new locals

RETURN

Returns from a function.

Stack: ..., result → ... (in caller), result
Effect:
  1. Pop frame
  2. Restore caller state
  3. Push result on caller's stack

Tensor operations

TENSOR_CREATE shape

Creates a tensor with the given shape.

Stack: ..., elem₀, ..., elemₙ₋₁ → ..., VTensor

The instruction pops as many elements as the shape holds, and it creates the tensor. The number of elements is the product of the dimensions.

TENSOR_DOT

Computes the dot product of two matrices or vectors.

Stack: ..., a, b → ..., result
Types: VTensor, VTensor → VTensor or VFloat

The result is a float when both inputs are vectors, and a tensor otherwise.

TENSOR_TRANSPOSE

Transposes a 2D tensor.

Stack: ..., t → ..., transposed
Types: VTensor → VTensor

TENSOR_RESHAPE shape

Reshapes a tensor.

Stack: ..., t → ..., reshaped
Types: VTensor → VTensor
Error: if sizes don't match

Debugging and control

PRINT

Prints the top of the stack, for debugging.

Stack: ..., v → ..., v
Effect: prints string_of_value(v)

HALT

Stops execution.

Effect: terminates VM execution loop

NOP

Does nothing.

Effect: none

9.4 Bytecode examples

The listings in section 9.4 are simplified by hand to show how each construct works. The real compiler stores each intermediate result in a local, and its optimizer folds constant expressions such as 1 + 2 into one value. The --dump-bytecode option prints the real code, which is often longer or different.

Simple expression

Source: 1 + 2

PUSH_INT 1
PUSH_INT 2
IADD
RETURN

Let binding

Source: let x = 5 in x + 1

PUSH_INT 5
STORE_LOCAL 0
LOAD_LOCAL 0
PUSH_INT 1
IADD
RETURN

Function definition

Source: fn x -> x + 1

// Main function
MAKE_CLOSURE 1 0    // func_id=1, no captures
RETURN

// Function 1: x -> x + 1
LOAD_LOCAL 0        // x
PUSH_INT 1
IADD
RETURN

Function call

Source: (fn x -> x + 1) 5

PUSH_INT 5          // Argument
MAKE_CLOSURE 1 0    // Create closure
CALL 1              // Call with 1 arg
RETURN

Conditional

Source: if true then 1 else 2

PUSH_BOOL true
JUMP_IF_FALSE 4     // Jump to else
PUSH_INT 1
JUMP 5              // Skip else
PUSH_INT 2          // else branch
RETURN

Closure with capture

Source: let y = 5 in fn x -> x + y

// Main
PUSH_INT 5
STORE_LOCAL 0       // y = 5
LOAD_LOCAL 0        // Push y for capture
MAKE_CLOSURE 1 1    // func_id=1, 1 capture
RETURN

// Function 1
LOAD_LOCAL 0        // x
LOAD_CAPTURE 0      // y
IADD
RETURN

9.5 Serialization format

A compiled program is saved as a .baglc file with the layout below. Every integer in the file, including each uint32 in the layout, is written as 4 bytes in big-endian order, which puts the most significant byte first.

Header:
  magic: "BAGL" (4 bytes)
  version: uint32
  entry: uint32
  num_chunks: uint32

For each chunk:
  num_locals: uint32
  num_params: uint32
  num_captures: uint32
  code_length: uint32
  code: opcode[] (variable)

Each opcode is written in two parts:

• One tag byte, which says which opcode it is.
• The operands, which depend on the opcode.