google / heir

A compiler for homomorphic encryption
https://heir.dev/
Apache License 2.0
350 stars 49 forks source link

Reduce reliance on main "canonicalize" pattern #1069

Open AlexanderViand-Intel opened 1 month ago

AlexanderViand-Intel commented 1 month ago

Image

I think the correct thing here is to do the following:

j2kun commented 1 month ago

The only place I know that we do this involves HEIR-owned patterns being added to the canonicalizer (for the HECO-ported pipeline). We could easily extract those into their own pass, but the risk in this issue is nonexistent since we control them. This could also fail if patterns removed from the canonicalizer result in code that we haven't code-genned things for yet (e.g., a leftover tensor.extract that isn't yet code-genned as a manual mask+rotate), but those also seem like problems that would go away on their own, though they may create some slight interim churn.

For other cases, the proposed solution (an automated test to check if they work when canonicalize is replaced by a no-op) would be ideal, but in my opinion a relatively low priority considering everything else we have to work on. By the same token, should we check that the pipelines still work when all folders are replaced by no-ops (in every single application of the greedy rewrite engine)?

If we had to do it, I'd say the best approach would be to add a CLI flag for no-op-canonicalize, replace the heir-opt.cpp invocations of createCanonicalizerPass with a helper function that branches on "no-op" vs actually creating the canonicalizer pass, and then be deliberate about the chosen end-to-end pipeline tests to replicate with the flag enabled. Mainly the tests whose assertions only involve running the final code and asserting an expected output.