From d4aace7094619154e8a93a55c2e201bc388d24e6 Mon Sep 17 00:00:00 2001 From: Thomas Alexander Date: Wed, 26 Aug 2026 14:45:51 -0300 Subject: [PATCH 1/3] Accelerate commutation-aware endpoint search Build a lazy block-local interaction index keyed by proven logical qubit identity. Segmented streams let find_nearest skip unrelated quantum operations while preserving the existing SSA frontier, endpoint predicates, conservative barriers, and physical fallback for unresolved anchors. Keep index state private to CommutationAnalysis and maintain or discard it through the existing rewrite listener. Reuse normalized anchor support after construction and avoid indexing work on adjacent or otherwise inexpensive searches. Verified with the focused CommutationAwareRewrite unit suite, QuakeSimplify commutation lit coverage, the verifier-enabled linear-value pipeline, public issue reproducers, FTQC hotspots, and NISQ regression guards. Signed-off-by: Thomas Alexander --- .../Optimizer/Analysis/CommutationAnalysis.h | 37 +- .../Transforms/CommutationAwareRewrite.h | 10 +- .../Analysis/CommutationAnalysis.cpp | 450 +++++++++++++++++- .../Transforms/CommutationAwareRewrite.cpp | 171 +++---- 4 files changed, 547 insertions(+), 121 deletions(-) diff --git a/cudaq/include/cudaq/Optimizer/Analysis/CommutationAnalysis.h b/cudaq/include/cudaq/Optimizer/Analysis/CommutationAnalysis.h index e0357f4de92..5467afb79aa 100644 --- a/cudaq/include/cudaq/Optimizer/Analysis/CommutationAnalysis.h +++ b/cudaq/include/cudaq/Optimizer/Analysis/CommutationAnalysis.h @@ -9,6 +9,8 @@ #pragma once #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLFunctionalExtras.h" +#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" #include "mlir/IR/ValueRange.h" #include @@ -29,6 +31,7 @@ class CommutationAwareRewriteListener; namespace cudaq::quake::detail { class QubitIdentityAnalysis; +class LogicalQubitInteractionIndex; /// The outcome of a commutation query. enum class commutation_status { commutes, does_not_commute, indeterminate }; @@ -169,6 +172,8 @@ class CommutationAnalysis { bool canCommute(mlir::Operation *lhs, mlir::Operation *rhs); private: + enum class prior_interaction_walk_result { unavailable, conclusive }; + using OperationPair = std::pair; /// Return true when every control and target role has a distinct known @@ -183,25 +188,49 @@ class CommutationAnalysis { /// identities. bool haveSameOrderedQuantumOperands(mlir::Operation *lhs, mlir::Operation *rhs) const; + /// Return true only when the operation's types, structure, and effects prove + /// that it cannot access or redirect an indexed qubit. + static bool isIgnorableNonQuantumOperation(mlir::Operation *operation); + /// Return the scalar wire established by a supported identity boundary. + static mlir::Value getIdentityBoundaryWire(mlir::Operation *operation); + /// Collect complete scalar-wire captures for a supported ordinary scope. + /// Return false when every search whose anchor wires have known identities + /// must stop at the scope. + static bool + collectScopeWireCaptures(mlir::Operation *operation, + llvm::SmallVectorImpl &captures); + /// Visit indexed interactions at or before `inclusiveUpperBound` in + /// descending block order. `unavailable` instructs the caller to continue + /// the block-order scan. `conclusive` means traversal exhausted the segment, + /// reached its boundary, or the visitor ended the search. + prior_interaction_walk_result + walkPriorInteractions(mlir::Operation *anchor, + mlir::Operation *inclusiveUpperBound, + llvm::function_ref visitor); /// Register a newly inserted scalar-wire operation only when every input /// identity is known. A classical-only insertion succeeds without changing /// identity state only when it is not call-like, owns no regions, and is /// memory-effect-free. Return false for every other insertion. bool registerIdentityPreservingOperation(mlir::Operation *operation); - /// Validate an identity-preserving replacement and clear cached relations. + /// Validate an identity-preserving replacement, clear cached relations, and + /// maintain or discard the interaction index. bool prepareIdentityPreservingReplacement(mlir::Operation *operation, - mlir::ValueRange replacement); - /// Clear cached relations without changing proved qubit identities. + mlir::ValueRange replacement, + mlir::Operation *replacementOp); + /// Clear cached pairwise relations without changing ordered search state. void clearCachedRelations(); - /// Clear cached relations, then remove an operation's result identities. + /// Clear cached relations, remove the operation from the interaction index, + /// then erase its result identities. void eraseOperation(mlir::Operation *operation); mlir::Block *block; std::unique_ptr qubitIdentity; + std::unique_ptr interactionIndex; llvm::DenseMap cache; friend class cudaq::opt::CommutationAwareRewriteMatcher; friend class cudaq::opt::detail::CommutationAwareRewriteListener; + friend class LogicalQubitInteractionIndex; }; } // namespace cudaq::quake::detail diff --git a/cudaq/include/cudaq/Optimizer/Transforms/CommutationAwareRewrite.h b/cudaq/include/cudaq/Optimizer/Transforms/CommutationAwareRewrite.h index db39780b305..204c4e7729a 100644 --- a/cudaq/include/cudaq/Optimizer/Transforms/CommutationAwareRewrite.h +++ b/cudaq/include/cudaq/Optimizer/Transforms/CommutationAwareRewrite.h @@ -45,11 +45,13 @@ struct CommutationAwareRewriteStatistics { /// branched. /// /// The search expects block-local linear-wire Quake. Candidate endpoints are -/// use-def frontier heads on the anchor's own wires. Physical block order -/// audits every intervening operation and selects the latest head when the -/// frontier is split. A frontier head that the consumer declines must have a +/// use-def frontier heads on the anchor's own wires. The search begins with a +/// block-order scan, including the latest head when the frontier is split. Once +/// analysis is required, it uses ordered per-qubit interaction streams when +/// every anchor wire has a known identity. Otherwise it continues the +/// block-order scan. A frontier head that the consumer declines must have a /// pairwise commutation proof with the anchor before the frontier advances. -/// Every other intervening scalar-wire operation requires either that pairwise +/// Every other enumerated scalar-wire operation requires either that pairwise /// proof or a disjoint-support proof. Fresh local identity sources may be /// crossed structurally because they cannot alias an existing logical qubit. /// Other identity boundaries require a disjoint-support proof. diff --git a/cudaq/lib/Optimizer/Analysis/CommutationAnalysis.cpp b/cudaq/lib/Optimizer/Analysis/CommutationAnalysis.cpp index c5d82cce80c..b2a3adb930a 100644 --- a/cudaq/lib/Optimizer/Analysis/CommutationAnalysis.cpp +++ b/cudaq/lib/Optimizer/Analysis/CommutationAnalysis.cpp @@ -8,12 +8,18 @@ #include "cudaq/Optimizer/Analysis/CommutationAnalysis.h" #include "QubitIdentityAnalysis.h" +#include "cudaq/Optimizer/Dialect/CC/CCOps.h" #include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h" +#include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/ErrorHandling.h" #include "mlir/IR/Matchers.h" +#include "mlir/Interfaces/CallInterfaces.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include #include #include #include @@ -677,6 +683,403 @@ static CommutationResult evaluate(Operation *lhs, Operation *rhs, return dispatchRules(lhsView, rhsView); } +bool CommutationAnalysis::isIgnorableNonQuantumOperation(Operation *operation) { + return !isa(operation) && operation->getNumRegions() == 0 && + operation->getNumSuccessors() == 0 && isMemoryEffectFree(operation) && + llvm::none_of(operation->getOperandTypes(), + cudaq::quake::isQuantumType) && + llvm::none_of(operation->getResultTypes(), + cudaq::quake::isQuantumType); +} + +Value CommutationAnalysis::getIdentityBoundaryWire(Operation *operation) { + if (auto unwrap = dyn_cast(operation)) + return unwrap.getResult(); + if (auto borrow = dyn_cast(operation)) + return borrow.getResult(); + if (auto wrapNew = dyn_cast(operation)) + return wrapNew.getWireValue(); + return {}; +} + +bool CommutationAnalysis::collectScopeWireCaptures( + Operation *operation, llvm::SmallVectorImpl &captures) { + auto scope = dyn_cast(operation); + if (!scope || scope.getAtomicQuantumRegionAttr() || + !scope.getInitRegion().hasOneBlock()) + return false; + + WalkResult result = operation->walk([&](Operation *nested) -> WalkResult { + if (nested == operation) + return WalkResult::advance(); + if (isa(nested) || nested->getNumRegions() != 0) + return WalkResult::interrupt(); + if (isIgnorableNonQuantumOperation(nested)) + return WalkResult::advance(); + + auto isScalarWire = [](Type type) { + return !cudaq::quake::isQuantumType(type) || + isa(type); + }; + if (!llvm::all_of(nested->getOperandTypes(), isScalarWire) || + !llvm::all_of(nested->getResultTypes(), isScalarWire)) + return WalkResult::interrupt(); + + for (Value operand : nested->getOperands()) { + if (!cudaq::quake::isQuantumType(operand.getType())) + continue; + Operation *definition = operand.getDefiningOp(); + if (!definition || !operation->isAncestor(definition)) + captures.push_back(operand); + } + + if (cudaq::quake::detail::getScalarWireFlow(nested) || + isa(nested)) + return WalkResult::advance(); + return WalkResult::interrupt(); + }); + return !result.wasInterrupted(); +} + +static std::optional> +resolveQubitSupport(ValueRange values, + const QubitIdentityAnalysis &qubitIdentity) { + llvm::SmallVector support; + llvm::DenseSet seen; + for (Value value : values) { + if (!isa(value.getType())) + return std::nullopt; + auto qubitId = qubitIdentity.getQubitId(value); + if (!qubitId) + return std::nullopt; + if (seen.insert(*qubitId).second) + support.push_back(*qubitId); + } + return support; +} + +// Build one ordered interaction stream for each logical qubit. A search on `q0` +// can then skip an `X(q1)` and visit only `S(q0), H(q0)`. Hard barriers split +// the streams into segments, which keeps indexed searches from crossing an +// operation that the block-order search would reject. +// +// Stable ordinals preserve block order without using MLIR's mutable order +// cache. Erasures leave gaps. A replacement inherits an ordinal only when it +// takes the erased operation's position. The rewrite listener removes erased +// pointers and rebuilds the index after insertions that cannot be placed +// safely. +class cudaq::quake::detail::LogicalQubitInteractionIndex { + // Describe where an operation belongs in the index. Omitted operations have + // empty support. A boundary stays in the preceding segment, while the next + // operation starts a new one, so a segment change ends the search. + struct OperationPosition { + unsigned ordinal; + unsigned segment; + bool isBoundary; + llvm::SmallVector support; + }; + + // Place an operation in one logical-qubit stream. Keeping the ordinal beside + // the pointer lets the streams be searched and merged without looking up the + // operation's position each time. + struct Interaction { + Operation *operation; + unsigned ordinal; + }; + + // Group the per-qubit streams between two hard barriers. A multi-qubit + // operation appears in each stream it affects. + struct Segment { + llvm::DenseMap> interactions; + }; + + struct WalkBounds { + const OperationPosition *anchor; + const OperationPosition *upper; + }; + + struct StreamCursor { + llvm::ArrayRef interactions; + std::ptrdiff_t index; + + const Interaction *current() const { + return index < 0 ? nullptr : &interactions[index]; + } + + void advancePast(unsigned ordinal) { + if (const Interaction *interaction = current(); + interaction && interaction->ordinal == ordinal) + --index; + } + }; + +public: + using PriorInteractionWalkResult = + CommutationAnalysis::prior_interaction_walk_result; + + LogicalQubitInteractionIndex(Block &block, + const QubitIdentityAnalysis &qubitIdentity) { + segments.emplace_back(); + unsigned ordinal = 0; + for (Operation &operation : block) { + unsigned segment = segments.size() - 1; + auto support = classifyIndexedOperation(&operation, qubitIdentity); + OperationPosition indexedPosition{ordinal, segment, !support, {}}; + if (support) + indexedPosition.support = std::move(*support); + auto [position, inserted] = + positions.try_emplace(&operation, std::move(indexedPosition)); + assert(inserted && "block operation already indexed"); + if (!support) { + segments.emplace_back(); + } else { + addInteractions(&operation, position->second); + } + ++ordinal; + } + } + + void noteInsertion(Operation *operation, + const QubitIdentityAnalysis &qubitIdentity) { + auto support = classifyIndexedOperation(operation, qubitIdentity); + // A new barrier or interaction changes the search order and requires a + // rebuild. Ignorable operations and fresh local wires do not. + if (!support || !support->empty()) + pendingInsertions.insert(operation); + } + + bool replaceOrEraseOperation(Operation *operation, Operation *replacement, + const QubitIdentityAnalysis &qubitIdentity) { + // Only a newly inserted replacement needs the old operation's position. + // An existing or value-only replacement adds no new position, so removing + // the old operation is enough. + if (!replacement || !pendingInsertions.contains(replacement)) + return eraseOperation(operation); + auto position = positions.find(operation); + if (position == positions.end()) { + // The rewriter may replace a new operation before the index is rebuilt. + // Drop the old pending operation, but keep a relevant replacement so the + // next rebuild includes it. + pendingInsertions.erase(operation); + return true; + } + pendingInsertions.erase(replacement); + auto replacementSupport = + classifyIndexedOperation(replacement, qubitIdentity); + // The greedy rewriter inserts a replacement beside the operation it will + // erase. Only that adjacency proves the old ordinal still describes the + // replacement's block position. + bool canInheritOrdinal = replacement->getBlock() == operation->getBlock() && + (replacement->getNextNode() == operation || + operation->getNextNode() == replacement); + if (!canInheritOrdinal || position->second.isBoundary || + !replacementSupport) + return false; + + OperationPosition newPosition{position->second.ordinal, + position->second.segment, false, + std::move(*replacementSupport)}; + removeInteractions(operation, position->second); + positions.erase(position); + auto [inserted, didInsert] = + positions.try_emplace(replacement, std::move(newPosition)); + assert(didInsert && "replacement operation already indexed"); + addInteractions(replacement, inserted->second); + return true; + } + + bool eraseOperation(Operation *operation) { + pendingInsertions.erase(operation); + auto position = positions.find(operation); + if (position == positions.end()) + return true; + if (position->second.isBoundary) + return false; + removeInteractions(operation, position->second); + positions.erase(position); + return true; + } + + bool hasPendingInsertions() const { return !pendingInsertions.empty(); } + + PriorInteractionWalkResult + walk(Operation *anchor, Operation *inclusiveUpperBound, + llvm::function_ref visitor) const { + auto bounds = findWalkBounds(anchor, inclusiveUpperBound); + if (!bounds) + return PriorInteractionWalkResult::unavailable; + + // Reuse the anchor support stored with its index position instead of + // resolving the same identities for every search. This also keeps QubitId + // storage private to the index. + llvm::ArrayRef anchorSupport = bounds->anchor->support; + if (bounds->anchor->segment != bounds->upper->segment) + return PriorInteractionWalkResult::conclusive; + + const Segment &segment = segments[bounds->anchor->segment]; + unsigned upperOrdinal = bounds->upper->ordinal; + if (anchorSupport.size() == 1) + return walkSingleQubitStream(segment, anchorSupport.front(), upperOrdinal, + visitor); + walkMergedStreams(segment, anchorSupport, upperOrdinal, visitor); + return PriorInteractionWalkResult::conclusive; + } + +private: + std::optional + findWalkBounds(Operation *anchor, Operation *inclusiveUpperBound) const { + auto anchorPosition = positions.find(anchor); + auto upperPosition = positions.find(inclusiveUpperBound); + if (anchorPosition == positions.end() || upperPosition == positions.end()) + return std::nullopt; + if (anchorPosition->second.support.empty() || + upperPosition->second.ordinal >= anchorPosition->second.ordinal) + return std::nullopt; + return WalkBounds{&anchorPosition->second, &upperPosition->second}; + } + + static std::ptrdiff_t findUpperIndex(llvm::ArrayRef interactions, + unsigned upperOrdinal) { + auto end = + std::upper_bound(interactions.begin(), interactions.end(), upperOrdinal, + [](unsigned ordinal, const Interaction &interaction) { + return ordinal < interaction.ordinal; + }); + return std::distance(interactions.begin(), end) - 1; + } + + PriorInteractionWalkResult + walkSingleQubitStream(const Segment &segment, QubitId qubitId, + unsigned upperOrdinal, + llvm::function_ref visitor) const { + auto stream = segment.interactions.find(qubitId); + if (stream == segment.interactions.end()) + return PriorInteractionWalkResult::conclusive; + for (std::ptrdiff_t index = findUpperIndex(stream->second, upperOrdinal); + index >= 0; --index) + if (!visitor(stream->second[index].operation)) + break; + return PriorInteractionWalkResult::conclusive; + } + + void walkMergedStreams(const Segment &segment, + llvm::ArrayRef anchorSupport, + unsigned upperOrdinal, + llvm::function_ref visitor) const { + auto cursors = makeStreamCursors(segment, anchorSupport, upperOrdinal); + while (const Interaction *next = findLatestInteraction(cursors)) { + // Advance every stream at this ordinal because a multi-qubit operation + // appears in several streams but should be visited only once. + for (StreamCursor &cursor : cursors) + cursor.advancePast(next->ordinal); + if (!visitor(next->operation)) + return; + } + } + + llvm::SmallVector + makeStreamCursors(const Segment &segment, + llvm::ArrayRef anchorSupport, + unsigned upperOrdinal) const { + llvm::SmallVector cursors; + for (QubitId qubitId : anchorSupport) { + auto stream = segment.interactions.find(qubitId); + if (stream == segment.interactions.end()) + continue; + std::ptrdiff_t index = findUpperIndex(stream->second, upperOrdinal); + if (index >= 0) + cursors.push_back({stream->second, index}); + } + return cursors; + } + + static const Interaction * + findLatestInteraction(llvm::ArrayRef cursors) { + const Interaction *latest = nullptr; + for (const StreamCursor &cursor : cursors) { + const Interaction *current = cursor.current(); + if (current && (!latest || current->ordinal > latest->ordinal)) + latest = current; + } + return latest; + } + // Return empty support to omit an operation, non-empty support to add it to + // every affected stream, and no result to end the segment. The index may stop + // more often than the matcher, but it must never omit an operation that the + // block-order search would inspect. + static std::optional> + classifyIndexedOperation(Operation *operation, + const QubitIdentityAnalysis &qubitIdentity) { + if (CommutationAnalysis::isIgnorableNonQuantumOperation(operation) || + isa(operation)) + return llvm::SmallVector{}; + + if (operation->getNumRegions() != 0) { + llvm::SmallVector captures; + if (!CommutationAnalysis::collectScopeWireCaptures(operation, captures)) + return std::nullopt; + return resolveQubitSupport(captures, qubitIdentity); + } + + if (auto flow = cudaq::quake::detail::getScalarWireFlow(operation)) + return resolveQubitSupport(flow->inputs, qubitIdentity); + + llvm::SmallVector support; + if (Value wire = CommutationAnalysis::getIdentityBoundaryWire(operation)) + support.push_back(wire); + else if (auto sink = dyn_cast(operation)) + support.push_back(sink.getTarget()); + else + return std::nullopt; + return resolveQubitSupport(support, qubitIdentity); + } + + void removeInteractions(Operation *operation, + const OperationPosition &position) { + if (position.support.empty()) + return; + Segment &segment = segments[position.segment]; + for (QubitId qubitId : position.support) { + auto stream = segment.interactions.find(qubitId); + assert(stream != segment.interactions.end() && + "operation is missing its interaction stream"); + auto interaction = std::lower_bound( + stream->second.begin(), stream->second.end(), position.ordinal, + [](const Interaction &candidate, unsigned ordinal) { + return candidate.ordinal < ordinal; + }); + assert(interaction != stream->second.end() && + interaction->ordinal == position.ordinal && + interaction->operation == operation && + "operation is missing from its interaction stream"); + stream->second.erase(interaction); + } + } + + void addInteractions(Operation *operation, + const OperationPosition &position) { + if (position.support.empty()) + return; + Segment &segment = segments[position.segment]; + for (QubitId qubitId : position.support) { + auto &stream = segment.interactions[qubitId]; + auto insertion = std::lower_bound( + stream.begin(), stream.end(), position.ordinal, + [](const Interaction &interaction, unsigned ordinal) { + return interaction.ordinal < ordinal; + }); + stream.insert(insertion, {operation, position.ordinal}); + } + } + + llvm::DenseMap positions; + llvm::SmallVector segments; + // A relevant insertion has no index position. Remember it so the next walk + // rebuilds the index before using the old order. + llvm::DenseSet pendingInsertions; +}; + llvm::StringRef cudaq::quake::detail::getCommutationReasonId(commutation_reason reason) { switch (reason) { @@ -769,18 +1172,53 @@ bool CommutationAnalysis::haveSameOrderedQuantumOperands(Operation *lhs, lhsInterface.getTargets(), rhsInterface.getTargets()); } +CommutationAnalysis::prior_interaction_walk_result +CommutationAnalysis::walkPriorInteractions( + Operation *anchor, Operation *inclusiveUpperBound, + llvm::function_ref visitor) { + if (!anchor || !inclusiveUpperBound || anchor->getBlock() != block || + inclusiveUpperBound->getBlock() != block) + return prior_interaction_walk_result::unavailable; + if (interactionIndex && interactionIndex->hasPendingInsertions()) + interactionIndex.reset(); + + if (!interactionIndex) { + // Avoid a block-wide build until the anchor has known support that can use + // the index. + auto flow = cudaq::quake::detail::getScalarWireFlow(anchor); + if (!flow) + return prior_interaction_walk_result::unavailable; + auto support = resolveQubitSupport(flow->inputs, *qubitIdentity); + if (!support || support->empty()) + return prior_interaction_walk_result::unavailable; + interactionIndex = + std::make_unique(*block, *qubitIdentity); + } + return interactionIndex->walk(anchor, inclusiveUpperBound, visitor); +} + bool CommutationAnalysis::registerIdentityPreservingOperation( Operation *operation) { - return operation && operation->getBlock() == block && - qubitIdentity->registerOperation(*operation); + if (!operation || operation->getBlock() != block || + !qubitIdentity->registerOperation(*operation)) + return false; + if (interactionIndex) + interactionIndex->noteInsertion(operation, *qubitIdentity); + return true; } bool CommutationAnalysis::prepareIdentityPreservingReplacement( - Operation *operation, ValueRange replacement) { + Operation *operation, ValueRange replacement, Operation *replacementOp) { if (!operation || operation->getBlock() != block || !qubitIdentity->replacementPreservesIdentities(*operation, replacement)) return false; - clearCachedRelations(); + cache.clear(); + if (interactionIndex) { + bool updated = interactionIndex->replaceOrEraseOperation( + operation, replacementOp, *qubitIdentity); + if (!updated) + interactionIndex.reset(); + } return true; } @@ -789,7 +1227,9 @@ void CommutationAnalysis::clearCachedRelations() { cache.clear(); } void CommutationAnalysis::eraseOperation(Operation *operation) { if (!operation || operation->getBlock() != block) return; - clearCachedRelations(); + cache.clear(); + if (interactionIndex && !interactionIndex->eraseOperation(operation)) + interactionIndex.reset(); qubitIdentity->eraseOperation(*operation); } diff --git a/cudaq/lib/Optimizer/Transforms/CommutationAwareRewrite.cpp b/cudaq/lib/Optimizer/Transforms/CommutationAwareRewrite.cpp index ec0b1dd130c..2b2bdc22cb1 100644 --- a/cudaq/lib/Optimizer/Transforms/CommutationAwareRewrite.cpp +++ b/cudaq/lib/Optimizer/Transforms/CommutationAwareRewrite.cpp @@ -8,15 +8,12 @@ #include "cudaq/Optimizer/Transforms/CommutationAwareRewrite.h" #include "cudaq/Optimizer/Analysis/CommutationAnalysis.h" -#include "cudaq/Optimizer/Dialect/CC/CCOps.h" #include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" -#include "mlir/Interfaces/CallInterfaces.h" -#include "mlir/Interfaces/SideEffectInterfaces.h" #include #include @@ -76,70 +73,6 @@ static bool collectFrontierHeads(llvm::ArrayRef frontier, return true; } -// Operations without quantum values can be ignored only when their structure -// and effects cannot hide access to an anchor qubit. Calls, region owners, and -// effectful operations remain conservative barriers. -static bool isKnownNonQuantumOperation(Operation *operation) { - return !isa(operation) && operation->getNumRegions() == 0 && - operation->getNumSuccessors() == 0 && isMemoryEffectFree(operation) && - llvm::none_of(operation->getOperandTypes(), - cudaq::quake::isQuantumType) && - llvm::none_of(operation->getResultTypes(), - cudaq::quake::isQuantumType); -} - -// Collect scalar wires captured by an ordinary single-block scope. Marked -// atomic scopes remain barriers even when their captures are disjoint. Every -// quantum operation inside must expose supported scalar-wire flow or be a local -// wire source, sink, or scope terminator. -static bool collectScopeWireCaptures(Operation *operation, - llvm::SmallVectorImpl &captures) { - auto scope = dyn_cast(operation); - if (!scope || scope.getAtomicQuantumRegionAttr() || - !scope.getInitRegion().hasOneBlock()) - return false; - - for (Value operand : operation->getOperands()) { - if (!cudaq::quake::isQuantumType(operand.getType())) - continue; - if (!isa(operand.getType())) - return false; - captures.push_back(operand); - } - - WalkResult result = operation->walk([&](Operation *nested) -> WalkResult { - if (nested == operation) - return WalkResult::advance(); - if (isa(nested) || nested->getNumRegions() != 0) - return WalkResult::interrupt(); - if (isKnownNonQuantumOperation(nested)) - return WalkResult::advance(); - - auto isScalarWire = [](Type type) { - return !cudaq::quake::isQuantumType(type) || - isa(type); - }; - if (!llvm::all_of(nested->getOperandTypes(), isScalarWire) || - !llvm::all_of(nested->getResultTypes(), isScalarWire)) - return WalkResult::interrupt(); - - for (Value operand : nested->getOperands()) { - if (!cudaq::quake::isQuantumType(operand.getType())) - continue; - Operation *definition = operand.getDefiningOp(); - if (!definition || !operation->isAncestor(definition)) - captures.push_back(operand); - } - - if (cudaq::quake::detail::getScalarWireFlow(nested) || - isa(nested)) - return WalkResult::advance(); - return WalkResult::interrupt(); - }); - return !result.wasInterrupted(); -} - // `QubitIdentityAnalysis` identifies logical qubits, not SSA paths. For // example, an endpoint can consume a second `borrow_wire` for the same wire-set // slot while the frontier lane from the anchor reaches another defining @@ -345,14 +278,20 @@ class CommutationAwareRewriteListener Operation *replacement) override { RewriterBase::ForwardingListener::notifyOperationReplaced(operation, replacement); - updateReplacement(operation, replacement->getResults()); + updateReplacement(operation, replacement->getResults(), replacement); } void notifyOperationReplaced(Operation *operation, ValueRange replacement) override { RewriterBase::ForwardingListener::notifyOperationReplaced(operation, replacement); - updateReplacement(operation, replacement); + Operation *replacementOp = nullptr; + if (!replacement.empty()) { + Operation *definition = replacement.front().getDefiningOp(); + if (definition && llvm::equal(definition->getResults(), replacement)) + replacementOp = definition; + } + updateReplacement(operation, replacement, replacementOp); } void notifyOperationErased(Operation *operation) override { @@ -403,7 +342,8 @@ class CommutationAwareRewriteListener discardBlock(block); } - void updateReplacement(Operation *operation, ValueRange replacement) { + void updateReplacement(Operation *operation, ValueRange replacement, + Operation *replacementOp) { // Replacement callbacks and their per-use modification callbacks are // synchronous. Starting another replacement or falling back must never // leave counts that could suppress a later genuine modification. @@ -412,11 +352,12 @@ class CommutationAwareRewriteListener if (analysis == matcher.impl->analyses.end()) return; - // A validated replacement preserves qubit identity state and clears cached - // relations. Modification notifications cover every rewired quantum or - // classical use. Failure requires block fallback. - if (!analysis->second->prepareIdentityPreservingReplacement(operation, - replacement)) { + // A valid replacement preserves qubit identities, clears cached relations, + // and either updates the old index position or discards the index. + // Modification notifications cover every rewired quantum or classical use. + // Failure requires block fallback. + if (!analysis->second->prepareIdentityPreservingReplacement( + operation, replacement, replacementOp)) { discardBlock(operation->getBlock()); return; } @@ -487,9 +428,6 @@ Operation *cudaq::opt::CommutationAwareRewriteMatcher::find_nearest( return nullptr; Block *block = anchor->getBlock(); - // A later anchor aligns the search with the greedy driver's bottom-up - // schedule. The physical scan audits the interval while selecting use-def - // frontier heads, including the latest head when the frontier is split. auto frontier = openFrontier(anchor, *anchorFlow); cudaq::quake::detail::CommutationAnalysis *analysis = nullptr; @@ -498,54 +436,47 @@ Operation *cudaq::opt::CommutationAwareRewriteMatcher::find_nearest( analysis = &impl->getAnalysis(block); return *analysis; }; - auto canCrossRegion = [&](Operation *operation) { + auto canCrossScope = [&](Operation *operation) { llvm::SmallVector captures; - return collectScopeWireCaptures(operation, captures) && + return cudaq::quake::detail::CommutationAnalysis::collectScopeWireCaptures( + operation, captures) && requireAnalysis().hasDisjointQuantumSupport(anchor, captures); }; auto canCrossIdentityBoundary = [&](Operation *operation) { if (isa(operation)) return true; - llvm::SmallVector wires; - if (auto unwrap = dyn_cast(operation)) - wires.push_back(unwrap.getResult()); - else if (auto borrow = dyn_cast(operation)) - wires.push_back(borrow.getResult()); - else if (auto wrapNew = dyn_cast(operation)) - wires.push_back(wrapNew.getWireValue()); - else + Value wire = + cudaq::quake::detail::CommutationAnalysis::getIdentityBoundaryWire( + operation); + if (!wire) return false; + llvm::SmallVector wires{wire}; return requireAnalysis().hasDisjointQuantumSupport(anchor, wires); }; llvm::DenseSet frontierHeads; if (!collectFrontierHeads(frontier, block, frontierHeads)) return nullptr; - // Audit physical block order while selecting use-def frontier heads. This - // catches operations on aliases that are absent from the anchor's SSA - // chains. The monotone scan visits each intervening operation at most once - // per search without consulting MLIR's mutable operation-order cache. - for (Operation *candidate = anchor->getPrevNode(); candidate; - candidate = candidate->getPrevNode()) { + Operation *match = nullptr; + auto processCandidate = [&](Operation *candidate) { bool isFrontierHead = frontierHeads.contains(candidate); - if (!isFrontierHead && isKnownNonQuantumOperation(candidate)) - continue; + if (!isFrontierHead && cudaq::quake::detail::CommutationAnalysis:: + isIgnorableNonQuantumOperation(candidate)) + return true; if (!isFrontierHead && candidate->getNumRegions() != 0) { - if (canCrossRegion(candidate)) - continue; - return nullptr; + return canCrossScope(candidate); } auto candidateFlow = cudaq::quake::detail::getScalarWireFlow(candidate); if (!isFrontierHead && !candidateFlow) { if (canCrossIdentityBoundary(candidate) || requireAnalysis().canCommute(anchor, candidate)) - continue; - return nullptr; + return true; + return false; } if (!candidateFlow) - return nullptr; + return false; auto candidateInterface = dyn_cast(candidate); @@ -553,11 +484,12 @@ Operation *cudaq::opt::CommutationAwareRewriteMatcher::find_nearest( if (isFrontierHead && candidateInterface && (!has_distinct_quantum_operands(anchor) || !has_distinct_quantum_operands(candidate))) - return nullptr; + return false; if (isFrontierHead && candidateInterface && isEndpoint(candidate)) { if (!doesCompleteFrontierReach(frontier, candidate)) - return nullptr; - return candidate; + return false; + match = candidate; + return false; } // Every frontier head the consumer declines and every other crossed @@ -565,15 +497,38 @@ Operation *cudaq::opt::CommutationAwareRewriteMatcher::find_nearest( auto &blockAnalysis = requireAnalysis(); if (!isFrontierHead && blockAnalysis.hasDisjointQuantumSupport(anchor, candidateFlow->inputs)) - continue; + return true; if (!blockAnalysis.canCommute(anchor, candidate)) - return nullptr; + return false; if (isFrontierHead) { stepFrontierBackward(frontier, candidate, *candidateFlow); if (!collectFrontierHeads(frontier, block, frontierHeads)) - return nullptr; + return false; + } + return true; + }; + + // Check nearby frontier heads and ignorable operations before building the + // index. Once analysis is needed, indexed traversal skips operations on + // unrelated known qubits when it is available. + bool indexUnavailable = false; + for (Operation *candidate = anchor->getPrevNode(); candidate;) { + if (!processCandidate(candidate)) + return match; + + Operation *inclusiveUpperBound = candidate->getPrevNode(); + if (!inclusiveUpperBound) + return nullptr; + if (analysis && !indexUnavailable) { + auto walkResult = analysis->walkPriorInteractions( + anchor, inclusiveUpperBound, processCandidate); + if (walkResult == cudaq::quake::detail::CommutationAnalysis:: + prior_interaction_walk_result::conclusive) + return match; + indexUnavailable = true; } + candidate = inclusiveUpperBound; } return nullptr; } From e80d6e78528ebde59fd219f8af861998ee317b3c Mon Sep 17 00:00:00 2001 From: Thomas Alexander Date: Thu, 27 Aug 2026 11:34:11 -0300 Subject: [PATCH 2/3] Avoid rebuilding commutation index during legalization Run rotation legalization between two commutation-aware simplification phases. This preserves pre-legalization rotation combining and post-legalization cleanup without invalidating the interaction index after every expanding rewrite. Add focused coverage for cleanup between a legalized rotation sequence and an existing adjacent gate. Verified with the QuakeSimplify and commutation-aware focused suites on the combined FTQC candidate. Signed-off-by: Thomas Alexander --- .../Optimizer/Transforms/QuakeSimplify.cpp | 89 ++++++++++++------- .../Transforms/quake_simplify_clifford_t.qke | 28 +++++- 2 files changed, 81 insertions(+), 36 deletions(-) diff --git a/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp b/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp index 4d5f9648282..a276d1e96b5 100644 --- a/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp +++ b/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp @@ -835,38 +835,63 @@ class QuakeSimplifyPass GreedyRewriteConfig config; config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled); auto *ctx = &getContext(); - cudaq::opt::CommutationAwareRewriteDriver driver(*ctx, config); - auto &patterns = driver.get_patterns(); - patterns.add(ctx, numResetsErased); - patterns.add(ctx, numReduceYSXRewrites); - auto &matcher = driver.get_matcher(); - - patterns.add, - RotationCombine, - RotationCombine, - RotationCombine>( - ctx, matcher, threshold, numZeroRotationsEliminated, - numRotationsCombined); - patterns.add>( - ctx, matcher, numDoubleSRewrites); - patterns.add>( - ctx, matcher, numDoubleTRewrites); - patterns.add, - InverseElimination, - InverseElimination, - InverseElimination, - InverseElimination>( - ctx, matcher, numHermitianEliminations); - patterns.add, - InverseElimination>( - ctx, matcher, numAdjointEliminations); - if (rotationsToCliffordT) - populateRotationsToCliffordTPatterns(patterns, cliffordTEpsilon, - numCliffordTRotations); - - if (failed(driver.run(getOperation()->getRegion(0)))) + Region ®ion = getOperation()->getRegion(0); + // Each simplification run gets a fresh matcher and interaction index. + auto simplify = [&]() { + cudaq::opt::CommutationAwareRewriteDriver driver(*ctx, config); + auto &patterns = driver.get_patterns(); + patterns.add(ctx, numResetsErased); + patterns.add(ctx, numReduceYSXRewrites); + auto &matcher = driver.get_matcher(); + + patterns.add, + RotationCombine, + RotationCombine, + RotationCombine>( + ctx, matcher, threshold, numZeroRotationsEliminated, + numRotationsCombined); + patterns.add>( + ctx, matcher, numDoubleSRewrites); + patterns.add>( + ctx, matcher, numDoubleTRewrites); + patterns.add, + InverseElimination, + InverseElimination, + InverseElimination, + InverseElimination>( + ctx, matcher, numHermitianEliminations); + patterns.add, + InverseElimination>( + ctx, matcher, numAdjointEliminations); + return driver.run(region); + }; + + if (rotationsToCliffordT) { + // Legalization expands one rotation into several gates. Keeping it out + // of the commutation-aware simplification loop avoids rebuilding the + // interaction index after every expansion. + + // First combine rotations so legalization only sees those that remain. + if (failed(simplify())) { + signalPassFailure(); + return; + } + + RewritePatternSet legalizationPatterns(ctx); + populateRotationsToCliffordTPatterns( + legalizationPatterns, cliffordTEpsilon, numCliffordTRotations); + if (failed(applyPatternsGreedily(region, std::move(legalizationPatterns), + config))) { + signalPassFailure(); + return; + } + } + + // Clean up gates created by legalization with a fresh simplification run. + // Without legalization, this is the pass's only simplification run. + if (failed(simplify())) signalPassFailure(); } }; diff --git a/cudaq/test/Transforms/quake_simplify_clifford_t.qke b/cudaq/test/Transforms/quake_simplify_clifford_t.qke index 6b11250dc47..1ded393634c 100644 --- a/cudaq/test/Transforms/quake_simplify_clifford_t.qke +++ b/cudaq/test/Transforms/quake_simplify_clifford_t.qke @@ -26,10 +26,10 @@ // `--symbol-dce` can remove unrelated kernels and `CircuitCheck` compares only // the selected test group. // RUN: cudaq-opt \ -// RUN: --symbol-privatize='exclude=r1_residues,r1_modulo,rx_residues,rx_modulo,ry_residues,ry_modulo,rz_residues,rz_modulo' \ +// RUN: --symbol-privatize='exclude=r1_residues,r1_modulo,rx_residues,rx_modulo,ry_residues,ry_modulo,rz_residues,rz_modulo,cleanup_after_legalization' \ // RUN: --symbol-dce %s -o %t.exact.qke // RUN: cudaq-opt --quake-simplify='rotations-to-clifford-t=true clifford-t-epsilon=0' \ -// RUN: --symbol-privatize='exclude=r1_residues,r1_modulo,rx_residues,rx_modulo,ry_residues,ry_modulo,rz_residues,rz_modulo' \ +// RUN: --symbol-privatize='exclude=r1_residues,r1_modulo,rx_residues,rx_modulo,ry_residues,ry_modulo,rz_residues,rz_modulo,cleanup_after_legalization' \ // RUN: --symbol-dce %s | \ // RUN: CircuitCheck %t.exact.qke @@ -500,8 +500,8 @@ func.func @unsupported_phased_rx() { // EXACT-LABEL: func.func @unsupported_phased_rx // EXACT: quake.phased_rx -// Constant folding and rotation merging participate in the same greedy -// fixpoint, so both ways of producing pi/4 reach the exact family. +// Pre-legalization simplification folds and merges both ways of producing +// pi/4, allowing the subsequent legalization phase to recognize them. func.func @fold_and_merge() { %eighth = arith.constant 0.39269908169872414 : f64 %sum = arith.addf %eighth, %eighth : f64 @@ -519,6 +519,26 @@ func.func @fold_and_merge() { // EXACT-COUNT-2: quake.t // EXACT-NOT: quake.r1 +// Legalizing Rx(pi/4) emits phase, H, T, H. The fresh simplification driver +// must cancel the final emitted H with the existing H. +func.func @cleanup_after_legalization() { + %quarter = arith.constant 0.7853981633974483 : f64 + %q = quake.null_wire + %r = quake.rx (%quarter) %q : (f64, !quake.wire) -> !quake.wire + %h = quake.h %r : (!quake.wire) -> !quake.wire + quake.sink %h : !quake.wire + return +} + +// EXACT-LABEL: func.func @cleanup_after_legalization +// EXACT: %[[PHASE_ANGLE:.*]] = arith.constant -0.39269908169872414 : f64 +// EXACT-NEXT: %[[Q:.*]] = quake.null_wire +// EXACT-NEXT: %[[PHASE:.*]] = quake.phase (%[[PHASE_ANGLE]]) %[[Q]] : (f64, !quake.wire) -> !quake.wire +// EXACT-NEXT: %[[H:.*]] = quake.h %[[PHASE]] : (!quake.wire) -> !quake.wire +// EXACT-NEXT: %[[T:.*]] = quake.t %[[H]] : (!quake.wire) -> !quake.wire +// EXACT-NEXT: quake.sink %[[T]] : !quake.wire +// EXACT-NOT: quake.rx + // The later rotation must combine backward before the exact Clifford+T family // can legalize it. The sum is not a pi/4 multiple, so the combined rotation // remains at the later location. From cab6a4d3084687a2f757a824ee9aad0c732abf9a Mon Sep 17 00:00:00 2001 From: Thomas Alexander Date: Thu, 27 Aug 2026 20:46:10 -0300 Subject: [PATCH 3/3] Clarify commutation search index contracts Rename the index around the operations stored for each logical qubit and define segment boundaries as unconditional indexed-search stops. Replace the two-value walk result with a boolean completion contract, and document traversal order, fallback behavior, query-specific stops, and ordinal maintenance. Verified by building cudaq-opt and OptimizerUnitTests. All 41 optimizer unit tests and both focused lit tests pass. The issue reproducer output hashes remain unchanged. Signed-off-by: Thomas Alexander --- .../Optimizer/Analysis/CommutationAnalysis.h | 30 +- .../Transforms/CommutationAwareRewrite.h | 8 +- .../Analysis/CommutationAnalysis.cpp | 315 +++++++++--------- .../Transforms/CommutationAwareRewrite.cpp | 11 +- .../Optimizer/Transforms/QuakeSimplify.cpp | 5 +- 5 files changed, 189 insertions(+), 180 deletions(-) diff --git a/cudaq/include/cudaq/Optimizer/Analysis/CommutationAnalysis.h b/cudaq/include/cudaq/Optimizer/Analysis/CommutationAnalysis.h index 5467afb79aa..dfe039dbb32 100644 --- a/cudaq/include/cudaq/Optimizer/Analysis/CommutationAnalysis.h +++ b/cudaq/include/cudaq/Optimizer/Analysis/CommutationAnalysis.h @@ -31,7 +31,7 @@ class CommutationAwareRewriteListener; namespace cudaq::quake::detail { class QubitIdentityAnalysis; -class LogicalQubitInteractionIndex; +class LogicalQubitOperationIndex; /// The outcome of a commutation query. enum class commutation_status { commutes, does_not_commute, indeterminate }; @@ -172,8 +172,6 @@ class CommutationAnalysis { bool canCommute(mlir::Operation *lhs, mlir::Operation *rhs); private: - enum class prior_interaction_walk_result { unavailable, conclusive }; - using OperationPair = std::pair; /// Return true when every control and target role has a distinct known @@ -199,38 +197,40 @@ class CommutationAnalysis { static bool collectScopeWireCaptures(mlir::Operation *operation, llvm::SmallVectorImpl &captures); - /// Visit indexed interactions at or before `inclusiveUpperBound` in - /// descending block order. `unavailable` instructs the caller to continue - /// the block-order scan. `conclusive` means traversal exhausted the segment, - /// reached its boundary, or the visitor ended the search. - prior_interaction_walk_result - walkPriorInteractions(mlir::Operation *anchor, - mlir::Operation *inclusiveUpperBound, - llvm::function_ref visitor); + /// Try to finish the remaining search using operations indexed by the + /// anchor's logical qubits. Visit operations at or before + /// `inclusiveUpperBound` in descending block order. Return true when the + /// visitor ends the search, the current segment is exhausted, or a segment + /// boundary prevents further traversal. Return false when the caller must + /// continue with the block-order scan. + bool + tryWalkPriorOperations(mlir::Operation *anchor, + mlir::Operation *inclusiveUpperBound, + llvm::function_ref visitor); /// Register a newly inserted scalar-wire operation only when every input /// identity is known. A classical-only insertion succeeds without changing /// identity state only when it is not call-like, owns no regions, and is /// memory-effect-free. Return false for every other insertion. bool registerIdentityPreservingOperation(mlir::Operation *operation); /// Validate an identity-preserving replacement, clear cached relations, and - /// maintain or discard the interaction index. + /// maintain or discard the operation index. bool prepareIdentityPreservingReplacement(mlir::Operation *operation, mlir::ValueRange replacement, mlir::Operation *replacementOp); /// Clear cached pairwise relations without changing ordered search state. void clearCachedRelations(); - /// Clear cached relations, remove the operation from the interaction index, + /// Clear cached relations, remove the operation from the operation index, /// then erase its result identities. void eraseOperation(mlir::Operation *operation); mlir::Block *block; std::unique_ptr qubitIdentity; - std::unique_ptr interactionIndex; + std::unique_ptr operationIndex; llvm::DenseMap cache; friend class cudaq::opt::CommutationAwareRewriteMatcher; friend class cudaq::opt::detail::CommutationAwareRewriteListener; - friend class LogicalQubitInteractionIndex; + friend class LogicalQubitOperationIndex; }; } // namespace cudaq::quake::detail diff --git a/cudaq/include/cudaq/Optimizer/Transforms/CommutationAwareRewrite.h b/cudaq/include/cudaq/Optimizer/Transforms/CommutationAwareRewrite.h index 204c4e7729a..b1282616bfa 100644 --- a/cudaq/include/cudaq/Optimizer/Transforms/CommutationAwareRewrite.h +++ b/cudaq/include/cudaq/Optimizer/Transforms/CommutationAwareRewrite.h @@ -47,10 +47,10 @@ struct CommutationAwareRewriteStatistics { /// The search expects block-local linear-wire Quake. Candidate endpoints are /// use-def frontier heads on the anchor's own wires. The search begins with a /// block-order scan, including the latest head when the frontier is split. Once -/// analysis is required, it uses ordered per-qubit interaction streams when -/// every anchor wire has a known identity. Otherwise it continues the -/// block-order scan. A frontier head that the consumer declines must have a -/// pairwise commutation proof with the anchor before the frontier advances. +/// analysis is required, it uses a per-qubit operation index when every anchor +/// wire has a known identity. Otherwise it continues the block-order scan. A +/// frontier head that the consumer declines must have a pairwise commutation +/// proof with the anchor before the frontier advances. /// Every other enumerated scalar-wire operation requires either that pairwise /// proof or a disjoint-support proof. Fresh local identity sources may be /// crossed structurally because they cannot alias an existing logical qubit. diff --git a/cudaq/lib/Optimizer/Analysis/CommutationAnalysis.cpp b/cudaq/lib/Optimizer/Analysis/CommutationAnalysis.cpp index b2a3adb930a..98d54fd4a28 100644 --- a/cudaq/lib/Optimizer/Analysis/CommutationAnalysis.cpp +++ b/cudaq/lib/Optimizer/Analysis/CommutationAnalysis.cpp @@ -759,82 +759,86 @@ resolveQubitSupport(ValueRange values, return support; } -// Build one ordered interaction stream for each logical qubit. A search on `q0` -// can then skip an `X(q1)` and visit only `S(q0), H(q0)`. Hard barriers split -// the streams into segments, which keeps indexed searches from crossing an -// operation that the block-order search would reject. +// Index block operations by the known logical qubits they may affect. For +// example, a search on `q0` reads the ordered operations for `q0` and skips an +// `X(q1)`. A multi-qubit search merges the ordered operations for its qubits. +// +// Calls, unsupported regions, and operations with unknown qubit support divide +// the index into segments. A segment boundary is an unconditional stop for an +// indexed search. Query-specific stops, such as a noncommuting operation on +// `q0`, remain in `q0`'s list and are decided by the visitor. // // Stable ordinals preserve block order without using MLIR's mutable order // cache. Erasures leave gaps. A replacement inherits an ordinal only when it // takes the erased operation's position. The rewrite listener removes erased // pointers and rebuilds the index after insertions that cannot be placed // safely. -class cudaq::quake::detail::LogicalQubitInteractionIndex { - // Describe where an operation belongs in the index. Omitted operations have - // empty support. A boundary stays in the preceding segment, while the next - // operation starts a new one, so a segment change ends the search. +class cudaq::quake::detail::LogicalQubitOperationIndex { + // Record an operation's immutable block position and index membership. A + // segment boundary belongs to the preceding segment, so a backward search + // from the next segment stops before crossing it. struct OperationPosition { unsigned ordinal; unsigned segment; - bool isBoundary; + bool isSegmentBoundary; llvm::SmallVector support; }; - // Place an operation in one logical-qubit stream. Keeping the ordinal beside - // the pointer lets the streams be searched and merged without looking up the - // operation's position each time. - struct Interaction { + // Store an operation in block order under each logical qubit it may affect. + // The copied ordinal makes those ordered lists directly searchable. + struct IndexedOperation { Operation *operation; unsigned ordinal; }; - // Group the per-qubit streams between two hard barriers. A multi-qubit - // operation appears in each stream it affects. + // Group the per-qubit operation lists within one searchable block interval. + // A multi-qubit operation appears in every list it affects. struct Segment { - llvm::DenseMap> interactions; + llvm::DenseMap> + operationsByQubit; }; - struct WalkBounds { + // Hold the indexed positions that bound one backward walk. + struct WalkPositions { const OperationPosition *anchor; const OperationPosition *upper; }; - struct StreamCursor { - llvm::ArrayRef interactions; + // Track one per-qubit operation list during a multi-qubit backward merge. + struct OperationCursor { + llvm::ArrayRef operations; std::ptrdiff_t index; - const Interaction *current() const { - return index < 0 ? nullptr : &interactions[index]; + const IndexedOperation *current() const { + return index < 0 ? nullptr : &operations[index]; } void advancePast(unsigned ordinal) { - if (const Interaction *interaction = current(); - interaction && interaction->ordinal == ordinal) + if (const IndexedOperation *operation = current(); + operation && operation->ordinal == ordinal) --index; } }; public: - using PriorInteractionWalkResult = - CommutationAnalysis::prior_interaction_walk_result; - - LogicalQubitInteractionIndex(Block &block, - const QubitIdentityAnalysis &qubitIdentity) { + LogicalQubitOperationIndex(Block &block, + const QubitIdentityAnalysis &qubitIdentity) { segments.emplace_back(); unsigned ordinal = 0; for (Operation &operation : block) { unsigned segment = segments.size() - 1; - auto support = classifyIndexedOperation(&operation, qubitIdentity); - OperationPosition indexedPosition{ordinal, segment, !support, {}}; - if (support) - indexedPosition.support = std::move(*support); - auto [position, inserted] = + auto indexedSupport = + classifyOperationForIndex(&operation, qubitIdentity); + OperationPosition indexedPosition{ordinal, segment, !indexedSupport, {}}; + if (indexedSupport) + indexedPosition.support = std::move(*indexedSupport); + auto [position, didInsert] = positions.try_emplace(&operation, std::move(indexedPosition)); - assert(inserted && "block operation already indexed"); - if (!support) { + assert(didInsert && "block operation already indexed"); + if (!indexedSupport) { segments.emplace_back(); } else { - addInteractions(&operation, position->second); + addOperation(&operation, position->second); } ++ordinal; } @@ -842,10 +846,10 @@ class cudaq::quake::detail::LogicalQubitInteractionIndex { void noteInsertion(Operation *operation, const QubitIdentityAnalysis &qubitIdentity) { - auto support = classifyIndexedOperation(operation, qubitIdentity); - // A new barrier or interaction changes the search order and requires a - // rebuild. Ignorable operations and fresh local wires do not. - if (!support || !support->empty()) + auto indexedSupport = classifyOperationForIndex(operation, qubitIdentity); + // A new segment boundary or indexed operation changes the search order and + // requires a rebuild. Ignorable operations and fresh local wires do not. + if (!indexedSupport || !indexedSupport->empty()) pendingInsertions.insert(operation); } @@ -866,26 +870,26 @@ class cudaq::quake::detail::LogicalQubitInteractionIndex { } pendingInsertions.erase(replacement); auto replacementSupport = - classifyIndexedOperation(replacement, qubitIdentity); + classifyOperationForIndex(replacement, qubitIdentity); // The greedy rewriter inserts a replacement beside the operation it will // erase. Only that adjacency proves the old ordinal still describes the // replacement's block position. bool canInheritOrdinal = replacement->getBlock() == operation->getBlock() && (replacement->getNextNode() == operation || operation->getNextNode() == replacement); - if (!canInheritOrdinal || position->second.isBoundary || + if (!canInheritOrdinal || position->second.isSegmentBoundary || !replacementSupport) return false; OperationPosition newPosition{position->second.ordinal, position->second.segment, false, std::move(*replacementSupport)}; - removeInteractions(operation, position->second); + removeOperation(operation, position->second); positions.erase(position); - auto [inserted, didInsert] = + auto [replacementPosition, didInsert] = positions.try_emplace(replacement, std::move(newPosition)); assert(didInsert && "replacement operation already indexed"); - addInteractions(replacement, inserted->second); + addOperation(replacement, replacementPosition->second); return true; } @@ -894,41 +898,42 @@ class cudaq::quake::detail::LogicalQubitInteractionIndex { auto position = positions.find(operation); if (position == positions.end()) return true; - if (position->second.isBoundary) + if (position->second.isSegmentBoundary) return false; - removeInteractions(operation, position->second); + removeOperation(operation, position->second); positions.erase(position); return true; } bool hasPendingInsertions() const { return !pendingInsertions.empty(); } - PriorInteractionWalkResult - walk(Operation *anchor, Operation *inclusiveUpperBound, - llvm::function_ref visitor) const { - auto bounds = findWalkBounds(anchor, inclusiveUpperBound); - if (!bounds) - return PriorInteractionWalkResult::unavailable; + bool + tryWalkPriorOperations(Operation *anchor, Operation *inclusiveUpperBound, + llvm::function_ref visitor) const { + auto walkPositions = findWalkPositions(anchor, inclusiveUpperBound); + if (!walkPositions) + return false; // Reuse the anchor support stored with its index position instead of // resolving the same identities for every search. This also keeps QubitId // storage private to the index. - llvm::ArrayRef anchorSupport = bounds->anchor->support; - if (bounds->anchor->segment != bounds->upper->segment) - return PriorInteractionWalkResult::conclusive; + llvm::ArrayRef anchorSupport = walkPositions->anchor->support; + if (walkPositions->anchor->segment != walkPositions->upper->segment) + return true; - const Segment &segment = segments[bounds->anchor->segment]; - unsigned upperOrdinal = bounds->upper->ordinal; + const Segment &segment = segments[walkPositions->anchor->segment]; + unsigned upperOrdinal = walkPositions->upper->ordinal; if (anchorSupport.size() == 1) - return walkSingleQubitStream(segment, anchorSupport.front(), upperOrdinal, - visitor); - walkMergedStreams(segment, anchorSupport, upperOrdinal, visitor); - return PriorInteractionWalkResult::conclusive; + walkQubitOperations(segment, anchorSupport.front(), upperOrdinal, + visitor); + else + walkMergedOperations(segment, anchorSupport, upperOrdinal, visitor); + return true; } private: - std::optional - findWalkBounds(Operation *anchor, Operation *inclusiveUpperBound) const { + std::optional + findWalkPositions(Operation *anchor, Operation *inclusiveUpperBound) const { auto anchorPosition = positions.find(anchor); auto upperPosition = positions.find(inclusiveUpperBound); if (anchorPosition == positions.end() || upperPosition == positions.end()) @@ -936,81 +941,85 @@ class cudaq::quake::detail::LogicalQubitInteractionIndex { if (anchorPosition->second.support.empty() || upperPosition->second.ordinal >= anchorPosition->second.ordinal) return std::nullopt; - return WalkBounds{&anchorPosition->second, &upperPosition->second}; + return WalkPositions{&anchorPosition->second, &upperPosition->second}; } - static std::ptrdiff_t findUpperIndex(llvm::ArrayRef interactions, - unsigned upperOrdinal) { - auto end = - std::upper_bound(interactions.begin(), interactions.end(), upperOrdinal, - [](unsigned ordinal, const Interaction &interaction) { - return ordinal < interaction.ordinal; - }); - return std::distance(interactions.begin(), end) - 1; + static std::ptrdiff_t + findLastOperationAtOrBefore(llvm::ArrayRef operations, + unsigned upperOrdinal) { + auto end = std::upper_bound( + operations.begin(), operations.end(), upperOrdinal, + [](unsigned ordinal, const IndexedOperation &operation) { + return ordinal < operation.ordinal; + }); + return std::distance(operations.begin(), end) - 1; } - PriorInteractionWalkResult - walkSingleQubitStream(const Segment &segment, QubitId qubitId, - unsigned upperOrdinal, - llvm::function_ref visitor) const { - auto stream = segment.interactions.find(qubitId); - if (stream == segment.interactions.end()) - return PriorInteractionWalkResult::conclusive; - for (std::ptrdiff_t index = findUpperIndex(stream->second, upperOrdinal); + void + walkQubitOperations(const Segment &segment, QubitId qubitId, + unsigned upperOrdinal, + llvm::function_ref visitor) const { + auto operations = segment.operationsByQubit.find(qubitId); + if (operations == segment.operationsByQubit.end()) + return; + for (std::ptrdiff_t index = + findLastOperationAtOrBefore(operations->second, upperOrdinal); index >= 0; --index) - if (!visitor(stream->second[index].operation)) + if (!visitor(operations->second[index].operation)) break; - return PriorInteractionWalkResult::conclusive; } - void walkMergedStreams(const Segment &segment, - llvm::ArrayRef anchorSupport, - unsigned upperOrdinal, - llvm::function_ref visitor) const { - auto cursors = makeStreamCursors(segment, anchorSupport, upperOrdinal); - while (const Interaction *next = findLatestInteraction(cursors)) { - // Advance every stream at this ordinal because a multi-qubit operation - // appears in several streams but should be visited only once. - for (StreamCursor &cursor : cursors) + void + walkMergedOperations(const Segment &segment, + llvm::ArrayRef anchorSupport, + unsigned upperOrdinal, + llvm::function_ref visitor) const { + auto cursors = makeOperationCursors(segment, anchorSupport, upperOrdinal); + while (const IndexedOperation *next = findNextOperation(cursors)) { + // Advance every list at this ordinal because a multi-qubit operation + // appears in several lists but should be visited only once. + for (OperationCursor &cursor : cursors) cursor.advancePast(next->ordinal); if (!visitor(next->operation)) return; } } - llvm::SmallVector - makeStreamCursors(const Segment &segment, - llvm::ArrayRef anchorSupport, - unsigned upperOrdinal) const { - llvm::SmallVector cursors; + llvm::SmallVector + makeOperationCursors(const Segment &segment, + llvm::ArrayRef anchorSupport, + unsigned upperOrdinal) const { + llvm::SmallVector cursors; for (QubitId qubitId : anchorSupport) { - auto stream = segment.interactions.find(qubitId); - if (stream == segment.interactions.end()) + auto operations = segment.operationsByQubit.find(qubitId); + if (operations == segment.operationsByQubit.end()) continue; - std::ptrdiff_t index = findUpperIndex(stream->second, upperOrdinal); + std::ptrdiff_t index = + findLastOperationAtOrBefore(operations->second, upperOrdinal); if (index >= 0) - cursors.push_back({stream->second, index}); + cursors.push_back({operations->second, index}); } return cursors; } - static const Interaction * - findLatestInteraction(llvm::ArrayRef cursors) { - const Interaction *latest = nullptr; - for (const StreamCursor &cursor : cursors) { - const Interaction *current = cursor.current(); - if (current && (!latest || current->ordinal > latest->ordinal)) - latest = current; + static const IndexedOperation * + findNextOperation(llvm::ArrayRef cursors) { + const IndexedOperation *next = nullptr; + for (const OperationCursor &cursor : cursors) { + const IndexedOperation *current = cursor.current(); + if (current && (!next || current->ordinal > next->ordinal)) + next = current; } - return latest; + return next; } - // Return empty support to omit an operation, non-empty support to add it to - // every affected stream, and no result to end the segment. The index may stop - // more often than the matcher, but it must never omit an operation that the - // block-order search would inspect. + + // Empty support omits an operation that every search can cross without + // analysis. Known support adds the operation under every affected qubit. + // Unknown support starts a new segment because the matcher cannot prove it + // safe to cross for any anchor with known support. static std::optional> - classifyIndexedOperation(Operation *operation, - const QubitIdentityAnalysis &qubitIdentity) { + classifyOperationForIndex(Operation *operation, + const QubitIdentityAnalysis &qubitIdentity) { if (CommutationAnalysis::isIgnorableNonQuantumOperation(operation) || isa(operation)) return llvm::SmallVector{}; @@ -1035,41 +1044,41 @@ class cudaq::quake::detail::LogicalQubitInteractionIndex { return resolveQubitSupport(support, qubitIdentity); } - void removeInteractions(Operation *operation, - const OperationPosition &position) { + void removeOperation(Operation *operation, + const OperationPosition &position) { if (position.support.empty()) return; Segment &segment = segments[position.segment]; for (QubitId qubitId : position.support) { - auto stream = segment.interactions.find(qubitId); - assert(stream != segment.interactions.end() && - "operation is missing its interaction stream"); - auto interaction = std::lower_bound( - stream->second.begin(), stream->second.end(), position.ordinal, - [](const Interaction &candidate, unsigned ordinal) { + auto operations = segment.operationsByQubit.find(qubitId); + assert(operations != segment.operationsByQubit.end() && + "operation is missing from its logical-qubit index"); + auto indexedOperation = std::lower_bound( + operations->second.begin(), operations->second.end(), + position.ordinal, + [](const IndexedOperation &candidate, unsigned ordinal) { return candidate.ordinal < ordinal; }); - assert(interaction != stream->second.end() && - interaction->ordinal == position.ordinal && - interaction->operation == operation && - "operation is missing from its interaction stream"); - stream->second.erase(interaction); + assert(indexedOperation != operations->second.end() && + indexedOperation->ordinal == position.ordinal && + indexedOperation->operation == operation && + "operation is missing from its logical-qubit index"); + operations->second.erase(indexedOperation); } } - void addInteractions(Operation *operation, - const OperationPosition &position) { + void addOperation(Operation *operation, const OperationPosition &position) { if (position.support.empty()) return; Segment &segment = segments[position.segment]; for (QubitId qubitId : position.support) { - auto &stream = segment.interactions[qubitId]; + auto &operations = segment.operationsByQubit[qubitId]; auto insertion = std::lower_bound( - stream.begin(), stream.end(), position.ordinal, - [](const Interaction &interaction, unsigned ordinal) { - return interaction.ordinal < ordinal; + operations.begin(), operations.end(), position.ordinal, + [](const IndexedOperation &operation, unsigned ordinal) { + return operation.ordinal < ordinal; }); - stream.insert(insertion, {operation, position.ordinal}); + operations.insert(insertion, {operation, position.ordinal}); } } @@ -1172,29 +1181,29 @@ bool CommutationAnalysis::haveSameOrderedQuantumOperands(Operation *lhs, lhsInterface.getTargets(), rhsInterface.getTargets()); } -CommutationAnalysis::prior_interaction_walk_result -CommutationAnalysis::walkPriorInteractions( +bool CommutationAnalysis::tryWalkPriorOperations( Operation *anchor, Operation *inclusiveUpperBound, llvm::function_ref visitor) { if (!anchor || !inclusiveUpperBound || anchor->getBlock() != block || inclusiveUpperBound->getBlock() != block) - return prior_interaction_walk_result::unavailable; - if (interactionIndex && interactionIndex->hasPendingInsertions()) - interactionIndex.reset(); + return false; + if (operationIndex && operationIndex->hasPendingInsertions()) + operationIndex.reset(); - if (!interactionIndex) { + if (!operationIndex) { // Avoid a block-wide build until the anchor has known support that can use // the index. auto flow = cudaq::quake::detail::getScalarWireFlow(anchor); if (!flow) - return prior_interaction_walk_result::unavailable; + return false; auto support = resolveQubitSupport(flow->inputs, *qubitIdentity); if (!support || support->empty()) - return prior_interaction_walk_result::unavailable; - interactionIndex = - std::make_unique(*block, *qubitIdentity); + return false; + operationIndex = + std::make_unique(*block, *qubitIdentity); } - return interactionIndex->walk(anchor, inclusiveUpperBound, visitor); + return operationIndex->tryWalkPriorOperations(anchor, inclusiveUpperBound, + visitor); } bool CommutationAnalysis::registerIdentityPreservingOperation( @@ -1202,8 +1211,8 @@ bool CommutationAnalysis::registerIdentityPreservingOperation( if (!operation || operation->getBlock() != block || !qubitIdentity->registerOperation(*operation)) return false; - if (interactionIndex) - interactionIndex->noteInsertion(operation, *qubitIdentity); + if (operationIndex) + operationIndex->noteInsertion(operation, *qubitIdentity); return true; } @@ -1213,11 +1222,11 @@ bool CommutationAnalysis::prepareIdentityPreservingReplacement( !qubitIdentity->replacementPreservesIdentities(*operation, replacement)) return false; cache.clear(); - if (interactionIndex) { - bool updated = interactionIndex->replaceOrEraseOperation( + if (operationIndex) { + bool updated = operationIndex->replaceOrEraseOperation( operation, replacementOp, *qubitIdentity); if (!updated) - interactionIndex.reset(); + operationIndex.reset(); } return true; } @@ -1228,8 +1237,8 @@ void CommutationAnalysis::eraseOperation(Operation *operation) { if (!operation || operation->getBlock() != block) return; cache.clear(); - if (interactionIndex && !interactionIndex->eraseOperation(operation)) - interactionIndex.reset(); + if (operationIndex && !operationIndex->eraseOperation(operation)) + operationIndex.reset(); qubitIdentity->eraseOperation(*operation); } diff --git a/cudaq/lib/Optimizer/Transforms/CommutationAwareRewrite.cpp b/cudaq/lib/Optimizer/Transforms/CommutationAwareRewrite.cpp index 2b2bdc22cb1..d9bb46945e9 100644 --- a/cudaq/lib/Optimizer/Transforms/CommutationAwareRewrite.cpp +++ b/cudaq/lib/Optimizer/Transforms/CommutationAwareRewrite.cpp @@ -512,7 +512,7 @@ Operation *cudaq::opt::CommutationAwareRewriteMatcher::find_nearest( // Check nearby frontier heads and ignorable operations before building the // index. Once analysis is needed, indexed traversal skips operations on // unrelated known qubits when it is available. - bool indexUnavailable = false; + bool useBlockOrderScan = false; for (Operation *candidate = anchor->getPrevNode(); candidate;) { if (!processCandidate(candidate)) return match; @@ -520,13 +520,12 @@ Operation *cudaq::opt::CommutationAwareRewriteMatcher::find_nearest( Operation *inclusiveUpperBound = candidate->getPrevNode(); if (!inclusiveUpperBound) return nullptr; - if (analysis && !indexUnavailable) { - auto walkResult = analysis->walkPriorInteractions( + if (analysis && !useBlockOrderScan) { + bool searchFinished = analysis->tryWalkPriorOperations( anchor, inclusiveUpperBound, processCandidate); - if (walkResult == cudaq::quake::detail::CommutationAnalysis:: - prior_interaction_walk_result::conclusive) + if (searchFinished) return match; - indexUnavailable = true; + useBlockOrderScan = true; } candidate = inclusiveUpperBound; } diff --git a/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp b/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp index a276d1e96b5..fd2643ccb31 100644 --- a/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp +++ b/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp @@ -836,7 +836,8 @@ class QuakeSimplifyPass config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled); auto *ctx = &getContext(); Region ®ion = getOperation()->getRegion(0); - // Each simplification run gets a fresh matcher and interaction index. + // Each simplification run gets a fresh matcher and logical-qubit operation + // index. auto simplify = [&]() { cudaq::opt::CommutationAwareRewriteDriver driver(*ctx, config); auto &patterns = driver.get_patterns(); @@ -871,7 +872,7 @@ class QuakeSimplifyPass if (rotationsToCliffordT) { // Legalization expands one rotation into several gates. Keeping it out // of the commutation-aware simplification loop avoids rebuilding the - // interaction index after every expansion. + // logical-qubit operation index after every expansion. // First combine rotations so legalization only sees those that remain. if (failed(simplify())) {