Skip to content

Tail Recursion Optimization #218

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 62 commits into from
Jun 20, 2024
Merged

Tail Recursion Optimization #218

merged 62 commits into from
Jun 20, 2024

Conversation

CAG2Mark
Copy link
Contributor

@CAG2Mark CAG2Mark commented Mar 29, 2024

Adds a pass to the IR which optimizes mutually tail-recursive and mutually tail-recursive modulo cons functions.

The full pass first optimizes modulo cons functions to be tail recursive, then optimizes those functions to be tail recursive.

This PR consists of the following features

IR Field Assignment

Adds an assignment primitive to the IR. Given x: CtorApp and val: TrvialExpr, we now have the primitive

x.field = val in ...

Also modified the interpreter to no longer assume references to constructors are referentially transparent.

Tail Recursion Modulo Cons

Extends the idea in this paper to optimize strongly connected components of tail recursive functions, of which some calls may be within a constructor (guarded recursion).

Supports both modulo-cons calls and regular tail calls within the same strongly connected component.

Rewrite mutually tail recursive functions

Suppose we have functions:

def f(n) = e1
def g(m) = e2

where e1 and e2 tail-call f or g. This function will be rewritten a single function:

def f_g(branch, n, m) =
  let join j(branch, m, n) = 
    if branch == 0 then e1'
    else e2'
  j(branch, m, n)

def f(n) = f_g(0, n, dummy)
def g(n) = f_g(1, dummy, m)

where e1' and e2' are e1 and e2 with their tail calls to f and g replaced with jumps to j.

In general, strongly connected components of mutually tail recursive functions will be optimized in this way. All join points used in each strongly connected component will be rewritten and merged into the same join point.

@CAG2Mark CAG2Mark self-assigned this Mar 29, 2024
@CAG2Mark CAG2Mark added the enhancement New feature or request label Mar 29, 2024
@CAG2Mark
Copy link
Contributor Author

CAG2Mark commented Mar 29, 2024

TODO:

  • Make sure variable names in the stack frame do not clash
  • Make sure the merged function's name does not clash with other functions
  • Substitute dummy variables (may need extension to the IR). See here
  • Propagage @tailrec annotations to the IR.

@CAG2Mark
Copy link
Contributor Author

OK, should be good to merge. There's just one trivial merge conflict in DiffTests.scala to handle

@CAG2Mark CAG2Mark marked this pull request as ready for review June 10, 2024 08:31
@LPTK
Copy link
Contributor

LPTK commented Jun 13, 2024

OK, should be good to merge.

Great!

There's just one trivial merge conflict in DiffTests.scala to handle

So... do you intend to ever fix these conflicts? 🤔

@CAG2Mark
Copy link
Contributor Author

CAG2Mark commented Jun 13, 2024

So... do you intend to ever fix these conflicts? 🤔

I figured doing it now will make the diff a bit harder to review. Once everything else is sorted, I'll fix the conflicts.

Copy link
Contributor

@LPTK LPTK left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Amazing change, thanks! Should we merge now? (I'll do it.)

@CAG2Mark
Copy link
Contributor Author

Go ahead

@LPTK LPTK merged commit 13ba521 into hkust-taco:mlscript Jun 20, 2024
1 check passed
@LPTK LPTK deleted the tailrec branch June 20, 2024 03:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants