You are getting early access to this article as a subscriber. Your support makes articles like this possible. Thank you.
OxCaml, Jane Street’s superset of OCaml, allows you to assert that entire functions do not allocate (on the heap). And if you start allocating within that call tree, the compiler will fail and tell you that you allocated. While you might be able to achieve this with static analysis, there are few mainstream languages that let you do this in the compiler itself. Swift and Clang are the only other exceptions I know of.1
In most languages (Java, Go, C#, Rust, Zig, OCaml, etc.) the process is reversed: you take a profiler to try and find allocations (usually in loops that happen millions of times). Then you go and eliminate or minimize the allocations. But as soon as you edit a line of code in the hot path, you might forget the context and start allocating again, and you’re back to square one with a profiler.
In D you can mark a function as @nogc but it doesn't stop you from heap allocating without the garbage collector. In Zig (and maybe more recent Rust) you might be able to minimize regression through convention by not passing an allocator to a function. But convention can be ignored and bypassed. Why not let the compiler do the work?
In this article we'll take a look at how OxCaml's [@zero_alloc] assertion works.
Prerequisites#
First install OxCaml (this will take a while).
sudo apt-get -y install unzip bubblewrap hyperfine build-essential autoconf
bash -c "sh <(curl -fsSL https://opam.ocaml.org/install.sh)"
opam init
opam switch create 5.2.0+ox --repos ox=git+https://github.com/oxcaml/opam-repository.git,default
eval $(opam env --switch 5.2.0+ox)
Now let's try it out.
A naive program to sum CSV data#
Let’s write a naive OCaml program (no OxCaml features yet) to sum up rows of a CSV. It might not be idiomatic but the point is to show off this zero allocation assertion feature.
First we open the file at the path in the first command line argument. Then we call sum_column with the open file and pass along the second command line argument to indicate the column to sum.
let () =
let path = Sys.argv.(1) in
let col = int_of_string Sys.argv.(2) in
let ic = open_in_bin path in
let sum = sum_column ic col in
close_in ic;
Printf.printf "sum=%d\n" sum
main_alloc.ml
Remember that parentheses for function calls in OCaml are not necessary. The .(X) syntax is array access. And open_in_bin and close_in are standard library functions for files.
Next we’ll iterate through the file in 64KB chunks.
let chunk = 1 lsl 16
let sum_column (ic: in_channel) (col: int) =
let buf = Bytes.create chunk in
let st: state = { column = 0; field = Buffer.create 16; sum = 0 } in
let continue = ref true in
while !continue do
let got = input ic buf 0 chunk in
if got > 0 then sum_column_incremental st buf got col
else continue := false
done;
st.sum
main_alloc.ml
Remember that ref is a wrapper for mutable values in OCaml. ! is the dereference operator for ref types and := is assignment for ref types.
Now within sum_column_incremental we’ll iterate through the buffer, character at a time. We naively accumulate the field characters into a buffer that we can convert into a string once we’ve parsed past the column the user asked us to sum.
type state = { mutable column : int; field : Buffer.t; mutable sum : int }
let sum_column_incremental (st: state) (buf: Bytes.t) (len: int) (col: int): unit =
for i = 0 to len - 1 do
let c = Bytes.get buf i in
if (c = ',' || c = '\n') && st.column = col then begin
let value = int_of_string (Buffer.contents st.field) in
st.sum <- st.sum + value;
Buffer.clear st.field;
end;
match c with
| '\n' ->
st.column <- 0
| ',' ->
st.column <- st.column + 1
| c ->
if st.column = col then
Buffer.add_char st.field c
done
main_alloc.ml
Remember that <- is record field assignment in OCaml.
This is not a great CSV parser. We don’t try to handle quotes or anything.
Let’s create a large CSV dataset (2 columns, 100 million rows) to run it against.
awk -v n=100000000 'BEGIN { srand(1); for (i = 1; i <= n; i++) print i "," int(rand() * 1000) }' > data.csv
And try it out, summing up the second column.
$ ocamlopt main_alloc.ml -o main_alloc
$ ./main_alloc data.csv 1
sum=49947827758
Not bad! But we’re unhappy with the performance. Hyperfine tells us it usually takes around 13 seconds to complete on this Digital Ocean virtual machine with 4vCPUs and 8GB RAM.
$ hyperfine './main_alloc data.csv 1'
Benchmark 1: ./main_alloc data.csv 1
Time (mean ± σ): 13.637 s ± 0.572 s [User: 13.263 s, System: 0.373 s]
Range (min … max): 12.579 s … 14.411 s 10 runs
Debugging the program#
At this point you indeed need to go in with a profiler and find out where we’re spending our time. Let's recompile the program with debug symbols and run perf against it.
$ ocamlopt -g main_alloc.ml -o main_alloc
$ perf record -g --call-graph dwarf ./main_alloc data.csv 1
$ perf report
Samples: 50K of event 'task-clock:ppp', Event count (approx.): 12669750000
Children Self Command Shared Object Symbol
+ 70.68% 62.15% main_alloc main_alloc [.] Main_alloc.sum_column_incremental_0_2_code
+ 65.11% 0.00% main_alloc main_alloc [.] caml_start_program
+ 65.11% 0.00% main_alloc main_alloc [.] caml_program
+ 65.11% 0.00% main_alloc main_alloc [.] Main_alloc.entry
+ 65.11% 0.01% main_alloc main_alloc [.] Main_alloc.sum_column_1_3_code
+ 33.02% 4.82% main_alloc main_alloc [.] caml_c_call
+ 16.26% 1.81% main_alloc main_alloc [.] caml_int_of_string
+ 14.45% 10.62% main_alloc main_alloc [.] parse_intnat
+ 7.83% 7.76% main_alloc main_alloc [.] caml_alloc_string
+ 5.58% 0.00% main_alloc [unknown] [.] 0xffffffffffffffff
+ 4.78% 1.40% main_alloc main_alloc [.] caml_blit_bytes
+ 3.60% 0.00% main_alloc main_alloc [.] memmove (inlined)
Inside the perf report we notice a lot of time spent in in sum_column_incremental which makes sense. But we also see a surprisingly large chunk of time allocating strings. Can we... stop?
Assertion-guided optimizing#
In any other language we'd have to continue down the path with perf and with some close peering at code to discover allocations. But with OxCaml we can simply label a function as [@zero_alloc] and the compiler will fail and tell us every single place we allocate.
That's amazing.
Let's copy main_alloc.ml into main_noalloc.ml and let's mark the sum_column function as [@zero_alloc].
let[@zero_alloc] sum_column (ic: in_channel) (col: int) =
let buf = Bytes.create chunk in
let st = { column = 0; field = Buffer.create 16; sum = 0 } in
let continue = ref true in
while !continue do
let read = input ic buf 0 chunk in
if read > 0 then sum_column_incremental st buf read col
else continue := false
done;
st.sum
main_noalloc.ml
And run the OxCaml compiler.
$ File "main_noalloc.ml", line 23, characters 5-15:
23 | let[@zero_alloc] sum_column (ic: in_channel) (col: int) =
^^^^^^^^^^
Error: Annotation check for zero_alloc failed on function Main_noalloc.sum_column (camlMain_noalloc__sum_column_1_3_code).
File "main_noalloc.ml", line 24, characters 12-30:
24 | let buf = Bytes.create chunk in
^^^^^^^^^^^^^^^^^^
Error: called function may allocate (external call to caml_create_bytes)
File "main_noalloc.ml", line 25, characters 40-56:
25 | let st: state = { column = 0; field = Buffer.create 16; sum = 0 } in
^^^^^^^^^^^^^^^^
Error: called function may allocate (external call to caml_create_bytes) (stdlib/buffer.ml:50,9--23;main_noalloc.ml:25,40--56)
File "main_noalloc.ml", line 25, characters 40-56:
25 | let st: state = { column = 0; field = Buffer.create 16; sum = 0 } in
^^^^^^^^^^^^^^^^
Error: allocation of 24 bytes (stdlib/buffer.ml:51,11--36;main_noalloc.ml:25,40--56)
File "main_noalloc.ml", line 25, characters 40-56:
25 | let st: state = { column = 0; field = Buffer.create 16; sum = 0 } in
^^^^^^^^^^^^^^^^
Error: allocation of 32 bytes (stdlib/buffer.ml:51,1--71;main_noalloc.ml:25,40--56)
File "main_noalloc.ml", line 25, characters 18-67:
25 | let st: state = { column = 0; field = Buffer.create 16; sum = 0 } in
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Error: allocation of 32 bytes
File "main_noalloc.ml", line 28, characters 14-34:
28 | let got = input ic buf 0 chunk in
^^^^^^^^^^^^^^^^^^^^
Error: called function may allocate (external call to caml_ml_input) (stdlib/stdlib.ml:441,7--32;main_noalloc.ml:28,14--34)
File "main_noalloc.ml", line 29, characters 20-57:
29 | if got > 0 then sum_column_incremental st buf got col
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Error: called function may allocate (direct call camlMain_noalloc__sum_column_incremental_0_2_code)
And we get a ton of errors pointing to where we allocated! (Or at least areas it couldn’t prove we didn’t allocate.) How helpful.
Let’s work through these. First, let's move the creation of our 64KB buffer outside of sum_column. We can create it in the main function.
let[@zero_alloc] sum_column (ic: in_channel) (col: int) (buf: Bytes.t) =
let st = { column = 0; field = Buffer.create 16; sum = 0 } in
let continue = ref true in
while !continue do
let read = input ic buf 0 chunk in
if read > 0 then sum_column_incremental st buf read col
else continue := false
done;
st.sum
let () =
let path = Sys.argv.(1) in
let col = int_of_string Sys.argv.(2) in
let ic = open_in_bin path in
let buf = Bytes.create chunk in
let sum = sum_column ic col buf in
close_in ic;
Printf.printf "sum=%d\n" sum
main_noalloc.ml
The main function passes buf to sum_column.
Let’s try compiling again.
$ ocamlopt main_noalloc.ml -o main_noalloc
File "main_noalloc.ml", line 23, characters 5-15:
23 | let[@zero_alloc] sum_column (ic: in_channel) (col: int) (buf: Bytes.t) =
^^^^^^^^^^
Error: Annotation check for zero_alloc failed on function Main_noalloc.sum_column (camlMain_noalloc__sum_column_1_3_code).
File "main_noalloc.ml", line 24, characters 33-49:
24 | let st = { column = 0; field = Buffer.create 16; sum = 0 } in
^^^^^^^^^^^^^^^^
Error: called function may allocate (external call to caml_create_bytes) (stdlib/buffer.ml:50,9--23;main_noalloc.ml:24,33--49)
File "main_noalloc.ml", line 24, characters 33-49:
24 | let st = { column = 0; field = Buffer.create 16; sum = 0 } in
^^^^^^^^^^^^^^^^
Error: allocation of 24 bytes (stdlib/buffer.ml:51,11--36;main_noalloc.ml:24,33--49)
File "main_noalloc.ml", line 24, characters 33-49:
24 | let st = { column = 0; field = Buffer.create 16; sum = 0 } in
^^^^^^^^^^^^^^^^
Error: allocation of 32 bytes (stdlib/buffer.ml:51,1--71;main_noalloc.ml:24,33--49)
File "main_noalloc.ml", line 24, characters 11-60:
24 | let st = { column = 0; field = Buffer.create 16; sum = 0 } in
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Error: allocation of 32 bytes
File "main_noalloc.ml", line 27, characters 15-35:
27 | let read = input ic buf 0 chunk in
^^^^^^^^^^^^^^^^^^^^
Error: called function may allocate (external call to caml_ml_input) (stdlib/stdlib.ml:441,7--32;main_noalloc.ml:27,15--35)
File "main_noalloc.ml", line 28, characters 21-59:
28 | if read > 0 then sum_column_incremental st buf read col
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Error: called function may allocate (direct call camlMain_noalloc__sum_column_incremental_0_2_code)
main_noalloc.ml
Ok! One down, a few more to go.
The next issue is this Buffer we use to store a column as we go so we can convert it to an integer with int_of_string. But instead of duplicating the characters and converting characters to integer after we pass the column we want to sum, we can do integer conversion ourselves, one character at a time.
type state = { mutable column : int; mutable value : int; mutable sum : int }
let sum_column_incremental (st: state) (buf: Bytes.t) (len: int) (col: int): unit =
for i = 0 to len - 1 do
let c = Bytes.get buf i in
if (c = ',' || c = '\n') && st.column = col then begin
st.sum <- st.sum + st.value;
st.value <- 0;
end;
match c with
| '\n' ->
st.column <- 0
| ',' ->
st.column <- st.column + 1
| c ->
if st.column = col then
st.value <- st.value * 10 + (Char.code c - Char.code '0')
done
let chunk = 1 lsl 16
let[@zero_alloc] sum_column (ic: in_channel) (col: int) (buf: Bytes.t) =
let st = { column = 0; value = 0; sum = 0 } in
let continue = ref true in
while !continue do
let read = input ic buf 0 chunk in
if read > 0 then sum_column_incremental st buf read col
else continue := false
done;
st.sum
main_noalloc.ml
Let’s see what the compiler says.
$ ocamlopt main_noalloc.ml -o main_noalloc
File "main_noalloc.ml", line 22, characters 5-15:
22 | let[@zero_alloc] sum_column (ic: in_channel) (col: int) (buf: Bytes.t) =
^^^^^^^^^^
Error: Annotation check for zero_alloc failed on function Main_noalloc.sum_column (camlMain_noalloc__sum_column_1_3_code).
File "main_noalloc.ml", line 23, characters 11-45:
23 | let st = { column = 0; value = 0; sum = 0 } in
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Error: allocation of 32 bytes
File "main_noalloc.ml", line 26, characters 15-35:
26 | let read = input ic buf 0 chunk in
^^^^^^^^^^^^^^^^^^^^
Error: called function may allocate (external call to caml_ml_input) (stdlib/stdlib.ml:441,7--32;main_noalloc.ml:26,15--35)
And we’re almost there!
The next challenge is this state record that gets heap allocated. But OxCaml gives us a tool for this too. We can mark a value as stack_ to have it stack allocated and to assert it will never escape the stack (which might otherwise happen if we tried to return it or assign it to a global mutable value, etc).
let[@zero_alloc] sum_column (ic: in_channel) (col: int) (buf: Bytes.t) =
let st = stack_ { column = 0; value = 0; sum = 0 } in
let continue = ref true in
while !continue do
let read = input ic buf 0 chunk in
if read > 0 then sum_column_incremental st buf read col
else continue := false
done;
st.sum
main_noalloc.ml
We could have equivalently marked st as local_ and OxCaml’s type inference would have deduced that st can be stack allocated.
let[@zero_alloc] sum_column (ic: in_channel) (col: int) (buf: Bytes.t) =
let local_ st = { column = 0; value = 0; sum = 0 } in
let continue = ref true in
while !continue do
let read = input ic buf 0 chunk in
if read > 0 then sum_column_incremental st buf read col
else continue := false
done;
st.sum
main_noalloc.ml
Let’s see what the compiler thinks now.
$ ocamlopt main_noalloc.ml -o main_noalloc
File "main_noalloc.ml", line 22, characters 5-15:
22 | let[@zero_alloc] sum_column (ic: in_channel) (col: int) (buf: Bytes.t) =
^^^^^^^^^^
Error: Annotation check for zero_alloc failed on function Main_noalloc.sum_column (camlMain_noalloc__sum_column_1_3_code).
File "main_noalloc.ml", line 26, characters 15-35:
26 | let read = input ic buf 0 chunk in
^^^^^^^^^^^^^^^^^^^^
Error: called function may allocate (external call to caml_ml_input) (stdlib/stdlib.ml:441,7--32;main_noalloc.ml:26,15--35)
Now we’ve culled every allocation except for the standard library function input when we copy bytes from the file on disk into our buffer.
We have a few options. First we could give up reading the file incrementally; reading it all at once in the beginning and then processing it in its entirety.
Second, we could wrap input in our own function and assert (i.e. lie) that it is zero_alloc.
let[@zero_alloc assume] input_unsafe ic buf pos len = input ic buf pos len
main_noalloc.ml
One area where input might allocate is on the exception path. Exceptions, being variant constructors, might need to allocate on the heap if they contain something bigger than an int (e.g. a string error message). So if there’s ever an issue accessing the file it might allocate. In this example program that’s certainly fine. We don’t expect it to fail. And even if it did, it would be fine to just run it again.
Third, we could use the C FFI to create a version of input that copies bytes into our buffer but never raises an exception.
#include <unistd.h>
#include <caml/mlvalues.h>
#include <caml/io.h>
CAMLprim value caml_read_noalloc(value ic, value buf, value pos, value len) {
int fd = Channel(ic)->fd;
ssize_t n = read(fd, Bytes_val(buf) + Long_val(pos), Long_val(len));
return Val_long(n);
}
read_noalloc.c
And in the OCaml stubs we’d mark it as [@@noalloc] which satisfies [@zero_alloc].
external read_noalloc : in_channel -> bytes -> int -> int -> int
= "caml_read_noalloc" [@@noalloc]
read_noalloc.ml
Fourth, we could have a separate reader thread and a shared buffer pool. Using noalloc Atomics and spinlocks so the reader and process threads can know when it’s safe to modify and read from buffers.
And there are probably other options still.
Let's go with the [@zero_alloc assume] path. Here's main_noalloc.ml in full.
type state = { mutable column : int; mutable value : int; mutable sum : int }
let sum_column_incremental (st: state) (buf: Bytes.t) (len: int) (col: int): unit =
for i = 0 to len - 1 do
let c = Bytes.get buf i in
if (c = ',' || c = '\n') && st.column = col then begin
st.sum <- st.sum + st.value;
st.value <- 0;
end;
match c with
| '\n' ->
st.column <- 0
| ',' ->
st.column <- st.column + 1
| c ->
if st.column = col then
st.value <- st.value * 10 + (Char.code c - Char.code '0')
done
let chunk = 1 lsl 16
let[@zero_alloc assume] input_unsafe ic buf pos limit = input ic buf pos limit
let[@zero_alloc] sum_column (ic: in_channel) (col: int) (buf: Bytes.t) =
let local_ st = { column = 0; value = 0; sum = 0 } in
let continue = ref true in
while !continue do
let read = input_unsafe ic buf 0 chunk in
if read > 0 then sum_column_incremental st buf read col
else continue := false
done;
st.sum
let () =
let path = Sys.argv.(1) in
let col = int_of_string Sys.argv.(2) in
let ic = open_in_bin path in
let buf = Bytes.create chunk in
let sum = sum_column ic col buf in
close_in ic;
Printf.printf "sum=%d\n" sum
main_noalloc.ml
And give it a go.
$ ocamlopt main_noalloc.ml -o main_noalloc
$ ./main_noalloc data.csv 1
sum=49947827758
Great! Let's see how it compares to main_alloc.ml with Hyperfine.
$ hyperfine --warmup 3 './main_alloc data.csv 1' './main_noalloc data.csv 1'
Benchmark 1: ./main_alloc data.csv 1
Time (mean ± σ): 12.626 s ± 0.564 s [User: 12.287 s, System: 0.337 s]
Range (min … max): 11.337 s … 13.356 s 10 runs
Benchmark 2: ./main_noalloc data.csv 1
Time (mean ± σ): 5.962 s ± 0.251 s [User: 5.648 s, System: 0.313 s]
Range (min … max): 5.659 s … 6.424 s 10 runs
Summary
./main_noalloc data.csv 1 ran
2.12 ± 0.13 times faster than ./main_alloc data.csv 1
Not bad!
This ability to mark call trees as not allocating is incredibly useful. It’s a great feature the Jane Street team added and I’d love to see more languages build this into their compilers.
An earlier version of this article said there were no other mainstream languages that allow you to annotate zero allocations.
weissi on Hacker News pointed out that Swift has a
@_noAllocationattribute that behaves similarly. This flag requires compiling a Swift program with the-experimental-performance-annotationsflag.And while looking into this Swift feature I also noticed that Clang has a nonallocating attribute that also behaves like OxCaml's
[@zero_alloc]. You must compile your C/C++ program with-Wfunction-effectsfor the compiler to warn you.In both cases, this is excellent news. Let's hope more languages continue down this path.↩