openqasm / qe-compiler

An MLIR based compiler dynamic circuit compiler for real-time control systems supporting OpenQASM 3
Other
38 stars 12 forks source link

Update parameter/constant op handling for performance #335

Open taalexander opened 2 weeks ago

taalexander commented 2 weeks ago

This PR updates how parameters are handled. They are now directly read from "parameter_load" ops which removes the need for QUIR variable analysis and casting which significantly improves performance of large parameter programs by reducing the total number of operations in the program. This also cuts down on the total number of arguments to each circuit/sequence by storing the parameter loads/constants directly in the sequence itself

For example before for the program below

OPENQASM 3;

qubit $0;
qubit $1;
qubit $2;

gate sx q { }
gate rz(phi) q { }

input float[64] theta = 3.14159265358979;
input angle phi;

sx $0;
rz(theta) $0;

sx $0;

rz(phi) $1;
rz(3.141592) $1;

rz(theta) $2;
rz(phi) $2;

bit b;

b = measure $0;

The MLIR generated is now:

module {
  func.func @sx(%arg0: !quir.qubit<1>) attributes {quir.classicalOnly = false} {
    return
  }
  func.func @rz(%arg0: !quir.qubit<1>, %arg1: !quir.angle<64>) attributes {quir.classicalOnly = false} {
    return
  }
  quir.circuit @circuit_0(%arg0: !quir.qubit<1> {quir.physicalId = 0 : i32}, %arg1: !quir.qubit<1> {quir.physicalId = 1 : i32}, %arg2: !quir.qubit<1> {quir.physicalId = 2 : i32}) -> i1 attributes {quir.classicalOnly = false, quir.physicalIds = [0 : i32, 1 : i32, 2 : i32]} {
    quir.call_gate @sx(%arg0) : (!quir.qubit<1>) -> ()
    %0 = qcs.parameter_load "theta" : !quir.angle<64> {initialValue = 3.14159265358979 : f64}
    quir.call_gate @rz(%arg0, %0) : (!quir.qubit<1>, !quir.angle<64>) -> ()
    quir.call_gate @sx(%arg0) : (!quir.qubit<1>) -> ()
    %1 = qcs.parameter_load "phi" : !quir.angle<64>
    quir.call_gate @rz(%arg1, %1) : (!quir.qubit<1>, !quir.angle<64>) -> ()
    %angle = quir.constant #quir.angle<3.1415920000000002> : !quir.angle<64>
    quir.call_gate @rz(%arg1, %angle) : (!quir.qubit<1>, !quir.angle<64>) -> ()
    %2 = qcs.parameter_load "theta" : !quir.angle<64> {initialValue = 3.14159265358979 : f64}
    quir.call_gate @rz(%arg2, %2) : (!quir.qubit<1>, !quir.angle<64>) -> ()
    %3 = qcs.parameter_load "phi" : !quir.angle<64>
    quir.call_gate @rz(%arg2, %3) : (!quir.qubit<1>, !quir.angle<64>) -> ()
    %4 = quir.measure(%arg0) {quir.noFastPathComm, quir.noJunoComm, quir.noJunoUse} : (!quir.qubit<1>) -> i1
    quir.return %4 : i1
  }
  func.func @main() -> i32 attributes {quir.classicalOnly = false} {
    %c0_i32 = arith.constant 0 : i32
    %dur = quir.constant #quir.duration<4.500000e+06> : !quir.duration<dt>
    %c1 = arith.constant 1 : index
    %c1000 = arith.constant 1000 : index
    %c0 = arith.constant 0 : index
    qcs.init
    scf.for %arg0 = %c0 to %c1000 step %c1 {
      quir.delay %dur, () : !quir.duration<dt>, () -> ()
      qcs.shot_init {qcs.num_shots = 1000 : i32}
      %0 = quir.declare_qubit {id = 0 : i32} : !quir.qubit<1>
      %1 = quir.declare_qubit {id = 1 : i32} : !quir.qubit<1>
      %2 = quir.declare_qubit {id = 2 : i32} : !quir.qubit<1>
      %3 = quir.call_circuit @circuit_0(%0, %1, %2) : (!quir.qubit<1>, !quir.qubit<1>, !quir.qubit<1>) -> i1
    } {qcs.shot_loop, quir.classicalOnly = false, quir.physicalIds = [0 : i32, 1 : i32, 2 : i32]}
    qcs.finalize
    return %c0_i32 : i32
  }
}