TRMC, reloaded

[gasche]

This PR is a reboot of @let-def's RFC #181, which implements Tail Recursion Modulo Cons (TRMC) for the OCaml compiler through an opt-in [@tailrec_mod_constr] attribute. It is a near-complete reimplementation of the original prototype, coming out of user-interface discussions with @let-def and a different implementation approach (more science, more fun!). It is joint work with @let-def, of course, and also @ksromanov who helped with testing, benchmarks and motivation.

TRMC in a hurry

Today you have to choose between writing

let rec map f = function
| [] -> []
| x :: xs ->
  let y = f x in
  y :: map f xs

which is fast, but fails on large inputs (Stack Overflow), and

let map f xs =
  let rec map f acc xs = function
  | [] -> List.rev acc
  | x :: xs -> map f (f x :: acc) xs
  in map f [] xs

which is tail-recursive, but slower as it allocates twice as many cons cells.

With this PR you can just write

let[@tailrec_mod_constr] rec map f = function
| [] -> []
| x :: xs -> f x :: map f xs

and be done with it: you have a version of map that is almost as fast as the non-tail-recursive map (and allocates exactly as much), yet is tail-recursive.

This solves a tension we have in stdlib design (the compiler stdlib goes for "fast and incorrect", most other stdlibs use sensibly more complex code to try to be tail-recursive at a smaller price than List.rev), and it can also apply to many other functions (not just on lists, any constructor application will do), including in user code. Note that TRMC is not an implicit/automatic optimization, it is an explicitly-requested transformation -- in particular, it duplicates code and changes evaluation order (the TRMC call is evaluated last).

Performance summary for List.map ((+) 1) on my machine:

• On lists of size 3, the naive-tailrec version is 55% slower than the non-tailrec version, TRMC is only 12% slower.

• The TRMC and stdlib start having indinstinguishable performance starting at size 15 or so, until size 3000 or so.
  TRMC starts again being 10% slower on larger lists (5000 or more) due to promotion effects (much more on this below), but this is also the range at which many systems will start to Stack Overflow on the non-tailrec implementation.
  At size 10_000, the naive tail-recursive version is 63% slower than the TRMC version.

Alternative stdlibs also perform manual unrolling of List.map, which is an independent/orthogonal optimization: if you TRMC-transform an unrolled version, you get an implementation that is both faster and noticeably simpler than the List.map of Base and Containers (of course, complexity moved into the compiler).

Plan for the rest of this document:

• Summary of the PR discussion so far.
• What TRMC does precisely.
• Design choices in the current PR.
• PR structure (for reviewing).
• Remaining TODOs.
• Benchmarking details.

Summary of the PR discussion so far

• We discussed the naming of the attribute; the original proposal
  [@trmc] is judged too obscure. The consensus was to go for
  [@tailrec_mod_constr] instead.

What TRMC does precisely

If you consider a simple functional language with function application and pattern-matching, a function call is in tail-call position with respect to an expression if occurs within a hole of the following family of "tail contexts" (those that correspond to an empty continuation):

• □
• let x = t in □
• match t with ... | p → □ | ...

The TRMC code transformation adds the following contexts to the grammar of tail-position contexts, where K is a datatype constructor:

• K(..., □, ...)

More precisely, TRMC is a destination-passing-style translation such that contexts in this extended family in the source program are in tail-recursive position in the transformed program.

For example:

let[@tailrec_mod_constr] rec map f = function
| [] -> []
| x :: xs -> f x :: map f xs

is transformed roughly as follows. This is not well-typed OCaml code (we use in-place mutation on a constructor block), the transformation actually happens on the Lambda intermediate representation rather at source level.

let rec map f = function
| [] -> []
| x :: xs ->
  let block = f x ::{mutable} Uninitialized in
  map_trmc block 1 f xs;
  block

and map_trmc dst offset f = function
| [] -> dst.(offset) <-{initialize} []
| x :: xs ->
  let block = f x ::{mutable} Unitialized in
  dst.(offset) <-{initialize} block;
  map_trmc block 1 f xs

A call to map f xs will now perform exactly one non-tail-recursive call (the first call to map_trmc in the map body), and then all following recursive calls are in tail-position.

The proper use of "mutable" blocks and "initializing" writes in the generated code makes the TRMC transformation perfectly compatible with Flambda (and Flambda 2, 3, etc.) and Multicore.

Note: the "destination" of the TRMC call is a couple (dst, offset) of a mutable block and an offset into the block, passed as separate arguments. In #181, @let-def implemented a version where offset is not a dynamic argument: one DPS version of the function is generated for each static offset used in the code. This static-offset version is more complex, generates more code, cannot be extended higher-order or cross-module TRMC (we don't support it but we could), and our performance testing showed that it is not noticeably faster.

