pydata
[WIP] Draft autoscheduler
CodSpeed Performance Report
Merging #831 will not alter performance
Comparing ms/draft-autoscheduler (3b49c64) with main (6722b73)
Summary
✅ 340 untouched benchmarks
![codspeed-hq[bot] avatar](https://avatars.githubusercontent.com/in/257293?v=4 "Published by a pub.dev verified publisher"){kind=small}
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.
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).
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.
as for defining AST's, we might consider using the approach used in EXO: https://github.com/ChezJrk/asdl
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.
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.
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.
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.
@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.