Add functionality to obtain some specs/info on quantum operations

frontend/catalyst/compiler.py

@@ -21,6 +21,7 @@
 import pathlib
 import platform
 import shutil
+import signal
 import subprocess
 import sys
 import tempfile
@@ -405,8 +406,8 @@
         )
         return cmd
 
     @debug_logger
     def run_from_ir(self, ir: str, module_name: str, workspace: Directory):
         """Compile a shared object from a textual IR (MLIR or LLVM).
 
         Args:
@@ -424,9 +425,9 @@
             workspace, Directory
         ), f"Compiler expects a Directory type, got {type(workspace)}."
         assert workspace.is_dir(), f"Compiler expects an existing directory, got {workspace}."
-        assert (
-            self.options.lower_to_llvm
-        ), "lower_to_llvm must be set to True in order to compile to a shared object"
+        # assert (
+        #     self.options.lower_to_llvm
+        # ), "lower_to_llvm must be set to True in order to compile to a shared object"
 
         if self.options.verbose:
             print(f"[LIB] Running compiler driver in {workspace}", file=self.options.logfile)
@@ -444,12 +445,47 @@
         try:
             if self.options.verbose:
                 print(f"[SYSTEM] {' '.join(cmd)}", file=self.options.logfile)
-            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+            # result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+
+            if self.options.debug_compiler != True:
+                result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+            else:  # self.options.debug_compiler == True:
+                with subprocess.Popen(
+                    cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
+                ) as p:
+                    # Ensure process creation succeeds
+                    if p.returncode not in {0, None}:
+                        raise subprocess.CalledProcessError(p.returncode, cmd)
+                    if self.options.debug_compiler:
+                        print(f"Compiler PID={p.pid}")
+                        print(
+                            f"""Ensure C++ debugger is attached and running before continuing with:
+                            kill -s SIGCONT {p.pid}"""
+                        )
+                        p.send_signal(signal.SIGSTOP)
+                        res_stdout, res_stderr = p.communicate()
+                        # Ensure process execution succeeds
+                        if p.returncode not in {0, None}:
+                            raise subprocess.CalledProcessError(
+                                p.returncode, cmd, res_stdout, res_stderr
+                            )
+
             if self.options.verbose or os.getenv("ENABLE_DIAGNOSTICS"):
-                if result.stdout:
-                    print(result.stdout.strip(), file=self.options.logfile)
-                if result.stderr:
-                    print(result.stderr.strip(), file=self.options.logfile)
+                # if result.stdout:
+                #     print(result.stdout.strip(), file=self.options.logfile)
+                # if result.stderr:
+                #     print(result.stderr.strip(), file=self.options.logfile)
+                if self.options.debug_compiler != True:
+                    if result.stdout:
+                        print(result.stdout.strip(), file=self.options.logfile)
+                    if result.stderr:
+                        print(result.stderr.strip(), file=self.options.logfile)
+                else:  # self.options.debug_compiler == True:
+                    if res_stdout:
+                        print(res_stdout.strip(), file=self.options.logfile)
+                    if res_stderr:
+                        print(res_stderr.strip(), file=self.options.logfile)
+
         except subprocess.CalledProcessError as e:  # pragma: nocover
             raise CompileError(f"catalyst failed with error code {e.returncode}: {e.stderr}") from e
 

frontend/catalyst/debug/__init__.py

@@ -22,6 +22,7 @@
     get_compilation_stages_groups,
     replace_ir,
 )
+from catalyst.debug.debugger import is_debugger_active
 from catalyst.debug.instruments import instrumentation
 from catalyst.debug.printing import print, print_memref  # pylint: disable=redefined-builtin
 

frontend/catalyst/jit.py

@@ -90,6 +90,7 @@ def qjit(
     circuit_transform_pipeline=None,
     pass_plugins=None,
     dialect_plugins=None,
+    debug_compiler=False,
 ):  # pylint: disable=too-many-arguments,unused-argument
     """A just-in-time decorator for PennyLane and JAX programs using Catalyst.
 
