The Flambda2 Snippets, Episode 3

---

Welcome, dear reader, to a new episode in the Flambda2 Snippets series.

Today, we will cover important and high level aspects of the algorithm of Flambda2. We will do our best to explain the fundamental design decisions pertaining to the architecture of the compiler.

At OCamlPro, the main ongoing task on the OCaml Compiler is to improve the high-level optimisation. This is something that we have been doing for quite some time now. Indeed, we are the authors behind the Flambda optimisation pass and today we would like to introduce the series of blog snippets showcasing the direct successor to it, the creatively named Flambda2.

This work was developed in collaboration with, and funded by Jane Street. Warm thanks to Mark Shinwell for shepherding the Flambda project and to Ron Minsky for his support.

Inlining in general

Given the way people write functional programs, inlining is an important part of the optimisation pipeline of such functional langages.

What we call inlining in this series is the process of duplicating some code to specialise it to a specific context.

Usually, this can be thought as copy-pasting the body of a function at its call site. A common misunderstanding is to think that the main benefit of this optimisation is to remove the cost of the function call. However with modern computer architectures, this has become less and less relevant in the last decades. The actual benefit is to use the specific context to trigger further optimisations.

Suppose we have the following option_map and double functions:

let option_map f x =
  match x with
  | None -> None
  | Some x -> Some (f x)

let double i =
  i + i

Additionally, suppose we are currently considering the following function:

let stuff () =
  option_map double (Some 21)

In this short example, inlining the option_map function would perform the following transformation:

let stuff () =
  let f = double in
  let x = Some 21 in
  match x with
  | None -> None
  | Some x -> Some (f x)

Now we can inline the double function.

let stuff () =
  let x = Some 21 in
  match x with
  | None -> None
  | Some x ->
    Some (let i = x in i + i)

As you can see, inlining alone isn't that useful of an optimisation per se. In this context, the application of the Constant Propagation optimisation simplifies it to the following:

let stuff () = Some 42

Although this is a toy example, combining small functions is a quite common pattern in functional programs. It's very useful that using combinators is not significantly than writing this function by hand.

When inlining is detrimental

We cannot just go around and inline everything, everywhere... all at once. The reason for that, as we said, inlining is mainly code duplication and that would be detrimental and blow the size of the compiled code drastically. However, there is a sweet spot to be found, between both absolute inlining and no inlining at all, but it is hard to find.

Here's an example of exploding code at inlining time:

(* val h : int -> int *)
let h n = (* Some non constant expression *)

(* val f : (int -> int) -> int -> int *)
let f g x = g (g x)

(* 4 calls to f -> 2^4 calls to h *)
let n = f (f (f (f h))) 42

Following through with the inlining process will produce a very large binary relative to its source code. This contrived example highlights potential problems that might arise in ordinary codebases in the wild, even if this one is tailored to be quite nasty for inlining: notice the exponential blowup in the number of nested calls, every additional call to f doubles the number of calls to h after inlining.

How to decide when inlining is beneficial

Most compilers use a collection of heuristics to guide them in the decision making. A good collection of heuristics is hard to both design, and fine-tune. They also can be quite specific to a programming style and unfit for other compilers to integrate. The take away is: there is no best way.

TODO: Trouver des exemples d'heuristiques d'autres langages pour l'inlining ? Pour documenter que c'est pas facile, qu'il pas de solution générique etc. Trouver une liste d'articles qui vont dans ce sens. TODO: provide numbers and links to argue about the higher cost of having a large codesize against having unspecified function calls

The following pattern is quite common in OCaml and other functional languages.

Example 1: Notice the higher-order function f:

(*
  val f :
    (condition:bool -> int -> unit) 
    -> condition:bool
    -> int
    -> unit
 *)
let f g ~condition n =
  for i = 0 to n do
    g ~condition i
  done

let g_real ~condition x =
  if condition then
    (* small operation *)
  else
    (* big piece of code *)

let condition = true

let foo x =
  f g_real ~condition x

Even for such a small example could we easily see the heuristics involved to find the right solution become quite complex.

Keeping in mind the fact that condition is always true, the best set of inlining decisions would yield the following code:

(* All the code before [foo] is kept as is, from the previous codeblock *)
let foo x = 
  for i = 0 to x do
    (* small operation *)
  done

But if condition was instead always false, instead of small operation, we would have had a big chunk of g_real duplicated again in foo, but it would have spared the cost of running only a few function calls. Thus, we would have preferred to have kept ourselves from inlining anything.

Specifically, we would have liked to have stopped from inlining g, as well as to have avoided inlining f because it would have needlessly increased the size of the code with no substantial benefit.

