Improve compiler driver tool stages [1/3]

doc/catalyst-cli/catalyst-cli.rst

@@ -5,30 +5,45 @@ Catalyst includes a standalone command-line tool, ``catalyst``, that
 compiles quantum programs written in our MLIR dialects into an object file,
 independent of the Catalyst Python frontend.
 
-This compiler tool combines three stages of compilation:
+This compiler tool combines four stages of compilation:
 
 #. ``quantum-opt``: Performs the MLIR-level optimizations, including quantum optimizations, and
    translates the input dialect to the `LLVM dialect <https://mlir.llvm.org/docs/Dialects/LLVM/>`_.
 #. ``mlir-translate``: Translates the input LLVM dialect into
    `LLVM IR <https://llvm.org/docs/LangRef.html>`_.
-#. ``llc``: Performs lower-level optimizations on the LLVM IR input and creates the object file.
+#. ``llvm-opt``: Performs lower-level optimization and compilation on the LLVM IR needed for certain
+   features.
+#. ``llc``: Runs LLVM backend code generation and creates the object file.
 
-The ``catalyst`` tool will run all three stages under the hood by default, but it also has the
+The ``catalyst`` tool will run all stages under the hood by default, but it also has the
 ability to run each stage individually. For example:
 
 .. code-block:: console
 
-    # Creates both the optimized IR and an object file
-    catalyst input.mlir -o output.o
+    # Run MLIR passes, returning LLVM-dialect MLIR
+    catalyst --tool=opt input.mlir > llvm-dialect.mlir
 
-    # Only performs MLIR optimizations and translates to LLVM dialect
-    catalyst --tool=opt input.mlir -o llvm-dialect.mlir
+    # Translate LLVM-dialect MLIR to LLVM IR
+    catalyst --tool=translate llvm-dialect.mlir > llvm-ir.ll
 
-    # Only lowers LLVM dialect input to LLVM IR
-    catalyst --tool=translate llvm-dialect.mlir -o llvm-ir.ll
+    # Run Catalyst's LLVM transformations
+    catalyst --tool=llvm-opt llvm-ir.ll > optimized.ll
 
-    # Only performs lower-level optimizations and creates object file (object.o)
-    catalyst --tool=llc llvm-ir.ll -o output.ll --module-name object
+    # Generate object.o from backend-ready LLVM IR
+    catalyst --tool=llc optimized.ll --module-name object
+
+Stages can be combined for a partial or full execution of the pipeline:
+
+.. code-block:: console
+
+    # Run MLIR passes and translation in one invocation, without LLVM transformations
+    catalyst --tool=opt-translate input.mlir > llvm-ir.ll
+
+    # Run through LLVM transformations in one invocation, without object generation
+    catalyst --tool=opt-translate-llvm input.mlir > optimized.ll
+
+    # Full compilation from mlir to object code
+    catalyst input.mlir
 
 .. note::
 
@@ -45,10 +60,10 @@ Usage
 
     catalyst [options] <input file>
 
-Calling ``catalyst`` without any options runs the three compilation stages (``quantum-opt``,
-``mlir-translate`` and ``llc``) using all default configurations, and outputs by default an object
-file named ``catalyst_module.o``. The name of the output file can be set by changing the output
-module name using the ``--module-name`` option (the default module name is ``catalyst_module``).
+Calling ``catalyst`` without any options runs the full compilation pipeline using all default
+configurations, and outputs by default an object file named ``catalyst_module.o``.
+The name of the output file can be set by changing the output module name using the
+``--module-name`` option.
 
 Command line options
 ^^^^^^^^^^^^^^^^^^^^
@@ -65,22 +80,28 @@ Show available command-line options and exit.
 ``--verbose``
 """""""""""""
 
-Emit verbose messages.
+Emit verbose progress messages to stderr.
 
 ``-o <filename>``
 """""""""""""""""
 
 Output IR filename. If no output filename is provided, the resulting IR is output to stdout.
 
-``--tool=<opt|translate|llc|all>``
-""""""""""""""""""""""""""""""""""
+``--tool=<opt|translate|llvm-opt|llc>``
+"""""""""""""""""""""""""""""""""""""""
+
+Select the tool to run individually. Each stage requires appropriate IR as input.
+
+* ``opt``: Run MLIR passes on MLIR input.
+* ``translate``: Translate LLVM-dialect MLIR to LLVM IR.
+* ``llvm-opt``: Run Catalyst's applicable LLVM transformations on LLVM IR input.
+* ``llc``: Run backend code generation on backend-ready LLVM IR input.
 
-Select the tool to run individually. The default is ``all``.
+A combination of stages (starting from MLIR) can also be run. ``all`` is the default.
 
-* ``opt``: Run ``quantum-opt`` on the MLIR input.
-* ``translate``: Run ``mlir-translate`` on the LLVM dialect input.
-* ``llc``: Run ``llc`` on the LLVM IR input.
-* ``all``: Run all of ``opt``, ``translate`` and ``llc`` on the MLIR input.
+* ``opt-translate``: Run ``opt`` and ``translate`` on MLIR input.
+* ``opt-translate-llvm``: Run ``opt``, ``translate``, and ``llvm-opt`` on MLIR input.
+* ``all``: Run all stages. LLVM IR input is also accepted, skipping the MLIR stages.
 
 ``--save-ir-after-each=<pass|pipeline|changed>``
 """""""""""""""""""""""""""""""""""""""""""""""

