[WIP] Draft autoscheduler

[mtsokol]

WIP

[codspeed-hq]

CodSpeed Performance Report

Merging #831 will not alter performance

_{Comparing ms/draft-autoscheduler (3b49c64) with main (6722b73)}

Summary

✅ 340 untouched benchmarks

[hameerabbasi]

I feel we're implementing a lot of IR pass infrastructure which we don't really need -- IMHO it'd be better to define an IR and do it via egglog or MLIR.

[mtsokol]

I think the only part of pass infrastructure rewritten here is rewrite_tools.py taken from Julia Finch. The rest is scheduler logic (passes, IR).
I used egglog for the first try but using pure Python and its structural pattern matching ended up more convenient for me (I could directly map Julia implementation, I can have side-effects inside matches).

[willow-ahrens]

egglog is overkill for a lot of Finch rewriting (especially since e-graph behavior is unbounded for the rules Finch uses, and deterministic matching works fine). The only time we might want to consider using Egglog would be for the theorem-proving rules. In the past, Jaeyeon has used Z3 for this to great success. However, the use of e-graphs should be considered a research topic for our initial rewriting.

[willow-ahrens]

as for defining AST's, we might consider using the approach used in EXO: https://github.com/ChezJrk/asdl

[willow-ahrens]

one last thing: I think the deferred node may change soon, all of the rest of this ast is stable but we may need to change deferred node. I'll try to remind y'all when I make the update.

[willow-ahrens]

I personally really like that you can use the default python pattern matching here, it's a big benefit for readability and maintainability. I struggled for a long time to teach new users how to use the custom julia rewriting library because julia had no inbuilt pattern matching syntax.

[mtsokol]

as for defining AST's, we might consider using the approach used in EXO: https://github.com/ChezJrk/asdl

I will take a look at it! Here the requirement is that each AST class must be a Python @dataclass to be used in the pattern matching. Right, pattern matching well suited here.

[willow-ahrens]

I don't want to push too hard for ASDL,having each ast node be it's own class is simpler and easy to understand.

[hameerabbasi]

@mtsokol Is this still relevant?

[mtsokol]

@mtsokol Is this still relevant?

I think it is. Once we have Finch in MLIR we might want to first use Python version of the autoscheduler if porting it to MLIR is more time-consuming.