Add State Outline Evolution xDSL Pass

[rniczh]

Context:

Hackweek project

Description of the Change:

This change isolates the quantum state evolution logic into a separate function. Previously, the operations for applying gates and evolving the quantum state (e.g., lines %4 to %16 in the original code) were inline within the main function. This PR moves them to a new @state_evolution function.

before outline-state-evolution
builtin.module @module_main {
  func.func public @main(%arg0 : tensor<f64>) -> (tensor<f64>) attributes {diff_method = "parameter-shift", llvm.linkage = #llvm.linkage<internal>, qnode} {
    %0 = "stablehlo.constant"() <{value = dense<10> : tensor<i64>}> : () -> tensor<i64>
    %1 = tensor.extract %0[] : tensor<i64>
    quantum.device shots(%1) ["....dylib", "NullQubit", "{'track_resources': False}"]
    %2 = "stablehlo.constant"() <{value = dense<2> : tensor<i64>}> : () -> tensor<i64>
    %3 = quantum.alloc(2) : !quantum.reg
    %4 = "stablehlo.constant"() <{value = dense<0> : tensor<i64>}> : () -> tensor<i64>
    %5 = tensor.extract %4[] : tensor<i64>
    %6 = quantum.extract %3[%5] : !quantum.reg -> !quantum.bit
    %7 = quantum.custom "Hadamard"() %6 : !quantum.bit
    %8 = "stablehlo.constant"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
    %9 = tensor.extract %8[] : tensor<i64>
    %10 = quantum.extract %3[%9] : !quantum.reg -> !quantum.bit
    %11 = quantum.custom "Hadamard"() %10 : !quantum.bit
    %12, %13 = quantum.measure %7 : i1, !quantum.bit
    %14 = tensor.from_elements %12 : tensor<i1>
    %15 = tensor.extract %arg0[] : tensor<f64>
    %16 = quantum.custom "RX"(%15) %13 : !quantum.bit
    %17 = quantum.namedobs %16[PauliX] : !quantum.obs
    %18 = quantum.expval %17 : f64
    %19 = tensor.from_elements %18 : tensor<f64>
    %20 = tensor.extract %4[] : tensor<i64>
    %21 = quantum.insert %3[%20], %16 : !quantum.reg, !quantum.bit
    %22 = tensor.extract %8[] : tensor<i64>
    %23 = quantum.insert %21[%22], %11 : !quantum.reg, !quantum.bit
    quantum.dealloc %23 : !quantum.reg
    quantum.device_release
    func.return %19 : tensor<f64>
  }
}

Result:

builtin.module @module_main {
  func.func public @main(%arg0 : tensor<f64>) -> (tensor<f64>) attributes {diff_method = "parameter-shift", llvm.linkage = #llvm.linkage<internal>, qnode} {
    %0 = "stablehlo.constant"() <{value = dense<10> : tensor<i64>}> : () -> tensor<i64>
    %1 = tensor.extract %0[] : tensor<i64>
    quantum.device shots(%1) ["....dylib", "NullQubit", "{'track_resources': False}"]
    %2 = "stablehlo.constant"() <{value = dense<2> : tensor<i64>}> : () -> tensor<i64>
    %3 = quantum.alloc(2) : !quantum.reg
    %4, %5, %6, %7, %8 = func.call @state_evolution(%3, %arg0) : (!quantum.reg, tensor<f64>) -> (!quantum.reg, tensor<i64>, i64, i64, tensor<i64>)
    %9 = quantum.extract %4[%6] : !quantum.reg -> !quantum.bit
    %10 = quantum.extract %4[%7] : !quantum.reg -> !quantum.bit
    %11 = quantum.namedobs %10[PauliX] : !quantum.obs
    %12 = quantum.expval %11 : f64
    %13 = tensor.from_elements %12 : tensor<f64>
    %14 = tensor.extract %8[] : tensor<i64>
    %15 = quantum.insert %4[%14], %10 : !quantum.reg, !quantum.bit
    %16 = tensor.extract %5[] : tensor<i64>
    %17 = quantum.insert %15[%16], %9 : !quantum.reg, !quantum.bit
    quantum.dealloc %17 : !quantum.reg
    quantum.device_release
    func.return %13 : tensor<f64>
  }
  func.func @state_evolution(%0 : !quantum.reg, %1 : tensor<f64>) -> (!quantum.reg, tensor<i64>, i64, i64, tensor<i64>) {
    %2 = "stablehlo.constant"() <{value = dense<0> : tensor<i64>}> : () -> tensor<i64>
    %3 = tensor.extract %2[] : tensor<i64>
    %4 = quantum.extract %0[%3] : !quantum.reg -> !quantum.bit
    %5 = quantum.custom "Hadamard"() %4 : !quantum.bit
    %6 = "stablehlo.constant"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
    %7 = tensor.extract %6[] : tensor<i64>
    %8 = quantum.extract %0[%7] : !quantum.reg -> !quantum.bit
    %9 = quantum.custom "Hadamard"() %8 : !quantum.bit
    %10, %11 = quantum.measure %5 : i1, !quantum.bit
    %12 = tensor.from_elements %10 : tensor<i1>
    %13 = tensor.extract %1[] : tensor<f64>
    %14 = quantum.custom "RX"(%13) %11 : !quantum.bit
    %15 = quantum.insert %0[%7], %9 : !quantum.reg, !quantum.bit
    %16 = quantum.insert %15[%3], %14 {terminal_boundary} : !quantum.reg, !quantum.bit
    func.return %16, %6, %7, %3, %2 : !quantum.reg, tensor<i64>, i64, i64, tensor<i64>
  }
}

Benefits:

Possible Drawbacks:

Related GitHub Issues:
[sc-96060]