diff --git a/CHANGELOG.md b/CHANGELOG.md index 3232b2171f..ef7be25d24 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,9 +42,9 @@ with the exception that minor releases may include breaking changes. [#1569], [#1570], [#1572], [#1573], [#1580], [#1602], [#1620], [#1623], [#1624], [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700], [#1710], [#1717], [#1728], [#1730], [#1749], [#1751], [#1762], [#1765], - [#1774], [#1780], [#1781], [#1782], [#1787], [#1802], [#1806], [#1807]) - ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], - [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], + [#1774], [#1780], [#1781], [#1782], [#1787], [#1802], [#1806], [#1807], + [#1808]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**], + [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**]) ### Changed @@ -598,6 +598,7 @@ changelogs._ +[#1808]: https://github.com/munich-quantum-toolkit/core/pull/1808 [#1807]: https://github.com/munich-quantum-toolkit/core/pull/1807 [#1806]: https://github.com/munich-quantum-toolkit/core/pull/1806 [#1802]: https://github.com/munich-quantum-toolkit/core/pull/1802 diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h index 0a59e7d83e..3d360ebb78 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h @@ -10,19 +10,23 @@ #pragma once -#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/IR/QCOInterfaces.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QCO/Utils/Qubits.h" #include "mlir/Dialect/QCO/Utils/WireIterator.h" #include "mlir/Dialect/QTensor/IR/QTensorOps.h" +#include #include +#include +#include #include #include #include #include +#include +#include #include #include #include @@ -88,13 +92,37 @@ LogicalResult walkProgram(Region& region, const WalkProgramFn& fn) { return success(); } -using ReleasedOps = SmallVector; -using PendingWiresMap = - DenseMap>; +using ReleasedOps = SmallVector; +using PendingWiresMap = DenseMap>; + +namespace impl { + +/// Return the number of qubit arguments of unitary-like operation. +inline size_t getNumQubitArgs(Operation* op) { + return TypeSwitch(op) + .Case( + [&](UnitaryOpInterface op) { return op.getNumQubits(); }) + .Case([&](auto op) { + return llvm::count_if( + op.getInits(), [](Value v) { return isa(v.getType()); }); + }) + .Case([&](qco::IfOp op) { + return llvm::count_if(op.getQubits(), [](Value v) { + return isa(v.getType()); + }); + }) + .Default([&](Operation* op) { + const auto name = op->getName().getStringRef(); + reportFatalInternalError("unknown op: " + name); + return 0; + }); +} +} // namespace impl struct IsReady { bool operator()(PendingWiresMap::value_type& kv) const { - return kv.second.size() == kv.first.getNumQubits(); + const auto npending = kv.second.size(); + return impl::getNumQubitArgs(kv.first) == npending; } }; @@ -134,67 +162,66 @@ LogicalResult walkProgramGraph(MutableArrayRef wires, using Traits = WireTraversalTraits; ReleasedOps released; - PendingWiresMap pending; pending.reserve(wires.size()); - SmallVector curr(wires.size()); + SmallVector curr(wires.size()); std::iota(curr.begin(), curr.end(), 0UL); - SmallVector next; + SmallVector next; next.reserve(wires.size()); while (!curr.empty()) { - for (std::size_t i : curr) { + for (size_t i : curr) { auto& it = wires[i]; + + if (it.operation() == nullptr) { // isa + std::ranges::advance(it, Traits::stride()); + } + while (Traits::isActive(it)) { - const auto res = - TypeSwitch(it.operation()) - .template Case([&](UnitaryOpInterface& op) { - // If there are fewer wires than the qubit requires inputs, - // it's impossible to release the operation. Hence, fail. - if (op.getNumQubits() > wires.size()) { - return WalkResult::interrupt(); - } - - if (op.getNumQubits() == 1) { - std::ranges::advance(it, Traits::stride()); - return WalkResult::advance(); - } - - // Insert the unitary to the pending map. - // The caller decides if this op should be released. - const auto [mapIt, inserted] = pending.try_emplace(op); - auto& indices = mapIt->second; - - if (inserted) { - indices.reserve(op.getNumQubits()); - } - - indices.emplace_back(i); - - return WalkResult::skip(); // Stop at multi-qubit gate. - }) - // AllocOp, StaticOp, and qtensor::ExtractOp are only reachable - // on the forward path; backward isActive() halts before - // reaching them (decrementing at a source op is a no-op). - .template Case([&](auto) { - std::ranges::advance(it, Traits::stride()); - return WalkResult::advance(); - }) - .Default([&](Operation* op) -> WalkResult { - const auto name = op->getName().getStringRef(); - report_fatal_error("unknown op encountered: " + name); - }); - - if (res.wasSkipped()) { - break; + + // For source-like (AllocOp, StaticOp, qtensor::ExtractOp), + // sink-like (SinkOp, YieldOp, qtensor::InsertOp, scf::YieldOp), and + // one-qubit non-unitary (ResetOp, MeasureOp) operations, simply advance + // the iterator. + + if (isa( + it.operation())) { + std::ranges::advance(it, Traits::stride()); + continue; + } + + const auto nqubits = impl::getNumQubitArgs(it.operation()); + + // Advance past one-qubit operations. + + if (nqubits == 1) { + std::ranges::advance(it, Traits::stride()); + continue; } - if (res.wasInterrupted()) { + // If there are fewer wires than the operation requires inputs, + // it's impossible to release the operation. Hence, fail. + + if (nqubits > wires.size()) { return failure(); } + + // Insert the unitary to the pending map. + // The caller decides if this op should be released. + + const auto [mapIt, inserted] = pending.try_emplace(it.operation()); + auto& indices = mapIt->second; + + if (inserted) { + indices.reserve(nqubits); + } + + indices.emplace_back(i); + + break; // Stop at multi-qubit unitary. } } @@ -212,11 +239,11 @@ LogicalResult walkProgramGraph(MutableArrayRef wires, } } - for (const UnitaryOpInterface& op : released) { + for (Operation* op : released) { const auto mapIt = pending.find(op); assert(mapIt != pending.end()); - for (std::size_t i : mapIt->second) { + for (size_t i : mapIt->second) { std::ranges::advance(wires[i], Traits::stride()); next.emplace_back(i); } diff --git a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp index 616e08d66d..c9f84fefd0 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp @@ -20,9 +20,11 @@ #include #include #include +#include #include #include #include +#include #include #include @@ -36,7 +38,8 @@ class DriversTest : public testing::Test { protected: void SetUp() override { DialectRegistry registry; - registry.insert(); + registry.insert(); context = std::make_unique(); context->appendDialectRegistry(registry); @@ -99,6 +102,27 @@ TEST_F(DriversTest, ProgramWalk) { ASSERT_EQ(ex3, q31); } +TEST_F(DriversTest, ProgramGraphWalkTooFewWires) { + qco::QCOProgramBuilder builder(context.get()); + builder.initialize(); + + const auto q00 = builder.allocQubit(); + const auto q10 = builder.allocQubit(); + const auto [q01, q11] = builder.cx(q00, q10); + + [[maybe_unused]] auto mod = builder.finalize(); + + // Collect just one wire. + SmallVector wires; + wires.emplace_back(q00); + + auto res = qco::walkProgramGraph( + wires, [&](const qco::ReadyRange&, qco::ReleasedOps&) { + return WalkResult::skip(); + }); + ASSERT_TRUE(res.failed()); +} + TEST_F(DriversTest, ProgramGraphWalk) { qco::QCOProgramBuilder builder(context.get()); builder.initialize(); @@ -121,10 +145,27 @@ TEST_F(DriversTest, ProgramGraphWalk) { const auto [q04, q13] = builder.cx(q03, q12); const auto q14 = builder.h(q13); - builder.measure(q04); - builder.measure(q14); - builder.measure(q23); - builder.measure(q31); + Value iterQ0; + Value iterQ1; + ValueRange blockArgs; + const auto forResults = builder.scfFor( + 0, 3, 1, {q04, q14, q23, q31}, [&](Value, ValueRange args) { + blockArgs = args; + std::tie(iterQ0, iterQ1) = builder.cx(args[0], args[1]); + return SmallVector{iterQ0, iterQ1, args[2], args[3]}; + }); + + const auto q05 = builder.qcoIf( + false, forResults[0], + [&](ValueRange args) { return SmallVector{builder.h(args[0])}; }, + [&](ValueRange args) { + return SmallVector{builder.id(args[0])}; + })[0]; + + builder.measure(q05); + builder.measure(forResults[1]); + builder.measure(forResults[2]); + builder.measure(forResults[3]); auto mod = builder.finalize(); auto func = *(mod->getOps().begin()); @@ -151,11 +192,12 @@ TEST_F(DriversTest, ProgramGraphWalk) { }); ASSERT_TRUE(res.succeeded()); - ASSERT_GE(readyPerLayer.size(), 3); + ASSERT_GE(readyPerLayer.size(), 4); ASSERT_TRUE(readyPerLayer[0].contains(q02.getDefiningOp())); ASSERT_TRUE(readyPerLayer[0].contains(q21.getDefiningOp())); ASSERT_TRUE(readyPerLayer[1].contains(q12.getDefiningOp())); ASSERT_TRUE(readyPerLayer[2].contains(q04.getDefiningOp())); + ASSERT_TRUE(readyPerLayer[3].contains(forResults[0].getDefiningOp())); // Backward pass. readyPerLayer.clear(); @@ -171,11 +213,12 @@ TEST_F(DriversTest, ProgramGraphWalk) { }); ASSERT_TRUE(res.succeeded()); - ASSERT_GE(readyPerLayer.size(), 3); - ASSERT_TRUE(readyPerLayer[0].contains(q04.getDefiningOp())); - ASSERT_TRUE(readyPerLayer[1].contains(q12.getDefiningOp())); - ASSERT_TRUE(readyPerLayer[2].contains(q02.getDefiningOp())); - ASSERT_TRUE(readyPerLayer[2].contains(q21.getDefiningOp())); + ASSERT_GE(readyPerLayer.size(), 4); + ASSERT_TRUE(readyPerLayer[0].contains(forResults[0].getDefiningOp())); + ASSERT_TRUE(readyPerLayer[1].contains(q04.getDefiningOp())); + ASSERT_TRUE(readyPerLayer[2].contains(q12.getDefiningOp())); + ASSERT_TRUE(readyPerLayer[3].contains(q02.getDefiningOp())); + ASSERT_TRUE(readyPerLayer[3].contains(q21.getDefiningOp())); // Forward, but instead of releasing all, we use ::skip(). readyPerLayer.clear(); @@ -186,16 +229,16 @@ TEST_F(DriversTest, ProgramGraphWalk) { layer.insert(op); } readyPerLayer.emplace_back(layer); - return WalkResult::skip(); }); ASSERT_TRUE(res.succeeded()); - ASSERT_GE(readyPerLayer.size(), 3); + ASSERT_GE(readyPerLayer.size(), 4); ASSERT_TRUE(readyPerLayer[0].contains(q02.getDefiningOp())); ASSERT_TRUE(readyPerLayer[0].contains(q21.getDefiningOp())); ASSERT_TRUE(readyPerLayer[1].contains(q12.getDefiningOp())); ASSERT_TRUE(readyPerLayer[2].contains(q04.getDefiningOp())); + ASSERT_TRUE(readyPerLayer[3].contains(forResults[0].getDefiningOp())); // Backward, but stop after first layer. readyPerLayer.clear(); @@ -212,5 +255,27 @@ TEST_F(DriversTest, ProgramGraphWalk) { ASSERT_TRUE(res.failed()); ASSERT_EQ(readyPerLayer.size(), 1); - ASSERT_TRUE(readyPerLayer[0].contains(q04.getDefiningOp())); + ASSERT_TRUE(readyPerLayer[0].contains(forResults[0].getDefiningOp())); + + // Forward, but start at block arguments. + wires.clear(); + for (Value arg : blockArgs) { + wires.emplace_back(arg); + } + + readyPerLayer.clear(); + res = qco::walkProgramGraph( + wires, [&](const qco::ReadyRange& ready, qco::ReleasedOps& released) { + DenseSet layer; + for (const auto& [op, progs] : ready) { + layer.insert(op); + released.emplace_back(op); + } + readyPerLayer.emplace_back(layer); + return WalkResult::advance(); + }); + + ASSERT_TRUE(res.succeeded()); + ASSERT_GE(readyPerLayer.size(), 1); + ASSERT_TRUE(readyPerLayer[0].contains(iterQ0.getDefiningOp())); }