diff --git a/CHANGELOG.md b/CHANGELOG.md index 31a7e5e259..487e69676c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -86,7 +86,7 @@ releases may include breaking changes. circuits to compiler-target topologies while preserving target site IDs and materializing routing workspace on demand ([#1537], [#1547], [#1568], [#1581], [#1583], [#1588], [#1600], [#1664], [#1709], [#1716], [#1748], [#1805], - [#1870], [#1904], [#1911], [#1951], [#1997]) ([**@MatthiasReumann**], + [#1870], [#1904], [#1911], [#1951], [#1997], [#2016]) ([**@MatthiasReumann**], [**@burgholzer**]) - ✨ Add a pass for qubit reuse in quantum programs, as well as related auxiliary passes and patterns ([#1705], [#1755], [#1756], [#1923], [#1924]) @@ -728,6 +728,7 @@ for previous changelogs._ [#2028]: https://github.com/munich-quantum-toolkit/core/pull/2028 [#2026]: https://github.com/munich-quantum-toolkit/core/pull/2026 [#2017]: https://github.com/munich-quantum-toolkit/core/pull/2017 +[#2016]: https://github.com/munich-quantum-toolkit/core/pull/2016 [#2014]: https://github.com/munich-quantum-toolkit/core/pull/2014 [#2011]: https://github.com/munich-quantum-toolkit/core/pull/2011 [#2010]: https://github.com/munich-quantum-toolkit/core/pull/2010 diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index 33dfb4c517..b258c5b1ec 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -22,14 +22,12 @@ #include "mlir/Dialect/QTensor/Utils/TensorIterator.h" #include "mlir/Dialect/Utils/Utils.h" -#include #include #include #include #include #include #include -#include #include #include #include @@ -78,11 +76,16 @@ struct MappingPass : impl::MappingPassBase { using IndexPairType = std::pair; using Window = SmallVector; using Wires = SmallVector; - using RecursiveRoutingStackItem = std::pair>; - using RecursiveRoutingStack = SmallVector; enum class RoutingMode : bool { Cold, Hot }; + struct CompositeUnitary { + /// The composite op (e.g. SCF). + Operation* op = nullptr; + /// Indices into a wire vector, where the order of indices has no meaning. + SmallVector indices; + }; + struct WireInfos { /// Return the mapped wire index of a program index. [[nodiscard]] size_t lookupIndex(const size_t prog) const { @@ -149,12 +152,17 @@ struct MappingPass : impl::MappingPassBase { /// Statistics collected while routing. struct Statistics { + /// The number of inserted swaps. size_t nswaps{0}; + /// Merge another statistics object into this one. + void merge(const Statistics& other) { nswaps += other.nswaps; } }; /// Parameters influencing the behavior of the A* search algorithm. struct Parameters { + /// The path weight. float alpha; + /// The lookahead decay factor. float lambda; }; @@ -163,6 +171,25 @@ struct MappingPass : impl::MappingPassBase { Wires wires; WireInfos infos; Layout layout; + + struct Patch { + std::optional layout; + std::optional infos; + std::optional wires; + }; + + void applyPatch(Patch&& patch) { + Patch p = std::move(patch); + if (p.layout) { + layout = std::move(*p.layout); + } + if (p.infos) { + infos = std::move(*p.infos); + } + if (p.wires) { + wires = std::move(*p.wires); + } + } }; /// Describes a node in the A* search graph. @@ -366,20 +393,20 @@ struct MappingPass : impl::MappingPassBase { std::tie(wires, infos) = std::move(place(body, *layout, *comp, rewriter)); - Statistics stats; RoutingBundle bundle{.wires = std::move(wires), .infos = std::move(infos), .layout = std::move(*layout)}; - const auto res = route( - bundle, stats, &rewriter); - if (res.failed()) { + const auto routeRes = + route(bundle, &rewriter); + if (failed(routeRes)) { func.emitError() << "failed to map the function"; signalPassFailure(); return; } // Collect statistics. + const auto stats = *routeRes; numSwaps += stats.nswaps; // Fix SSA Dominance issues. @@ -698,126 +725,6 @@ struct MappingPass : impl::MappingPassBase { SinkOp::create(rewriter, body.getLoc(), qubit); } - // Finally, update the SCF operations such that they take all static qubits - // as input. To handle recursively nested SCF operations, use a stack of - // (region, mapping) pairs. - - SmallVector>> stack; - stack.emplace_back(body, DenseSet{}); - - while (!stack.empty()) { - for (auto [region, qubits] = stack.pop_back_val(); - Operation& op : make_early_inc_range(region.getOps())) { - TypeSwitch(&op) - .Case( - [&](StaticOp staticOp) { qubits.insert(staticOp.getQubit()); }) - .Case([&](UnitaryOpInterface& uOp) { - for (const auto [pred, succ] : llvm::zip_equal( - uOp.getInputQubits(), uOp.getOutputQubits())) { - qubits.insert(succ); - qubits.erase(pred); - } - }) - .Case([&](scf::ForOp forOp) { - assert(qubits.size() == layout.nqubits()); - - llvm::for_each(getQubitValues(forOp.getInits()), - [&](Value v) { qubits.erase(v); }); - - auto newForOp = extend(forOp, to_vector(qubits), rewriter); - for (const auto [init, result] : llvm::zip_equal( - newForOp.getInits(), *newForOp.getLoopResults())) { - if (isa(init.getType())) { - qubits.insert(result); - qubits.erase(init); - } - } - - const auto regionQubits = - getQubitValues(newForOp.getRegionIterArgs()); - stack.emplace_back( - newForOp.getRegion(), - DenseSet(regionQubits.begin(), regionQubits.end())); - }) - .Case([&](scf::WhileOp whileOp) { - assert(qubits.size() == layout.nqubits()); - - llvm::for_each(getQubitValues(whileOp.getInits()), - [&](Value v) { qubits.erase(v); }); - - auto newWhileOp = extend(whileOp, to_vector(qubits), rewriter); - for (const auto [init, result] : llvm::zip_equal( - newWhileOp.getInits(), newWhileOp.getResults())) { - if (isa(init.getType())) { - qubits.insert(result); - qubits.erase(init); - } - } - - const auto beforeArgs = - getQubitValues(newWhileOp.getBeforeArguments()); - const auto afterArgs = - getQubitValues(newWhileOp.getAfterArguments()); - stack.emplace_back( - newWhileOp.getBefore(), - DenseSet(beforeArgs.begin(), beforeArgs.end())); - stack.emplace_back( - newWhileOp.getAfter(), - DenseSet(afterArgs.begin(), afterArgs.end())); - }) - .Case([&](IfOp ifOp) { - assert(qubits.size() == layout.nqubits()); - - llvm::for_each(ifOp.getQubits(), - [&](Value v) { qubits.erase(v); }); - - auto newIfOp = extend(ifOp, to_vector(qubits), rewriter); - - for (const auto [qubit, result] : llvm::zip_equal( - newIfOp.getQubits(), newIfOp.getLinearResults())) { - qubits.insert(result); - qubits.erase(qubit); - } - - const auto thenArgs = newIfOp.getThenRegion().getArguments(); - const auto elseArgs = newIfOp.getElseRegion().getArguments(); - stack.emplace_back( - newIfOp.getThenRegion(), - DenseSet(thenArgs.begin(), thenArgs.end())); - stack.emplace_back( - newIfOp.getElseRegion(), - DenseSet(elseArgs.begin(), elseArgs.end())); - }) - .Case([&](IndexSwitchOp switchOp) { - assert(qubits.size() == layout.nqubits()); - - llvm::for_each(switchOp.getTargets(), - [&](Value value) { qubits.erase(value); }); - - auto newSwitchOp = extend(switchOp, to_vector(qubits), rewriter); - for (const auto [target, result] : - llvm::zip_equal(newSwitchOp.getTargets(), - newSwitchOp.getLinearResults())) { - qubits.insert(result); - qubits.erase(target); - } - - for (Region* region : newSwitchOp.getRegions()) { - const auto args = region->getArguments(); - stack.emplace_back(*region, - DenseSet(args.begin(), args.end())); - } - }) - .Case([&](auto resetOp) { - qubits.insert(resetOp.getQubitOut()); - qubits.erase(resetOp.getQubitIn()); - }) - .Case([&](auto) { - llvm::reportFatalInternalError("unexpected dynamic qubit alloc"); - }); - } - } - return {wires, infos}; } @@ -854,13 +761,17 @@ struct MappingPass : impl::MappingPassBase { parallelForEach(&getContext(), trials, [&, this](Trial& t) { for (size_t i = 0; i < niterations; ++i) { - if (route(t.bundle, t.stats).failed()) { + const auto fwRouteRes = route(t.bundle); + if (failed(fwRouteRes)) { return; } - t.stats.nswaps = 0; - if (route(t.bundle, t.stats).failed()) { + + const auto bwRouteRes = route(t.bundle); + if (failed(bwRouteRes)) { return; } + + t.stats = *bwRouteRes; } t.success = true; @@ -1224,14 +1135,13 @@ struct MappingPass : impl::MappingPassBase { /// gates are found. After the function returns, the wires point at the /// results of non-executable gates or operations with nested regions. template - RecursiveRoutingStack advance(Wires& wires, const WireInfos& infos, - const Layout& layout) { + SmallVector advance(Wires& wires, const WireInfos& infos, + const Layout& layout) { DenseSet visited; - RecursiveRoutingStack stack; + SmallVector composites; - // Advance wires past all executable gates and push operations with - // nested regions and the respective wire indices of their inputs onto the - // result stack. + // Advance wires past all executable gates and push composite unitaries and + // the respective wire indices of their inputs onto the vector. walkProgramGraph(wires, [&](const ReadyMap& ready, ReleasedOps& released) { @@ -1257,7 +1167,7 @@ struct MappingPass : impl::MappingPassBase { if (op->getNumRegions() > 0 && visited.insert(op).second) { assert((isa(op))); - stack.emplace_back(op, indices); + composites.emplace_back(op, indices); continue; } } @@ -1269,7 +1179,81 @@ struct MappingPass : impl::MappingPassBase { return WalkResult::advance(); }); - return stack; + // Preserve the block order when multiple independent composite operations + // become ready at once. Hot routing threads every qubit through each + // composite, so processing a later operation first could introduce a + // use-before-definition for an earlier operation. + llvm::sort(composites, + [](const CompositeUnitary& lhs, const CompositeUnitary& rhs) { + assert(lhs.op->getBlock() == rhs.op->getBlock()); + return lhs.op->isBeforeInBlock(rhs.op); + }); + + return composites; + } + + /// Extends the composite unitary's operation to cover all target qubits by + /// adding operands for indices not in the composite's index set. Returns a + /// patch with the updated wire mapping which preserves the parent's wire + /// infos and layout. + RoutingBundle::Patch place(CompositeUnitary& composite, + const RoutingBundle& parent, + IRRewriter& rewriter) { + DenseSet included; // Already included indices. + included.reserve(composite.indices.size()); + + // Maps the i-th included index to its result number. + DenseMap indexToResultNum; + indexToResultNum.reserve(composite.indices.size()); + + for (const auto index : composite.indices) { + const WireIterator& it = parent.wires[index]; + indexToResultNum.try_emplace( + index, cast(it.qubit()).getResultNumber()); + included.insert(index); + } + + const auto allIndices = to_vector(llvm::seq(target->numQubits())); + + const SmallVector excluded(llvm::make_filter_range( + allIndices, [&](const size_t i) { return !included.contains(i); })); + + const SmallVector addons(map_range(excluded, [&](const size_t i) { + // Make sure the qubits point to an already processed operation. + const auto& it = std::prev( + parent.wires[i], parent.wires[i] == std::default_sentinel ? 2 : 1); + return it.qubit(); + })); + + composite = CompositeUnitary{ + .op = TypeSwitch(composite.op) + .Case( + [&](auto cfOp) { return extend(cfOp, addons, rewriter); }) + .Default([](Operation* op) { + report_fatal_error("place: unhandled op: " + + op->getName().getStringRef()); + return nullptr; + }), + .indices = allIndices}; + + const auto results = composite.op->getResults(); + + Wires wires(allIndices.size()); + for (size_t index : included) { + wires[index] = WireIterator(results[indexToResultNum.at(index)]); + } + for (const auto [index, res] : + llvm::zip_equal(excluded, results.take_back(excluded.size()))) { + wires[index] = WireIterator(res); + } + + assert(llvm::all_of(wires, [&](WireIterator& it) { + return it.operation() == composite.op; + })); + + return RoutingBundle::Patch{.layout = std::nullopt, + .infos = std::nullopt, + .wires = std::move(wires)}; } /// Return `values` with only the qubit entries realigned according to the @@ -1296,14 +1280,16 @@ struct MappingPass : impl::MappingPassBase { return realigned; } - /// Processes the recursive stack item by routing the nested operation and - /// inserting epilogue SWAPs. + /// Processes the composite unitary by routing the nested operation and + /// inserting a SWAP appendix. Returns a pair of the patch to apply to the + /// parent bundle and the accumulated statistics, or `failure` if routing + /// fails. template requires(Mode != RoutingMode::Hot || Direction == WireDirection::Forward) - LogicalResult dispatch(const RecursiveRoutingStackItem& item, - RoutingBundle& parent, Statistics& stats, - IRRewriter* rewriter = nullptr) { - const auto& [op, indices] = item; + FailureOr> + dispatch(const CompositeUnitary& composite, const RoutingBundle& parent, + IRRewriter* rewriter = nullptr) { + const auto& [op, indices] = composite; SmallVector permutation(indices.size()); SmallVector children = @@ -1320,10 +1306,6 @@ struct MappingPass : impl::MappingPassBase { return SmallVector( switchOp.getNumRegions(), RoutingBundle{.layout = parent.layout}); - }) - .Default([](Operation* op) -> SmallVector { - report_fatal_error("unhandled region op in dispatch: " + - op->getName().getStringRef()); }); SmallVector> resultToQubitIndex(op->getNumResults()); @@ -1408,11 +1390,16 @@ struct MappingPass : impl::MappingPassBase { // qubit op (note: might be a measurement) before the yield. // TODO: Parallelize multiple children, if possible. + Statistics totalStats; + for (auto& child : children) { - if (failed(route(child, stats, rewriter))) { + const auto stats = route(child, rewriter); + if (failed(stats)) { return failure(); } + totalStats.merge(*stats); + if constexpr (Mode == RoutingMode::Hot) { for_each(child.wires, [](auto& it) { std::advance(it, -2); }); } @@ -1441,10 +1428,13 @@ struct MappingPass : impl::MappingPassBase { children[1].infos.insertOrUpdate(i, prog); } - if (failed(route(children[1], stats, rewriter))) { + const auto stats = route(children[1], rewriter); + if (failed(stats)) { return failure(); } + totalStats.merge(*stats); + if constexpr (Mode == RoutingMode::Hot) { for_each(children[1].wires, [](auto& it) { std::advance(it, -2); }); } @@ -1454,16 +1444,16 @@ struct MappingPass : impl::MappingPassBase { // using the restore (scf::ForOp, scf::While), converge (IfOp), and vote // and restore (IndexSwitchOp) strategies. - const Layout exit = + Layout exit = TypeSwitch(op) .Case([&](scf::ForOp) { const auto swaps = restore(children[0].layout, parent.layout); - insertSWAPs(swaps, children[0], stats, rewriter); + insertSWAPs(swaps, children[0], totalStats, rewriter); return parent.layout; }) .template Case([&](scf::WhileOp) { const auto swaps = restore(children[1].layout, parent.layout); - insertSWAPs(swaps, children[1], stats, rewriter); + insertSWAPs(swaps, children[1], totalStats, rewriter); // The scf::YieldOp is the terminator in the before region and // thus determines the final output layout. return children[0].layout; @@ -1471,8 +1461,8 @@ struct MappingPass : impl::MappingPassBase { .template Case([&](IfOp) { const auto [convergedLayout, fst, snd] = converge(children[0].layout, children[1].layout); - insertSWAPs(fst, children[0], stats, rewriter); - insertSWAPs(snd, children[1], stats, rewriter); + insertSWAPs(fst, children[0], totalStats, rewriter); + insertSWAPs(snd, children[1], totalStats, rewriter); return convergedLayout; }) .template Case([&](IndexSwitchOp) { @@ -1482,7 +1472,7 @@ struct MappingPass : impl::MappingPassBase { })); for (RoutingBundle& child : children) { const auto swaps = restore(child.layout, winner); - insertSWAPs(swaps, child, stats, rewriter); + insertSWAPs(swaps, child, totalStats, rewriter); } return winner; }); @@ -1524,56 +1514,67 @@ struct MappingPass : impl::MappingPassBase { } } - // If the operation is a scf::ForOp, where the parent.layout = - // child.layout, we are done. Otherwise, propagate the final layout and - // index-to-program mapping to the parent. - - if (!isa(op)) { - WireInfos realigendInfos; - for (size_t i = 0; i < parent.wires.size(); ++i) { - const auto oldProg = parent.infos.lookupProgram(i); - const auto oldHw = parent.layout.getHardwareIndex(oldProg); - const auto newProg = exit.getProgramIndex(oldHw); - realigendInfos.insertOrUpdate(i, newProg); - } + // If the operation is a scf::ForOp, where the parent.layout = child.layout, + // we are done. Otherwise, propagate a patch with the final layout and + // index-to-program mapping. - parent.layout = exit; - parent.infos = std::move(realigendInfos); + if (isa(op)) { + return std::make_pair(RoutingBundle::Patch{}, totalStats); } - // Finally, move past the operation with nested regions by - // incrementing the respective global wires. - - for_each(indices, [&](size_t i) { - std::advance(parent.wires[i], WireTraversalTraits::stride()); - }); + RoutingBundle::Patch patch{.layout = std::nullopt, .infos = WireInfos{}}; + for (size_t i = 0; i < parent.wires.size(); ++i) { + const auto oldProg = parent.infos.lookupProgram(i); + const auto oldHw = parent.layout.getHardwareIndex(oldProg); + const auto newProg = exit.getProgramIndex(oldHw); + patch.infos->insertOrUpdate(i, newProg); + } + patch.layout = std::move(exit); - return success(); + return std::make_pair(std::move(patch), totalStats); } /// Iterates over a dynamically computed window of layers and uses A* search /// to find a SWAP sequence that makes each layer executable. Depending on /// the template parameter, this function only updates the layout or also - /// inserts the SWAPs into the IR. The function returns `failure` if A* is - /// unable to find a solution. + /// inserts the SWAPs into the IR. Returns `FailureOr` containing + /// the accumulated statistics on success, or `failure` if A* is unable to + /// find a solution. template requires(Mode != RoutingMode::Hot || Direction == WireDirection::Forward) - LogicalResult route(RoutingBundle& bundle, Statistics& stats, - IRRewriter* rewriter = nullptr) { + FailureOr route(RoutingBundle& bundle, + IRRewriter* rewriter = nullptr) { auto& [wires, infos, layout] = bundle; - while (true) { + Statistics stats; + while (true) { while (true) { - const auto stack = advance(wires, infos, layout); - if (stack.empty()) { + auto composites = advance(wires, infos, layout); + if (composites.empty()) { break; } - for (const auto& item : stack) { - if (dispatch(item, bundle, stats, rewriter) - .failed()) { + + for (auto& composite : composites) { + if constexpr (Mode == RoutingMode::Hot) { + auto patch = place(composite, bundle, *rewriter); + bundle.applyPatch(std::move(patch)); + } + + auto res = dispatch(composite, bundle, rewriter); + if (failed(res)) { return failure(); } + + bundle.applyPatch(std::move(res->first)); + stats.merge(res->second); + + // Once the composite is mapped, move past this op by incrementing the + // respective wires. + + for_each(composite.indices, [&](size_t i) { + std::advance(wires[i], WireTraversalTraits::stride()); + }); } } @@ -1616,7 +1617,7 @@ struct MappingPass : impl::MappingPassBase { } } - return success(); + return stats; } std::optional target; diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 9db4d637be..76a797190c 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -940,6 +940,71 @@ TEST_P(MappingPassTest, MapParallelLoops) { EXPECT_TRUE(isExecutable(getEntryPoint(m.get()), target)); } +TEST_P(MappingPassTest, MapParallelLoopsWithClassicalDependencies) { + const auto& target = GetParam(); + constexpr StringLiteral source = R"mlir( + module { + func.func @main() attributes {passthrough = ["entry_point"]} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %q0 = qco.alloc : !qco.qubit + %q1 = qco.alloc : !qco.qubit + %q2 = qco.alloc : !qco.qubit + %q3 = qco.alloc : !qco.qubit + %q4 = qco.alloc : !qco.qubit + %q5 = qco.alloc : !qco.qubit + %q6 = qco.alloc : !qco.qubit + %q7 = qco.alloc : !qco.qubit + %a0, %a1, %s1 = scf.for %i = %c0 to %c1 step %c1 + iter_args(%x = %q0, %y = %q1, %s = %c1) + -> (!qco.qubit, !qco.qubit, index) { + %nx, %ny = qco.swap %x, %y + : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit + scf.yield %nx, %ny, %s : !qco.qubit, !qco.qubit, index + } + %b0, %b1, %s2 = scf.for %i = %c0 to %s1 step %c1 + iter_args(%x = %q2, %y = %q3, %s = %s1) + -> (!qco.qubit, !qco.qubit, index) { + %nx, %ny = qco.swap %x, %y + : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit + scf.yield %nx, %ny, %s : !qco.qubit, !qco.qubit, index + } + %d0, %d1, %s3 = scf.for %i = %c0 to %s2 step %c1 + iter_args(%x = %q4, %y = %q5, %s = %s2) + -> (!qco.qubit, !qco.qubit, index) { + %nx, %ny = qco.swap %x, %y + : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit + scf.yield %nx, %ny, %s : !qco.qubit, !qco.qubit, index + } + %e0, %e1, %s4 = scf.for %i = %c0 to %s3 step %c1 + iter_args(%x = %q6, %y = %q7, %s = %s3) + -> (!qco.qubit, !qco.qubit, index) { + %nx, %ny = qco.swap %x, %y + : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit + scf.yield %nx, %ny, %s : !qco.qubit, !qco.qubit, index + } + qco.sink %a0 : !qco.qubit + qco.sink %a1 : !qco.qubit + qco.sink %b0 : !qco.qubit + qco.sink %b1 : !qco.qubit + qco.sink %d0 : !qco.qubit + qco.sink %d1 : !qco.qubit + qco.sink %e0 : !qco.qubit + qco.sink %e1 : !qco.qubit + return + } + } + )mlir"; + + auto m = parseSourceString(source, context.get()); + ASSERT_TRUE(m); + ASSERT_TRUE(succeeded(verify(*m))); + ASSERT_TRUE( + runPass(m.get(), target, MappingPassOptions{.ntrials = 1}).succeeded()); + EXPECT_TRUE(succeeded(verify(*m))); + EXPECT_TRUE(isExecutable(getEntryPoint(m.get()), target)); +} + TEST_P(MappingPassTest, MapForWithClassicalIterArg) { const auto& target = GetParam(); constexpr StringLiteral source = R"mlir(