This post is about our use of OCaml inside LexiFi and how we managed a large refactoring of our codebase that we enterprised in order to take advantage of a new feature of OCaml.
Namely, in OCaml 4.08 a specific type floatarray for arrays of floating-point
numbers was introduced which has a number of advantages relative to the usual
type float array (see below for some details or this blog
post to learn more about the
motivation for this change).
At LexiFi we had been wanting to switch to the new type for a long time to profit from the advantages it offers, but the magnitude of the task had kept us from doing so until now. Luckily for us, Nicolas Chataing, who interned at LexiFi from Sep 2020 to Feb 2021 came along with the right skill set at the right time to tackle this task and was able to make the migration happen (thank you Nicolas!).
In the sections that follow we will recall the background for this migration, the challenges we faced and the solutions we came up with. We hope this will be useful to others who are thinking of doing the same.
Before going any further, note that the issues discussed here are only relevant
if you have a medium-to-large codebase making heavy use of float array. If
your codebase is small or if your use of float array is light, then migration
to floatarray is easy: just change the signatures in your code to use
floatarray instead of float array, and let the compiler guide you in fixing
the rest of the code.
However, for larger codebases (LexiFi’s codebase is ~600k LOC, including third-party dependencies), this manual approach is not feasible. To give an idea of the size of the migration, the final patch represented ~10k LOC, most of it produced by an automatic tool.
The tool is available online at https://github.com/nchataing/caml-migrate-floatarray.
Let us explain the “why” of the migration. For more details, head to this previous post where you will find an in-depth discussion.
Recall that at runtime every OCaml value is represented either by a word-sized integer, or as a pointer to a memory block:
│
┌─────┬──▼──┬─────┬─────┬──────┬─────┐
│ TAG │ │ │ │ ... │ │
└─────┴─────┴─────┴─────┴──────┴─────┘
In order for the GC to differentiate between the two kinds of representations, values encoded as integers have their least significant bit set to one (pointers are always aligned, so their least significant bit is zero). In particular, integer-encoded values must fit in 31 or 63 bits (depending on machine architecture). Otherwise, they must be encoded as pointers to a memory block.
So bool (1 bit), char (8 bits) and int (31/63 bits) can all be represented
as integers. On the other hand, float (64 bits) cannot, and must be
represented as pointer:
│
┌─────┬──▼────────┐
│ 253 │ 1.92e12 │
└─────┴───────────┘
An array a of length n must also be represented as a pointer:
│
┌─────┬──▼──┬─────┬─────┬──────┬───────┐
│ 0 │a.(0)│a.(1)│a.(2)│ ... │a.(n-1)│
└─────┴─────┴─────┴─────┴──────┴───────┘
Putting these two representations together, it would mean that a value of type
float array would be represented thus:
│
┌─────┬──▼──┬─────┬─────┬──────┬───────┐
│ 0 │a.(0)│a.(1)│a.(2)│ ... │a.(n-1)│
└─────┴──┬──┴──┬──┴──┬──┴──────┴──┬────┘
│ │ │ │
┌─────┬──▼─────┴──┐ │ │
│ 253 │ 1.92e12 │ │ │
└─────┴────────┬──┘ │ │
┌─────┬──▼─────┴──┐ │
│ 253 │ 1.92e12 │ │
└─────┴────────┬──┘ │
┌─────┬──▼────────┐ │
│ 253 │ 1.92e12 │ │
└─────┴───────────┘ │
┌─────┬──▼────────┐
│ 253 │ 1.92e12 │
└─────┴───────────┘
This is inefficient in terms of memory use and locality, but that is not the
only issue. For example, this memory layout does not match the C memory layout
for an array of floats. This means that interacting with C libraries (a common
use-case in numerical code) requires converting float array back-and-forth
into the C representation.
To mitigate this problem, OCaml has had an optimization in place since the
beginning of time: values of type float array are represented specially by
embedding the float values directly in the memory block:
│
┌────────┬──▼──────┬─────────┬─────────┬──────┬──────────┐
│ 254 │ 1.92e12 │ 1.92e12 │ 1.92e12 │ ... │ 1.92e12 │
└────────┴─────────┴─────────┴─────────┴──────┴──────────┘
(we refer to this representation as “flat”). This solves the most of the issues mentioned above, but introduces some problems of its own.
Every runtime operation on arrays is less efficient, as it must check to see which representation it is dealing with.
The compiler is able to optimize away this check in some cases: for example,
if it is able to determine that the type of the array elements is float or
that it is not float. But this still leaves many other cases which are not
covered: for example, polymorphic functions are never able to opitimize the
check away, or if the type of elements is abstract, etc.
Reading an element from a float array involves an allocation of a float
value, not just a memory read.
The different runtime representation can inhibit certain other optimizations (see the blog post referred to above for more on that).
floatarrayGiven these issues, an alternative that would still allow to profit from the optimized, “flat” representation for arrays of floats, without incurring in the downsides was desired.
With this objective in mind, three changes were made to the compiler:
a new type floatarray was introduced in the standard library. Values of this
type are always represented using the “flat” layout.
a module Float.Array was added to the standard library which mirrors the
Array module but that acts of values of type floatarray
a configure-time option was added to OCaml to disable the “flat float
array” mode.
Incidentally, note that code compiled with the mode disabled is incompatible with code compiled with it enabled, so if you are thinking about disabling it, make sure all your code is compiled the same way!
floatarrayThe objective for us was to migrate all or most of the code that used float array into similar code using floatarray, and to disable the “flat float
array” mode.
As a starting point, we considered the following naïve approach to carry out the migration:
replace every instance float array by floatarray (in type annotations,
type signatures, type declarations, etc).
replace every array operation (get, set, length, map, etc) on a value
of type float array by the corresponding operation on floatarray.
The expectation was that after doing these two changes, the resulting code would be close to compilable and we would be able to fix the remaining issues by hand.
One sees right away that while the the first step can be done mechanically, the
second step is harder. Indeed, most of the time, OCaml code has few or no type
annotations, and finding expressions of type float array requires typechecking
information, which has to be extracted from the compiler.
With this in mind, we refined our initial approach into an iterative workflow:
An automatic tool will handle the mechanical part of the migration:
replacing every instance of float array in type annotations, signatures or
type definitions by floatarray, replacing calls to Array module acting on
float array values by the corresponding functions Float.Array, etc. It
will do this in two passes.
An annotation pass tasked with adding type annotations to make functions
which are polymorphic on 'a array but are only ever used on float array
values effectively monomorphic. The objective is to maximize the code that can
be handled by the following pass.
A rewriting pass that rewrites code that is statically known to be acting or
returning values of type float array into similar code involving
floatarray. Note that polymorphic functions cannot be rewritten in this way,
as the resulting code would no longer be capable of being used with values of
type 'a array.
Whenever the tool cames across instances of such polymorphic functions or any other case it cannot handle, the issue will be logged into a file, together with location information that would allow to locate the issue in the codebase.
After the run, the log file is investigated manually. Once an issue is identified, changes are made to the codebase to put it into a form that will allow the tool to handle the code in question the next time around.
A key point is that the tool requires typing information to work, so it can only be run in a codebase that can be compiled. In particular, it cannot be run in the codebase that results from applying the tool itself during an intermediate stage of the migration. Indeed, this code is not typically in a compilable state due to the spots in the code that the tool is not yet able to handle.
Because of this, we organized the above iterative workflow around two copies of the codebase, “source” and “test”. The tool is always run on the “source” repo and modifies files in the “test” repo. The manual fixes done by the developer are always done in the “source” repo, in a way that the code remains compilable, so that the process can be iterated. In a diagram:
adapt "test"
┌──────────────────┐
┌──────┴─────┐ ┌─────┴─────┐ ┌─────────┐
│ run │ │ run │ │ │
│ │ │ │ │ examine │
┌──► "annotate" ├──────► "rewrite" ├───────► ├──┐
│ │ │ │ │ │ log │ │
│ │ tool │ │ tool │ │ │ │
│ └────────────┘ └───────────┘ └─────────┘ │
└───────────────────────────────────────────────────────┘
adapt "source"
As mentioned, the goal of this pass is to add type annotations to functions
which are only ever used with float array values but which have a definition
polymorphic on 'a array. By constraining their signature to float array, the
effectiveness of the following rewriting pass is increased considerably.
For example, this pass will transform the following polymorphic function
(assumed to only ever be used with x of type float array),
let get_re x i = x.(2 * i)
into
let get_re (x : float array) i = x.(2 * i)
The annotation logic also works on the return type: if the function
let create x = Array.make 10 x
is ever only used in contexts where x is of type float, then the this pass
will annotate the function as
let create x : float array = Array.make 10 x
Both this and the rewriting pass need access to the typing information
produced by the compiler. This information is available in the .cmt[i] files
produced by the compiler during compilation. These files contain a version of
the source AST of every module where each node is annotated with typing
information.
Armed with this information, the tool scans the AST looking for function calls
the signature of which is polymorphic on 'a array. It considers two cases:
If the function is local to the module (ie not exposed in the interface), all
callsites in the same module can be found and the tool can decide whether the
function is only used at float array or not. If this is the case then it
syntactically modifies the source code to add a type annotation constraining
the signature of the function.
If the function is used both at float array and at some other type 'a array, an entry is made in the log for manual examination.
If the function is exposed in the module interface, then it is not possible to
know if it is used at some other type other than float array without doing a
global analysis. To keep things simple, in this case we simply logged the
function for manual examination.
An interesting aspect of this transformation is that by “monomorphising”
functions which were polymorphic on 'a array, this pass is effectively
optimizing the code relative to the original version, independently of the
migration to floatarray. Indeed, as we mentioned in the background section, by
removing the polymorphism, the compiler would then be able to optimize away the
representation check that is required for polymorphic functions.
For the rewriting pass, the tool needed to identify expressions of type float array in the code. Again it made use of the “typed” AST stored in .cmt[i]
files to do so.
This pass performs the following syntactic rewritings:
For calls to functions such as Array.get, Array.length, Array.map, etc
the type of the argument is examined. If it is float array, then the
function call is rewritten into Float.Array.get, Float.Array.length,
Float.Array.map, etc.
For example, this pass will rewrite the function get_re from the previous
section to
let get_re x = Float.Array.get x (2 * i)
(The tool has some logic to remove the type annotations it inserts in the annotation pass during rewriting, in order to keep the code clean.)
For functions such as Array.make or Array.init, the type of the result
is examined. If it is found to be float array then the function is
rewritten into Float.Array.make, Float.Array.init, etc.
For example, this pass will rewrite the function create from the previous
section to
let create x = Float.Array.make 10 x
Sometimes the corresponding function in the Float.Array module did not
exist. For example, Array.iter2 applied to a float array in the first
argument and a (non-float) 'a array in the second argument. In this case,
the missing function was written, and the tool adapted to use it during
rewriting. After the migration was complete, a PR was submitted
upstream to add the missing
functions to the standard library.
Instances of float array appearing in type signatures and type definitions
are changed to floatarray. Note that this includes the signature of
externals. Normally this should not cause problems if the C bindings are
written using the usual Double_field and Store_double_field macros (see
here for more on this).
float array literals are rewritten. As there is no syntax for floatarray
literals, we added functions
val make1: float -> floatarray
val make2: float -> float -> floatarray
val make3: float -> float -> float -> floatarray
val make4: float -> float -> float -> float -> floatarray
that the tool uses to rewrite small literals. For maximum efficiency, these functions are defined as follows:
type float2 = {mutable x1: float; mutable x2: float}
let make2 x1 x2 : floatarray = Obj.magic {x1; x2}
taking advantange of the fact that floatarray values share the same
representation at runtime as values of record types where all the fields have
type float. For larger literals, the tool rewrites it into a call to
Float.Array.map_from_array Fun.id applied to the float array literal.
Note that previous annotation pass increases the chance that an expression may
be given the type float array, thereby increasing the chances that it could be
handled by this pass.
The previous passes still left the case of really polymorphic functions (ie functions that could not be made monomorphic by adding type annotations) unhandled. This case was treated manually. We did so in two different ways.
Duplicating code
For small library functions, we simply duplicated the code. For example, the function:
let map_sum f xs =
let len = Array.length xs in
let sum = ref 0.0 in
for i = 0 to len - 1 do
sum := !sum +. f (Array.unsafe_get xs i);
done;
!sum
would be duplicated to a specific function acting on float array:
let map_sum_float f xs =
let len = Array.length xs in
let sum = ref 0.0 in
for i = 0 to len - 1 do
sum := !sum +. f (Array.unsafe_get xs i);
done;
!sum
and all places where map_sum was called with arguments of type float array
would be changed to calls to map_sum_float. Note that map_sum_float is
identical to map_sum, which made the change easy to review. Once this no-op
change was done, the annotation/rewriting passes of the tool would detect that
map_sum_float could be made monomorphic and could then proceed to rewrite the
function to its final form:
let map_sum_float f xs =
let len = Float.Array.length xs in
let sum = ref 0.0 in
for i = 0 to len - 1 do
sum := !sum +. f (Float.Array.unsafe_get xs i);
done;
!sum
Note that it was important that map_sum_float act on float array and not
floatarray so that the “source” repo continued to be in a compilable state. As
already mentioned, this was necessary in order to run the tool, since it depends
on typing information produced during compilation.
Ad-hoc polymorphism
For more complicated functions, or functions polymorphic on complex types 'a t
that included in their structure 'a array, we opted to maintain the
polymorphism. To do so, we utilized a GADT type to express polymorphism over
both 'a array and floatarray, as follows:
type (_, _) array_kind =
| Array : ('a, 'a array) array_kind
| Floatarray : (float, floatarray) array_kind
This allowed to define array operations which could act on both 'a array and
floatarray, eg:
let get : type a arr. (a, arr) kind -> arr -> int -> a =
fun kind arr n ->
match kind with
| Array -> Array.get arr n
| Floatarray -> Float.Array.get arr n
We would then simply thread an extra argument of type (a, arr) kind through
the polymorphic function in question allowing the same code to be used both for
'a array and for floatarray.
floatarray pattern matchingThere is no syntax to pattern match on floatarray values, so instances of
pattern matching on values of type float array were flagged by the automatic
tool. As an example, consider the following function found in our codebase:
let add_float xs_ ys_ =
match xs_, ys_ with
| [| x |], ys
| ys, [| x |] ->
let zs = Array.create_float (Array.length ys) in
A.shift ys x ~dst:zs;
zs
| xs, ys ->
map2_add xs ys
We systematically rewrote this by a pattern match on the length of the arrays:
let swap_args xs ys =
if Array.length ys = 1 then ys, xs else xs, ys
let add_float xs_ ys_ =
let xs, ys = swap_args xs_ ys_ in
match Array.length xs with
| 1 ->
let x = A.unsafe_get xs 0 in
let zs = Array.create_float (Array.length ys) in
A.shift ys x ~dst:zs;
zs
| _ ->
map2_add xs_ ys_
ObjFor low-level code that uses Obj, one needs to be careful when disabling the
“flat float array” mode. Indeed, Obj.field and Obj.set_field do the right
thing for both flat and normal blocks when “flat float array” mode is enabled.
However, when the mode is disabled, these two functions assume that the argument
is a normal block, and will segfault if passed a flat block. In order to get or
set the field of a “flat” block, one must use Obj.double_field and
Obj.set_double_field.
Another instance of this problem is the idiom (Obj.magic x : float array),
used to access the elements of x, a block representing a flat block. After
disabling the “flat float array” mode, this will lead to a segfault. Instead,
the code should be changed to (Obj.magic x : floatarray).
See https://github.com/ocsigen/js_of_ocaml/pull/1083 for an example of such a
bug in the js_of_ocaml compiler, discovered during our migration.
There are three sets of macros to read and write to arrays of floating-point numbers that can be used in C bindings:
Double_flat_field and Store_double_flat_field. These macros are to be used
for floatarray values only.
Double_array_field and Store_double_array_field. These macros are to be
used for float array values only. They will work correctly regardless of
whether the “flat float array” mode is enabled or not.
Double_field and Store_double_field. These macros are for backwards
compatibility and will work correctly on both floatarray and float array
values, regardless of whether the “flat float array” mode is enabled or
not. However, since they perform check at runtime to branch according to the
representation of the array, they are a bit less efficient that the other
macros.
In pratice, LexiFi’s code did not use any of these macros. Instead, array values
are casted on the C side to (double *) and accessed directly. This amounts to
assuming a “flat” representation for these values.
This will work correctly with float array values if the “flat float array”
mode is enabled, and with floatarray values regardless of whether the mode is
anabled or not. However, this kind of code will not work with float array
values if the mode is disabled (since the representation is no longer “flat” in
that case).
In LexiFi’s codebase, all the C externals were rewritten to use floatarray
instead of float array, so this problem did not come up.
In order to increase our confidence on the migration tool, a runtime warning was
also implemented that would print a backtrace to standard output whenever a
float array was allocated by the OCaml runtime.
By running our testsuite and looking for such runtime warnings we were able to
to gain confidence that the tool was working well and not missing any instances
where float array was being used.
The runtime warning has also been submitted for inclusion upstream https://github.com/ocaml/ocaml/pull/10236.
In some places of our code, we decided to stick with float array instead of
floatarray. We did this in places that were not performance-sensitive and
where switching to floatarray would have compexified the code too much. A
feature was added to the tool to allow “whitelisting” such places so that they
would not be automatically rewritten.
Globally, we did not observe a significant efficiency increase, but disabling the “flat float array” mode of course makes some GC tag operations a bit quicker, and the simplification in runtime semantics was more than enough justification for us to consider this project.
We disabled the “flat float array” mode shortly after the migration. However,
the code obtained by the migration process could have been used with that mode
enabled as well. Indeed, the memory layout of floatarray values is the same
in both cases. The only difference is that if we had kept the mode enabled we
would not have profited from the simplification in runtime semantics.
For floatarray values, there is no built-in indexing operators similar to
the ones for usual arrays: a.(n) to read an element from an array, and
a.(n) <- x to write an element to it.
The corresponding expressions for floatarray would be Float.Array.get a n
and Float.Array.set a n x. As this notation was a bit too heavy, we defined
a pair of custom infix operators:
let (.!()) = Float.Array.get
let (.!()<-) = Float.Array.set
This makes it possible to write a.!(n), and a.!(n) <- x when dealing with
floatarray, which is almost as good as the notation for usual arrays.
We migrated LexiFi’s codebase to use floatarray in place of float array with
an iterative workflow and the help of a semi-automatic tool:
adapt "test"
┌──────────────────┐
┌──────┴─────┐ ┌─────┴─────┐ ┌─────────┐
│ run │ │ run │ │ │
│ │ │ │ │ examine │
┌──► "annotate" ├──────► "rewrite" ├───────► ├──┐
│ │ │ │ │ │ log │ │
│ │ tool │ │ tool │ │ │ │
│ └────────────┘ └───────────┘ └─────────┘ │
└───────────────────────────────────────────────────────┘
adapt "source"
The annotation pass adds type annotations to those functions which can be
guaranteed to be monorphic on float array, thus increasing the applicability
of the rewriting pass.
The rewriting pass uses type information to mechanically rewrite most
instances of uses of float array and functions operating on values of this
type to the corresponding code in terms of floatarray.
Any unhandled problem is then investigated by the tool operator, who adapts the code in question into a form that can be handled by the tool. The process is then repeated.
Additionally, a runtime warning was implemented to track allocations of float array values. By using it, we could be confident that the tool had not missed
any rewriting opportunities.
While the tool was developed with LexiFi’s needs in mind, if you have a big code base to migrate, you may be interested in checking it out, either to use it directly or to draw inspiration from it for your own solution.
If you have any questions about this blog post, you can get in touch at nicolas.ojeda.bar AT lexifi.com.
Thanks for reading!