[WIP] experimental memref_stream unroll-and-jam pass using interpreter-based cost model

[superlopuh]

Most of the juicy stuff here is in the autotune.py marimo notebook, best experienced by checking out this branch and running uv run marimo edit docs/marimo.

The main idea here is to have a system that proposes a bunch of rewrites, and then evaluates each of these rewrites based on an additional lowering + interpreter tracing.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 91.28%. Comparing base (9c53363) to head (45d57ed).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3724      +/-   ##
==========================================
- Coverage   91.30%   91.28%   -0.02%     
==========================================
  Files         468      471       +3     
  Lines       58636    58697      +61     
  Branches     5656     5661       +5     
==========================================
+ Hits        53535    53581      +46     
- Misses       3650     3662      +12     
- Partials     1451     1454       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[qaco]

In the marimo notebook, when I click on autotune.py, I have the following error : "f-string expression part cannot include a backslash (autotune.py, line 516)"

[superlopuh]

uh oh

[superlopuh]

I have a feeling it was a python 3.12 thing, should work now

[superlopuh]

OK despite the CI still failing it should work locally for you

[qaco]

It's nice. I like the "lazy pass" pattern. And I like the fact that the unroll and jam is confined at linalg level (rather than materializing loops).

In the case of this cost model, it's minimizing the number of jumps/blt (thanks to unrolling) that reduces the cost, right?

[superlopuh]

Yep, basically. It should be relatively easy to make the cost model a bit smarter and to add latency of instructions.