@@ -161,6 +162,8 @@ def qjit(
             If not specified, the default pass pipeline will be applied.
         pass_plugins (Optional[List[Path]]): List of paths to pass plugins.
         dialect_plugins (Optional[List[Path]]): List of paths to dialect plugins.
+        debug_compiler (Optional[bool]): Enable external debugger attachment to the compiler
+            driver when launching from an active Python debugging environment.
 
     Returns:
         QJIT object.

frontend/catalyst/pipelines.py

@@ -26,7 +26,7 @@
 """
 
 import enum
 import sys
 from copy import deepcopy
 from dataclasses import dataclass
 from functools import partial
@@ -109,10 +109,12 @@
             Default is None.
         pass_plugins (Optional[Iterable[Path]]): List of paths to pass plugins.
         dialect_plugins (Optional[Iterable[Path]]): List of paths to dialect plugins.
+        debug_compiler (Optional[bool]): Enable external debugger attachment to the compiler
+            driver when launching from an active Python debugging environment.
     """
 
     verbose: Optional[bool] = False
-    logfile: Optional[TextIOWrapper] = sys.stderr
+    logfile: Optional[TextIOWrapper] = "test1.txt"  # sys.stderr
     target: Optional[str] = "binary"
     keep_intermediate: Optional[Union[str, int, bool, KeepIntermediateLevel]] = False
     pipelines: Optional[List[Any]] = None
@@ -122,13 +124,14 @@
     static_argnums: Optional[Union[int, Iterable[int]]] = None
     static_argnames: Optional[Union[str, Iterable[str]]] = None
     abstracted_axes: Optional[Union[Iterable[Iterable[str]], Dict[int, str]]] = None
-    lower_to_llvm: Optional[bool] = True
+    lower_to_llvm: Optional[bool] = False  # "True"
     checkpoint_stage: Optional[str] = ""
-    disable_assertions: Optional[bool] = False
+    disable_assertions: Optional[bool] = True
     seed: Optional[int] = None
     circuit_transform_pipeline: Optional[dict[str, dict[str, str]]] = None
     pass_plugins: Optional[Set[Path]] = None
     dialect_plugins: Optional[Set[Path]] = None
+    debug_compiler: Optional[bool] = False
 
     def __post_init__(self):
         # Convert keep_intermediate to Enum

frontend/catalyst/utils/filesystem.py

@@ -121,12 +121,12 @@
         preferred_name = name.name
 
         # TODO: Maybe just look for the last one?
-        while curr_preferred_abspath.exists():
+        while curr_preferred_abspath.exists() and False:
             curr_preferred_name_str = preferred_name + "_" + str(count)
             curr_preferred_name = pathlib.Path(curr_preferred_name_str)
             curr_preferred_abspath = path / curr_preferred_name
             count += 1
 
         free_preferred_abspath = pathlib.Path(curr_preferred_abspath)
-        free_preferred_abspath.mkdir()
+        free_preferred_abspath.mkdir(exist_ok=True)
         return free_preferred_abspath

mlir/cmake/modules/CMakeLists.txt

@@ -26,7 +26,7 @@ set(llvm_cmake_builddir "${LLVM_BINARY_DIR}/${LLVM_INSTALL_PACKAGE_DIR}")
 
 # Targets with third party dependencies are not exported.
 get_property(MLIR_EXPORTS GLOBAL PROPERTY MLIR_EXPORTS)
-set(TARGETS_TO_REMOVE nlohmann_json tomlplusplus_tomlplusplus ion-transforms CatalystCompilerDriver QECUtils QuantumCAPI qec-transforms)
+set(TARGETS_TO_REMOVE nlohmann_json tomlplusplus_tomlplusplus ion-transforms CatalystCompilerDriver QECUtils QuantumCAPI qec-transforms quantum-transforms)
 list(REMOVE_ITEM MLIR_EXPORTS ${TARGETS_TO_REMOVE})
 export(TARGETS ${MLIR_EXPORTS} ExternalStablehloLib FILE ${catalyst_cmake_builddir}/CatalystTargets.cmake)
 

mlir/include/Quantum/Transforms/Passes.h

@@ -34,5 +34,6 @@ std::unique_ptr<mlir::Pass> createDisentangleCNOTPass();
 std::unique_ptr<mlir::Pass> createDisentangleSWAPPass();
 std::unique_ptr<mlir::Pass> createIonsDecompositionPass();
 std::unique_ptr<mlir::Pass> createLoopBoundaryOptimizationPass();
+std::unique_ptr<mlir::Pass> createQuantumSpecsInfoPass();
 
 } // namespace catalyst

mlir/include/Quantum/Transforms/Passes.td