Design choices in the current PR

[@tailrec_mod_constr] is opt-in: only functions that are explicitly marked are transformed (this is good because, as you can see in the example, the transformation duplicates code).

More generally, @let-def and myself wanted to have something very explicit, where programmers have to indicate their intent if something is undetermined or subtle.

The current PR warns if a [@tailrec_mod_constr]-annotated function does not actually benefit from TRMC (no additional tail-calls). It warns in two additional situations.

If several calls under a constructor could be optimized, the implementation will not make an arbitrary choice but request that the user be explicit, either marking the call to be optimized with [@tailcall], or the other calls with [@tailcall false]. For example, the first version below is rejected with an error, while the three following versions are accepted:

let rec map f = function
| Leaf v -> Leaf (f v)
| Node (left, right) -> Node (map f left, map f right)

let rec map f = function
| Leaf v -> Leaf (f v)
| Node (left, right) -> Node (map f left, (map [@tailcall]) f right)

let rec map f = function
| Leaf v -> Leaf (f v)
| Node (left, right) -> Node ((map [@tailcall false]) f left, map f right)

let rec map f = function
| Leaf v -> Leaf (f v)
| Node (left, right) ->
  let left = map f left in
  Node (left, map f right)

If a function call appears in tail-call position in the source definition, but the function itself is not marked [@tailrec_mod_constr] (it does not expose a destination-passing-style version), the TRMC translation will not preserve tail-callness. In this case a warning is emitted, the user has to use [@tailcall false] to explicit mark this intent, or mark the called function with [@tailrec_mod_constr] as well.

TRMC only applies to let-rec blocks (not simple let-bindings), and calls to mutually-recursive functions can be optimized (if all relevant functions are marked [@tailrec_mod_constr]). On the other hand, this is a first-order transformation: uses of the function that are not direct applications are not transformed. In particular, a [@tailrec_mod_constr]-transformed function may still have non-constant stack usage, due to direct self-calls that are not in TRMC-tailcall position, higher-order usage, or calling another function that consumes stack space. (We discussed emitting more warnings if we cannot guarantee that stack usage is constant, but this is mildly annoying to implement and also irritating to use, so I do not plan to implement it for now.)

Finally, in this PR, the "scope" of the TRMC declaration on a definition is local to its block of mutually-recursive definitions: when you write

let rec f1 = e1 and ... and[@tailrec_mod_constr] fi = ei and ... in body

then only calls to fi in one of the e1, ... may be turned into destination-passing-style calls. In particular, body and the rest of the program only call the direct version of fi, even under data constructors. (In particular, there is no point in cross-module TRMC information.) Note that, except for higher-order use-cases, such "outer" calls will only call the function a constant number of time, so optimizing them would not change call-stack-usage asymptotics.

PR structure (for reviewing)

The PR is designed to be reviewed commit-by-commit, starting at the commit "preliminary implementation of TRMC". The previous commits correspond to merges of buildup PRs that are to reviewed independently, or testing code that we left for now but need not be part of the final implementation.

The current PR implements a very compositional approach to TRMC: the translation of a term depends only the translation of its subterms. The first versions are very simple, but produce disappointing code when the source program nests data constructor, for example:

let[@tailrec_mod_constr] rec dup = function
| [] -> []
| x :: xs -> x :: x :: dup xs

(Nested TRMC constructors are very common when in recursive traversals of ASTs, for desugaring for example, and nesting depth can be relatively high in this case.)

Further commits add more sophistication which result in nice code being produced. The final result is more complex than #181 but more principled (it optimizes a more general class of functions; #181 only supports a tail-context followed by constructor applications, instead of arbitrary composition of tail-contexts and constructors; not that the difference matters terribly in practice), and is free of any quadratic compile-time behavior in the implementation.

Finally, the end of the PR implements the various errors and warnings, which are also a source of implementation complexity.

Remaining TODOs

The PR is already ready for review, and starting review sooner rather than later would be very helpful. But some aspects are not complete yet:

• The PR does not yet include documentation in the manual. (blocker)

• The PR does not change standard library functions to use [@tailrec_mod_constr] attributes. (It can come as a later PR.)

