address comments

Signed-off-by: Luca Mondada <[email protected]>

Author: lmondada

lib/Optimizer/Transforms/DecompositionPatterns.cpp

@@ -30,6 +30,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Rewrite/FrozenRewritePatternSet.h"
 #include <llvm/ADT/SmallVector.h>
+#include <llvm/ADT/StringMap.h>
 #include <llvm/ADT/StringRef.h>
 #include <llvm/Support/Casting.h>
 #include <llvm/Support/Error.h>
@@ -333,8 +334,9 @@ LogicalResult checkAndExtractControls(quake::OperatorInterface op,
   };                                                                           \
   CUDAQ_REGISTER_TYPE(cudaq::DecompositionPatternType, PATTERN##Type, PATTERN)
 
-// TODO: "SToR1", "TToR1", "R1ToU3", "U3ToRotations" can be generalised
-// arbitrary number of controls, but we would need to reason over n HOp patterns
+// TODO: The decomposition patterns "SToR1", "TToR1", "R1ToU3", "U3ToRotations"
+// can handle arbitrary number of controls, but currently metadata cannot
+// capture this. The pattern types therefore only advertise them for 0 controls.
 
 //===----------------------------------------------------------------------===//
 // HOp decompositions
@@ -1812,8 +1814,22 @@ REGISTER_DECOMPOSITION_PATTERN(U3ToRotations, "u3", "rz", "rx");
 
 void cudaq::populateWithAllDecompositionPatterns(
     mlir::RewritePatternSet &patterns) {
-  for (auto &patternType :
-       cudaq::DecompositionPatternType::RegistryType::entries()) {
-    patterns.add(patternType.instantiate()->create(patterns.getContext()));
+  // For deterministic ordering, sort the registered pattern types by name
+  // Note that this assumes that no additional patterns are registered at
+  // runtime.
+  static std::map<std::string, std::unique_ptr<cudaq::DecompositionPatternType>>
+      patternTypes = []() {
+        std::map<std::string, std::unique_ptr<cudaq::DecompositionPatternType>>
+            map;
+        for (auto &patternType :
+             cudaq::DecompositionPatternType::RegistryType::entries()) {
+          map[patternType.getName().str()] = patternType.instantiate();
+        }
+        return map;
+      }();
+
+  for (auto it = patternTypes.begin(), ie = patternTypes.end(); it != ie;
+       ++it) {
+    patterns.add(it->second->create(patterns.getContext()));
   }
 }

unittests/Optimizer/DecompositionPatternsTest.cpp

@@ -25,9 +25,9 @@
 #include <iterator>
 #include <llvm/ADT/APFloat.h>
 #include <llvm/ADT/STLExtras.h>
+#include <llvm/ADT/StringMap.h>
 #include <memory>
 #include <mlir/IR/BuiltinOps.h>
-#include <ranges>
 
 using namespace mlir;
 
@@ -203,8 +203,10 @@ void stripNamespace(std::string &debugName) {
 
 // Test 1: Verify the total number of registered decomposition patterns
 TEST_F(DecompositionPatternsTest, TotalPatternCount) {
-  auto patternNames = cudaq::DecompositionPatternType::RegistryType::entries();
-  unsigned int size = std::distance(patternNames.begin(), patternNames.end());
+  auto patternEntries =
+      cudaq::DecompositionPatternType::RegistryType::entries();
+  unsigned int size =
+      std::distance(patternEntries.begin(), patternEntries.end());
   EXPECT_EQ(size, 31) << "Expected 31 decomposition patterns, but found "
                       << size;
 }
@@ -297,18 +299,29 @@ TEST_F(DecompositionPatternsTest, DecompositionProducesOnlyTargetGates) {
     // Collect all gates in the output
     auto outputGates = collectGateTypesInModule(module);
 
+    // Map from gate prefix to allowed number of controls
+    llvm::StringMap<llvm::SmallVector<size_t>> allowedGates;
+    for (auto targetGate : targetGates) {
+      auto [tPrefix, tNum] = splitGateAndControls(targetGate);
+      allowedGates[tPrefix].push_back(tNum);
+    }
     auto isAllowedGate = [&](StringRef gate) {
       // Split gate into prefix and number (e.g., "h(1)" -> "h", 1) using
       // utility function
       auto [gatePrefix, gateNum] = splitGateAndControls(gate);
 
-      for (auto targetGate : targetGates) {
-        auto [tPrefix, tNum] = splitGateAndControls(targetGate);
-        if (gatePrefix == tPrefix && gateNum <= tNum) {
-          return true;
-        }
+      auto it = allowedGates.find(gatePrefix);
+      if (it == allowedGates.end()) {
+        return false;
       }
-      return false;
+      auto allowedNumControls = it->second;
+      // Check if the number of controls is in the allowed list (or if any
+      // number is allowed)
+      auto isEqOrMax = [gateNum](size_t num) {
+        return num == gateNum || num == std::numeric_limits<size_t>::max();
+      };
+      return std::find_if(allowedNumControls.begin(), allowedNumControls.end(),
+                          isEqOrMax) != allowedNumControls.end();
     };
 
     std::vector<std::string> unexpectedGates;