@@ -139,4 +139,11 @@ def LoopBoundaryOptimizationPass : Pass<"loop-boundary"> {
 }
 // ----- Quantum circuit transformation passes end ----- //
 
+
+def QuantumSpecsInfoPass : Pass<"quantum-specs-info"> {
+    let summary = "Informative pass to provide quantum specs, like gate counts.";
+
+    let constructor = "catalyst::createQuantumSpecsInfoPass()";
+}
+
 #endif // QUANTUM_PASSES

mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp

@@ -75,4 +75,5 @@ void catalyst::registerAllCatalystPasses()
     mlir::registerPass(catalyst::createSplitMultipleTapesPass);
     mlir::registerPass(catalyst::createTestPass);
     mlir::registerPass(catalyst::createTLayerReductionPass);
+    mlir::registerPass(catalyst::createQuantumSpecsInfoPass);
 }

mlir/lib/QEC/Transforms/CountPPMSpecs.cpp

@@ -19,13 +19,17 @@
 
 #include <nlohmann/json.hpp>
 
+// #include "llvm/Support/raw_ostream.h"
+#include "llvm/Support/FileSystem.h"
+
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
 #include "Catalyst/Utils/SCFUtils.h"
 #include "QEC/IR/QECDialect.h"
 #include "QEC/IR/QECOpInterfaces.h"
+// #include "QEC/Utils/PauliStringWrapper.h"
 #include "QEC/Utils/QECLayer.h"
 #include "QEC/Utils/QECOpUtils.h"
 #include "Quantum/IR/QuantumOps.h"
@@ -229,9 +233,18 @@ struct CountPPMSpecsPass : public impl::CountPPMSpecsPassBase<CountPPMSpecsPass>
 
         countDepths(layers, PPMSpecs, stringAllocator);
 
+        std::error_code EC;
+        llvm::raw_fd_ostream fileOutputStream("test.txt", EC, llvm::sys::fs::OF_Append);
+        if (EC) {
+            llvm::errs() << "Error opening file: " << EC.message() << "\n";
+            return failure(); // Handle error
+        }
         json PPMSpecsJson = PPMSpecs;
         llvm::outs() << PPMSpecsJson.dump(4)
                      << "\n"; // dump(4) makes an indent with 4 spaces when printing JSON
+        fileOutputStream << PPMSpecsJson.dump(4)
+                         << "\n"; // dump(4) makes an indent with 4 spaces when printing JSON
+        fileOutputStream.flush();
         return success();
     }
 

mlir/lib/Quantum/Transforms/CMakeLists.txt