• For some [@tailrec_mod_constr] functions, only the destination-style version is used, not the direct-style version. We make no attempt to avoid generating the direct-style version in this case (this is something that RFC: tail recursion modulo constructors #181 did). (No immediate plans to do it.)

• One could decide to rewrite TRMC calls in a more global way, outside the mutually-recursive block; to do this in a modular way, one should use a [@tailrec_mod_constr] attribute on function declarations. (No immediate plans to do it.)

• One could consider a higher-order TRMC transformation, where a function argument can be marked [@tailrec_mod_constr] to indicates that it supports the destination-passing-style protocol (this would be desugared into two arguments, for the direct-style and DPS versions). (No plans to do it.)

• The TRMC transformation breaks usages of non-linear continuations; TRMC-annotated code will change semantics when using the delimcc library in a way that duplicates continuations, in particular the HANSEI probabilistic-programming library.

  The approach we recommend to solve this problem is to implement a fallback mode for TRMC that uses (a restricted form of) a general continuation-passing-style transformation, instead of destination-passing-style with mutation. (No immediate plans to do it.)

Benchmarking details

Note: all performance measurements were produced with a version on this PR rebased on the 4.10 release branch, which may be of interest to others for testing purposes:

https://github.com/gasche/ocaml/tree/4.10-new-trmc

Here are good-looking performance graphs produced by the faster-map benchmarks ( https://github.com/aantron/faster-map/ , adapted by @ksromanov as https://github.com/ksromanov/faster-map/ ) for various graph functions, including TRMC maps. In these tests, the result of List.map lives long after the call (more details on this below).

All graphs are shown relative to the naive tail-recursion line, the flat gray at 100%. There are two TRMC functions:

• The "trmc map" (orange) which is the example above, which should be compared to the "stdlib map" which is the non-recursive version (using ulimit -s unlimited to test on larger stack size), and Batteries' List.map which use a simple implementation in destination-passing-style (using Obj.magic).

• The "trmc unrolled 4x" (yellow) which is the trmc version of a map function unrolled four times. It competes with "stdlib unrolled 5x" (I only did 4x as I was too bored to unroll a last time), and the Base and Containers map that use unrolling (plus some extra logic to start being tail-recursive after some cutoff).

You can see that "trmc" has mostly the same performance as (non-tail-recursive) List.map, except on very small lists (they get even at n=17), where List.map is somewhat faster, and on very large lists, where trmc is much faster.

The "trmc unrolled 4x" version is faster than all other implementations, which is nice as its implementation is sensibly simpler than the competition. Note: we are not specifically advertising for unrolling, but demonstrating that applying this independent optimization makes trmc-map competitive with highly-tuned implementations.

Note: these graphs correspond to benchmark runs without Flambda, but I also ran benchmarks with Flambda and the result look very similar -- no qualitative difference.

Details

On very small lists, I believe that the difference comes from the cost of the write, which calls the non-GC C function 'caml_initialize'; trmc map has a noticeably higher instruction count than standard map, which is in large part compensated by a higher instruction-per-cycles ratio.

On very large lists, the call stack of the standard List.map trashes the cache, traversing the a large portion of memory twice. This explains why the performance get closer to the naive tail-recursive version (which has the same behavior, except it traverses a large heap region rather than a large stack region).

Batteries' implementation (also in destination-passing style) has the same performance as trmc map on large enough inputs. It is noticeably slower on very small inputs, because it always allocates a word as the initial location. Our implementation of TRMC does not allocate on empty lists, at the cost of code duplication. (We could change easily this to use Batteries' approach of always allocating an initial word.)

Promotion

When the mapped function has very little compute and no allocation (here ((+) 1)), the performance of List.map is dominated by allocation costs and GC behaviors, and locality effects on very large inputs.

The promotion behavior of a micro-benchmark depends on whether the result of List.map is long-lived or short-lived. When it is very short-lived (typically, in benchmarks where it is computed and immediately ignored), the numbers are less good for TRMC in a "middle regime" where input lists are close to the minor-heap size (between 10^4 and 10^5 elements on my machine).

This comes from a subtle difference in the way the output is constructed. With a non-tail-recursive implementation, the output list is constructed from the end to the beginning. If a minor GC happens in the middle of the computation, the suffix of the output that is already built is promoted, but the prefix remains in the minor heap. With a TRMC implementation, the output list is constructed from the beginning to the end (by mutation). If a minor GC happens in the middle, the prefix of the output is promoted; the next write creates a pointer from the major heap to the minor heap, which results in the whole output being eventually promoted to the major heap.

For example, if the amount of allocation in one List.map call exactly fills the minor heap, a minor collection will happen at some arbitrary point (depending on what happened before in the program) during the List.map call. On average it happens in the middle of the list; with the non-tail-recursive version, half the output will get promoted; with the TRMC version, the full output will be promoted, so we see 100% more promotion to the major heap, making the TRMC version noticeably slower -- I call this the "promotion effect". Of course, if the output is used afterwards, and another collection happens before it becomes dead, the whole output gets promoted in both cases and there is no difference.

Discaring the result of List.map right after it was computed happens exactly never in real-world scenarios, but almost always is benchmarking code, so you will see this effect in microbenchmarks. One scenario in real-world code that is somewhat similar is data pipelines where the output of List.map is only traversed once -- it will exhibit those promotion effects.

Here is a graph of the same benchmark where the result of List.map is ignored, instead of being long-lived:

There are three different performance regimes:

• on small lists (smaller than minor-heap size, up to 10^3 here), promotion almost never happens so promotion effects are neglectible.

• on lists of the order of the minor-heap size (10^4 to 10^5 here), promotion effects can result in 100% more promotion, so TRMC versions are noticeably slower than their non-tail-recursive counterparts

• on lists much larger than the minor-heap size (10^6), promotion happens K times per List.map call, so only 1/K of the list cells remain in the minor heap in the non-tail-recursive version: promotion effects become neglectible again.

Promotion effects can be observed in the core_bench output. For example, on lists of size 1000, when the List.map result is ignored, we see

With lists of size 1000
  Name                   Time/Run   mWd/Run   mjWd/Run   Prom/Run   Percentage  
 ---------------------- ---------- --------- ---------- ---------- ------------ 
  naive-tail-recursive    10.07us    6.01kw     51.63w     51.63w       99.78%  
  base                     6.43us    3.01kw     17.01w     17.01w       63.71%  
  batteries                7.50us    3.01kw     34.53w     34.53w       74.29%  
  containers               6.66us    3.00kw     17.10w     17.10w       65.95%  
  trmc                     8.07us    3.00kw     34.23w     34.23w       79.99%  
  trmc_unrolled4           5.93us    3.00kw     34.02w     34.02w       58.71%  
  stdlib                   7.96us    3.00kw     17.13w     17.13w       78.88%  
  unrolled5                6.45us    3.00kw     16.96w     16.96w       63.92%  

We see that the naive tail-recursive implementation allocates twice as much as all other solutions. TRMC solutions allocate as much as non-TRMC solutions, but they promote twice as much to the major heap on average -- this is the promotion effect. On the other hand, the amount of promotion is very small (34 words promoted for 3000 words allocated), so the performance impact is neglectible; in particular, "trmc" and "stdlib" have the same performance.

With lists of size around 5.10^5, promotion happens on most function calls (but at most once per call), so promotion effects are at their strongest (when the result is short-lived):

With lists of size 56234
  Name                   Time/Run    mWd/Run   mjWd/Run   Prom/Run   Percentage  
 ---------------------- ---------- ---------- ---------- ---------- ------------ 
  naive-tail-recursive     3.94ms   337.41kw   159.60kw   159.60kw       99.90%  
  base                     3.64ms   322.38kw   144.33kw   144.33kw       92.34%  
  batteries                2.69ms   168.71kw   109.83kw   109.83kw       68.24%  
  containers               2.22ms   244.15kw    87.07kw    87.07kw       56.15%  
  trmc                     2.73ms   168.71kw   110.52kw   110.52kw       69.16%  
  trmc_unrolled4           2.64ms   168.71kw   110.24kw   110.24kw       66.93%  
  stdlib                   2.06ms   168.71kw    54.77kw    54.77kw       52.32%  
  unrolled5                1.68ms   168.71kw    55.50kw    55.50kw       42.66%  

The TRMC versions still promote twice as much as the non-tail-recursive versions, but now a large portion of the allocated cons cells end up promoted, so the performance impact is large: indeed, trmc (2.73ms) is noticeably slower than naive List.map (2.06ms) at this size. Notice that optimized implementations (base, containers) are out of their "short lists" regime, so they start behaving less like List.map and more like manual tail-recursive versions.

Finally, for very large lists, most allocated cells are promoted for all versions, so the performance impact becomes small again, dominated by locality effects.

With lists of size 1000000
  Name                   Time/Run   mWd/Run   mjWd/Run   Prom/Run   Percentage  
 ---------------------- ---------- --------- ---------- ---------- ------------ 
  naive-tail-recursive   132.50ms    6.00Mw     5.79Mw     5.79Mw       95.69%  
  base                   135.66ms    5.98Mw     5.74Mw     5.74Mw       97.97%  
  batteries               69.11ms    3.00Mw     3.01Mw     3.01Mw       49.91%  
  containers              88.12ms    4.44Mw     4.27Mw     4.27Mw       63.64%  
  trmc                    71.23ms    3.00Mw     3.00Mw     3.00Mw       51.44%  
  trmc_unrolled4          68.02ms    3.00Mw     3.01Mw     3.01Mw       49.12%  
  stdlib                 138.47ms    3.00Mw     2.85Mw     2.85Mw      100.00%  
  unrolled5               86.85ms    3.00Mw     2.85Mw     2.85Mw       62.72%  

[lpw25]

This all looks great, but can I request that the bike-shed not be luminous pink: please can we find a less cryptic annotation than [@trmc].

[gasche]

What's a better name? Users have to know of the feature to use it, and know when it applies (the structure of optimizable functions). "Tail-recursion modulo constructor" is the canonical name for this feature. Personally I think that @trmc is fine.

With @let-def we discussed having a generic attribute [@tailcall] on declarations that would perform a CPS transformation to make the function tail-recursive, or use the more efficient TRMC transformation when we are in the optimizable subset. This could make sense, but we haven't implemented the general CPS transformation (and it is not completely clear how to specify first-order CPS and whether it's a really good idea to restrict oneself to first-order uses), so we can't do this right now.

