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

frontend/catalyst/passes/__init__.py

@@ -42,6 +42,7 @@
     ppm_compilation,
     ppr_to_ppm,
     to_ppr,
+    get_ppm_spec,
 )
 from catalyst.passes.pass_api import Pass, PassPlugin, apply_pass, apply_pass_plugin
 
@@ -58,4 +59,5 @@
     "merge_ppr_ppm",
     "ppr_to_ppm",
     "ppm_compilation",
+    "get_ppm_spec",
 )

frontend/catalyst/passes/builtin_passes.py

@@ -15,10 +15,29 @@
 """This module exposes built-in Catalyst MLIR passes to the frontend."""
 
 import functools
+import re
 
+from catalyst.compiler import _quantum_opt
 from catalyst.passes.pass_api import PassPipelineWrapper
 
 
+def get_ppm_spec(QJIT):
+    """TODO: docstring"""
+
+    if QJIT.mlir is not None:
+        # aot mode
+        raw_result = _quantum_opt(
+            ("--pass-pipeline", "builtin.module(ppm_specs)"), [], stdin=QJIT.mlir
+        )
+
+        # breakpoint()
+        regex_search_for_json = re.search(r"\{.*?\}", raw_result, re.DOTALL)
+        return regex_search_for_json.group(0)
+
+    else:
+        raise "ppm passes only support aot mode, please make sure ..."
+
+
 ## API ##
 def cancel_inverses(qnode):
     """

frontend/catalyst/passes/pass_api.py

@@ -382,4 +382,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,10 @@ 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 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
@@ -31,7 +32,27 @@ set(DEPENDS
 
 add_mlir_library(${LIBRARY_NAME} STATIC ${SRC} LINK_LIBS PRIVATE ${LIBS} DEPENDS ${DEPENDS})
 target_compile_features(${LIBRARY_NAME} PUBLIC cxx_std_20)
+
+include(FetchContent)
+FetchContent_Declare(json
+    URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz
+    URL_HASH  SHA256=d6c65aca6b1ed68e7a182f4757257b107ae403032760ed6ef121c9d55e81757d
+    DOWNLOAD_EXTRACT_TIMESTAMP true
+    SYSTEM
+)
+FetchContent_MakeAvailable(json)
+target_link_libraries(${LIBRARY_NAME} 
+    PRIVATE nlohmann_json::nlohmann_json
+)
+
+if(CMAKE_CXX_COMPILER_ID MATCHES ".*Clang")
+set_source_files_properties(
+    CountPPMSpecs.cpp
+    COMPILE_FLAGS "-Wno-covered-switch-default"
+)
+endif()
+
 target_include_directories(${LIBRARY_NAME} PUBLIC
                            .
                            ${PROJECT_SOURCE_DIR}/include
-                           ${CMAKE_BINARY_DIR}/include)
+                           ${CMAKE_BINARY_DIR}/include)

mlir/lib/QEC/Transforms/CountPPMSpecs.cpp

@@ -0,0 +1,109 @@
+// 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 <algorithm>
+#include <string>
+
+#include <nlohmann/json.hpp>
+
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+#include "QEC/IR/QECDialect.h"
+
+using namespace llvm;
+using namespace mlir;
+using namespace catalyst;
+using namespace catalyst::qec;
+using json = nlohmann::json;
+
+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;
+
+            StringRef gate_name = op->getName().getStringRef();
+
+            if (gate_name == "quantum.alloc") {
+                auto num_qubits_attr = op->getAttrOfType<mlir::IntegerAttr>("nqubits_attr");
+                assert(num_qubits_attr == nullptr && "num_qubits_attr is a dynamic SSA value");
+                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"]++;
+            }
+
+            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[max_weight_pi_key] =
+                            std::max(PPM_Specs[max_weight_pi_key],
+                                     static_cast<int>(pauli_product_attr.size()));
+                    }
+                }
+            }
+        });
+
+        json PPM_Specs_Json = PPM_Specs;
+        llvm::outs() << PPM_Specs_Json.dump(4) << "\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