@@ -20,6 +20,7 @@ file(GLOB SRC
     IonsDecompositionPatterns.cpp
     loop_boundary_optimization.cpp
     LoopBoundaryOptimizationPatterns.cpp
+    QuantumSpecsInfo.cpp
 )
 
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
@@ -36,6 +37,19 @@ set(DEPENDS
 
 add_mlir_library(${LIBRARY_NAME} STATIC ${SRC} LINK_LIBS PRIVATE ${LIBS} DEPENDS ${DEPENDS})
 target_compile_features(${LIBRARY_NAME} PUBLIC cxx_std_20)
+
+target_link_libraries(${LIBRARY_NAME} 
+    PRIVATE nlohmann_json::nlohmann_json
+    )
+
+# Disable warning from json library
+if(CMAKE_CXX_COMPILER_ID MATCHES ".*Clang")
+set_source_files_properties(
+    QuantumSpecsInfo.cpp
+    COMPILE_FLAGS "-Wno-covered-switch-default"
+)
+endif()
+
 target_include_directories(${LIBRARY_NAME} PUBLIC
                            .
                            ${PROJECT_SOURCE_DIR}/include

mlir/lib/Quantum/Transforms/QuantumSpecsInfo.cpp

@@ -0,0 +1,132 @@
+// 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 "quantum-specs-info"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/FileSystem.h"
+#include <nlohmann/json.hpp>
+
+#include "QEC/IR/QECDialect.h"
+#include "Quantum/IR/QuantumDialect.h"
+#include "Quantum/IR/QuantumInterfaces.h"
+#include "Quantum/IR/QuantumOps.h"
+
+using namespace llvm;
+using namespace mlir;
+using namespace catalyst;
+using namespace catalyst::quantum;
+using namespace catalyst::qec;
+using json = nlohmann::json;
+
+namespace catalyst {
+namespace quantum {
+
+#define GEN_PASS_DEF_QUANTUMSPECSINFOPASS
+#define GEN_PASS_DECL_QUANTUMSPECSINFOPASS
+
+#include "Quantum/Transforms/Passes.h.inc"
+
+struct QuantumSpecsInfoPass : public impl::QuantumSpecsInfoPassBase<QuantumSpecsInfoPass> {
+    using impl::QuantumSpecsInfoPassBase<QuantumSpecsInfoPass>::QuantumSpecsInfoPassBase;
+
+    LogicalResult printQuantumSpecs()
+    {
+        llvm::BumpPtrAllocator stringAllocator;
+        llvm::DenseMap<StringRef, llvm::DenseMap<StringRef, int>> PPMSpecs;
+
+        int numQuantumGates = 0;
+        int numMeasureOp = 0;
+        int numPPRotation = 0;
+        int numPPMeasurement = 0;
+        int numTotal = 0;
+
+        // Walk over all operations in the IR (could be ModuleOp or FuncOp)
+        WalkResult wr = getOperation()->walk([&](Operation *op) {
+            if (isa<quantum::QuantumGate>(op)) {
+                auto parentFuncOp = op->getParentOfType<func::FuncOp>();
+                StringRef funcName = parentFuncOp.getName();
+                // llvm::errs() << funcName;
+                PPMSpecs[funcName]["QuantumGate_Count"] = ++numQuantumGates;
+                PPMSpecs[funcName]["Total_Count"] = ++numTotal;
+                return WalkResult::advance();
+            }
+
+            if (isa<quantum::MeasureOp>(op)) {
+                auto parentFuncOp = op->getParentOfType<func::FuncOp>();
+                StringRef funcName = parentFuncOp.getName();
+                // llvm::errs() << funcName;
+                PPMSpecs[funcName]["MeasureOp_Count"] = ++numMeasureOp;
+                PPMSpecs[funcName]["Total_Count"] = ++numTotal;
+                return WalkResult::advance();
+            }
+            // Count PPMs
+            else if (isa<qec::PPMeasurementOp>(op)) {
+                auto parentFuncOp = op->getParentOfType<func::FuncOp>();
+                StringRef funcName = parentFuncOp.getName();
+                PPMSpecs[funcName]["PPMeasurement_Count"] = ++numPPMeasurement;
+                PPMSpecs[funcName]["Total_Count"] = ++numTotal;
+                return WalkResult::advance();
+            }
+            // Count PPRs
+            else if (isa<qec::PPRotationOp>(op)) {
+                auto parentFuncOp = op->getParentOfType<func::FuncOp>();
+                StringRef funcName = parentFuncOp.getName();
+                PPMSpecs[funcName]["PPRotation_Count"] = ++numPPRotation;
+                PPMSpecs[funcName]["Total_Count"] = ++numTotal;
+                return WalkResult::advance();
+            }
+            // Skip other ops
+            else {
+                return WalkResult::skip();
+            }
+        });
+        if (wr.wasInterrupted()) {
+            return failure();
+        }
+
+        std::error_code EC;
+        // llvm::raw_fd_ostream fileOutputStream("test.txt", EC, llvm::sys::fs::OF_None);
+        llvm::raw_fd_ostream fileOutputStream("test.txt", EC, llvm::sys::fs::OF_Append);
+        if (EC) {
+            llvm::errs() << "Error opening file: " << EC.message() << "\n";
+            return failure(); // Handle error
+        }
+        json PPMSpecsJson = PPMSpecs;
+        llvm::outs() << PPMSpecsJson.dump(4)
+                     << "\n"; // dump(4) makes an indent with 4 spaces when printing JSON
+        fileOutputStream << PPMSpecsJson.dump(4)
+                         << "\n"; // dump(4) makes an indent with 4 spaces when printing JSON
+        fileOutputStream.flush();
+        return success();
+    }
+
+    void runOnOperation() final
+    {
+        if (failed(printQuantumSpecs())) {
+            signalPassFailure();
+        }
+    }
+};
+
+} // namespace quantum
+
+std::unique_ptr<Pass> createQuantumSpecsInfoPass()
+{
+    return std::make_unique<QuantumSpecsInfoPass>();
+}
+
+} // namespace catalyst