[lpw25]

Anything which uses words would be fine. Even just [@tailrec_mod_cons] would be an improvement. It would be a bit easier if the annotation ensured some property of the function rather than just transforming it. For example, if it checked that all recursive calls were in tail-position, optimised, or marked [@tailcall false] then you could simply use [@tailrec] since the annotation would ensure that the function was tail recursive.

[xavierleroy]

Lovely bikeshedding discussion, I feel compelled to participate!

@TRMC would be better than @trmc because it's an acronym and it is pronounced letter by letter (Tee-Are-Em-Cee).

A Web search for "TRMC" doesn't find "tail-recursion modulo cons". The best it finds is https://acronyms.thefreedictionary.com/TRMC (in English) or "taux de réjection du mode commun" (in French).

@tail_rec_mod_cons is a bit verbose, but at least it is clear that it's about tail recursion optimization.

[gasche]

Ok, let's go for @tailrec_mod_cons.

Note: Leo's [@tailrec] suggestion is something we discussed with @let-def (and I think in #181 ?). I mentioned it briefly above, it works well if we consider TRMC as an implicit transformation that is obtained to obtain "fully tail-recursive functions" (whatever that means) in certain cases where functions are not naturally so. This is a distinct feature from having an explicit attribute to say "I know what TRMC is and I want it here". This PR implements the "explicit transformation" approach, not the "implicit optimization" approach. I think that we can provide both features, but that it is important to consider the two separately (the implicit way, in a separate PR!). In particular:

