update QPD test

Author: kipawaa

frontend/test/lit/GraphDecomposition/TestQPD.mlir

@@ -14,25 +14,17 @@
 
 // Test that graph-decomposition is able to use compile-time python decompositions
 
-// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=RX=1.0,Hadamard=1.0,MultiRZ=1.0,GlobalPhase=1.0})' %s | FileCheck %s
+// RUN: catalyst --tool=opt --pass-pipeline='builtin.module(graph-decomposition{gate-set=PhaseShift=1.0})' %s | FileCheck %s
 
 
-func.func @circuit() -> (!quantum.bit, !quantum.bit, !quantum.bit) {
-    %0 = quantum.alloc(3) : !quantum.reg
-    %q0 = quantum.extract %0[0] : !quantum.reg -> !quantum.bit
-    %q1 = quantum.extract %0[1] : !quantum.reg -> !quantum.bit
-    %q2 = quantum.extract %0[2] : !quantum.reg -> !quantum.bit
+func.func @circuit() -> !quantum.bit {
+  %reg = quantum.alloc(1): !quantum.reg
+  %0 = quantum.extract %reg[0] : !quantum.reg -> !quantum.bit
 
-    %pi = arith.constant 3.1 : f64
-    %qout:3 = quantum.paulirot ["Y", "X", "Z"](%pi) %q0, %q1, %q2 : !quantum.bit, !quantum.bit, !quantum.bit
-
-    // CHECK-NOT: quantum.paulirot
-
-    // CHECK-DAG: RX
-    // CHECK-DAG: Hadamard
-    // CHECK: multirz
-    // CHECK-DAG: RX
-    // CHECK-DAG: Hadamard
-
-    return %qout#0, %qout#1, %qout#2 : !quantum.bit, !quantum.bit, !quantum.bit
+  // CHECK-NOT: quantum.custom "S"
+  // CHECK: quantum.custom "PhaseShift"
+  %out_qubits = quantum.custom "S"() %0 : !quantum.bit
+  return %out_qubits : !quantum.bit
 }
+
+// TODO: add a tests for more complicated operators once they're migrated

mlir/lib/Quantum/Transforms/GraphDecomposition/graph_decomposition.cpp

@@ -395,7 +395,7 @@ struct GraphDecompositionPass : public impl::GraphDecompositionPassBase<GraphDec
             }
 
             moduleOp->walk([&](mlir::func::FuncOp func) {
-                if (func.getName().starts_with(opId)) {
+                if (func.getName().ends_with(opId)) {
                     mlir::OwningOpRef<mlir::func::FuncOp> outOp;
                     func->remove();
                     outOp = mlir::OwningOpRef<mlir::func::FuncOp>(func);