Fix AggressiveUnroll again

[cduck]

Do not merge until we understand the issue being fixed here. This became an issue again while upgrading my code for bloqade-circuit v0.9.

This pass has been unreliable before (see #593) and I'm not sure why adding a second copy of the ilist.rewrite.Unroll pass fixes the issue.

Here's the error it is fixing:

  File "test.py", line 59, in compile_to_stim
    SquinToStimPass(main.dialects, no_raise=no_raise)(main)
  File "kirin/passes/abc.py", line 30, in __call__
    result = self.unsafe_run(mt)
  File "bloqade/stim/passes/squin_to_stim.py", line 71, in unsafe_run
    .rewrite(mt.code)
  File "kirin/rewrite/chain.py", line 29, in rewrite
    result = rule.rewrite(node)
  File "kirin/rewrite/walk.py", line 40, in rewrite
    result = self.rule.rewrite(subnode)
  File "kirin/rewrite/abc.py", line 38, in rewrite
    return self.rewrite_Statement(cast(Statement, node))
  File "bloqade/stim/rewrite/set_observable_to_stim.py", line 26, in rewrite_Statement
    return self.rewrite_SetObservable(node)
  File "bloqade/stim/rewrite/set_observable_to_stim.py", line 39, in rewrite_SetObservable
    measure_ids = self.measure_id_frame.entries[node.measurements]
KeyError: <ResultValue[IList[MeasurementResult, AnyType()]] inputs, uses: 1>

Tagging: @kaihsin

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

☂️ Python Coverage

current status: ✅

Overall Coverage

Lines	Covered	Coverage	Threshold	Status
10143	8897	88%	0%	🟢

New Files

No new covered files...

Modified Files

File	Coverage	Status
src/bloqade/rewrite/passes/aggressive_unroll.py	100%	🟢
TOTAL	100%	🟢

updated for commit: 462b0cd by action🐍

[david-pl]

@cduck The problem here seems to be that the AggressiveUnroll fails to inline some lambda functions. The MeasureIdAnalysis then fails since it has no method for lambdas. We could add a method for that, but that's not easy and it will probably cause a failure further down the line.

Somehow your change causes the inline to succeed, not sure why.

[kaihsin]

Is there an issue to track this? I can dig a bit more on this since it's affecting downstream

[david-pl]

Oh, I think it might be as simple as this:

@squin.kernel
def main(n: int):
    def map_func(i):
        return n + 1
    
    return ilist.map(map_func, range(4))

main.print()

AggressiveUnroll(main.dialects).fixpoint(main)

main.print()

AggressiveUnroll just doesn't rewrite ilist.map, I think. @cduck's added rewrite does. I'm not quite sure how that worked on previous releases though.

Edit: okay, this is really odd. If I change the map_func to return n + 1 it works as expected on main.

[cduck]

The problem here seems to be that the AggressiveUnroll fails to inline some lambda functions.

This makes sense. Let's open an issue to fix AggressiveUnroll. I am confused how the line I added here changes the behavior, because it is a duplicate copy of the ilist.rewrite.Unroll pass but run earlier in the sequence.

[weinbe58]

@cduck are you using scf in the kernels?

[kaihsin]

@cduck are you using scf in the kernels?

Yes. Squin is using scf!

[kaihsin]

For some reason typeinfer marked this as bottom:

%0 = py.constant.constant IList(range(0, 4)) : !Bottom

For some reason the final typeinfer rewrite in the UnrollScf pass modify that to Bottom

[weinbe58]

My bad, accidentally closed the PR

[weinbe58]

For some reason typeinfer marked this as bottom:

%0 = py.constant.constant IList(range(0, 4)) : !Bottom

For some reason the final typeinfer rewrite in the UnrollScf pass modify that to Bottom

I see