1. [@tailrec] only makes sense for a subset of the functions we accept to TRMC (see the examples above for map on binary trees).
2. the [@trmc] transformation changes evaluation order, so seeing it as a more neutral optimization is going to face design challenges.

[dbuenzli]

Two other non-serious proposal: @jump_the_cons and @guarded_tailrec (by analogy with guarded variables on the right hand-size of lazy values).

[hcarty]

I was going to suggest @tailrec.mod_cons since . is used elsewhere in attribute names but tailrec_mod_cons might be more search engine friendly when an unsuspecting user comes across this attribute in the wild.

[gasche]

I renamed the attribute to [@tailrec_mod_constr] in the PR description, and I am in the process of updating the implementation as well. I added a section "Summary of the PR discussion" to the original post that I will update to keep a synthesis of the PR discussion.

(I like "constr" better than "cons" to avoid the misunderstanding that this feature is limited to list "cons" cells, which was the original idea when the feature was proposed in LISP implementations.)

[Elarnon]

@gasche suggested I look at this PR as a first foray into the compiler's internal, which I am doing now. As stated I am not familiar with the compiler, so I might make remarks which should be obvious or show a lack of understanding, please be forgiving :)

Here are some general remarks after reading only the description of the PR:

• This is a well-written PR with lots of explanations and context, very pleasant to read 👍
• In the documentation I would remind the reader that, if TRMC changes the evaluation order, said order was unspecified to begin with and should not be relied upon
• I will remark that (unless I am mistaken) the Batteries-style version with an extra allocation can be implemented today (and would synergise well with the ability to not generate unused direct versions, especially re: compile times):

type 'a box = Box of 'a

let map =
  let rec box_map f xs =
    Box (map f xs)
  and[@tailcall_mod_constr] map f = function
  | [] -> []
  | x :: xs -> f x :: map f xs
  in fun f xs -> match box_map f xs with Box ys -> ys

It might be worth mentioning in the documentation as a trick to avoid code duplication, if that is ever implemented.

