Implement ppm_specs pass

frontend/catalyst/compiler.py

@@ -16,10 +16,12 @@
 """
 import glob
 import importlib
+import json
 import logging
 import os
 import pathlib
 import platform
+import re
 import shutil
 import subprocess
 import sys
@@ -371,7 +373,32 @@
         return _quantum_opt(*args, stdin=stdin)
 
     opts = _options_to_cli_flags(options)
-    return _quantum_opt(*opts, *args, stdin=stdin)
+    raw_result = _quantum_opt(*opts, *args, stdin=stdin)
+    regex_search_for_json = re.search(r"\{[a-zA-Z0-9_\":\{\},\n]+\}", raw_result)
+    raw_result = raw_result.replace(regex_search_for_json.group(0), "")
+    return raw_result
+
+
+def to_ppm_spec(*args, stdin=None, options: Optional[CompileOptions] = None):
+    """echo ${input} | catalyst --tool=opt *args *opts -"""
+    # These are the options that may affect compilation
+    if not options:
+        return _quantum_opt(*args, stdin=stdin)
+
+    opts = _options_to_cli_flags(options)
+
+    raw_json_format = _quantum_opt(*opts, *args, stdin=stdin)
+    regex_search_for_json = re.search(r"\{[a-zA-Z0-9_\":\{\},\n]+\}", raw_json_format)
+
+    # No ppm_specs json is found
+    if regex_search_for_json is None:
+        return raw_json_format
+
+    json_ppm_specs = regex_search_for_json.group(0)
+    json_ppm_specs = json_ppm_specs.replace(",\n}", "\n}")
+    json_ppm_specs = json.loads(json_ppm_specs)
+    return json_ppm_specs
+
 
 
 class Compiler:

frontend/catalyst/jit.py

@@ -39,6 +39,7 @@
     canonicalize,
     to_llvmir,
     to_mlir_opt,
+    to_ppm_spec,
 )
 from catalyst.debug.instruments import instrument
 from catalyst.from_plxpr import trace_from_pennylane
@@ -578,6 +579,13 @@ def mlir_opt(self):
 
         return to_mlir_opt(stdin=str(self.mlir_module), options=self.compile_options)
 
+    def get_ppm_spec(self):
+        """obtain the PPM specs after optimization"""
+        if not self.mlir_module:
+            return None
+
+        return to_ppm_spec(stdin=str(self.mlir_module), options=self.compile_options)
+
     @debug_logger
     def __call__(self, *args, **kwargs):
         # Transparantly call Python function in case of nested QJIT calls.

frontend/catalyst/passes/__init__.py

@@ -40,6 +40,7 @@
     merge_ppr_ppm,
     merge_rotations,
     ppm_compilation,
+    ppm_specs,
     ppr_to_ppm,
     to_ppr,
 )
@@ -58,4 +59,5 @@
     "merge_ppr_ppm",
     "ppr_to_ppm",
     "ppm_compilation",
+    "ppm_specs",
 )

frontend/catalyst/passes/builtin_passes.py

@@ -781,3 +781,8 @@
         )
 
     return PassPipelineWrapper(qnode, passes)
+
+
+def ppm_specs(qnode):
+    # TODO: Add docstring
+    return PassPipelineWrapper(qnode, "ppm_specs")

frontend/catalyst/passes/pass_api.py

@@ -379,4 +379,5 @@ def _API_name_to_pass_name():
         "merge_ppr_ppm": "merge_ppr_ppm",
         "ppr_to_ppm": "ppr_to_ppm",
         "ppm_compilation": "ppm_compilation",
+        "ppm_specs": "ppm_specs",
     }

mlir/include/QEC/Transforms/Passes.h

@@ -29,5 +29,6 @@ std::unique_ptr<mlir::Pass> createCommuteCliffordPastPPMPass();
 std::unique_ptr<mlir::Pass> createDecomposeNonCliffordPPRPass();
 std::unique_ptr<mlir::Pass> createDecomposeCliffordPPRPass();
 std::unique_ptr<mlir::Pass> createCliffordTToPPMPass();
+std::unique_ptr<mlir::Pass> createCountPPMSpecsPass();
 
 } // namespace catalyst

mlir/include/QEC/Transforms/Passes.td

@@ -121,4 +121,14 @@ def CliffordTToPPMPass : Pass<"ppm_compilation"> {
     let options = [MaxPauliSizeOption, DecomposeMethodOption, AvoidYMeasureOption];
 }
 
+def CountPPMSpecsPass : Pass<"ppm_specs"> {
+    let summary = "Count specs in Pauli Product Measurement operations.";
+
+    let dependentDialects = [
+        "catalyst::qec::QECDialect",
+    ];
+
+    let constructor = "catalyst::createCountPPMSpecsPass()";
+}
+
 #endif // QEC_PASSES

mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp

@@ -38,6 +38,7 @@ void catalyst::registerAllCatalystPasses()
     mlir::registerPass(catalyst::createCliffordTToPPMPass);
     mlir::registerPass(catalyst::createDecomposeNonCliffordPPRPass);
     mlir::registerPass(catalyst::createDecomposeCliffordPPRPass);
+    mlir::registerPass(catalyst::createCountPPMSpecsPass);
     mlir::registerPass(catalyst::createDetensorizeSCFPass);
     mlir::registerPass(catalyst::createDisableAssertionPass);
     mlir::registerPass(catalyst::createDisentangleCNOTPass);

mlir/lib/QEC/Transforms/CMakeLists.txt

@@ -8,6 +8,7 @@ file(GLOB SRC
     commute_clifford_t_ppr.cpp
     CommuteCliffordPastPPM.cpp
     commute_clifford_past_ppm.cpp
+    CountPPMSpecs.cpp
     decompose_non_clifford_ppr.cpp
     DecomposeNonCliffordPPR.cpp
     decompose_clifford_ppr.cpp

mlir/lib/QEC/Transforms/CountPPMSpecs.cpp

@@ -0,0 +1,126 @@
+// Copyright 2025 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#define DEBUG_TYPE "ppm_specs"
+
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+
+#include "QEC/IR/QECDialect.h"
+#include "QEC/Transforms/Patterns.h"
+#include "Quantum/IR/QuantumOps.h"
+#include "QEC/Utils/PauliStringWrapper.h"
+#include <algorithm>
+#include <string>
+
+using namespace llvm;
+using namespace mlir;
+using namespace catalyst;
+using namespace catalyst::qec;
+
+namespace catalyst {
+namespace qec {
+
+#define GEN_PASS_DEF_COUNTPPMSPECSPASS
+#define GEN_PASS_DECL_COUNTPPMSPECSPASS
+#include "QEC/Transforms/Passes.h.inc"
+
+
+struct CountPPMSpecsPass : public impl::CountPPMSpecsPassBase<CountPPMSpecsPass> {
+    using CountPPMSpecsPassBase::CountPPMSpecsPassBase;
+
+    void print_specs()
+    {
+        llvm::BumpPtrAllocator string_allocator;
+        llvm::DenseMap<StringRef, int> PPM_Specs;
+        PPM_Specs["num_logical_qubits"] = 0;
+        PPM_Specs["num_of_ppm"] = 0;
+
+        // Walk over all operations in the IR (could be ModuleOp or FuncOp)
+        getOperation()->walk([&](Operation *op) {
+            // Skip top-level container ops if desired
+            if (isa<ModuleOp>(op)) return;
+
+            // TODO: Remove debug in future
+            // llvm::outs()<<"\n-----------------------------MLIR------------------------------\n";
+            // op->print(llvm::outs());
+            // llvm::outs()<<"\n-----------------------------MLIR------------------------------\n";
+
+            StringRef gate_name = op->getName().getStringRef();
+
+            if (gate_name == "quantum.alloc") {
+                auto num_qubits_attr = op->getAttrOfType<mlir::IntegerAttr>("nqubits_attr");
+                u_int64_t num_qubits = num_qubits_attr ? static_cast<u_int64_t>(num_qubits_attr.getInt()) : 0;
+                PPM_Specs["num_logical_qubits"] = num_qubits;
+            }
+
+            if (gate_name == "qec.ppm") {
+                PPM_Specs["num_of_ppm"] = PPM_Specs["num_of_ppm"] + 1;
+            }
+
+            if (gate_name == "qec.ppr") {
+                auto rotation_attr = op->getAttrOfType<mlir::IntegerAttr>("rotation_kind");
+                auto pauli_product_attr = op->getAttrOfType<mlir::ArrayAttr>("pauli_product");
+                int16_t rotation_kind = rotation_attr ? static_cast<int16_t>(rotation_attr.getInt()) : 0;
+                if (rotation_kind) {
+                    llvm::StringSaver saver(string_allocator);
+                    StringRef num_pi_key = saver.save("num_pi"+std::to_string(abs(rotation_kind))+"_gates");
+                    StringRef max_weight_pi_key = saver.save("max_weight_pi"+std::to_string(abs(rotation_kind)));
+
+                    if (PPM_Specs.find(llvm::StringRef(num_pi_key)) == PPM_Specs.end()) {
+                        PPM_Specs[num_pi_key] = 1;
+                        PPM_Specs[max_weight_pi_key] = static_cast<int>(pauli_product_attr.size());
+                    }
+                    else {
+                        PPM_Specs[num_pi_key] = PPM_Specs[num_pi_key] + 1;
+                        PPM_Specs[max_weight_pi_key] = std::max(PPM_Specs[max_weight_pi_key], static_cast<int>(pauli_product_attr.size()));
+                    }
+                }
+            }
+            // TODO: Implement depth using slicing 
+            // mlir::SetVector <Operation *> backwardSlice;
+            // getBackwardSlice(op, &backwardSlice);
+            // llvm::outs()<<"\n-----------------------------SLICE-----------------------------\n";
+            // llvm::outs()<<"Backward slicing\n";
+            // for (Operation *o : backwardSlice) {
+            //     if (o->getName().getStringRef() == "quantum.extract") {
+            //         llvm::outs() << *o << "\n"; 
+            //     }
+            // }
+            // llvm::outs()<<"\n-----------------------------SLICE------------------------------\n";
+        });
+
+        llvm::outs() << "{\n";
+        for (const auto &entry : PPM_Specs) {
+            llvm::outs() << '"' << entry.first << '"' << ":" << entry.second << ",\n";
+        }
+        llvm::outs() << "}\n";
+        return;
+    }
+
+    void runOnOperation() final
+    {
+        print_specs();
+    }
+};
+
+} // namespace qec
+
+/// Create a pass for lowering operations in the `QECDialect`.
+std::unique_ptr<Pass> createCountPPMSpecsPass() { return std::make_unique<CountPPMSpecsPass>(); }
+
+} // namespace catalyst