SquinToStimPass fails to lower loops with non-pure inputs #484
Description
The below valid squin kernel fails to lower to stim. I found this bug while lowering another dialect to stim through squin.
Side note: Technically len(result) can be replaced with 5 during constant-propagation but I don't think the const prop pass does this currently. This bug is the same if result is a function argument of test_squin_kernel instead.
Package versions:
bloqade-circuit==0.6.8
kirin-toolchain==0.17.22
Minimum code to reproduce:
from bloqade.squin import kernel, qubit
from bloqade.stim.passes import SquinToStimPass
@kernel
def test_squin_kernel():
q = qubit.New(5)
result = qubit.measure(q)
outputs = []
for rnd in range(len(result)): # <----- Non-pure loop iterator
outputs += [] # Doesn't matter what is in the body
return outputs
main = test_squin_kernel.similar()
SquinToStimPass(main.dialects)(main) # Fails here with KeyError
main.print() # Should print the kernel rewritten into stimTrackback:
---------------------------------------------------------------------------
KeyError Traceback (most recent call last)
Cell In[1], line 17
14 return outputs
16 main = test_squin_kernel.similar()
---> 17 SquinToStimPass(main.dialects)(main) # Fails here with KeyError
18 main.print() # Should print the kernel rewritten into stim
File .../kirin/passes/abc.py:30, in Pass.__call__(self, mt)
29 def __call__(self, mt: Method) -> RewriteResult:
---> 30 result = self.unsafe_run(mt)
31 mt.code.verify()
32 return result
File .../bloqade/stim/passes/squin_to_stim.py:130, in SquinToStimPass.unsafe_run(self, mt)
125 address_analysis_frame, _ = aa.run_analysis(mt, no_raise=self.no_raise)
127 # wrap the address analysis result
128 rewrite_result = (
129 Walk(WrapAddressAnalysis(address_analysis=address_analysis_frame.entries))
--> 130 .rewrite(mt.code)
131 .join(rewrite_result)
132 )
134 # 2. rewrite
135 ## Invoke DCE afterwards to eliminate any GetItems
136 ## that are no longer being used. This allows for
137 ## SquinMeasureToStim to safely eliminate
138 ## unused measure statements.
139 rewrite_result = (
140 Chain(
141 Walk(IfToStim(measure_frame=meas_analysis_frame)),
(...)
145 .join(rewrite_result)
146 )
File .../kirin/rewrite/walk.py:38, in Walk.rewrite(self, node)
36 subnode = self.worklist.pop()
37 while subnode is not None:
---> 38 result = self.rule.rewrite(subnode)
39 if result.terminated:
40 return result
File .../kirin/rewrite/abc.py:35, in RewriteRule.rewrite(self, node)
33 return self.rewrite_Region(node)
34 elif isinstance(node, Block):
---> 35 return self.rewrite_Block(node)
36 elif isinstance(node, Statement):
37 return self.rewrite_Statement(node)
File .../bloqade/squin/rewrite/wrap_analysis.py:51, in WrapAnalysis.rewrite_Block(self, node)
50 def rewrite_Block(self, node: ir.Block) -> RewriteResult:
---> 51 has_done_something = any(self.wrap(arg) for arg in node.args)
52 return RewriteResult(has_done_something=has_done_something)
File .../bloqade/squin/rewrite/wrap_analysis.py:51, in <genexpr>(.0)
50 def rewrite_Block(self, node: ir.Block) -> RewriteResult:
---> 51 has_done_something = any(self.wrap(arg) for arg in node.args)
52 return RewriteResult(has_done_something=has_done_something)
File .../bloqade/squin/rewrite/wrap_analysis.py:64, in WrapAddressAnalysis.wrap(self, value)
63 def wrap(self, value: ir.SSAValue) -> bool:
---> 64 address_analysis_result = self.address_analysis[value]
66 if value.hints.get("address") is not None:
67 return False
KeyError: <BlockArgument[int] rnd, uses: 1>