[WIP] Draft autoscheduler

[mtsokol]

WIP

[codspeed-hq]

CodSpeed Performance Report

Merging #831 will degrade performances by 30.19%

_{Comparing ms/draft-autoscheduler (219b11f) with main (ada7871)}

Summary

❌ 1 regressions
✅ 339 untouched benchmarks

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Benchmarks breakdown

	Benchmark	main	ms/draft-autoscheduler	Change
❌	test_index_slice[side=100-rank=2-format='gcxs']	2.4 ms	3.4 ms	-30.19%

[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.