• The PR talks multiple times about unrolling and the boredom of doing so manually. I will admit this is something I am incredibly reluctant to do myself, because it is tedious and leads to code that is hard to maintain, but is sometimes necessary for performance. This is a bit of a tangent and not directly related to the PR but it makes me think about the usefulness of an [@unroll n] annotation, to unroll recursive calls (inline up to n nested calls but only inside the defining let rec block). Does something like this exist?

[Elarnon]

Oh, is this used to define functions local to the module which will not be exported even though there is no .mli? If so that is a neat trick I haven't seen used yet.

[gasche]

Yes, it's a recent idiom given that open M didn't accept struct .. end module expressions before 4.08.

[gasche]

Thanks for the comments!

In the documentation I would remind the reader that, if TRMC changes the evaluation order, said order was unspecified to begin with and should not be relied upon.

I had scenarios in mind such as K1(left1, K2(left2, trmc_call e, right2), right1), where I had the impression that we could get the evaluation order left1; left2; right2; right1; trmc_call e in the transformed program -- note that this order is invalid for the source program according to my understanding of what is specified of the evaluation order (each argument is evaluated entirely, in some order). But in fact I think that the implementation is careful to always evaluate (left1, right1) first (in some unspecified order) before (left2, right2) and then trmc_call e, so the order we can observe is a valid order for the original source program.

the Batteries-style version with an extra allocation can be implemented today

Yes, but I think that it is a bit too much sophistication for the user to write this version "just" to avoid code duplication. We could of course also use this strategy in the code generator, which eliminates the code duplication, but at the cost of an extra allocation in the empty-list case, which we thought could be fairly controversial.

I don't expect code duplication to be an issue in practice, given that TRMC is opt-in: you would only convert a handful of functions in your program, so code blowup is extremely limited. And if you compare to the code people write today to get fast tail-recursion equivalent, we probably generate less code.

(One aspect that I find mildly irritating is the idea of nested TRMC definitions leading to 2^n versions of the inner functions laying around. I think that in many case the "usage analysis" I mentioned would remove the direct version completely, removing the exponential aspect of TRMC nesting. But then again, nested TRMC annotations would be fairly rare in actual code.)

This is a bit of a tangent and not directly related to the PR but it makes me think about the usefulness of an [@unroll n] annotation, to unroll recursive calls (inline up to n nested calls but only inside the defining let rec block). Does something like this exist?

I believe that flambda does some unrolling. This could also be exposed to users as a (predictable) manually-controlled transformation, as we do here for TRMC. I prefer to keep it out of the scope of the current PR, because it's easy for discussions to be side-tracked by those tangents.

[gasche]

Note: The PR description mentions that the scope of TRMC is "local" in the sense that only calls in mutually-recursive definition bodies are transformed. This is in fact, I realized afterwards, not correct with the current implementation, and I am not sure that we want to keep this locality guarantee. I like the simple reasoning model it offers, and I used to think that it would never worsen stack-consumption complexity (if you neglect potential-TRMC calls that are outside the recursive body, they will not be part of the recursive loop so they only occur a constant number of time); but in fact this is not true in presence of nested TRMC defintions:

let[@tailrec_mod_constr] rec flatten xss = function
| [] -> []
| xs :: xss ->
  let[@tailrec_mod_constr] rec append_flatten xs xss =
    match xs with
    | [] -> flatten xss
    | x :: xs -> x :: append_flatten xs xss
  in append_flatten xs xss

In this example, the implementation will TRMC-optimize the call to append_flatten on the last line, while it is outside the "local" scope mentioned from the point of view of let rec append_flatten. This transformation is necessary to get constant-stack behavior for flatten.

The function could be rewritten this way by the user, and then the local-scope discipline would suffice, but I find the code less pleasant:

let[@tailrec_mod_constr] rec flatten xss = function
| [] -> []
| xs :: xss -> append_flatten xs xss

and[@tailrec_mod_constr] append_flatten xs xss =
    match xs with
    | [] -> flatten xss
    | x :: xs -> x :: append_flatten xs xss

I'm not completely sure whether I want to stick to the local-scope property, give up and implement a global (but not cross-module) transformation, or do a compromise where the scope is "global within @TRMC definition themselves but not outside" (meh).

[Elarnon]

This would be a good practice.

[Elarnon]

It looks like by reusing the ctx from traverse_letrec in traverse ctx body you allow performing destination-passing-style calls in the body of the let rec block. The text of the PR claims this is forbidden, so one of the code, the PR, or my understanding must be wrong.

[Elarnon]

Actually this does not allow destination-passing-style calls in the body of the let rec block (because traverse is called, not traverse_binding) but by reusing the context I believe in the following code:

let[@tailcall_mod_constr] rec f x = fdef in
let[@tailcall_mod_constr] rec g y = gdef in
body

