google-research / dex-lang

Research language for array processing in the Haskell/ML family
BSD 3-Clause "New" or "Revised" License
1.58k stars 107 forks source link

Coalesce code transformations across all three compilation paths #1214

Closed axch closed 1 year ago

axch commented 1 year ago

Now top-level blocks, standalone functions, and exported functions are all compiled to Imp with the same sequence of passes, namely simplification, inlining, dce, loop unrolling, dce again, loop invariant hoisting, abstract machine lowering (lowerFullySequential), dce again, loop invariant hoisting again (for destination allocs), vectorization, and finally imp lowering.

This required a little mumbo-jumbo with typeclasses, because the code being compiled is represented differently in the different paths: a top-level block is represented as a Block before abstract machine lowering and as a DestBlock after, whereas the other two paths represent the code as a LamExpr before abstract machine lowering and the newly introduced DestLamExpr after. However, all the passes I was dealing with are scope-relative, so running them under the binders in a LamExpr was easy.

I am a little unsure that I correctly adjusted Imp.toImpFunction to accept code that has already had a destination passing transform done on it, so I'd appreciate a second pair of eyes on that (though the test suite passes).

Also, now that this happened, is there anything we can delete (e.g., from Imp)?