New compiler driver interface (#1208)

**Context:**

The compiler driver previously had limited flexibility regarding
compiler passes, pipelines, and option management. There was also a need
for a standalone command-line compiler tool to facilitate easier
compilation and testing of MLIR files without python dependancy.
Additionally, the handling of intermediate compiler outputs in the
Python frontend was complex and needed simplification.


**Description of the Change:**

1. Refactored Compiler Driver:
- Updated CompilerOptions and related functions to accommodate new
options and pipeline configurations.
- Enhanced the handling of compiler pipelines to allow for custom
pipelines and better integration with MLIR’s pass management.

2. New Standalone Compiler Tool (catalyst-cli):
- catalyst-cli enables users to compile MLIR files from the command
line.
- Ability to isolate different tools that are used in the compilation
process (quantum-opt, mlir-translate, llc)

3. Improved Intermediate File Handling
- Changed the keepIntermediate option, allowing users to save
intermediate IR after each pass or pipeline. Currently, this is only
available via catalyst-cli and the python frontend is still unchanged,
which can updated later.

4. Pipeline Configuration:
- Defined catalyst pipelines explicitly in Pipelines.h and Pipelines.cpp
and also registered them as individual passes. Therefore, user can
easily run an specific pipeline on an mlir file.

5. Removed self.last_compiler_output
- Modified the compiler to search the workspace directory for outputs
which is inline with the keep_intermediate option.

**Benefits:**
1. Standalone executable to run the catalyst compilation pipeline
without dependancy to python frontend
2. Easier Debugging capabilities 
3. mlir options are directly accessible from catalyst-cli e.g.
  -mlir-print-ir-after-failure
  -mlir-print-ir-before-all
  -mlir-print-op-generic
  -mlir-timing
  -pass-pipeline 
4. New options added for catalyst compiler e.g.
-tool=[opt|translate|llc]
-keep-intermediate=[true|false]
-save-ir-after-each=[pass|pipeline]
-catalyst-pipeline=pipeline1(pass1;pass2),pipeline2(pass3)
5. The options that are required for enabling MLIR plugins are now
available. e.g.
-load-pass-plugin ,and -load-dialect-plugin

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-73536], [sc-73353]

---------

Co-authored-by: paul0403 <[email protected]>

.github/workflows/build-wheel-linux-x86_64.yaml

@@ -379,7 +379,7 @@ jobs:
               -DLLVM_ENABLE_ZSTD=FORCE_ON \
               -DLLVM_ENABLE_LLD=ON
 
-        cmake --build quantum-build --target check-dialects compiler_driver
+        cmake --build quantum-build --target check-dialects compiler_driver catalyst-cli
 
     - name: Build wheel
       run: |

.github/workflows/build-wheel-macos-arm64.yaml

@@ -344,7 +344,7 @@ jobs:
               -DLLVM_ENABLE_LLD=OFF \
               -DLLVM_DIR=$GITHUB_WORKSPACE/llvm-build/lib/cmake/llvm
 
-        cmake --build quantum-build --target check-dialects compiler_driver
+        cmake --build quantum-build --target check-dialects compiler_driver catalyst-cli
 
     - name: Build wheel
       run: |

.github/workflows/build-wheel-macos-x86_64.yaml

@@ -334,7 +334,7 @@ jobs:
               -DLLVM_ENABLE_ZSTD=FORCE_ON \
               -DLLVM_ENABLE_LLD=OFF
 
-        cmake --build quantum-build --target check-dialects compiler_driver
+        cmake --build quantum-build --target check-dialects compiler_driver catalyst-cli
 
     - name: Build wheel
       run: |

.github/workflows/scripts/linux_arm64/rh8/build_catalyst.sh

@@ -75,7 +75,7 @@ cmake -S mlir -B quantum-build -G Ninja \
     -DLLVM_ENABLE_ZSTD=FORCE_ON \
     -DLLVM_ENABLE_LLD=ON \
     -DLLVM_DIR=/catalyst/llvm-build/lib/cmake/llvm
