diff --git a/cudaq/include/cudaq/Optimizer/Analysis/CommutationAnalysis.h b/cudaq/include/cudaq/Optimizer/Analysis/CommutationAnalysis.h index e0357f4de92..dfe039dbb32 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 LogicalQubitOperationIndex; /// The outcome of a commutation query. enum class commutation_status { commutes, does_not_commute, indeterminate }; @@ -183,25 +186,51 @@ 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); + /// 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 and clear cached relations. + /// Validate an identity-preserving replacement, clear cached relations, and + /// maintain or discard the operation 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 operation index, + /// then erase its result identities. void eraseOperation(mlir::Operation *operation); mlir::Block *block; std::unique_ptr qubitIdentity; + std::unique_ptr operationIndex; llvm::DenseMap cache; friend class cudaq::opt::CommutationAwareRewriteMatcher; friend class cudaq::opt::detail::CommutationAwareRewriteListener; + 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 db39780b305..b1282616bfa 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 -/// pairwise commutation proof with the anchor before the frontier advances. -/// Every other intervening scalar-wire operation requires either that pairwise +/// 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 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. /// 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..98d54fd4a28 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,412 @@ 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; +} + +// 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::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 isSegmentBoundary; + llvm::SmallVector support; + }; + + // 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 operation lists within one searchable block interval. + // A multi-qubit operation appears in every list it affects. + struct Segment { + llvm::DenseMap> + operationsByQubit; + }; + + // Hold the indexed positions that bound one backward walk. + struct WalkPositions { + const OperationPosition *anchor; + const OperationPosition *upper; + }; + + // Track one per-qubit operation list during a multi-qubit backward merge. + struct OperationCursor { + llvm::ArrayRef operations; + std::ptrdiff_t index; + + const IndexedOperation *current() const { + return index < 0 ? nullptr : &operations[index]; + } + + void advancePast(unsigned ordinal) { + if (const IndexedOperation *operation = current(); + operation && operation->ordinal == ordinal) + --index; + } + }; + +public: + LogicalQubitOperationIndex(Block &block, + const QubitIdentityAnalysis &qubitIdentity) { + segments.emplace_back(); + unsigned ordinal = 0; + for (Operation &operation : block) { + unsigned segment = segments.size() - 1; + 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(didInsert && "block operation already indexed"); + if (!indexedSupport) { + segments.emplace_back(); + } else { + addOperation(&operation, position->second); + } + ++ordinal; + } + } + + void noteInsertion(Operation *operation, + const QubitIdentityAnalysis &qubitIdentity) { + 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); + } + + 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 = + 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.isSegmentBoundary || + !replacementSupport) + return false; + + OperationPosition newPosition{position->second.ordinal, + position->second.segment, false, + std::move(*replacementSupport)}; + removeOperation(operation, position->second); + positions.erase(position); + auto [replacementPosition, didInsert] = + positions.try_emplace(replacement, std::move(newPosition)); + assert(didInsert && "replacement operation already indexed"); + addOperation(replacement, replacementPosition->second); + return true; + } + + bool eraseOperation(Operation *operation) { + pendingInsertions.erase(operation); + auto position = positions.find(operation); + if (position == positions.end()) + return true; + if (position->second.isSegmentBoundary) + return false; + removeOperation(operation, position->second); + positions.erase(position); + return true; + } + + bool hasPendingInsertions() const { return !pendingInsertions.empty(); } + + 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 = walkPositions->anchor->support; + if (walkPositions->anchor->segment != walkPositions->upper->segment) + return true; + + const Segment &segment = segments[walkPositions->anchor->segment]; + unsigned upperOrdinal = walkPositions->upper->ordinal; + if (anchorSupport.size() == 1) + walkQubitOperations(segment, anchorSupport.front(), upperOrdinal, + visitor); + else + walkMergedOperations(segment, anchorSupport, upperOrdinal, visitor); + return true; + } + +private: + 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()) + return std::nullopt; + if (anchorPosition->second.support.empty() || + upperPosition->second.ordinal >= anchorPosition->second.ordinal) + return std::nullopt; + return WalkPositions{&anchorPosition->second, &upperPosition->second}; + } + + 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; + } + + 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(operations->second[index].operation)) + break; + } + + 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 + makeOperationCursors(const Segment &segment, + llvm::ArrayRef anchorSupport, + unsigned upperOrdinal) const { + llvm::SmallVector cursors; + for (QubitId qubitId : anchorSupport) { + auto operations = segment.operationsByQubit.find(qubitId); + if (operations == segment.operationsByQubit.end()) + continue; + std::ptrdiff_t index = + findLastOperationAtOrBefore(operations->second, upperOrdinal); + if (index >= 0) + cursors.push_back({operations->second, index}); + } + return cursors; + } + + 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 next; + } + + // 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> + classifyOperationForIndex(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 removeOperation(Operation *operation, + const OperationPosition &position) { + if (position.support.empty()) + return; + Segment &segment = segments[position.segment]; + for (QubitId qubitId : position.support) { + 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(indexedOperation != operations->second.end() && + indexedOperation->ordinal == position.ordinal && + indexedOperation->operation == operation && + "operation is missing from its logical-qubit index"); + operations->second.erase(indexedOperation); + } + } + + void addOperation(Operation *operation, const OperationPosition &position) { + if (position.support.empty()) + return; + Segment &segment = segments[position.segment]; + for (QubitId qubitId : position.support) { + auto &operations = segment.operationsByQubit[qubitId]; + auto insertion = std::lower_bound( + operations.begin(), operations.end(), position.ordinal, + [](const IndexedOperation &operation, unsigned ordinal) { + return operation.ordinal < ordinal; + }); + operations.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 +1181,53 @@ bool CommutationAnalysis::haveSameOrderedQuantumOperands(Operation *lhs, lhsInterface.getTargets(), rhsInterface.getTargets()); } +bool CommutationAnalysis::tryWalkPriorOperations( + Operation *anchor, Operation *inclusiveUpperBound, + llvm::function_ref visitor) { + if (!anchor || !inclusiveUpperBound || anchor->getBlock() != block || + inclusiveUpperBound->getBlock() != block) + return false; + if (operationIndex && operationIndex->hasPendingInsertions()) + operationIndex.reset(); + + 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 false; + auto support = resolveQubitSupport(flow->inputs, *qubitIdentity); + if (!support || support->empty()) + return false; + operationIndex = + std::make_unique(*block, *qubitIdentity); + } + return operationIndex->tryWalkPriorOperations(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 (operationIndex) + operationIndex->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 (operationIndex) { + bool updated = operationIndex->replaceOrEraseOperation( + operation, replacementOp, *qubitIdentity); + if (!updated) + operationIndex.reset(); + } return true; } @@ -789,7 +1236,9 @@ void CommutationAnalysis::clearCachedRelations() { cache.clear(); } void CommutationAnalysis::eraseOperation(Operation *operation) { if (!operation || operation->getBlock() != block) return; - clearCachedRelations(); + cache.clear(); + 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 ec0b1dd130c..d9bb46945e9 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,37 @@ 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 useBlockOrderScan = false; + for (Operation *candidate = anchor->getPrevNode(); candidate;) { + if (!processCandidate(candidate)) + return match; + + Operation *inclusiveUpperBound = candidate->getPrevNode(); + if (!inclusiveUpperBound) + return nullptr; + if (analysis && !useBlockOrderScan) { + bool searchFinished = analysis->tryWalkPriorOperations( + anchor, inclusiveUpperBound, processCandidate); + if (searchFinished) + return match; + useBlockOrderScan = true; } + candidate = inclusiveUpperBound; } return nullptr; } diff --git a/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp b/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp index 4d5f9648282..fd2643ccb31 100644 --- a/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp +++ b/cudaq/lib/Optimizer/Transforms/QuakeSimplify.cpp @@ -835,38 +835,64 @@ 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 logical-qubit operation + // 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 + // logical-qubit operation 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.