What comptime can do in Vx
TL;DR — Acomptimeblock runs ordinary Vx code inside the compiler: loops, recursion, structs, arrays and helpers that take&mut. Its value goes into the program as a constant, and the block itself is removed. A falseassertstops the build, and a block the compiler cannot finish is an error, so nothing quietly becomes run-time code. PR #858 rebuilt this on one interpreter that works out a block's value and tracks its writes at the same time. For facts about every input, contracts hand the proof to the z3 theorem prover.
A comptime block in Vx is ordinary Vx code that the compiler runs while it compiles your program. Loops, recursion, structs, arrays, helper functions that take &mut: all of it runs inside the compiler, and the answer goes into your program as a constant.
Vx PR #858 rebuilt this on a single interpreter. It works out the value of a block and, at the same time, tracks every write the block makes. This post shows what that makes possible. Every example below was compiled and run with vxc built from commit 4a6a58a3 on main.
There are two forms, and two rules.
let x = comptime { ... }; // the block's value becomes a constant
comptime { assert(...); } // checked while compiling, then removed- A false
assertstops the build. It never becomes a run-time check. - A block the compiler cannot finish is an error. It never quietly turns into run-time code, so a block that compiled really did run.
A whole program that becomes one number
fn is_prime(n : i64) -> bool {
if n < 2i64 {
return false;
}
let mut d : i64 = 2i64;
while d * d <= n {
if n % d == 0i64 {
return false;
}
d = d + 1i64;
}
return true;
}
fn count_primes_below(limit : i64) -> i64 {
let mut count : i64 = 0i64;
for n in 0i64..limit {
if is_prime(n) {
count = count + 1i64;
}
}
return count;
}
fn main() -> i32 {
let primes = comptime {
count_primes_below(10000i64)
};
print(primes);
return 0;
}This prints 1229. Here is main as the compiler generates it (vxc primes.vx --action emit-mlir):
func.func @main() -> i32 {
call @vx_init_signals() : () -> ()
%c1229_i32 = arith.constant 1229 : i32
%0 = call @print_i32(%c1229_i32) : (i32) -> i32
%c0_i32 = arith.constant 0 : i32
return %c0_i32 : i32
}The compiler tested ten thousand numbers for primality. The program tests none.
Lookup tables, built while compiling
A block can produce a whole tensor, which makes lookup tables easy. This one holds the first 16 primes, found by the same is_prime as above:
fn main() -> i32 {
let table : Tensor<i64, [16]> = comptime {
let mut t : Tensor<i64, [16]> = [ 0i64, 0i64, 0i64, 0i64, 0i64, 0i64, 0i64, 0i64,
0i64, 0i64, 0i64, 0i64, 0i64, 0i64, 0i64, 0i64 ];
let mut found : i64 = 0i64;
let mut n : i64 = 2i64;
while found < 16i64 {
if is_prime(n) {
t[found] = n;
found = found + 1i64;
}
n = n + 1i64;
}
t
};
print(table[15]);
return 0;
}It prints 53. In the generated code the table is sixteen constants, 2 through 53, stored into an array. No search runs at run time.
A small computer, inside the compiler
The code a comptime block calls is not limited to simple arithmetic. Here is a tiny register machine: a loop that fetches instructions from an array, decodes them, and jumps.
// A tiny register machine. Each instruction is three numbers: opcode, a, b.
// 0 halt 1 set r[a] = b 2 add r[a] += r[b] 3 mul r[a] *= r[b]
// 4 dec r[a] -= 1 5 jump to b if r[a] != 0
fn run(code : Tensor<i64, [18]>) -> i64 {
let mut r : Tensor<i64, [4]> = [ 0i64, 0i64, 0i64, 0i64 ];
let mut pc : i64 = 0i64;
loop {
let op : i64 = code[pc];
let a : i64 = code[pc + 1i64];
let b : i64 = code[pc + 2i64];
pc = pc + 3i64;
if op == 0i64 {
break;
} else if op == 1i64 {
r[a] = b;
} else if op == 2i64 {
r[a] = r[a] + r[b];
} else if op == 3i64 {
r[a] = r[a] * r[b];
} else if op == 4i64 {
r[a] = r[a] - 1i64;
} else if r[a] != 0i64 {
pc = b;
}
}
return r[0];
}
fn main() -> i32 {
let answer = comptime {
// r0 = 1; r1 = 10; do { r0 *= r1; r1 -= 1 } while r1 != 0
let factorial_10 : Tensor<i64, [18]> = [
1i64, 0i64, 1i64,
1i64, 1i64, 10i64,
3i64, 0i64, 1i64,
4i64, 1i64, 0i64,
5i64, 1i64, 6i64,
0i64, 0i64, 0i64
];
run(factorial_10)
};
print(answer);
return 0;
}The compiler runs one interpreter inside another and generates a program that prints 3628800, which is 10!. main contains that single constant.
Structs and &mut helpers
Compile-time code can be written the same way as run-time code. Helpers take &mut, change struct fields, and loop:
struct Account {
balance : i64,
deposits : i64,
}
fn deposit(acct : &mut Account, amount : i64) -> void {
acct.balance = acct.balance + amount;
acct.deposits = acct.deposits + 1i64;
}
fn add_interest(acct : &mut Account, percent : i64, years : i64) -> void {
for y in 0i64..years {
acct.balance = acct.balance + acct.balance * percent / 100i64;
}
}
fn main() -> i32 {
let final_balance = comptime {
let mut acct = Account {
balance : 0i64,
deposits : 0i64
};
deposit(&mut acct, 1000i64);
deposit(&mut acct, 500i64);
add_interest(&mut acct, 5i64, 10i64);
acct.balance
};
print(final_balance);
return 0;
}It prints 2438: 1500 at 5% for ten years, rounded down each year. Before #858 the compiler refused this block. The new interpreter follows each write through each &mut, so it can run it.
Tests that run in the compiler
A comptime block of asserts is a unit test that runs on every build:
fn gcd(a : i64, b : i64) -> i64 {
let mut x : i64 = a;
let mut y : i64 = b;
while y != 0i64 {
let t : i64 = x % y;
x = y;
y = t;
}
return x;
}
fn pow_mod(base : i64, exp : i64, m : i64) -> i64 {
let mut result : i64 = 1i64;
let mut b : i64 = base % m;
let mut e : i64 = exp;
while e > 0i64 {
if e % 2i64 == 1i64 {
result = result * b % m;
}
b = b * b % m;
e = e / 2i64;
}
return result;
}
fn main() -> i32 {
comptime {
assert(gcd(1071i64, 462i64) == 21i64, "gcd");
assert(pow_mod(2i64, 10i64, 1000i64) == 24i64, "2^10 mod 1000");
assert(pow_mod(7i64, 1000002i64, 1000003i64) == 1i64, "Fermat's little theorem");
}
return 0;
}1000003 is prime, so Fermat's little theorem says the last line holds. Change the expected value to 2i64 and the program no longer compiles:
Error[E8002] at 0:0: Comptime assert failed: Fermat's little theoremThe same idea scales to a real algorithm. This file checks, on every build, that insertion_sort returns its input in order and loses no numbers:
fn insertion_sort<const N : i32>(a : Tensor<i64, [N]>) -> Tensor<i64, [N]> {
let mut w : Tensor<i64, [N]> = a;
for i in 1..N {
let key : i64 = w[i];
let mut j : i32 = i - 1;
while j >= 0 && w[j] > key {
w[j + 1] = w[j];
j = j - 1;
}
w[j + 1] = key;
}
return w;
}
fn is_sorted<const N : i32>(a : &Tensor<i64, [N]>) -> bool {
for i in 1..N {
if a[i - 1] > a[i] {
return false;
}
}
return true;
}
// Two lists with the same sum and the same sum of squares almost always hold the same numbers.
fn checksum<const N : i32>(a : &Tensor<i64, [N]>) -> i64 {
let mut sum : i64 = 0i64;
let mut squares : i64 = 0i64;
for i in 0..N {
sum = sum + a[i];
squares = squares + a[i] * a[i];
}
return sum * 1000003i64 + squares;
}
fn main() -> i32 {
let sorted_ok = comptime {
let input : Tensor<i64, [10]> = [ 42i64, -7i64, 19i64, 0i64, 42i64, 3i64, -100i64, 8i64, 8i64, 1i64 ];
let expected : i64 = checksum(&input);
let out : Tensor<i64, [10]> = insertion_sort(input);
is_sorted(&out) && checksum(&out) == expected
};
comptime {
assert(sorted_ok, "insertion_sort sorts and keeps every number");
}
return 0;
}Flip w[j] > key to w[j] < key, or overwrite an element by mistake, and the build fails with Comptime assert failed: insertion_sort sorts and keeps every number. The sort is generic over its length N, and the compiler made the N = 10 copy and ran it.
Proofs for every input, with z3
A comptime test checks the inputs you write down. A contract checks every input. Vx functions can state what they need (requires) and what they promise (ensures), and the compiler hands both to the z3 theorem prover while it compiles.
Here is the index arithmetic for a 64 × 64 matrix stored row by row and cut into 16 × 16 tiles. The promise is that the index never leaves the matrix:
// The element (row, col) of tile (tr, tc) in a 64 x 64 matrix stored row by row,
// cut into 16 x 16 tiles.
fn tiled_index(tr : i32, tc : i32, row : i32, col : i32) -> i32
requires tr >= 0 && tr < 4 && tc >= 0 && tc < 4
requires row >= 0 && row < 16 && col >= 0 && col < 16
ensures return >= 0 && return < 4096 {
let r = tr * 16 + row;
let c = tc * 16 + col;
return r * 64 + c;
}
fn main() -> i32 {
print(tiled_index(3, 3, 15, 15));
return 0;
}The compiler turns the body into equations, adds the requires as facts, and asks z3 for any input that breaks the ensures. z3 finds none, for all 4,096 inputs, without trying them one by one. Each of these one-character mistakes stops the build:
row < 17in therequiresr * 65instead ofr * 64tc * 16 + col + 1return < 4095in theensures. The last element really is index 4095, and z3 finds the input that reaches it.
Error[E8001]: Function 'tiled_index' cannot prove postcondition (ensures) at compile timeWithout z3 installed, the build fails rather than skipping the proof.
The prover is young, and it is worth knowing where it stops today:
- It proves straight-line code built from
letbindings,+,-, multiplication by a constant, comparisons,&&and||. Code that reassigns a variable or returns from inside anifis refused, even when the promise is true. - z3 treats integers as unbounded, while an
i32wraps at 2^31. Bound every input inrequires, as above, so no step can overflow; otherwise a proof can hold for z3's integers and fail for the machine's. - The compiler does not yet check that callers meet
requires. The proof says "if the inputs are in range, the result is in range", and making sure they are in range is still the caller's job.
Generic code that recurses on its own parameter
A constant generic parameter can drive recursion, with if comptime choosing the base case:
fn fib<const N : i32>() -> i64 {
if comptime N < 2 {
return 1i64;
} else {
return fib<N - 1>() + fib<N - 2>();
}
}
fn main() -> i32 {
comptime {
assert(fib<20>() == 10946i64, "fib 20");
}
return 0;
}fib<20> makes the compiler create fib<19>, fib<18>, and so on down to fib<0>, then run them. Change 10946 to 10945 and the build fails with Comptime assert failed: fib 20.
Functions that only run while compiling
A closure whose body is a comptime block is a function that only exists at compile time, like C++'s consteval:
fn main() -> i32 {
let cube = |x : i32| comptime {
x * x * x
};
print(cube(3i32) + cube(4i32));
return 0;
}This prints 91. Each call became its answer (27 and 64), and cube itself is gone from the generated code. Call it with a value only known at run time, and the compiler says so:
fn main(argc : i32) -> i32 {
let cube = |x : i32| comptime {
x * x * x
};
return cube(argc);
}Error[E3033] at 0:0: this call to a `comptime` lambda cannot be evaluated: a comptime lambda runs while compiling, so every argument has to be known thenC++'s constexpr would quietly make this an ordinary run-time call. Vx stops instead.
Facts about the hardware
Vx knows which device a function runs on and which memory a tensor lives in, and both are known while compiling. So you can assert them:
fn on_gpu() on Topology::GPU -> void {
let data : Tensor<f32, [128]> = Tensor<f32, [128]>::new();
let moved = transfer(data, Memory::GPU_HBM);
comptime {
assert(Topology::Current == Topology::GPU, "this function runs on the GPU");
assert(moved.topology() == Some(Topology::GPU), "the tensor now lives in GPU memory");
}
}
fn main() -> i32 {
spawn on(Topology::GPU) {
on_gpu();
};
return 0;
}If a later change moves that tensor somewhere else, this file stops compiling, and the message says which assumption broke. The check costs nothing at run time.
It refuses rather than guess
A comptime block is removed from the program once it has run. So a block that writes to a variable declared outside it would lose that write. Vx makes this an error:
fn main() -> i32 {
let mut hits : i32 = 0i32;
let answer = comptime {
hits = hits + 1i32;
42i32
};
return answer + hits;
}Error[E3033] at 0:0: this `comptime` block cannot be evaluated: it writes to 'hits', which is declared outside it -- the block disappears, so the write would have to disappear with itThis is the bug class #858 was written to close. The old compiler had two separate passes: one worked out a block's value and another looked for writes. When they disagreed about which parts of the block ran, a write could vanish without a word. Now one interpreter does both jobs. It follows writes through &mut references, struct fields, array elements, closures, and branches it cannot decide. Every kind of expression in the language has to be handled explicitly, so new syntax cannot slip through unnoticed.
Integer overflow is refused too. 25! does not fit in an i64:
fn factorial(n : i64) -> i64 {
let mut acc : i64 = 1i64;
for i in 1i64..n + 1i64 {
acc = acc * i;
}
return acc;
}
fn main() -> i32 {
let f = comptime {
factorial(25i64)
};
return (f % 100i64) as i32;
}Error[E3033] at 11:5: this `comptime` block cannot be evaluated: its value cannot be worked outA wrapped-around number would be put into the program as if it were the answer. Vx does not produce one.
Fast enough to use
This block finds the start below 10,000 with the longest Collatz sequence. That is 849,637 loop steps in total:
fn collatz_steps(start : i64) -> i64 {
let mut n : i64 = start;
let mut steps : i64 = 0i64;
while n != 1i64 {
if n % 2i64 == 0i64 {
n = n / 2i64;
} else {
n = 3i64 * n + 1i64;
}
steps = steps + 1i64;
}
return steps;
}
fn main() -> i32 {
let best = comptime {
let mut best_start : i64 = 1i64;
let mut best_steps : i64 = 0i64;
for s in 1i64..10000i64 {
let steps : i64 = collatz_steps(s);
if steps > best_steps {
best_steps = steps;
best_start = s;
}
}
best_start
};
print(best);
return 0;
}It prints 6171, and the whole compile takes about 1.5 seconds on an Apple M4. The compiler from just before #858 printed 1161 for this program, which is wrong. The new interpreter gets it right.
What it cannot do yet
- An
assertthat comes after a&mutcall or a struct field write, in the same block, is not checked yet, and neither is one that calls a function with a borrowed tensor. Until that is fixed, compute the result as the block's value and assert on that value, as the sort example above does. - A block's value can be a number, a
boolor a tensor of numbers. It cannot be a struct yet. - Inside a block there are no
ascasts, strings orprint. Do the cast outside the block, as the overflow example above does. - A function call cannot yet be used as a tensor's size or as a constant generic argument (
total<square(2)>(xs)). Arithmetic likefib<N - 1>works.
Apart from the first, each of these is refused with an error, so a program that compiles does not rely on one of them without saying so.
The compile-time evaluation chapter of the docs covers the rules in full. The design is discussed in discussion #857, and the change itself is PR #858.