Use flat products instead of nested pairs in non-dependent telescopes.

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

Use flat products instead of nested pairs in non-dependent telescopes. #1227

Closed dougalm closed 1 year ago

dougalm commented 1 year ago

Partly this is just to make the IR more readable. But it might have some compile-time performance advantages too. Linearization can produce intermediate telescopes with hundreds of components. Indexing these as nested pairs means an "atom" can have hundreds ProjectElt layers. I wouldn't be surprised if some operations like type checking end up being quadratic in that depth.

dougalm commented 1 year ago

Running DCE after linearization reduced the amount of code that we produced, but it exposed a strange bug in AD: https://github.com/google-research/dex-lang/tree/flat-telescopes-dce-ad-bug. Reverting for now.