Fix BufferizationStage IR blowup for runtime-parameterized Hamiltonians (Scalarize)

[jk20342]

Context:
qml.TrotterProduct with runtime coefficients (via qml.dot) causes severe IR
amplification: tracing unrolls all Trotter steps, and every gate angle drags a chain of
single-element tensor ops through bufferization, each becoming an alloc plus copies.
For an H2 QPE workload this meant ~62x total IR growth and ~19 GB of compiler memory.

Description of the Change:

• New scalarize-tensor-extracts pass: folds tensor.extract through
  extract_slice/collapse_shape/small elementwise linalg.generic, turning gate-angle
  dataflow into scalar arithmetic.
• New reroll-loops pass: detects repeated op sequences via structural hashing and
  rebuilds them as scf.for loops.
• Default pipeline (HLO stage only): scalarize + linalg-fuse-elementwise-ops before
  reroll-loops, and a second scalarize after detensorization (which materializes the
  extract chains).

Benefits:
On the H2 QPE benchmark, before/after measured on the same machine, same build, same
input IR, differing only in the pipeline: compile time 1h41m -> ~25s (~240x), peak
memory 15.4 GB -> 1.2 GB (12.8x), final IR 18.4 MB -> 5.4 MB (3.4x), post-bufferization
IR 7.45 MB -> 2.64 MB (2.8x). Output probabilities are bit-identical to the old pipeline.

these baseline numbers differ from the issue report (306 s, 21.3 GB, 241 MB final
IR) because the issue predates current main, which already reduced the
baseline's bufferization blowup (60 MB -> 7.45 MB post-bufferization; input IR drifted
3.86 -> 4.69 MB), and these measurements use an assertions-enabled dev build, ~20x
slower in wall time than the release wheel the issue used. Ratios within this PR are
apples-to-apples; absolute times are not comparable to the issue's.

Possible Drawbacks:
Both passes are conservative (size caps, semantic verification, high min-savings
default) but rerolling may miss non-contiguous repeats; overhead on workloads without
the pattern is negligible.

Related GitHub Issues:
Fixes #2759. Related: #1464, #1507.

[jk20342]

here is some benchmark code

https://gist.github.com/jk20342/d6cc77c374522db0d59e646dfc4a04ae

and results through my docker container :]

config: mode=dyn n_est=4 n_trotter=10

compile time: 35.2 s
peak RSS: self=542 MB children=1206 MB

per-stage IR size:
  0_qpe_circuit.mlir                                          4.69 MB
  1_AfterQuantumCompilationStage.mlir                         4.59 MB  (1.0x)
  2_AfterHLOLoweringStage.mlir                                2.82 MB  (0.6x)
  3_AfterGradientLoweringStage.mlir                           2.82 MB  (1.0x)
  4_AfterBufferizationStage.mlir                              2.66 MB  (0.9x)
  5_AfterMLIRToLLVMDialectConversion.mlir                     7.32 MB  (2.8x)
  6_AfterLLVMIRTranslation.ll                                 5.37 MB  (0.7x)

[jk20342]

This could be amplified lower I think but there is upstream issue on LLVM that I bumped to try to get more traction on

[dime10]

Thanks for this contribution @jk20342, very nice results 💯 I love the scalarization pass, which is beneficial in more ways than performance!

I'd be curious to know the effect size of the scalarization vs the rerolling, as using loops is known to have massive impact on the compile time. Automatically detecting unrolled loops is definitely nice for programs that were suboptimally written, although we are continuously updating PL library code to make sure programs use loops to begin with as well.

[jk20342]

Thanks for this contribution @jk20342, very nice results 💯 I love the scalarization pass, which is beneficial in more ways than performance!

I'd be curious to know the effect size of the scalarization vs the rerolling, as using loops is known to have massive impact on the compile time. Automatically detecting unrolled loops is definitely nice for programs that were suboptimally written, although we are continuously updating PL library code to make sure programs use loops to begin with as well.

@dime10 when i tested it the reroll improved ~1.2x wall time ~1.35x memory and ~1.27x final IR and a 5x reduction in gate-op volume over just scalarize could remove it, if you think this will be fixed on PL directly though

[mehrdad2m]

Thanks @jk20342, this is a very impressive work! It seems to me that the reroll-loops is an independent work but still requires the cleanup scalarize-tensor-extracts provides to get the performance improvement, is that correct? If so, I would appreciate if we could split the work into two PRs (one for each pass). This would really help us in the review process to be focused on one pass and also be able to benchmark the effects of each independently.

[jk20342]

Thanks @jk20342, this is a very impressive work! It seems to me that the reroll-loops is an independent work but still requires the cleanup scalarize-tensor-extracts provides to get the performance improvement, is that correct? If so, I would appreciate if we could split the work into two PRs (one for each pass). This would really help us in the review process to be focused on one pass and also be able to benchmark the effects of each independently.

Yep I can split it up

[jk20342]

@mehrdad2m made this PR the scalarize #3036 is reroll now :)