However, if we want to be able to take an educated decision based on the value of condition, we will have to consider the entirety of the code relevant to that choice. Indeed, if we just look at the code for f, or its call site in foo, nothing would guide us to the right decision. In order to take the right decision, we need to understand that if the ~condition parameter to the g_real function is true, then we can remove a large piece of code, namely: the else branch and the condition check as well.

But to understand that the ~condition in g_real is always true, we need to see it in the context of f in foo. This implies again that, that choice of inlining is not based on a property of g_real but rather a property of the context of its call.

There exists a very large number of combinations of such difficult situations that would each require different heuristics which would be incredibly tedious to design, implement, and maintain.

Speculative inlining

We manage to circumvent the hurdle that this decision problem represents thanks to what we call Speculative Inlining. We will not describe it in detail in this blog post. For now, all you need to know is that it requires two properties from the compiler: the ability to inline and optimise at the same time, as well as being able to backtrack inlining decisions.

In order for you to grasp at the nature of the Speculative Inlining strategy, we will run it step-by-step on Example 1.

let f g ~condition n =
  for i = 0 to n do
    g ~condition i
  done

let g_real ~condition x =
  if condition then
    (* small operation *)
  else
    (* big piece of code *)

let condition = true

let foo x =
  f g_real ~condition x

We will focus only on the traversal of the foo function.

First, we need to keep in mind that the call to f in foo is not trivially a direct call to f. The reason for that is that functions are values in OCaml, we call them general functions in this context. Indeed, one can store functions in pairs, or lists, or even hashtables, retrieve them and apply them at will. On the other hand when e want to generate a direct call or inline, we need to know the body of the function, we call a function concrete when we have knowledge of its body. This means we need Constant Propagation to associate a concrete function to general function values and consequently be able to simplify it, either with the help of direct calls, or inlining.

To add an additional layer of fun, we must remember that all OCaml functions are of arity one, which means, that what the compiler may see in practice is the following:

let foo x =
  let f1 = f in
  let f2 = f1 g_real in 
  let f3 = f2 ~condition in
  f3 x

Morever, we can only inline a concrete function if all its arguments are provided, so we need to know the actual number or arguments for that function.

let foo_bar y =
  let pair = foo, y in
  (fst pair) (snd pair)

Here, we also have to look inside the pair in order to find the function, this demonstrates that we sometimes have to do some amount of value analysis in order to proceed. It's quite common to come across such cases in OCaml programs due to the module system. Other functional languages present similar caracteristics.

Other scenarios also require a decent amount of context to be sure about which function should be called. If a function passed as parameter is called, we need to know the context of the caller functions, and sometimes up to an arbitrarily large context. This will help us know which function is being called and thus help us make educated inlining decisions. This problem is specific to functional languages, functions in good old imperative languages are seldom ambiguous; such considerations are only relevant when function pointers are involved.

Expression traversal

on peut iagiiner une version naive de cet algo. on a un prog, on choisit une applicaiton de fonction, on tente l'ininling, on fait tputes les optimisations qui en découlent, et on mesure le resultat. Si c'est bien, on garde, sinon on jette, et on boucle jusqu'à satisfaction. C'est évidemment inacceptable d'un pdv perf, donc il a fallu qu'on fasse le schoses différemment.

Une meilleure strategie c'est: d'avoir un algo qui permet de faire les choix d'inlining en meme temps que les optimisations.

Naturellement, on ne veut pas explorer tous les choix potentiels, on veut faire ça raisonablement efficacement, et ne pas rater trop de choix, donc on exploite une des forces de CAML, la rpz pure des termes. La nature de cette representation nous permet d'écrire la spéculation ainsi:

(* Pseudo-code to rpz the actual speculation *)
let try_inlining f env args =
  let inlined_version_of_f = inline f env args in
  let benefit = compare inlined_version_of_f f in
  if benefit > 0 then
    inlined_version_of_f
  else
    f

Pseudocode pour une version juste traversée:

let rec traverse env acc expr = match expr with
  | Let (var, e, body) ->
    let env, acc = update_env_acc var expr in
    traverse env acc body
  | Let_cont (k, param, handler, body) ->
    let env_for_handler = record_cont env k param in
    let acc = traverse env_for_body acc body in
    let calls = find_apply_conts acc k in
    let env, acc = populate calls env acc arg in
    traverse env acc handler
  | Apply_cont (k, arg) ->
    register env acc k arg

Pour une version avec réecriture d'expression:

let rec traverse env acc expr = match expr with
  | Let (var, e, body) ->
    let env, acc = update_env_acc var expr in
    let new_e = optimize env acc e in
    let acc, new_body = traverse env acc body in
    acc, Let (var, new_e, new_body)
  | Let_cont (k, param, handler, body) ->
    let env_for_handler = record_cont env k param in
    let acc, new_body = traverse env_for_handler acc body in
    let calls = find_apply_conts acc k in
    let env, acc = populate calls env acc arg in
    let acc, new_handler = traverse env acc handler in
    acc, Let_cont (k, arg, new_handler, new_body)
  | Apply_cont (k, arg) ->
    register env acc k arg,
    Apply_cont (k, arg)
  | Apply (f, k_return, arg) ->
    let inline_acc, new_expr = try_inline env acc f k_return arg in
    if is_better new_expr expr then
      inline_acc, new_expr
    else
      let acc = register env acc k_return Unknown in
      acc, expr

On va expliquer deux opti qui permettent de démontrer les contraintes sur l'ordere de traversée, en fdescendant la propagation de constantes, et en remontant, l'élimination de code mort.

---

• Ordre de traversée

  • Dans un langage imp, les appels de fonctions ne sont pas ambigus.
  • Dans un langage fonctionnels, c'est l'inverse. Il faut généralement ajouter pas mal de contexte pour arriver a manipuler des fonctions moins génériques.

• Simplify

• Structure des termes

• Environnement vs accumulateur (Needed for our backtracking)

• Deux trois exemples de fonctions dure à inliner pour imager les problématiques et les design decisions fondamentales dans le fonctionnement de fl2.

• la ou bcp de langag fonc qui travillent sur du lambda calcul, les decision d'inlining se font sur des terms, et il existe pas mal de littérature pour faire la résolution de constante

  ex:

1. les langage fonctionnels ont besoin d'heuristiques marantes pour inliner
2. nous on a été feignant on voulait pas faire d'heuristiques. Ex: voir au dessus. Il faudrait une combinaison d'heuristiques compliquées pour faire le bon choix ici...
3. Quel sont les bons et mauvais choix ?
4. La bonne manière pour éviter les heuristiques c'est de connaitre directement le résultat optimal final dès le début.
5. C'est là qu'on dit qu'on a décidé que le le compilatuer ferait la chose suivante:
   • essayer d'inliner,
   • mesurer si le résultat est OK
   • choisire en fonction du résultat si on inline ou non
   • Speculative inlining, tel que l'on entend ici, est, à notre connaissance, propres à FL1 et FL2
6. La technique naive serait de faire le choix de l'inlining pour une fonction, de laisser le compilateur finir, et voir le résultat à la fin, sauf que c'est pas acceptable question perf.
7. Donc il faut faire tout le reste du compilateur pendant qu'on essaye de faire l'inlining.
8. On a un choix de ce qui est important ou non pour les interactions entre les optims et l'inlining. On va en donner quelques exemples genre la propagation de constantes. L'allocation de registre par contre ca l'est moins. On va en citer trois en tout, notamment sur la base des informations nécessaires à leur conduite: En avant. Puis, en arrière avec l'élimination de deadcode et enfin l'élimination de références qui fait un point fixe.
9. La propagation de constantes c'est une passe qui a beosin de propager des informations en avant. suivant l'ordre d'execution, et non pas l'ordre syntaxique
10. L'élimination de code mort, c'est quelque chose qui remonte la trace d'execution.
    • L'élimination de référence a besoin d'informations sur les boucles. Le format d'FL2 nous force (a mentioner ou pas).
11. Conclusion logique: il faut etre capable de faire toutes ces optims pendant un essai d'inlining, pour une fonction.
12. Explication technique de à quoi ça ressemble. Environnement + accumulateur
13. Name-dropping CPS, good.
14. Name-dropping Abstract Interpretation, good too.
15. La famine, bad.

---

Au sujet d'OCamlPro :

OCamlPro développe des applications à haute valeur ajoutée depuis plus de 10 ans, en utilisant les langages les plus avancés, tels que OCaml et Rust, visant aussi bien rapidité de développement que robustesse, et en ciblant les domaines les plus exigeants (méthodes formelles, cybersécurité, systèmes distribués/blockchain, conception de DSLs). Fort de plus de 20 ingénieurs R&D, avec une expertise unique sur les langages de programmation, aussi bien théorique (plus de 80% de nos ingénieurs ont une thèse en informatique) que pratique (participation active au développement de plusieurs compilateurs open-source, prototypage de la blockchain Tezos, etc.), diversifiée (OCaml, Rust, Cobol, Python, Scilab, C/C++, etc.) et appliquée à de multiples domaines. Nous dispensons également des [formations sur mesure certifiées Qualiopi sur OCaml, Rust, et les méthodes formelles] (https://training.ocamlpro.com/) Pour nous contacter : contact@ocamlpro.com.