`keep_intermediate=True` does not faithfully record initial IR received by first pass

[dime10]

catalyst/frontend/catalyst/compiler.py

Line 678 in e67b2d8

f.write(canonicalize(stdin=ir, options=self.options))

Can be misleading / source of confusion when debugging a pass run from the frontend, as the first recorded IR shows a different program than the pass pipeline is receiving. Specifically, the first recorded IR (0_circuit.mlir) is a canonicalized version of the program passed to the pass pipeline. For debugging purposes, which is the most common use case for keep_intermediate, it is important that the program is faithfully reproduced during inspection.

Reproducer:

import pennylane as qp

@qp.qjit(
    target="mlir",
    pipelines=[("short", ["builtin.module(apply-transform-sequence)"])],
    keep_intermediate=3,
)
@qp.qnode(qp.device("null.qubit", wires=1))
def circuit():
    @qp.for_loop(0, 10, 1)
    def loop(_):
        qp.S(0)
    loop()
    return qp.probs()

circuit.mlir_opt

Compare the output of circuit/0_circuit.mlir and circuit/1_short/1_ApplyTransformSequencePass.mlir (technically recorded after the first pass but in this case it doesn't do anything so is identical to the input):

    func.func public @circuit() -> tensor<2xf64> attributes {quantum.node} {
      %c1 = arith.constant 1 : index
      %c10 = arith.constant 10 : index
      %c0 = arith.constant 0 : index
      %c0_i64 = arith.constant 0 : i64
      quantum.device shots(%c0_i64) ["..."]
      %0 = quantum.alloc( 1) : !quantum.reg
      %1 = scf.for %arg0 = %c0 to %c10 step %c1 iter_args(%arg1 = %0) -> (!quantum.reg) {
        %4 = quantum.extract %arg1[ 0] : !quantum.reg -> !quantum.bit
        %out_qubits = quantum.custom "S"() %4 : !quantum.bit
        %5 = quantum.insert %arg1[ 0], %out_qubits : !quantum.reg, !quantum.bit
        scf.yield %5 : !quantum.reg
      }
      %2 = quantum.compbasis qreg %1 : !quantum.obs
      %3 = quantum.probs %2 : tensor<2xf64>
      quantum.dealloc %1 : !quantum.reg
      quantum.device_release
      return %3 : tensor<2xf64>
    }

and

  func.func public @circuit() -> tensor<2xf64> attributes {quantum.node} {
    %c = stablehlo.constant dense<0> : tensor<i64>
    %extracted = tensor.extract %c[] : tensor<i64>
    quantum.device shots(%extracted) ["..."]
    %c_0 = stablehlo.constant dense<1> : tensor<i64>
    %0 = quantum.alloc( 1) : !quantum.reg
    %c_1 = stablehlo.constant dense<10> : tensor<i64>
    %extracted_2 = tensor.extract %c[] : tensor<i64>
    %1 = arith.index_cast %extracted_2 : i64 to index
    %extracted_3 = tensor.extract %c_1[] : tensor<i64>
    %2 = arith.index_cast %extracted_3 : i64 to index
    %extracted_4 = tensor.extract %c_0[] : tensor<i64>
    %3 = arith.index_cast %extracted_4 : i64 to index
    %4 = scf.for %arg0 = %1 to %2 step %3 iter_args(%arg1 = %0) -> (!quantum.reg) {
      %7 = arith.index_cast %arg0 : index to i64
      %from_elements = tensor.from_elements %7 : tensor<i64>
      %c_5 = stablehlo.constant dense<0> : tensor<i64>
      %extracted_6 = tensor.extract %c_5[] : tensor<i64>
      %8 = quantum.extract %arg1[%extracted_6] : !quantum.reg -> !quantum.bit
      %out_qubits = quantum.custom "S"() %8 : !quantum.bit
      %extracted_7 = tensor.extract %c_5[] : tensor<i64>
      %9 = quantum.insert %arg1[%extracted_7], %out_qubits : !quantum.reg, !quantum.bit
      scf.yield %9 : !quantum.reg
    }
    %5 = quantum.compbasis qreg %4 : !quantum.obs
    %6 = quantum.probs %5 : tensor<2xf64>
    quantum.dealloc %4 : !quantum.reg
    quantum.device_release
    return %6 : tensor<2xf64>
  }