✨ Add QCO If operation (#1506)

Author: li-mingbao

CHANGELOG.md

@@ -12,7 +12,7 @@ This project adheres to [Semantic Versioning], with the exception that minor rel
 ### Added
 
 - ✨ Add initial infrastructure for new QC and QCO MLIR dialects
-  ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1510], [#1513], [#1521])
+  ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1506], [#1510], [#1513], [#1521])
   ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**])
 
 ### Changed
@@ -330,6 +330,7 @@ _📚 Refer to the [GitHub Release Notes](https://github.com/munich-quantum-tool
 [#1521]: https://github.com/munich-quantum-toolkit/core/pull/1521
 [#1513]: https://github.com/munich-quantum-toolkit/core/pull/1513
 [#1510]: https://github.com/munich-quantum-toolkit/core/pull/1510
+[#1506]: https://github.com/munich-quantum-toolkit/core/pull/1506
 [#1481]: https://github.com/munich-quantum-toolkit/core/pull/1481
 [#1475]: https://github.com/munich-quantum-toolkit/core/pull/1475
 [#1474]: https://github.com/munich-quantum-toolkit/core/pull/1474

mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h

@@ -11,7 +11,6 @@
 #pragma once
 
 #include <cstdint>
-#include <functional>
 #include <llvm/ADT/DenseSet.h>
 #include <llvm/ADT/STLFunctionalExtras.h>
 #include <llvm/ADT/SmallVector.h>
@@ -1071,6 +1070,51 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder {
    */
   QCOProgramBuilder& dealloc(Value qubit);
 
+  //===--------------------------------------------------------------------===//
+  // SCF operations
+  //===--------------------------------------------------------------------===//
+
+  /**
+   * @brief Construct an if operation for qubits with linear typing
+   *
+   * @details
+   * Constructs an if operation that takes a bool Value and a range of qubit
+   * values that are used in the then/else region of this operation. The qubit
+   * values are passed down as block arguments to each region.
+   *
+   * @param condition Bool condition
+   * @param qubits Input qubits
+   * @param thenBody Function that builds the then body of the if
+   * operation
+   * @param elseBody Function that builds the else body of the if
+   * operation
+   * @return ValueRange of the results (must be the same types as the input
+   * qubits)
+   *
+   * @par Example:
+   * ```c++
+   * auto result =
+   *   builder.qcoIf(condition, q0,
+   *     [&](ValueRange args) -> llvm::SmallVector<Value> {
+   *       auto q1 = builder.h(args[0]);
+   *       return {q1};
+   *     });
+   * ```
+   * ```mlir
+   * %q2 = qco.if %condition qubits(%arg0 = %q0) {
+   *      %q1 = qco.h %arg0 : !qco.qubit -> !qco.qubit
+   *      qco.yield %q1
+   * } else qubits(%arg0 = %q0) {
+   *      qco.yield %arg0
+   * } : {i1, !qco.qubit} -> {!qco.qubit}
+   * ```
+   */
+  ValueRange
+  qcoIf(const std::variant<bool, Value>& condition, ValueRange qubits,
+        llvm::function_ref<llvm::SmallVector<Value>(ValueRange)> thenBody,
+        llvm::function_ref<llvm::SmallVector<Value>(ValueRange)> elseBody =
+            nullptr);
+
   //===--------------------------------------------------------------------===//
   // Finalization
   //===--------------------------------------------------------------------===//

mlir/include/mlir/Dialect/QCO/IR/QCOOps.h

@@ -25,6 +25,7 @@
 #include "mlir/Dialect/QCO/IR/QCOInterfaces.h"
 
 #include <mlir/Bytecode/BytecodeOpInterface.h>
+#include <mlir/Interfaces/ControlFlowInterfaces.h>
 #include <mlir/Interfaces/SideEffectInterfaces.h>
 #include <variant>
 

mlir/include/mlir/Dialect/QCO/IR/QCOOps.td

@@ -14,6 +14,7 @@ include "mlir/IR/BuiltinTypeInterfaces.td"
 include "mlir/IR/DialectBase.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
@@ -1048,7 +1049,7 @@ def BarrierOp : QCOOp<"barrier", traits = [UnitaryOpInterface]> {
 // Modifiers
 //===----------------------------------------------------------------------===//
 
-def YieldOp : QCOOp<"yield", traits = [Terminator]> {
+def YieldOp : QCOOp<"yield", traits = [Terminator, ReturnLike]> {
     let summary = "Yield from a modifier region";
     let description = [{
         Terminates a modifier region, yielding the transformed target qubits back to the enclosing modifier operation.
@@ -1205,4 +1206,83 @@ def InvOp : QCOOp<"inv",
     let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// SCF operations
+//===----------------------------------------------------------------------===//
+
+def IfOp : QCOOp<"if", traits =
+            [
+                DeclareOpInterfaceMethods<RegionBranchOpInterface, [
+                "getNumRegionInvocations", "getRegionInvocationBounds",
+                "getEntrySuccessorRegions"]>,
+                SingleBlock,
+                SingleBlockImplicitTerminator<"::mlir::qco::YieldOp">,
+                RecursiveMemoryEffects
+            ]> {
+
+    let summary = "If-then-else operation for linear (qubit) types";
+    let description = [{
+        The `qco.if` operation is an if-then-else construct similar to the standard scf.if operation.
+        In addition to the condition, the operation takes a variadic number of qubits as inputs that are
+        required in the bodies of both branches. These qubits are passed down to the individual regions
+        as block arguments. The number of results and the type of the results must be equivalent to the
+        number and types of the input qubits.
+
+        Example:
+        ```mlir
+        %result = qco.if %condition qubits(%arg0 = %q0) {
+            %q1 = qco.h %arg0 : !qco.qubit -> !qco.qubit
+            qco.yield %q1
+         } else qubits(%arg0 = %q0) {
+            qco.yield %arg0
+        } : {i1, !qco.qubit} -> {!qco.qubit}
+        ```
+    }];
+
+    let arguments = (ins I1:$condition,  Variadic<QubitType>:$qubits);
+    let results = (outs Variadic<QubitType>:$results);
+    let regions = (region SizedRegion<1>:$thenRegion,
+                          SizedRegion<1>:$elseRegion);
+
+    let assemblyFormat = [{
+        $condition
+         custom<IfOpAliasing>($thenRegion, $elseRegion, $qubits)
+         attr-dict `:`
+        `{` type($condition) `,` type($qubits) `}`
+        `->`
+        `{` type($results) `}`
+    }];
+
+    let builders = [
+        OpBuilder<(ins "Value":$condition, "ValueRange":$qubits), [{
+           build($_builder, $_state, qubits.getTypes(), condition, qubits);
+        }]>,
+        OpBuilder<(ins "Value":$condition, "ValueRange":$qubits, "llvm::function_ref<llvm::SmallVector<Value>(ValueRange)>":$thenBuilder,
+        CArg<"llvm::function_ref<llvm::SmallVector<Value>(ValueRange)>","nullptr">:$elseBuilder)>
+    ];
+
+    let extraClassDeclaration = [{
+        Block *thenBlock() {
+           return &getThenRegion().back();
+        }
+        YieldOp thenYield();
+        Block* elseBlock() {
+           return &getElseRegion().back();
+        }
+        YieldOp elseYield();
+    }];
+
+    let extraClassDefinition = [{
+        YieldOp $cppClass::thenYield() {
+           return cast<YieldOp>(&thenBlock()->back());
+        }
+        YieldOp $cppClass::elseYield() {
+           return cast<YieldOp>(&elseBlock()->back());
+        }
+    }];
+
+    let hasCanonicalizer = 1;
+    let hasVerifier = 1;
+}
+
 #endif // QCOOPS

mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp

@@ -678,6 +678,71 @@ QCOProgramBuilder& QCOProgramBuilder::dealloc(Value qubit) {
   return *this;
 }
 
+//===----------------------------------------------------------------------===//
+// SCF Operations
+//===----------------------------------------------------------------------===//
+
+ValueRange QCOProgramBuilder::qcoIf(
+    const std::variant<bool, Value>& condition, ValueRange qubits,
+    llvm::function_ref<llvm::SmallVector<Value>(ValueRange)> thenBody,
+    llvm::function_ref<llvm::SmallVector<Value>(ValueRange)> elseBody) {
+  checkFinalized();
+
+  auto conditionValue = utils::variantToValue(*this, getLoc(), condition);
+
+  auto ifOp = IfOp::create(*this, conditionValue, qubits);
+  // Create the then and else block
+  auto& thenBlock = ifOp->getRegion(0).emplaceBlock();
+  auto& elseBlock = ifOp->getRegion(1).emplaceBlock();
+
+  // Create the block arguments and add them as valid qubits
+  for (auto qubitType : qubits.getTypes()) {
+    const auto thenArg = thenBlock.addArgument(qubitType, getLoc());
+    const auto elseArg = elseBlock.addArgument(qubitType, getLoc());
+    validQubits.insert(thenArg);
+    validQubits.insert(elseArg);
+  }
+
+  // Construct the bodies of the regions
+  const InsertionGuard guard(*this);
+  setInsertionPointToStart(&thenBlock);
+  const auto thenResult = thenBody(thenBlock.getArguments());
+  YieldOp::create(*this, thenResult);
+  setInsertionPointToStart(&elseBlock);
+  llvm::SmallVector<Value> elseResult;
+  if (elseBody) {
+    elseResult = elseBody(elseBlock.getArguments());
+    YieldOp::create(*this, elseResult);
+  } else {
+    elseResult.assign(elseBlock.getArguments().begin(),
+                      elseBlock.getArguments().end());
+    YieldOp::create(*this, elseBlock.getArguments());
+  }
+
+  if (thenResult.size() != qubits.size() ||
+      thenResult.size() != elseResult.size()) {
+    llvm::reportFatalUsageError(
+        "Then and else body must return the same amount of qubits as the "
+        "number of input qubits!");
+  }
+
+  // Update qubit tracking
+  const auto& ifResults = ifOp->getResults();
+  for (auto [input, output] : llvm::zip_equal(qubits, ifResults)) {
+    updateQubitTracking(input, output);
+  }
+
+  // Remove the inner qubits as valid qubits
+  for (auto thenOut : thenResult) {
+    validQubits.erase(thenOut);
+  }
+  for (auto elseOut : elseResult) {
+    validQubits.erase(elseOut);
+  }
+
+  return ifResults;
+}
+
 //===----------------------------------------------------------------------===//
 // Finalization
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/QCO/IR/CMakeLists.txt

@@ -9,13 +9,15 @@
 file(GLOB_RECURSE MODIFIERS "${CMAKE_CURRENT_SOURCE_DIR}/Modifiers/*.cpp")
 file(GLOB_RECURSE OPERATIONS "${CMAKE_CURRENT_SOURCE_DIR}/Operations/*.cpp")
 file(GLOB_RECURSE QUBIT_MANAGEMENT "${CMAKE_CURRENT_SOURCE_DIR}/QubitManagement/*.cpp")
+file(GLOB_RECURSE SCF "${CMAKE_CURRENT_SOURCE_DIR}/SCF/*.cpp")
 
 add_mlir_dialect_library(
   MLIRQCODialect
   QCOOps.cpp
   ${MODIFIERS}
   ${OPERATIONS}
   ${QUBIT_MANAGEMENT}
+  ${SCF}
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/mlir/include/mlir/Dialect/QCO
   DEPENDS

mlir/lib/Dialect/QCO/IR/QCOOps.cpp

@@ -12,6 +12,8 @@
 
 #include "mlir/Dialect/QCO/IR/QCODialect.h" // IWYU pragma: associated
 
+#include <llvm/ADT/DenseSet.h>
+#include <llvm/ADT/STLExtras.h>
 #include <mlir/IR/Block.h>
 #include <mlir/IR/OpImplementation.h>
 #include <mlir/IR/Operation.h>
@@ -113,6 +115,68 @@ static void printTargetAliasing(OpAsmPrinter& printer, Operation* /*op*/,
   printer.printRegion(region, false);
 }
 
+static ParseResult
+parseIfOpAliasing(OpAsmParser& parser, Region& thenRegion, Region& elseRegion,
+                  SmallVectorImpl<OpAsmParser::UnresolvedOperand>& operands) {
+  // Parse the qubits keyword
+  if (parser.parseKeyword("qubits")) {
+    return failure();
+  }
+
+  // Parse the then region
+  if (parseTargetAliasing(parser, thenRegion, operands)) {
+    return failure();
+  }
+  const auto thenCount = operands.size();
+
+  // Parse the else keyword
+  if (parser.parseKeyword("else")) {
+    return failure();
+  }
+
+  // Parse the qubits keyword
+  if (parser.parseKeyword("qubits")) {
+    return failure();
+  }
+
+  // Parse the else region
+  if (parseTargetAliasing(parser, elseRegion, operands)) {
+    return failure();
+  }
+
+  const auto elseCount = operands.size() - thenCount;
+
+  if (thenCount != elseCount) {
+    return parser.emitError(
+        parser.getCurrentLocation(),
+        "then/else qubit aliasing lists must be the same length");
+  }
+  for (unsigned i = 0; i < thenCount; ++i) {
+    if (operands[i].name != operands[thenCount + i].name) {
+      return parser.emitError(
+          parser.getCurrentLocation(),
+          "then/else qubit aliasing lists must match in order");
+    }
+  }
+
+  // Remove duplicate operands
+  llvm::DenseSet<llvm::StringRef> seen;
+  llvm::erase_if(operands,
+                 [&](const auto& op) { return !seen.insert(op.name).second; });
+
+  return success();
+}
+
+static void printIfOpAliasing(OpAsmPrinter& printer, Operation* op,
+                              Region& thenRegion, Region& elseRegion,
+                              OperandRange qubits) {
+  printer << "qubits";
+  printTargetAliasing(printer, op, thenRegion, qubits);
+  printer << " else ";
+  printer << "qubits";
+  printTargetAliasing(printer, op, elseRegion, qubits);
+}
+
 //===----------------------------------------------------------------------===//
 // Dialect
 //===----------------------------------------------------------------------===//