Add support for graphOpId to the Graph Solver

doc/releases/changelog-dev.md

@@ -264,8 +264,9 @@
   of the default else branch.
   [(#3018)](https://github.com/PennyLaneAI/catalyst/pull/3018)
 
-* Add the `DecomposableGate` op interface to allow generic handling of operations in the `graph-decomposition` pass. This allows arbitrary operations implementing the interface to be registered to and decomposed by the graph. This also allows the use of python-decompositions for any operator pre-registered in the frontend graph.
+* Add the `DecomposableGate` op interface to allow generic handling of operations in the `graph-decomposition` pass. This allows arbitrary operations implementing the interface to be registered to and decomposed by the graph. This also allows the use of python-decompositions for any operator pre-registered in the frontend graph. The graph solver now supports the new `graphOpId`s provided by the interface, as well as the legacy pathway with `name`, `numWires` etc.
   [(#2983)](https://github.com/PennyLaneAI/catalyst/pull/2983)
+  [(#3039)](https://github.com/PennyLaneAI/catalyst/pull/3039)
 
 * The `graph-decomposition` pass eliminates three redundant IR manipulations:
   the cloning, removal, and re-insertion of user rules. This optimization is particularly

mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGTypes.hpp

@@ -20,7 +20,9 @@
  */
 #pragma once
 
+#include <cstddef>
 #include <cstdint>
+#include <functional>
 #include <limits>
 #include <string>
 #include <unordered_map>
@@ -58,9 +60,16 @@ struct OperatorNode {
 
     // Optional static arguments for operators that require additional data.
     std::unordered_map<std::string, std::string> staticNamedArgs{};
+    std::string id{""};
 
     bool operator==(const OperatorNode &other) const
     {
+        // id match
+        if (!id.empty() && !other.id.empty() && id == other.id) {
+            return true;
+        }
+        // legacy fallback if either op is missing ID
+
         // For equality, we consider numWires and numParams conditionally equal
         // if they are not set to -1 (which indicates a wildcard that can match any value).
         const bool default_wires =
@@ -98,6 +107,10 @@ struct OperatorNode {
 struct OperatorNodeHash {
     std::size_t operator()(const OperatorNode &node) const
     {
+        // prefer id if available
+        if (!node.id.empty()) {
+            return std::hash<std::string>{}(node.id);
+        }
         return std::hash<std::string>{}(node.name);
     }
 };
@@ -108,7 +121,23 @@ struct OperatorNodeHash {
 struct WeightedGateset {
     std::unordered_map<OperatorNode, double, OperatorNodeHash> ops;
 
-    [[nodiscard]] bool contains(const OperatorNode &op) const { return ops.find(op) != ops.end(); }
+    [[nodiscard]] bool contains(const OperatorNode &op) const
+    {
+        // hash match
+        if (ops.find(op) != ops.end()) {
+            return true;
+        }
+
+        // use op-matching if hashes failed (could be id vs name)
+        for (auto [gatesetOp, cost] : ops) {
+            if (gatesetOp == op) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     [[nodiscard]] double getCost(const OperatorNode &op) const
     {
         auto it = ops.find(op);
@@ -137,7 +166,8 @@ struct RuleTerm {
  * @brief This represents the origin of a decomposition rule.
  *
  * This enum is used to categorize decomposition rules based on their source or type:
- * - Default: The default rule for decomposing an operator as defined in the decomposition graph.
+ * - Default: The default rule for decomposing an operator as defined in the decomposition
+ * graph.
  * - Fixed: A fixed rule that cannot be changed or overridden by the solver.
  * - Alternative: An alternative rule that can be used in place of the default rule.
  */
@@ -154,7 +184,8 @@ enum class RuleOrigin : uint8_t { Default = 0, Fixed = 1, Alternative = 2 };
  *
  * TODO:
  * - We can add a field for work_wires_required if we want to consider the number of ancillary
- * wires needed for the decomposition, which can be an important factor in resource optimization.
+ * wires needed for the decomposition, which can be an important factor in resource
+ * optimization.
  * - We can also consider adding a field for the decomposition function or a pointer to it,
  * which can be used to actually perform the decomposition after the graph solver selects
  * the rules.

mlir/lib/Quantum/Transforms/GraphDecomposition/DecompGraphSolver/DGUtils.hpp

@@ -33,6 +33,11 @@ namespace DecompGraph::Core {
 
 static inline auto print_op(const OperatorNode &op) -> std::string
 {
+    // id override
+    if (!op.id.empty()) {
+        return "id: " + op.id;
+    }
+
     std::ostringstream oss;
     oss << op.name;
     oss << "[w:" << op.numWires << "]";

mlir/unittests/DecompGraphSolver/Test_DecompGraphCore.cpp

@@ -146,3 +146,45 @@ TEST_CASE("Test ChosenDecompRule construction", "[DecompGraph::Core]")
     REQUIRE(chosenRule.basisCounts[rz] == 2);
     REQUIRE(chosenRule.basisCounts[rx] == 1);
 }
+
+TEST_CASE("Test graphOpId Support", "[DecompGraph::Core]")
+{
+    // comparing an op without an ID should fallback to legacy match
+    const OperatorNode h{"H"};
+    const OperatorNode hId{"H", -1, -1, false, {}, "H[][1]{}"};
+
+    REQUIRE(h == hId);
+
+    // id + legacy match with non-wildcard params
+    const OperatorNode x{"X", 1, 0};
+    const OperatorNode xId{"X", -1, -1, false, {}, "X[][1]{}"};
+
+    REQUIRE(x == xId);
+
+    // id nodes should fail legacy match if params differ
+    const OperatorNode op1{"op", 1, 1};
+    const OperatorNode op2{"op", 2, 2, false, {}, "op[f64,f64][2]{}"};
+
+    REQUIRE_FALSE(op1 == op2);
+
+    // nodes with same ids should match
+    const OperatorNode pr1{"PauliRot", -1, -1, false, {}, "PauliRot[f64][2]{pauli_word:XX}"};
+    const OperatorNode pr2{"PauliRot", -1, -1, false, {}, "PauliRot[f64][2]{pauli_word:XX}"};
+
+    REQUIRE(pr1 == pr2);
+
+    // nodes with matching ids should ignore other parameters (id is source of truth)
+    const OperatorNode id1{"name1", 1, 1, false, {}, "sameID"};
+    const OperatorNode id2{"name2", 2, 2, true, {}, "sameID"};
+
+    REQUIRE(id1 == id2);
+
+    // Unit test for `OperatorNodeHash`. Check that the hash function prefers ID over name
+    const OperatorNode hash1{"name", -1, -1, false, {}, "id"};
+    const OperatorNode hash2{"name2", 1, 1, true, {}, "id"};
+    const OperatorNode hash3{"name2", 1, 1, true, {}};
+
+    const OperatorNodeHash hashFunc;
+    REQUIRE(hashFunc(hash1) == hashFunc(hash2));
+    REQUIRE(hashFunc(hash2) != hashFunc(hash3));
+}