Refactor post-processing in measurements_from_samples

[lillian542]

Context:
One of several PRs for [sc-107541]. This one addresses refactoring to follow the new guidelines in the qnode transform guide
regarding transforms that add post-processing

Description of the Change:

We refactor to follow the new guidelines in the new qnode transform guide. The main difference is we create a separate function for each quantum.node that calls the quantum node and applies post-processing on the results, i.e. we go from including post-processing inside the qnode:

  func.func public @circuit() -> tensor<f64> attributes {quanutm.node} {
    %cst = arith.constant dense<0> : tensor<i64>
    %c1000_i64 = arith.constant 1000 : i64
    quantum.device shots(%c1000_i64)
    %0 = op.test —> !quantum.bit
    %1 = quantum.compbasis %0 : !quantum.obs
    %2 = quantum.sample %1 : tensor<1000x1xf64>
    %3 = call @expval_from_samples.tensor.1000x1xf64(%2, %cst) : (tensor<1000x1xf64>, tensor<i64>) -> tensor<f64>
    return %3 : tensor<f64>

to calling the qnode (which returns samples), and performing post-processing on the results:

  func.func public @circuit() -> tensor<f64> {
    %0 = call @circuit.from_samples() : () -> tensor<1000x1xf64>
    %1 = call @expval_from_samples.tensor.1000x1xf64(%0) : (tensor<1000x1xf64>) -> tensor<f64>
    return %1 : tensor<f64>
  }
  func.func public @circuit.from_samples() -> tensor<1000x1xf64> attributes {quantum.node} {
    %c1000_i64 = arith.constant 1000 : i64
    quantum.device shots(%c1000_i64) 
    %0 = op.test —> !quantum.bit
    %1 = quantum.compbasis %0 : !quantum.obs
    %2 = quantum.sample %1 : tensor<1000x1xf64>
    return %2 : tensor<1000x1xf64>
  }

Key changes are:

• We add a new Pattern, AddPostProcessingPattern, that inserts a classical function that calls the quantum_node for each in the module. This should be generic for re-use with other 1:1 passes that add post-processing (but won't work for passes that create multiple quantum.nodes meant to be recombined in post-processing)
• We update the match_and_rewrite functions to look for FuncOps instead of measurement processes, and apply the pattern to each quantum.node
• We add a method on the main MeasurementsFromSamples pattern that gets the module the quantum.node is nested in, walks through it, and finds the CallOp for the quantum.node, so that it is available as in insertion point for post-processing.
• We add an update_returns function that updates the operations and types returned in the outer function and the qnode to reflect the conversion to samples an the addition of a post-processing step
• We no longer assume that the shots are the same for all quantum.nodes - this is technically true currently, because weighted shot distribution for split-non-commuting isn't added yet, but its not a safe assumption to make.

Benefits:
More future-facing technical implementation (follows new convention, won't break when shot distribution is added). Will be easier to maintain and improves readability of the IR (especially when there are multiple measurements output). Maintains cleaner separation between hybrid/quantum execution and purely classical post-processing.

Possible Drawbacks:
Not sure if/how these changes will impact performance. We previously "matched" directly with the MPs in the pattern rewriter, and now we match with all FuncOps, check if they are quantum.nodes, and then walk through that to find the MPs. I'm not sure whether this will impact performance meaningfully.

Note that this will need to be moved from xDSL to MLIR eventually, so in terms of efficiency, "good enough for now" is fine.

Related GitHub Issues:
None

[multiphaseCFD]

LGTM! Thanks @lillian542 .

[mudit2812]

This looks great Lillian. A few final comments. I also wanted to ask, is there any plan to support dynamic shots? Seems like a pretty important feature if we want to link this pass to the pennylane transform

[lillian542]

I also wanted to ask, is there any plan to support dynamic shots? Seems like a pretty important feature if we want to link this pass to the pennylane transform

When we were scoping it early, my impression was we are alright with figuring out dynamic shots later, but this would definitely be a good thing to revisit before we link the tape transform.

[mudit2812]

Looks great. My only blocking comment is about removing the change to get_constant_from_ssa and updating any affected tests accordingly.

[mudit2812]

Nice. Customary 🚀