🐛 Fix conversions of jeff.switch and jeff.for

CHANGELOG.md

@@ -15,7 +15,7 @@ This project adheres to [Semantic Versioning], with the exception that minor rel
 - ✨ Add a `quantum-loop-unroll` pass for unrolling for-loop operations containing quantum operations ([#1718]) ([**@MatthiasReumann**])
 - ✨ Add a `hadamard-lifting` pass for lifting Hadamard gates above Pauli gates ([#1605]) ([**@lirem101**], [**@burgholzer**])
 - ✨ Add a `merge-single-qubit-rotation-gates` pass for merging consecutive rotation gates using quaternions ([#1407], [#1674]) ([**@J4MMlE**], [**@denialhaag**], [**@MatthiasReumann**])
-- ✨ Add conversions between `jeff` and QCO ([#1479], [#1548], [#1565], [#1637], [#1676], [#1706]) ([**@denialhaag**], [**@burgholzer**])
+- ✨ Add conversions between `jeff` and QCO ([#1479], [#1548], [#1565], [#1637], [#1676], [#1706], [#1776]) ([**@denialhaag**], [**@burgholzer**])
 - ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568], [#1581], [#1583], [#1588], [#1600], [#1664], [#1709], [#1716], [#1748]) ([**@MatthiasReumann**], [**@burgholzer**])
 - ✨ 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], [#1506], [#1510], [#1513], [#1521], [#1542], [#1548], [#1550], [#1554], [#1567], [#1569], [#1570], [#1572], [#1573], [#1580], [#1602], [#1620], [#1623], [#1624], [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700], [#1717], [#1728], [#1730], [#1749], [#1751], [#1762], [#1765], [#1774])
@@ -402,6 +402,7 @@ _📚 Refer to the [GitHub Release Notes](https://github.com/munich-quantum-tool
 
 <!-- PR links -->
 
+[#1776]: https://github.com/munich-quantum-toolkit/core/pull/1776
 [#1774]: https://github.com/munich-quantum-toolkit/core/pull/1774
 [#1765]: https://github.com/munich-quantum-toolkit/core/pull/1765
 [#1762]: https://github.com/munich-quantum-toolkit/core/pull/1762

cmake/ExternalDependencies.cmake

@@ -27,7 +27,9 @@ if(BUILD_MQT_CORE_MLIR)
   FetchContent_Declare(
     jeff-mlir
     GIT_REPOSITORY https://github.com/PennyLaneAI/jeff-mlir.git
-    GIT_TAG v0.2.0)
+    # Pinned to an unreleased commit until v0.3.0 is released. jeff-mlir's SCF operations are
+    # already marked as IsolatedFromAbove in the pinned version.
+    GIT_TAG 3f34dc3e2865ceaffb8003b2410404306a49f0ab)
   list(APPEND FETCH_PACKAGES jeff-mlir)
 endif()
 

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

@@ -105,6 +105,21 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder {
    */
   Value intConstant(int64_t value);
 
+  /**
+   * @brief Create a constant float value
+   * @param value The value to store in the constant
+   * @return The value produced by the constant operation
+   *
+   * @par Example:
+   * ```c++
+   * auto c = builder.floatConstant(0.123);
+   * ```
+   * ```mlir
+   * %c = arith.constant 0.123 : f64
+   * ```
+   */
+  Value floatConstant(double value);
+
   //===--------------------------------------------------------------------===//
   // Memory Management
   //===--------------------------------------------------------------------===//

mlir/lib/Conversion/JeffToQCO/JeffToQCO.cpp

@@ -861,18 +861,52 @@ struct ConvertJeffSwitchOpToQCO final : OpConversionPattern<jeff::SwitchOp> {
           op, "qco.if requires exactly two branches");
     }
 
-    auto qcoIf = IfOp::create(rewriter, op.getLoc(), adaptor.getSelection(),
-                              adaptor.getInValues());
+    auto isLinearType = [](Type t) {
+      return isa<jeff::QubitType, jeff::QuregType>(t);
+    };
+
+    auto inValues = adaptor.getInValues();
+
+    SmallVector<Value> qubits;
+    for (auto [value, adapted] : llvm::zip(op.getInValues(), inValues)) {
+      if (isLinearType(value.getType())) {
+        qubits.push_back(adapted);
+      }
+    }
+
+    auto qcoIf =
+        IfOp::create(rewriter, op.getLoc(), adaptor.getSelection(), qubits);
 
     auto moveRegion = [&](Region& source, Region& dest) -> LogicalResult {
-      rewriter.inlineRegionBefore(source, dest, dest.end());
-      Block* block = &dest.front();
-      TypeConverter::SignatureConversion sc(block->getNumArguments());
-      if (failed(getTypeConverter()->convertSignatureArgs(
-              block->getArgumentTypes(), sc))) {
-        return failure();
+      auto* oldBlock = &source.back();
+      auto* newBlock = &dest.emplaceBlock();
+      rewriter.setInsertionPointToEnd(newBlock);
+
+      IRMapping mapping;
+      for (auto [oldArg, adapted] :
+           llvm::zip(oldBlock->getArguments(), inValues)) {
+        if (isLinearType(oldArg.getType())) {
+          auto newArg = newBlock->addArgument(
+              typeConverter->convertType(oldArg.getType()), oldArg.getLoc());
+          mapping.map(oldArg, newArg);
+        } else {
+          mapping.map(oldArg, adapted);
+        }
+      }
+
+      for (auto& op : oldBlock->without_terminator()) {
+        rewriter.clone(op, mapping);
+      }
+
+      auto* oldTerminator = oldBlock->getTerminator();
+      SmallVector<Value> yields;
+      for (auto value : oldTerminator->getOperands()) {
+        if (isLinearType(value.getType())) {
+          yields.push_back(rewriter.getRemappedValue(mapping.lookup(value)));
+        }
       }
-      rewriter.applySignatureConversion(block, sc);
+      rewriter.replaceOpWithNewOp<YieldOp>(oldTerminator, yields);
+
       return success();
     };
 
@@ -883,7 +917,15 @@ struct ConvertJeffSwitchOpToQCO final : OpConversionPattern<jeff::SwitchOp> {
       return failure();
     }
 
-    rewriter.replaceOp(op, qcoIf.getResults());
+    SmallVector<Value> results;
+    size_t index = 0;
+    for (auto [value, adapted] : llvm::zip(op.getResults(), inValues)) {
+      results.push_back(isLinearType(value.getType())
+                            ? qcoIf.getResults()[index++]
+                            : adapted);
+    }
+    rewriter.replaceOp(op, results);
+
     return success();
   }
 };
@@ -934,8 +976,8 @@ struct ConvertJeffForOpToQCO final : OpConversionPattern<jeff::ForOp> {
     auto scfFor = scf::ForOp::create(rewriter, loc, start, stop, step,
                                      adaptor.getInValues());
 
-    Block* jeffBody = &op.getBody().front();
-    Block* scfBody = scfFor.getBody();
+    auto* jeffBody = &op.getBody().front();
+    auto* scfBody = scfFor.getBody();
 
     OpBuilder::InsertionGuard guard(rewriter);
     rewriter.setInsertionPointToStart(scfBody);
@@ -944,7 +986,7 @@ struct ConvertJeffForOpToQCO final : OpConversionPattern<jeff::ForOp> {
                                          jeffBody->getArgument(0).getType(),
                                          scfFor.getInductionVar());
     SmallVector<Value> args = {iv.getResult()};
-    for (Value arg : scfFor.getRegionIterArgs()) {
+    for (auto arg : scfFor.getRegionIterArgs()) {
       args.push_back(arg);
     }
 
@@ -964,11 +1006,6 @@ struct ConvertJeffYieldOpToQCO final : OpConversionPattern<jeff::YieldOp> {
   LogicalResult
   matchAndRewrite(jeff::YieldOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter& rewriter) const override {
-    if (isa<IfOp>(op->getParentOp())) {
-      rewriter.replaceOpWithNewOp<YieldOp>(op, adaptor.getOperands());
-      return success();
-    }
-
     rewriter.replaceOpWithNewOp<scf::YieldOp>(op, adaptor.getOperands());
     return success();
   }

mlir/lib/Conversion/QCOToJeff/QCOToJeff.cpp

@@ -10,7 +10,6 @@
 
 #include "mlir/Conversion/QCOToJeff/QCOToJeff.h"
 
-#include "mlir/Conversion/ConversionUtils.h"
 #include "mlir/Dialect/QCO/IR/QCODialect.h"
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
 #include "mlir/Dialect/QTensor/IR/QTensorDialect.h"
@@ -39,6 +38,7 @@
 #include <mlir/Support/LLVM.h>
 #include <mlir/Support/LogicalResult.h>
 #include <mlir/Transforms/DialectConversion.h>
+#include <mlir/Transforms/RegionUtils.h>
 
 #include <cassert>
 #include <cstddef>
@@ -342,6 +342,45 @@ static LogicalResult cleanUp(Operation* op, LoweringState& state) {
   return success();
 }
 
+/**
+ * @brief Move a region from QCO/SCF operation to a jeff operation
+ */
+static LogicalResult moveRegion(Region& source, Region& dest,
+                                ConversionPatternRewriter& rewriter,
+                                const TypeConverter* typeConverter,
+                                const SetVector<Value>& aboveValues) {
+  auto* oldBlock = &source.back();
+  auto* newBlock = &dest.emplaceBlock();
+  rewriter.setInsertionPointToEnd(newBlock);
+
+  IRMapping mapping;
+  for (auto oldArg : oldBlock->getArguments()) {
+    auto newArg = newBlock->addArgument(
+        typeConverter->convertType(oldArg.getType()), oldArg.getLoc());
+    mapping.map(oldArg, newArg);
+  }
+  for (auto value : aboveValues) {
+    auto newArg = newBlock->addArgument(
+        typeConverter->convertType(value.getType()), value.getLoc());
+    mapping.map(value, newArg);
+  }
+
+  for (auto& op : oldBlock->without_terminator()) {
+    rewriter.clone(op, mapping);
+  }
+
+  auto* oldTerminator = oldBlock->getTerminator();
+  SmallVector<Value> yields;
+  for (auto value : oldTerminator->getOperands()) {
+    yields.push_back(rewriter.getRemappedValue(mapping.lookup(value)));
+  }
+  llvm::append_range(yields,
+                     newBlock->getArguments().take_back(aboveValues.size()));
+  rewriter.replaceOpWithNewOp<jeff::YieldOp>(oldTerminator, yields);
+
+  return success();
+}
+
 namespace {
 
 /**
@@ -963,13 +1002,8 @@ struct ConvertQCOYieldOpToJeff final : StatefulOpConversionPattern<YieldOp> {
   using StatefulOpConversionPattern::StatefulOpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(YieldOp op, OpAdaptor adaptor,
+  matchAndRewrite(YieldOp op, OpAdaptor /*adaptor*/,
                   ConversionPatternRewriter& rewriter) const override {
-    if (isa<jeff::SwitchOp>(op->getParentOp())) {
-      rewriter.replaceOpWithNewOp<jeff::YieldOp>(op, adaptor.getOperands());
-      return success();
-    }
-
     auto& state = getState();
 
     if (state.inInvOp) {
@@ -1036,37 +1070,55 @@ struct ConvertQCOIfOpToJeff final : StatefulOpConversionPattern<IfOp> {
   matchAndRewrite(IfOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter& rewriter) const override {
     auto loc = op.getLoc();
+
+    SetVector<Value> aboveValues;
+    getUsedValuesDefinedAbove(op.getElseRegion(), aboveValues);
+    getUsedValuesDefinedAbove(op.getThenRegion(), aboveValues);
+
+    SmallVector<Value> initArgs;
+    llvm::append_range(initArgs, adaptor.getQubits());
+
     SmallVector<Type> outTypes;
     if (failed(
             getTypeConverter()->convertTypes(op.getResultTypes(), outTypes))) {
       return failure();
     }
 
-    auto jeffIf =
-        jeff::SwitchOp::create(rewriter, loc, outTypes, adaptor.getCondition(),
-                               adaptor.getQubits(), 2);
+    for (auto value : aboveValues) {
+      auto remappedValue = rewriter.getRemappedValue(value);
+      initArgs.push_back(remappedValue);
+      outTypes.push_back(remappedValue.getType());
+    }
+
+    auto jeffSwitch = jeff::SwitchOp::create(
+        rewriter, loc, outTypes, adaptor.getCondition(), initArgs, 2);
 
-    if (failed(moveRegion(op.getElseRegion(), jeffIf.getBranches()[0], rewriter,
-                          getTypeConverter()))) {
+    if (failed(moveRegion(op.getElseRegion(), jeffSwitch.getBranches()[0],
+                          rewriter, getTypeConverter(), aboveValues))) {
       return failure();
     }
-    if (failed(moveRegion(op.getThenRegion(), jeffIf.getBranches()[1], rewriter,
-                          getTypeConverter()))) {
+    if (failed(moveRegion(op.getThenRegion(), jeffSwitch.getBranches()[1],
+                          rewriter, getTypeConverter(), aboveValues))) {
       return failure();
     }
 
     // Add trivial default case
     {
-      auto* block = &jeffIf.getDefault().emplaceBlock();
+      auto* block = &jeffSwitch.getDefault().emplaceBlock();
       for (auto value : adaptor.getQubits()) {
         block->addArgument(value.getType(), loc);
       }
+      for (auto value : aboveValues) {
+        block->addArgument(typeConverter->convertType(value.getType()), loc);
+      }
       OpBuilder::InsertionGuard guard(rewriter);
       rewriter.setInsertionPointToStart(block);
       jeff::YieldOp::create(rewriter, loc, block->getArguments());
     }
 
-    rewriter.replaceOp(op, jeffIf.getResults());
+    rewriter.replaceOp(op,
+                       jeffSwitch.getResults().take_front(op.getNumResults()));
+
     return success();
   }
 };
@@ -1104,34 +1156,35 @@ struct ConvertSCFForOpToJeff final : StatefulOpConversionPattern<scf::ForOp> {
   LogicalResult
   matchAndRewrite(scf::ForOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter& rewriter) const override {
+    SetVector<Value> aboveValues;
+    getUsedValuesDefinedAbove(op.getRegion(), aboveValues);
+
+    SmallVector<Value> initArgs;
+    llvm::append_range(initArgs, adaptor.getInitArgs());
+
     SmallVector<Type> outTypes;
     if (failed(
             getTypeConverter()->convertTypes(op.getResultTypes(), outTypes))) {
       return failure();
     }
 
+    for (auto value : aboveValues) {
+      auto remappedValue = rewriter.getRemappedValue(value);
+      initArgs.push_back(remappedValue);
+      outTypes.push_back(remappedValue.getType());
+    }
+
     auto jeffFor = jeff::ForOp::create(
         rewriter, op.getLoc(), outTypes, adaptor.getLowerBound(),
-        adaptor.getUpperBound(), adaptor.getStep(), adaptor.getInitArgs());
+        adaptor.getUpperBound(), adaptor.getStep(), initArgs);
 
     if (failed(moveRegion(op.getRegion(), jeffFor.getRegion(), rewriter,
-                          getTypeConverter()))) {
+                          getTypeConverter(), aboveValues))) {
       return failure();
     }
 
-    rewriter.replaceOp(op, jeffFor.getResults());
-    return success();
-  }
-};
-
-struct ConvertSCFYieldOpToJeff final
-    : StatefulOpConversionPattern<scf::YieldOp> {
-  using StatefulOpConversionPattern::StatefulOpConversionPattern;
+    rewriter.replaceOp(op, jeffFor.getResults().take_front(op.getNumResults()));
 
-  LogicalResult
-  matchAndRewrite(scf::YieldOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter& rewriter) const override {
-    rewriter.replaceOpWithNewOp<jeff::YieldOp>(op, adaptor.getResults());
     return success();
   }
 };
@@ -1414,11 +1467,11 @@ struct QCOToJeff final : impl::QCOToJeffBase<QCOToJeff> {
     addQCOToJeffGatePattern<JK::Custom, 2, 2, XXMinusYYOp, void, false>(
         patterns, typeConverter, context, state, "xx_minus_yy");
 
-    patterns.add<ConvertQCOBarrierOpToJeff, ConvertQCOCtrlOpToJeff,
-                 ConvertQCOInvOpToJeff, ConvertQCOYieldOpToJeff,
-                 ConvertQCOIfOpToJeff, ConvertSCFForOpToJeff,
-                 ConvertSCFYieldOpToJeff, ConvertQCOMainToJeff>(
-        typeConverter, context, &state);
+    patterns
+        .add<ConvertQCOBarrierOpToJeff, ConvertQCOCtrlOpToJeff,
+             ConvertQCOInvOpToJeff, ConvertQCOYieldOpToJeff,
+             ConvertQCOIfOpToJeff, ConvertSCFForOpToJeff, ConvertQCOMainToJeff>(
+            typeConverter, context, &state);
 
     // Apply the conversion
     if (applyPartialConversion(module, target, std::move(patterns)).failed()) {

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

@@ -74,6 +74,11 @@ Value QCOProgramBuilder::intConstant(const int64_t value) {
   return arith::ConstantOp::create(*this, getI64IntegerAttr(value)).getResult();
 }
 
+Value QCOProgramBuilder::floatConstant(const double value) {
+  checkFinalized();
+  return arith::ConstantOp::create(*this, getF64FloatAttr(value)).getResult();
+}
+
 Value& QCOProgramBuilder::QubitRegister::operator[](const size_t index) {
   if (index >= qubits.size()) {
     llvm::reportFatalUsageError("Qubit index out of bounds");