you allow destination-passing-style calls to f inside gdef, which is still a violation of the "only in the defining let rec bindings".

[Elarnon]

Note: traverse is actually called for code outside tail position, so this piece of code is responsible of destination-passing-style calls to f inside gdef if the example code is not in tail position. Destination-passing-style calls to f inside both gdef and body are allowed by the similar behavior in choice I also pointed out.

See also this note from @gasche, which questions whether the locality specification is actually what we want, or if the behavior of the current implementation is better (if harder to describe).

[gasche]

Yes, currently I believe the description of the PR does not match the implementation, which implements the "global" scope by propagating the transformation context along the lexical scope.

I am thinking of relaxing the "local scope" specification into an "attribute scope" specification: a TRMC attribute scope over the recursive definitions, and calls are rewritten if they are within the scope of at least one TRMC attribute. This allows the flatten example, but it still provides a principled explanation for a strategy that does not require cross-module information for TRMC.

Once this aspect has been decided (I am starting to think that the "attribute scope" proposal is a good compromise, but other people may chime in with suggestions) is will push a separate commit to implement the consensual scoping strategy.

[Elarnon]

I don't know what is the local opinion on that sort of thing, but I would justify by a comment (if not an assertion), either here or at the definition of traverse_binding, why this find can not fail (because all the bindings which are traversed must have been declared beforehand).

[gasche]

Good point, will do.

[lthls]

This is a bit of a tangent and not directly related to the PR but it makes me think about the usefulness of an [@unroll n] annotation, to unroll recursive calls (inline up to n nested calls but only inside the defining let rec block). Does something like this exist?

I believe that flambda does some unrolling. This could also be exposed to users as a (predictable) manually-controlled transformation, as we do here for TRMC. I prefer to keep it out of the scope of the current PR, because it's easy for discussions to be side-tracked by those tangents.

The [@unrolled n] attribute is documented in section 21.6 of the manual. It is specific to flambda.

[Elarnon]

This warning is a bit surprising, because append is actually marked @tailrec_mod_constr. I would add a remark that both functions also need to be in the same TRMC "scope", whatever that ends up meaning.

[gasche]

Actually I think that seeing this warning is an artifact of using toplevel-style tests (with the current implementation, if the whole file was processed at once there would be no warning).

The actual warning we want to see is that this [@tailrec_mod_constr] definition is useless and wrong since the call to flatten itself is not TRMC-optimized. But in fact the current implementation does not provide this warning (in the case where the call to append itself is correctly TRMC-ed), as it only expects calls in TRMC position.

I don't know whether I want to have a warning on non-TRMC-position self-and-mutual calls that are not marked [@tailcall false] (this sounds potentially useful in some cases, including this one, but also annoying in other cases and also annoying to implement), or accept that we don't give a very useful warning in this example.

[lpw25]

On the design side: I think that it is worth considering the variation of this proposal that only extend the "tail contexts" with:

K(v, v, v... , □, ...)

i.e. the tail call must only be followed by values, not arbitrary expressions. The benefit here is that the transformation does not alter the evaluation order. Obviously the expressivity is the same, it just makes the user have their source code match the actual evaluation order.

[gasche]

But why would we require using A-normal forms for these specific definitions and not other parts of the source programs?

I think that the evaluation order resulting from the TRMC transformation is easy to specify, and to understand for users who already understand when TRMC can be used: in a constructor application, at most one argument is in TRMC position, and this argument will always be evaluated last.

[gasche]

(answering some inline comments)

[lthls]

But why would we require using A-normal forms for these specific definitions and not other parts of the source programs?

This is not exactly A-normal form, as we need the last block to contain a function call that we can turn into a tailrec call.
More importantly, the reason we can require a specific form here is that this optimisation only works if the recursive call is the last thing you evaluate, so it seems reasonable to ask programmers to make sure that it is their intent rather than magically transform their program.
This would also solve the problem of multiple recursive calls: you don't need a special warning for cases where you don't know which call should be turned into a TRMC call, because you've already required that the programmer makes it explicit.

Since you're already planning to make this feature opt-in, it won't be a major inconvenience if programmers have to change their code slightly to make sure the optimisation triggers. And you won't get strange bugs due to changes in the evaluation order if you start trying to apply it automatically either.

[gasche]

@lpw25 and @lthls : I'm willing to go ahead and implement your suggestion, but it's some amount of work and I expect that you will feel a moral duty to consider reviewing the PR in return :-)

[gasche]

