Ensure an ion.yield operation is returning an region output to all pa…

…rallal protocol ops
PennyLaneAI · Jan 20, 2025 · 82cabb8 · 82cabb8
1 parent 78c7251
commit 82cabb8
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/mlir/lib/Ion/IR/IonOps.cpp b/mlir/lib/Ion/IR/IonOps.cpp
@@ -42,18 +42,23 @@ void ParallelProtocolOp::build(OpBuilder &builder, OperationState &result, Value
                                BodyBuilderFn bodyBuilder)
 {
     OpBuilder::InsertionGuard guard(builder);
+    Location loc = result.location;
 
     result.addOperands(inQubits);
     for (Value v : inQubits)
         result.addTypes(v.getType());
 
     Region *bodyRegion = result.addRegion();
     Block *bodyBlock = builder.createBlock(bodyRegion);
-    for (Value v : inQubits)
+    for (Value v : inQubits) {
         bodyBlock->addArgument(v.getType(), v.getLoc());
+    }
 
     builder.setInsertionPointToStart(bodyBlock);
-    bodyBuilder(builder, result.location, bodyBlock->getArguments());
+    bodyBuilder(builder, loc, bodyBlock->getArguments());
+
+    builder.setInsertionPointToEnd(bodyBlock);
+    builder.create<ion::YieldOp>(loc, bodyBlock->getArguments());
 }
 
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/Ion/Transforms/QuantumToIonPatterns.cpp b/mlir/lib/Ion/Transforms/QuantumToIonPatterns.cpp
@@ -190,7 +190,6 @@ mlir::LogicalResult oneQubitGateToPulse(CustomOp op, mlir::PatternRewriter &rewr
                 auto qubit = qubits.front();
                 builder.create<ion::PulseOp>(loc, time, qubit, beam0toEAttr, phase1Attr);
                 builder.create<ion::PulseOp>(loc, time, qubit, beam1toEAttr, phase2Attr);
-                builder.create<ion::YieldOp>(loc);
             });
         rewriter.replaceOp(op, ppOp);
         return success();
@@ -408,8 +407,6 @@ mlir::LogicalResult MSGateToPulse(CustomOp op, mlir::PatternRewriter &rewriter,
                         rewriter.getI64VectorAttr(beam.polarization),
                         rewriter.getI64VectorAttr(flipSign(beam.wavevector)));
                     builder.create<ion::PulseOp>(loc, time, qubit1, beam6Attr, phase0Attr);
-
-                    builder.create<ion::YieldOp>(loc);
                 });
             rewriter.replaceOp(op, ppOp);
             return success();

diff --git a/mlir/test/Ion/QuantumToIon.mlir b/mlir/test/Ion/QuantumToIon.mlir
@@ -163,6 +163,7 @@ func.func @example_ion_two_qubit(%arg0: f64) -> !quantum.bit {
     // CHECK-SAME:         polarization = dense<[0, 1]> : vector<2xi64>,
     // CHECK-SAME:         wavevector = dense<[-2, 3]> : vector<2xi64>>,
     // CHECK-SAME:     phase = 0.000000e+00 : f64}
+    // CHECK-NEXT:   ion.yield %arg1 : !quantum.bit
     // CHECK-NEXT: }
     %4 = quantum.custom "RX"(%arg0) %2 : !quantum.bit
 
@@ -186,6 +187,7 @@ func.func @example_ion_two_qubit(%arg0: f64) -> !quantum.bit {
     // CHECK-SAME:         polarization = dense<[0, 1]> : vector<2xi64>,
     // CHECK-SAME:         wavevector = dense<[-2, 3]> : vector<2xi64>>,
     // CHECK-SAME:     phase = 3.1415926535{{[0-9]*}} : f64}
+    // CHECK-NEXT:   ion.yield %arg1 : !quantum.bit
     // CHECK-NEXT: }
     %5 = quantum.custom "RY"(%arg0) %4 : !quantum.bit
 
@@ -209,6 +211,7 @@ func.func @example_ion_two_qubit(%arg0: f64) -> !quantum.bit {
     // CHECK-SAME:         polarization = dense<[0, 1]> : vector<2xi64>,
     // CHECK-SAME:         wavevector = dense<[-2, 3]> : vector<2xi64>>,
     // CHECK-SAME:     phase = 0.000000e+00 : f64}
+    // CHECK-NEXT:   ion.yield %arg1 : !quantum.bit
     // CHECK-NEXT: }
     %6 = quantum.custom "RX"(%arg0) %5 : !quantum.bit
 
@@ -264,6 +267,7 @@ func.func @example_ion_two_qubit(%arg0: f64) -> !quantum.bit {
     // CHECK-SAME:         polarization = dense<[7, 8]> : vector<2xi64>,
     // CHECK-SAME:         wavevector = dense<[-9, -10]> : vector<2xi64>>,
     // CHECK-SAME:     phase = 0.000000e+00 : f64}
+    // CHECK-NEXT:   ion.yield %arg1, %arg2 : !quantum.bit, !quantum.bit
     // CHECK-NEXT: }
     %7:2 = quantum.custom "MS"(%arg0) %6, %3 : !quantum.bit, !quantum.bit
     return %7#0: !quantum.bit
@@ -340,6 +344,7 @@ func.func @example_ion_three_qubit(%arg0: f64) -> (!quantum.bit, !quantum.bit, !
     // CHECK-SAME:         polarization = dense<[7, 8]> : vector<2xi64>,
     // CHECK-SAME:         wavevector = dense<[-9, -10]> : vector<2xi64>>,
     // CHECK-SAME:     phase = 0.000000e+00 : f64}
+    // CHECK-NEXT:   ion.yield %arg1, %arg2 : !quantum.bit, !quantum.bit
     // CHECK-NEXT: }
     %5:2 = quantum.custom "MS"(%arg0) %2, %3 : !quantum.bit, !quantum.bit
 
@@ -395,6 +400,7 @@ func.func @example_ion_three_qubit(%arg0: f64) -> (!quantum.bit, !quantum.bit, !
     // CHECK-SAME:         polarization = dense<[1, 2]> : vector<2xi64>,
     // CHECK-SAME:         wavevector = dense<[3, -4]> : vector<2xi64>>,
     // CHECK-SAME:     phase = 0.000000e+00 : f64}
+    // CHECK-NEXT:   ion.yield %arg1, %arg2 : !quantum.bit, !quantum.bit
     // CHECK-NEXT: }
     %6:2 = quantum.custom "MS"(%arg0) %5#0, %4 : !quantum.bit, !quantum.bit
 
@@ -450,6 +456,7 @@ func.func @example_ion_three_qubit(%arg0: f64) -> (!quantum.bit, !quantum.bit, !
     // CHECK-SAME:         polarization = dense<[37, 42]> : vector<2xi64>,
     // CHECK-SAME:         wavevector = dense<[42, 37]> : vector<2xi64>>,
     // CHECK-SAME:     phase = 0.000000e+00 : f64}
+    // CHECK-NEXT:   ion.yield %arg1, %arg2 : !quantum.bit, !quantum.bit
     // CHECK-NEXT: }
     %7:2 = quantum.custom "MS"(%arg0) %5#1, %6#1 : !quantum.bit, !quantum.bit
     return %6#0, %7#0, %7#1: !quantum.bit, !quantum.bit, !quantum.bit