Speed up commutation-aware rewrite searches

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 <memory>
@@ -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<mlir::Operation *, mlir::Operation *>;
 
   /// 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<mlir::Value> &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<bool(mlir::Operation *)> 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<QubitIdentityAnalysis> qubitIdentity;
+  std::unique_ptr<LogicalQubitInteractionIndex> interactionIndex;
   llvm::DenseMap<OperationPair, CommutationResult> cache;
 
   friend class cudaq::opt::CommutationAwareRewriteMatcher;
   friend class cudaq::opt::detail::CommutationAwareRewriteListener;
+  friend class LogicalQubitInteractionIndex;
 };
 
 } // namespace cudaq::quake::detail

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.