I am only somewhat convinced by the proposal; I don't like the fact of making harder to convert a function to TRMC in the common case (where the code is robust to the unspecified order of argument evaluation), but I do agree that it is an improvement in that it removes the need for one of the warnings I implemented. One idea is to go look at the compiler codebase and see whether some interesting functions can be marked TRMC, to have more examples in mind before making these kind of user-interface decision.

On the other hand, I am starting to be convinced that implementing the analysis of which variants of each function (direct or TRMC) is actually used in the code would be useful. This gives the opportunity to not include unused variants in the generated code (eliminating code duplication in some cases), but it could also be useful to provide warnings (warn if the TRMC variant of a specialized function is never used).

[gasche]

Thanks @ksromanov for your eagle eyes again. I fixed the leftover 'trmc' -- and removed the force-trmc commit.

This PR is tickling my sense of "this could be more perfect": it was a lot of work to get it to this point, but I can't help feeling that it could be improved in various ways¹, and I've been hesitant to inflict it onto the world in its current state. However, it is ready by our usual standards, and I doubt that it would be able to improve without feedback from actual users, so I think it does need to get released.

I'm planning to wait for another round of CI, give a last opportunity for feedback, and merge late today or tomorrow if there are no objections.

¹ Some potential improvements:

1. With @Elarnon we discussed changing the code-generation strategy slightly to pass 3 parameters instead of 2, with the benefit that calls from non-DPS functions to DPS functions would then preserve tail-call-ness (currently they don't). This would be mostly invisible to users (now that there is not practical limit on tail-call parameter count anymore).
2. I think that the attribute-based control of the feature could still be refined, but this probably needs experience in real codebases before we decide to go one way or another.
3. There is no "force-tmc" mode to apply the transformation whenever beneficial (what that means is not completely clear), and implementing it well (I did a quick attempt yesterday) would require some refactoring in the treatment of warnings and errors.
4. Some warnings are produced after the TMC transformation, and will thus occur twice as TMC can duplicate code. I've only seen this happen for @tailcall-related warnings, and this is a minor UI glitch, but it's not obvious how to fix it.
5. The "driver" logic could get better at not generating the direct-style or DPS version of a function when it is not used, to reduce issues with code duplication in nested specialized let-bindings.

I believe that none of this is a problem for many beneficial uses of the transformation in practice, but some of those will show up as limitations if it gets used more heavily.

Note: I experimented with using the attribute in the compiler codebase (for example in CSEgen; it works), but this was not included in the present PR as it would require a bootstrap, which we prefer to avoid.

[dbuenzli]

@gasche congratulations ! Do you have an idea of which functions of List automagically become T.R. ? Ready to kill some List.{rev,rev_map} applications and rewire the brain for not using them.

[pmetzger]

This is cool enough to be worth a public announcement on Discuss, IMHO.

[hhugo]

@gasche, Have you tried to test this with jsoo by any chance ?

[gasche]

@pmetzger I have a long TODO list of urgent stuff, but I'll try to get to that, thanks for the suggestion.

@hhugo: no, because I didn't expect particular issues with it. (But then I'm not sure what assumptions jsoo is making about the incoming bytecode.) It's "almost" a source-to-source transformation, "almost" is that it creates mutable blocks that inhabit immutable types (eg. y ::{mutable} Hole that later gets written to to become a valid list cell), and uses the Psetfieldcomputed operation on these blocks. This is close in spirit to what some OCaml libraries are already doing with Obj tricks (Batteries and Extlib have versions of this), that I suppose jsoo already supports.
If you let me know of any issue, I would of course assist in finding a solution that works for jsoo as well.

[gasche]

@dbuenzli for now, we are not using the annotation in the stdlib or the compiler codebase. I think that we could eventually, but after the 5.0 transition. In the meantime, I would think of the feature as more experimental (it exists to be used, and we welcome feedback, but its usage may remain rare for a while.)

[xavierleroy]

@dbuenzli for now, we are not using the annotation in the stdlib or the compiler codebase. I think that we could eventually, but after the 5.0 transition.

Wise words. But if some brave soul wanted to experiment with using the "tmc" annotation in the standard library or the compiler code base, we'd be very interested in the results.

[gasche]

For the record, I tried to look at whether TMC is still useful with the Multicore runtime, where OCaml calls do not use the system stack, and which could thus free itself of the dreaded Stack_Overflow error. My experiments show that currently the Multicore runtime becomes very slow if the call stack grows very large; see #10947.

[hhugo]

TMC is still useful for jsoo which rely on the (small) Js stack

[yminsky]

This is a big deal for us too. We basically wouldn't be able to get rid of our carefully written, stack-avoiding list manipulation functions in the absence of TRMC, because those functions would then cease to function reasonably when compiling to Javascript.