-cmake --build quantum-build --target check-dialects compiler_driver
+cmake --build quantum-build --target check-dialects compiler_driver catalyst-cli
 
 # Copy files needed for the wheel where they are expected
 cp /catalyst/runtime-build/lib/*/*/*/*/librtd* /catalyst/runtime-build/lib

doc/releases/changelog-dev.md

@@ -186,6 +186,29 @@
   Array([2, 4, 6], dtype=int64)
   ```
 
+* Catalyst now has a standalone compiler tool called `catalyst-cli` which quantum
+  compiles MLIR input files into an object file without any dependancy to the python frontend.
+  [(#1208)](https://github.com/PennyLaneAI/catalyst/pull/1208)
+
+  This compiler tool combines three stages of compilation which are:
+
+  - qunatum-opt: Performs the mlir level optimizations and lowers the input dialect to LLVM dialect.
+  - mlir-translate: Translates the input in LLVM dialect into LLVM IR.
+  - llc: Performs lower level optimizations and creates the object file.
+
+  catalyst-cli runs all of the above stages under the hood, but it has the ability to isolate them on demand.
+  An example of usage whould look like below:
+
+  ```
+  // Creates both the optimized IR and an object file
+  catalyst-cli input.mlir -o output.o 
+  // Only performs MLIR optimizations
+  catalyst-cli --tool=opt input.mlir -o llvm-dialect.mlir
+  // Only lowers LLVM dialect MLIR input to LLVM IR
+  catalyst-cli --tool=translate llvm-dialect.mlir -o llvm-ir.ll
+  // Only performs lower-level optimizations and create object file
+  catalyst-cli --tool=llc llvm-ir.ll -o output.o
+
 * Static arguments of a qjit-compiled function can now be indicated by a `static_argnames`
   argument to `qjit`.
   [(#1158)](https://github.com/PennyLaneAI/catalyst/pull/1158)
@@ -301,6 +324,11 @@
 
   Please use `debug.replace_ir`.
 
+* Removes `compiler.last_compiler_output`.
+  [(#1208)](https://github.com/PennyLaneAI/catalyst/pull/1208)
+
+  Please use `compiler.get_output_of("last", workspace)`
+
 <h3>Bug fixes</h3>
 
 * Fix a bug in `catalyst.mitigate_with_zne` that would lead

frontend/catalyst/compiler.py

@@ -547,7 +547,6 @@ class Compiler:
     @debug_logger_init
     def __init__(self, options: Optional[CompileOptions] = None):
         self.options = options if options is not None else CompileOptions()
-        self.last_compiler_output = None
 
     @debug_logger
     def run_from_ir(self, ir: str, module_name: str, workspace: Directory):
@@ -604,7 +603,6 @@ def run_from_ir(self, ir: str, module_name: str, workspace: Directory):
         else:
             output_filename = filename
 
-        self.last_compiler_output = compiler_output
         return output_filename, out_IR
 
     @debug_logger
@@ -633,20 +631,49 @@ def run(self, mlir_module, *args, **kwargs):
         )
 
     @debug_logger
-    def get_output_of(self, pipeline) -> Optional[str]:
+    def get_output_of(self, pipeline, workspace) -> Optional[str]:
         """Get the output IR of a pipeline.
         Args:
             pipeline (str): name of pass class
 
         Returns
             (Optional[str]): output IR
         """
-        if not self.last_compiler_output or not self.last_compiler_output.get_pipeline_output(
-            pipeline
-        ):
+        file_content = None
+        for dirpath, _, filenames in os.walk(str(workspace)):
+            filenames = [f for f in filenames if f.endswith(".mlir") or f.endswith(".ll")]
+            if not filenames:
+                break
+            filenames_no_ext = [os.path.splitext(f)[0] for f in filenames]
+            if pipeline == "mlir":
+                # Sort files and pick the first one
+                selected_file = [
+                    sorted(filenames)[0],
+                ]
+            elif pipeline == "last":
+                # Sort files and pick the last one
+                selected_file = [
+                    sorted(filenames)[-1],
+                ]
+            else:
+                selected_file = [
+                    f
+                    for f, name_no_ext in zip(filenames, filenames_no_ext)
+                    if pipeline in name_no_ext
+                ]
+            if len(selected_file) != 1:
+                msg = f"Attempting to get output for pipeline: {pipeline},"
+                msg += " but no or more than one file was found.\n"
+                raise CompileError(msg)
+            filename = selected_file[0]
+
+            full_path = os.path.join(dirpath, filename)
+            with open(full_path, "r", encoding="utf-8") as file:
+                file_content = file.read()
+
+        if file_content is None:
             msg = f"Attempting to get output for pipeline: {pipeline},"
             msg += " but no file was found.\n"
             msg += "Are you sure the file exists?"
             raise CompileError(msg)
-
-        return self.last_compiler_output.get_pipeline_output(pipeline)
+        return file_content

frontend/catalyst/debug/compiler_functions.py

@@ -98,9 +98,7 @@ def func(x: float):
     if not isinstance(fn, catalyst.QJIT):
         raise TypeError(f"First argument needs to be a 'QJIT' object, got a {type(fn)}.")
 
-    if stage == "last":
-        return fn.compiler.last_compiler_output.get_output_ir()
-    return fn.compiler.get_output_of(stage)
+    return fn.compiler.get_output_of(stage, fn.workspace)
 
 
 @debug_logger

frontend/test/pytest/test_compiler.py

@@ -122,7 +122,7 @@ def test_attempts_to_get_inexistent_intermediate_file(self):
         """Test the return value if a user requests an intermediate file that doesn't exist."""
         compiler = Compiler()
         with pytest.raises(CompileError, match="Attempting to get output for pipeline"):
-            compiler.get_output_of("inexistent-file")
+            compiler.get_output_of("inexistent-file", ".")
 
     def test_runtime_error(self, backend):
         """Test with non-default flags."""
@@ -222,15 +222,15 @@ def workflow():
 
         compiler = workflow.compiler
         with pytest.raises(CompileError, match="Attempting to get output for pipeline"):
-            compiler.get_output_of("EmptyPipeline1")
-        assert compiler.get_output_of("HLOLoweringPass")
-        assert compiler.get_output_of("QuantumCompilationPass")
+            compiler.get_output_of("EmptyPipeline1", workflow.workspace)
+        assert compiler.get_output_of("HLOLoweringPass", workflow.workspace)
+        assert compiler.get_output_of("QuantumCompilationPass", workflow.workspace)
         with pytest.raises(CompileError, match="Attempting to get output for pipeline"):
-            compiler.get_output_of("EmptyPipeline2")
-        assert compiler.get_output_of("BufferizationPass")
-        assert compiler.get_output_of("MLIRToLLVMDialect")
+            compiler.get_output_of("EmptyPipeline2", workflow.workspace)
+        assert compiler.get_output_of("BufferizationPass", workflow.workspace)
+        assert compiler.get_output_of("MLIRToLLVMDialect", workflow.workspace)
         with pytest.raises(CompileError, match="Attempting to get output for pipeline"):
-            compiler.get_output_of("None-existing-pipeline")
+            compiler.get_output_of("None-existing-pipeline", workflow.workspace)
         workflow.workspace.cleanup()
 
     def test_print_nonexistent_stages(self, backend):
@@ -243,7 +243,7 @@ def workflow():
             return qml.state()
 
         with pytest.raises(CompileError, match="Attempting to get output for pipeline"):
-            workflow.compiler.get_output_of("None-existing-pipeline")
+            workflow.compiler.get_output_of("None-existing-pipeline", workflow.workspace)
         workflow.workspace.cleanup()
 
     def test_workspace(self):
@@ -305,10 +305,9 @@ def circuit():
             compiled.compile()
 
         assert "Failed to lower MLIR module" in e.value.args[0]
-        assert "While processing 'TestPass' pass of the 'PipelineB' pipeline" in e.value.args[0]
-        assert "PipelineA" not in e.value.args[0]
+        assert "While processing 'TestPass' pass " in e.value.args[0]
         assert "Trace" not in e.value.args[0]
-        assert isfile(os.path.join(str(compiled.workspace), "2_PipelineB_FAILED.mlir"))
+        assert isfile(os.path.join(str(compiled.workspace), "2_TestPass_FAILED.mlir"))
         compiled.workspace.cleanup()
 
         with pytest.raises(CompileError) as e:

frontend/test/pytest/test_debug.py

@@ -376,8 +376,6 @@ def f(x):
             """Square function."""
             return x**2
 
-        f.__name__ = f.__name__ + pass_name
-
         jit_f = qjit(f, keep_intermediate=True)
         data = 2.0
         old_result = jit_f(data)
@@ -400,8 +398,6 @@ def f(x: float):
             """Square function."""
             return x**2
 
-        f.__name__ = f.__name__ + pass_name
-
         jit_f = qjit(f)
         jit_grad_f = qjit(value_and_grad(jit_f), keep_intermediate=True)
         jit_grad_f(3.0)
@@ -418,7 +414,7 @@ def f(x: float):
         assert len(res) == 0
 
     def test_get_compilation_stage_without_keep_intermediate(self):
-        """Test if error is raised when using get_pipeline_output without keep_intermediate."""
+        """Test if error is raised when using get_compilation_stage without keep_intermediate."""
 
         @qjit
         def f(x: float):

mlir/Makefile

@@ -146,7 +146,7 @@ dialects:
 		-DLLVM_ENABLE_ZLIB=$(ENABLE_ZLIB) \
 		-DLLVM_ENABLE_ZSTD=$(ENABLE_ZSTD)
 
-	cmake --build $(DIALECTS_BUILD_DIR) --target check-dialects quantum-lsp-server compiler_driver
+	cmake --build $(DIALECTS_BUILD_DIR) --target check-dialects quantum-lsp-server compiler_driver catalyst-cli
 
 .PHONY: test
 test:

mlir/include/Driver/CompilerDriver.h

@@ -21,9 +21,12 @@
 
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Support/LogicalResult.h"
+#include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include "Driver/Pipelines.h"
+
 namespace catalyst {
 namespace driver {
 
@@ -32,6 +35,12 @@ namespace driver {
 // low-level messages, we might want to hide these.
 enum class Verbosity { Silent = 0, Urgent = 1, Debug = 2, All = 3 };
 
+enum SaveTemps { None, AfterPipeline, AfterPass };
+
+enum Action { OPT, Translate, LLC, All };
+
+enum InputType { MLIR, LLVMIR, OTHER };
+
 /// Helper verbose reporting macro.
 #define CO_MSG(opt, level, op)                                                                     \
     do {                                                                                           \
@@ -40,14 +49,6 @@ enum class Verbosity { Silent = 0, Urgent = 1, Debug = 2, All = 3 };
         }                                                                                          \
     } while (0)
 
-/// Pipeline descriptor
-struct Pipeline {
-    using Name = std::string;
-    using PassList = llvm::SmallVector<std::string>;
-    Name name;
-    PassList passes;
-};
-
 /// Optional parameters, for which we provide reasonable default values.
 struct CompilerOptions {
     /// The textual IR (MLIR or LLVM IR)
@@ -58,19 +59,21 @@ struct CompilerOptions {
     mlir::StringRef moduleName;
     /// The stream to output any error messages from MLIR/LLVM passes and translation.
     llvm::raw_ostream &diagnosticStream;
-    /// If true, the driver will output the module at intermediate points.
-    bool keepIntermediate;
+    /// If specified, the driver will output the module after each pipeline or each pass.
+    SaveTemps keepIntermediate;
     /// If true, the llvm.coroutine will be lowered.
     bool asyncQnodes;
     /// Sets the verbosity level to use when printing messages.
     Verbosity verbosity;
     /// Ordered list of named pipelines to execute, each pipeline is described by a list of MLIR
     /// passes it includes.
     std::vector<Pipeline> pipelinesCfg;
-    /// Whether to assume that the pipelines output is a valid LLVM dialect and lower it to LLVM IR
-    bool lowerToLLVM;
     /// Specify that the compiler should start after reaching the given pass.
     std::string checkpointStage;
+    /// Specify the loweting action to perform
+    Action loweringAction;
+    /// If true, the compiler will dump the pass pipeline that will be run.
+    bool dumpPassPipeline;
 
     /// Get the destination of the object file at the end of compilation.
     std::string getObjectFile() const
@@ -81,8 +84,8 @@ struct CompilerOptions {
 };
 
 struct CompilerOutput {
-    typedef std::unordered_map<Pipeline::Name, std::string> PipelineOutputs;
-    std::string objectFilename;
+    typedef std::unordered_map<std::string, std::string> PipelineOutputs;
+    std::string outputFilename;
     std::string outIR;
     std::string diagnosticMessages;
     PipelineOutputs pipelineOutputs;
@@ -91,7 +94,7 @@ struct CompilerOutput {
     bool isCheckpointFound;
 
     // Gets the next pipeline dump file name, prefixed with number.
-    std::string nextPipelineDumpFilename(Pipeline::Name pipelineName, std::string ext = ".mlir")
+    std::string nextPipelineDumpFilename(std::string pipelineName, std::string ext = ".mlir")
     {
         return std::filesystem::path(std::to_string(this->pipelineCounter++) + "_" + pipelineName)
             .replace_extension(ext);
@@ -103,15 +106,24 @@ struct CompilerOutput {
 
 /// Entry point to the MLIR portion of the compiler.
 mlir::LogicalResult QuantumDriverMain(const catalyst::driver::CompilerOptions &options,
-                                      catalyst::driver::CompilerOutput &output);
+                                      catalyst::driver::CompilerOutput &output,
+                                      mlir::DialectRegistry &registry);
+
+int QuantumDriverMainFromCL(int argc, char **argv);
+int QuantumDriverMainFromArgs(const std::string &source, const std::string &workspace,
+                              const std::string &moduleName, bool keepIntermediate,
+                              bool asyncQNodes, bool verbose, bool lowerToLLVM,
+                              const std::vector<catalyst::driver::Pipeline> &passPipelines,
+                              const std::string &checkpointStage,
+                              catalyst::driver::CompilerOutput &output);
 
 namespace llvm {
 
 inline raw_ostream &operator<<(raw_ostream &oss, const catalyst::driver::Pipeline &p)
 {
-    oss << "Pipeline('" << p.name << "', [";
+    oss << "Pipeline('" << p.getName() << "', [";
     bool first = true;
-    for (const auto &i : p.passes) {
+    for (const auto &i : p.getPasses()) {
         oss << (first ? "" : ", ") << i;
         first = false;
     }