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

Preliminary work on separating SimpIR from CoreIR. #1169

Closed axch closed 1 year ago

axch commented 1 year ago

Method: Make SimpIR a separate (but isomorphic) IR constructor to CoreIR and fix compilation errors.

I didn't finish, so I rolled back the impetus change. So, SimpIR is still an alias for CoreIR, but the correct one is referred to in more places.

The functional change is that attempting to substE at (AtomSubstVal CoreIR) into the LoweredTopFun constructor of TopFunBinding is now a runtime error, on the grounds that LoweredTopFun carries SimpIR (which I infer from its name and type-checking how it is constructed). This runtime error does not occur in the test suite.

I encountered three problems that stopped me from making further progress:

Custom Linearizations

The content of a custom linearization rule is typed as Atom CoreIR, and at runtime is presumably always a lambda. However, linearize just directly constructs an application thereof in SimpIR, using naryApp.

Is this just a bug? Should we call simplifyApp from linearizeExpr? (Maybe the only reason it isn't currently done that way is the circular module dependency it would cause between Simplify.hs and Linearize.hs.)

On the other hand, that Atom CoreIR is just the reconstruction of a definition that was already simplified (when the "create custom linearization" directive was processed). Is there perhaps some reason to expect that such a reconstruction can just be added to SimpIR without trouble? If so, should we express that reason to GHC?

Simplification substitutions

The simplifier's whole job is to change from CoreIR to SimpIR. However, the substitution it carries is typed AtomSubstVal CoreIR, so doesn't have a proof that it did its job! I am somewhat confused by the comment about top-level function names being an exception; but whatever the exception is, perhaps we can express it to GHC?

Reconstructions

What is the IR typology of reconstructions and ReconAbs? What information are they supposed to carry, and what are they even used for at the end of the day? In particular, applyRecon seems to be used a lot to emit SimpIR inside simplifyHof, but simplifyTopBlock seems to want to be able to use the recon to carry CoreIR.

Is there a chance that this problem might be solved by capturing the constraint on where Lam atoms can appear after simplification, and what we know about their bodies before and after? And maybe Lams that float all the way to the top level are allowed to be different from those that stop at a PrimHof?

axch commented 1 year ago

Subsumed by #1179.