doc/releases/changelog-dev.md

@@ -824,6 +824,13 @@
   :class:`~pennylane.transforms.core.CompilePipeline` on individual QNodes instead.
   [(#3324)](https://github.com/PennyLaneAI/catalyst/pull/3324)
 
+* The Catalyst CLI now exposes LLVM transformations separately from object code generation, via
+  the `--tool=llvm-opt` option. Originally this was part of the same tool that is responsible
+  for object code generation (`--tool=llc`).
+  Combined "tools" are now also exposed via `--tool=opt-translate` and `--tool=opt-translate-llvm`
+  for compilation from MLIR to a stage before object code generation.
+  [(#3402)](https://github.com/PennyLaneAI/catalyst/pull/3402)
+
 * A Backline per-message coprocessor function (`CoprocessorFn`, `CatalystCoprocessorFn` in
   `TransportABI.h`) is now called with the message's payload, its length and its fn_id,
   `fn(payload, len, fn_id, out, out_cap, ctx)`, on every transport, instead of with the raw
@@ -898,6 +905,10 @@
   conditional branches.
   [(#3342)](https://github.com/PennyLaneAI/catalyst/pull/3342)
 
+* In the Catalyst CLI, verbose compiler messages are now written to stderr instead of stdout, so
+  `--verbose` no longer contaminates IR returned through stdout.
+  [(#3402)](https://github.com/PennyLaneAI/catalyst/pull/3402)
+
 * Ahead-of-time compilation with ``target="llvmir"`` no longer links the program into a shared
   library. Calling the function still compiles it to a shared library.
   [(#3391)](https://github.com/PennyLaneAI/catalyst/pull/3391)

frontend/catalyst/compiler.py

@@ -382,9 +382,6 @@ def _options_to_cli_flags(options):
     if options.checkpoint_stage:
         extra_args += [("--checkpoint-stage", options.checkpoint_stage)]
 
-    if not options.lower_to_llvm:
-        extra_args += [("--tool", "opt")]
-
     if options.keep_intermediate:
         extra_args += ["--keep-intermediate"]
 
@@ -408,14 +405,25 @@ def _options_to_cli_flags(options):
     return extra_args
 
 
-def to_llvmir(*args, stdin=None, options: Optional[CompileOptions] = None):
-    """echo ${input} | catalyst *args -"""
-    # These are the options that may affect compilation
-    if not options:
-        return _catalyst(*args, stdin=stdin)
+def to_llvm(*args, stdin=None, options: Optional[CompileOptions] = None):
+    """Return LLVM IR after MLIR lowering and translation, before LLVM transformations.
 
-    opts = _options_to_cli_flags(options)
-    return _catalyst(*opts, *args, stdin=stdin)
+    Runs one CLI invocation with ``--tool=opt-translate``. Final IR is returned
+    independently of ``keep_intermediate``; no object file is generated.
+    """
+    opts = _options_to_cli_flags(options) if options is not None else []
+    return _catalyst(("--tool", "opt-translate"), *opts, *args, stdin=stdin)
+
+
+def to_llvm_opt(*args, stdin=None, options: Optional[CompileOptions] = None):
+    """Return LLVM IR after lowering, translation, and Catalyst's LLVM transformations.
+
+    Runs one CLI invocation with ``--tool=opt-translate-llvm``, including the
+    applicable coroutine, O2, and Enzyme passes. No object file is generated.
+    Final IR is returned independently of ``keep_intermediate``.
+    """
+    opts = _options_to_cli_flags(options) if options is not None else []
+    return _catalyst(("--tool", "opt-translate-llvm"), *opts, *args, stdin=stdin)
 
 
 def to_mlir_opt(
@@ -464,6 +472,8 @@ def get_cli_command(self, tmp_infile_name, output_ir_name, module_name, workspac
             cmd (str): The command to be executed.
         """
         opts = _options_to_cli_flags(self.options)
+        if not self.options.lower_to_llvm:
+            opts += [("--tool", "opt")]
         cmd = _get_catalyst_cli_cmd(
             ("-o", output_ir_name),
             ("--module-name", module_name),

frontend/catalyst/jit.py

@@ -45,7 +45,7 @@
     settle_executors,
 )
 from catalyst.compiled_functions import CompilationCache, CompiledFunction
-from catalyst.compiler import CompileOptions, Compiler, canonicalize, to_llvmir, to_mlir_opt
+from catalyst.compiler import CompileOptions, Compiler, canonicalize, to_llvm_opt, to_mlir_opt
 from catalyst.debug.instruments import instrument
 from catalyst.from_plxpr import trace_from_pennylane
 from catalyst.jax_tracer import lower_jaxpr_to_mlir, trace_to_jaxpr
@@ -778,7 +778,7 @@ def llvmir(self):
             return None
 
         _mlir = str(self.mlir_module)
-        self.llvm_ir = to_llvmir(stdin=_mlir, options=self.compile_options)
+        self.llvm_ir = to_llvm_opt(stdin=_mlir, options=self.compile_options)
         return self.llvm_ir
 
     @debug_logger

frontend/test/lit/test_opt.py

@@ -16,7 +16,7 @@
 
 # RUN: %PYTHON %s | FileCheck %s
 
-from catalyst.compiler import _quantum_opt, canonicalize, to_llvmir
+from catalyst.compiler import _quantum_opt, canonicalize, to_llvm_opt
 
 
 def test_opt_happy_path():
@@ -66,7 +66,7 @@ def test_to_llvmir():
         }
     }
     """
-    output = to_llvmir(stdin=mlir)
+    output = to_llvm_opt(stdin=mlir)
     print(output)
 
 

frontend/test/pytest/test_compiler.py

@@ -163,6 +163,20 @@ def test_options_to_cli_flags_keep_intermediate_pass(self):
 class TestCompilerDiagnostics:
     """Unit tests for diagnostics emitted by the Catalyst subprocess."""
 
+    def test_verbose_logging_preserves_ir(self, capsys):
+        """Verbose progress goes to stderr without contaminating the returned LLVM IR."""
+        tool = "opt-translate-llvm"
+        source = "func.func @foo() { return }"
+        quiet_ir = _catalyst(("--tool", tool), stdin=source)
+        verbose_ir = _catalyst(("--tool", tool), "--verbose", stdin=source, stderr_return=True)
+
+        assert "define void @foo" in quiet_ir
+        assert verbose_ir == quiet_ir
+        captured = capsys.readouterr()
+        assert captured.out == ""
+        assert "MLIR parsing successful" in captured.err
+        assert "Compilation successful:" in captured.err
+
     @pytest.mark.parametrize(
         "text,stdout,stderr",
         [

frontend/test/pytest/test_debug.py

@@ -24,7 +24,7 @@
 from pennylane import for_loop
 
 from catalyst import debug, qjit, value_and_grad
-from catalyst.compiler import _options_to_cli_flags, to_llvmir, to_mlir_opt
+from catalyst.compiler import Compiler, _options_to_cli_flags, to_llvm, to_llvm_opt, to_mlir_opt
 from catalyst.debug import compile_executable, get_cmain, get_compilation_stage, replace_ir
 from catalyst.pipelines import CompileOptions
 from catalyst.utils.exceptions import CompileError
@@ -581,12 +581,14 @@ def test_option_use_nameloc(self):
         assert "--use-nameloc-as-prefix" in flags
 
     def test_option_not_lower_to_llvm(self):
-        """Test not lower to llvm"""
+        """The compiler selects the MLIR-only tool when LLVM lowering is disabled."""
         options = CompileOptions(lower_to_llvm=False)
-        flags = _options_to_cli_flags(options)
-        assert ("--tool", "opt") in flags
+        command = Compiler(options).get_cli_command("input.mlir", "output.mlir", "module", ".")
+        assert "--tool" in command
+        assert command[command.index("--tool") + 1] == "opt"
 
-    def test_no_options_to_llvmir(self):
+    @pytest.mark.parametrize("lower", [to_llvm, to_llvm_opt])
+    def test_no_options_to_llvmir(self, lower):
         """Test that we can lower to llvmir without needing any options"""
         mlir = """
         module {
@@ -595,7 +597,7 @@ def test_no_options_to_llvmir(self):
             }
         }
         """
-        observed = to_llvmir(stdin=mlir)
+        observed = lower(stdin=mlir)
         # pylint: disable=line-too-long
         expected = textwrap.dedent("""
         define void @foo() {

mlir/include/Driver/CompilerDriver.h

@@ -57,7 +57,7 @@ enum SaveTemps { None, AfterPipeline, AfterPassChanged, AfterPass };
  * @brief Defines the different functional stages of the driver.
  *
  */
-enum Action { OPT, Translate, LLC, All };
+enum Action { OPT, Translate, LLVMOpt, LLC, OptTranslate, OptTranslateLLVM, All };
 
 /**
  * @brief Defines the format of the input data.

mlir/lib/Driver/CompilerDriver.cpp

@@ -534,11 +534,13 @@ llvm::LogicalResult catalyst::driver::verifyInputType(const CompilerOptions &opt
         CO_MSG(options, Verbosity::Urgent, "Wrong or unsupported input\n");
         return failure();
     }
-    if (options.loweringAction == Action::LLC && inType != InputType::LLVMIR) {
+    bool requiresLLVM =
+        options.loweringAction == Action::LLC || options.loweringAction == Action::LLVMOpt;
+    if (requiresLLVM && inType != InputType::LLVMIR) {
         CO_MSG(options, Verbosity::Urgent, "Expected LLVM IR input but received MLIR input.\n");
         return failure();
     }
-    if (options.loweringAction < Action::LLC && inType != InputType::MLIR) {
+    if (!requiresLLVM && options.loweringAction != Action::All && inType != InputType::MLIR) {
         CO_MSG(options, Verbosity::Urgent, "Expected MLIR input but received LLVM IR input.\n");
         return failure();
     }