Improve keep_intermediate to enable keeping IR after each pass

doc/releases/changelog-dev.md

@@ -95,6 +95,14 @@
   performance by eliminating indirect conversion.
   [(#1738)](https://github.com/PennyLaneAI/catalyst/pull/1738)
 
+* The `keep_intermediate` argument in the `qjit` decorator now accepts a new value that allows for
+  saving intermediate files after each pass. The updated possible options for this argument are:
+  * `False` or `0` or `"none"`: No intermediate files are kept.
+  * `True` or `1` or `"pipeline"`: Intermediate files are saved after each pipeline.
+  * `2` or `"pass"`: Intermediate files are saved after each pass.
+  The default value is `False`.
+  [(#1791)](https://github.com/PennyLaneAI/catalyst/pull/1791)
+
 <h3>Breaking changes 💔</h3>
 
 * (Device Developers Only) The `QuantumDevice` interface in the Catalyst Runtime plugin system

frontend/catalyst/compiler.py

@@ -29,7 +29,7 @@
 from typing import List, Optional
 
 from catalyst.logging import debug_logger, debug_logger_init
-from catalyst.pipelines import CompileOptions
+from catalyst.pipelines import CompileOptions, KeepIntermediateLevel
 from catalyst.utils.exceptions import CompileError
 from catalyst.utils.filesystem import Directory
 from catalyst.utils.runtime_environment import get_cli_path, get_lib_path
@@ -342,7 +342,7 @@ def _options_to_cli_flags(options):
     if not options.lower_to_llvm:
         extra_args += [("--tool", "opt")]
 
-    if options.keep_intermediate:
+    if options.keep_intermediate >= KeepIntermediateLevel.PIPELINE:
         extra_args += ["--keep-intermediate"]
 
     if options.verbose:
@@ -351,6 +351,9 @@ def _options_to_cli_flags(options):
     if options.async_qnodes:  # pragma: nocover
         extra_args += ["--async-qnodes"]
 
+    if options.keep_intermediate >= KeepIntermediateLevel.PASS:
+        extra_args += ["--save-ir-after-each=pass"]
+
     return extra_args
 
 

frontend/catalyst/jit.py

@@ -118,10 +118,15 @@ def qjit(
         async_qnodes (bool): Experimental support for automatically executing
             QNodes asynchronously, if supported by the device runtime.
         target (str): the compilation target
-        keep_intermediate (bool): Whether or not to store the intermediate files throughout the
-            compilation. If ``True``, intermediate representations are available via the
-            :attr:`~.QJIT.mlir`, :attr:`~.QJIT.mlir_opt`, :attr:`~.QJIT.jaxpr`,
-            and :attr:`~.QJIT.qir`, representing different stages in the optimization process.
+        keep_intermediate (Union[str, int, bool]): Level controlling intermediate file generation.
+            - ``False`` or ``0`` or ``"none"`` (default): No intermediate files are kept.
+            - ``True`` or ``1`` or ``"pipeline"``: Intermediate files are saved after each pipeline.
+            - ``2`` or ``"pass"``: Intermediate files are saved after each pass.
+            If enabled, intermediate representations are available via the following attributes:
+            - :attr:`~.QJIT.jaxpr`: JAX program representation
+            - :attr:`~.QJIT.mlir`: MLIR representation after canonicalization
+            - :attr:`~.QJIT.mlir_opt`: MLIR representation after optimization
+            - :attr:`~.QJIT.qir`: QIR in LLVM IR form
         verbose (bool): If ``True``, the tools and flags used by Catalyst behind the scenes are
             printed out.
         logfile (Optional[TextIOWrapper]): File object to write verbose messages to (default -

frontend/catalyst/pipelines.py

@@ -25,6 +25,7 @@
 
 """
 
+import enum
 import sys
 from copy import deepcopy
 from dataclasses import dataclass
@@ -37,6 +38,42 @@
 from catalyst.utils.exceptions import CompileError
 
 
+class KeepIntermediateLevel(enum.IntEnum):
+    """Enum to control the level of intermediate file keeping."""
+
+    NONE = 0  # No intermediate files are kept.
+    PIPELINE = 1  # Intermediate files are saved after each pipeline.
+    PASS = 2  # Intermediate files are saved after each pass.
+
+
+def _parse_keep_intermediate(
+    level: Union[str, int, bool, KeepIntermediateLevel, None],
+) -> KeepIntermediateLevel:
+    """Parse the keep_intermediate value into a KeepIntermediateLevel enum."""
+    if isinstance(level, KeepIntermediateLevel):
+        return level
+    elif level is None:
+        return KeepIntermediateLevel.NONE
+    elif isinstance(level, bool):
+        return KeepIntermediateLevel.PIPELINE if level else KeepIntermediateLevel.NONE
+    elif isinstance(level, int):
+        try:
+            return KeepIntermediateLevel(level)
+        except ValueError as e:
+            raise ValueError(
+                f"Invalid int for keep_intermediate: {level}. " "Valid integers are 0, 1, 2."
+            ) from e
+    elif isinstance(level, str):
+        try:
+            return KeepIntermediateLevel[level.upper()]
+        except KeyError as e:
+            raise ValueError(
+                f"Invalid string for keep_intermediate: {level}. "
+                "Valid strings are 'none', 'pipeline', 'pass'."
+            ) from e
+    raise TypeError(f"Invalid type for keep_intermediate: {type(level)}.")
+
+
 # pylint: disable=too-many-instance-attributes
 @dataclass
 class CompileOptions:
@@ -47,8 +84,11 @@ class CompileOptions:
             Default is ``False``
         logfile (Optional[TextIOWrapper]): the logfile to write output to.
             Default is ``sys.stderr``
-        keep_intermediate (Optional[bool]): flag indicating whether to keep intermediate results.
-            Default is ``False``
+        keep_intermediate (Optional[Union[str, int, bool]]): Level controlling intermediate file
+        generation.
+            - ``False`` or ``0`` or ``"none"`` (default): No intermediate files are kept.
+            - ``True`` or ``1`` or ``"pipeline"``: Intermediate files are saved after each pipeline.
+            - ``2`` or ``"pass"``: Intermediate files are saved after each pass.
         pipelines (Optional[List[Tuple[str,List[str]]]]): A list of tuples. The first entry of the
             tuple corresponds to the name of a pipeline. The second entry of the tuple corresponds
             to a list of MLIR passes.
@@ -79,7 +119,7 @@ class CompileOptions:
     verbose: Optional[bool] = False
     logfile: Optional[TextIOWrapper] = sys.stderr
     target: Optional[str] = "binary"
-    keep_intermediate: Optional[bool] = False
+    keep_intermediate: Optional[Union[str, int, bool, KeepIntermediateLevel]] = False
     pipelines: Optional[List[Any]] = None
     autograph: Optional[bool] = False
     autograph_include: Optional[Iterable[str]] = ()
@@ -96,6 +136,9 @@ class CompileOptions:
     dialect_plugins: Optional[Set[Path]] = None
 
     def __post_init__(self):
+        # Convert keep_intermediate to Enum
+        self.keep_intermediate = _parse_keep_intermediate(self.keep_intermediate)
+
         # Check that async runs must not be seeded
         if self.async_qnodes and self.seed is not None:
             raise CompileError(

frontend/test/pytest/test_compiler.py

@@ -30,8 +30,9 @@
 import pytest
 
 from catalyst import qjit
-from catalyst.compiler import CompileOptions, Compiler, LinkerDriver
+from catalyst.compiler import CompileOptions, Compiler, LinkerDriver, _options_to_cli_flags
 from catalyst.debug import instrumentation
+from catalyst.pipelines import KeepIntermediateLevel
 from catalyst.utils.exceptions import CompileError
 from catalyst.utils.filesystem import Directory
 
@@ -92,6 +93,74 @@ def circuit():
         capture = capture_result.out + capture_result.err
         assert "[DIAGNOSTICS]" in capture
 
+    @pytest.mark.parametrize(
+        "input_value, expected_level",
+        [
+            (None, KeepIntermediateLevel.NONE),
+            (False, KeepIntermediateLevel.NONE),
+            (True, KeepIntermediateLevel.PIPELINE),
+            (0, KeepIntermediateLevel.NONE),
+            (1, KeepIntermediateLevel.PIPELINE),
+            (2, KeepIntermediateLevel.PASS),
+            ("none", KeepIntermediateLevel.NONE),
+            ("NONE", KeepIntermediateLevel.NONE),
+            ("pipeline", KeepIntermediateLevel.PIPELINE),
+            ("PIPELINE", KeepIntermediateLevel.PIPELINE),
+            ("pass", KeepIntermediateLevel.PASS),
+            ("PASS", KeepIntermediateLevel.PASS),
+            (KeepIntermediateLevel.NONE, KeepIntermediateLevel.NONE),
+            (KeepIntermediateLevel.PIPELINE, KeepIntermediateLevel.PIPELINE),
+            (KeepIntermediateLevel.PASS, KeepIntermediateLevel.PASS),
+        ],
+    )
+    def test_keep_intermediate_levels_conversion(self, input_value, expected_level):
+        """Test that various inputs for keep_intermediate are correctly converted to Enum."""
+        options = CompileOptions(keep_intermediate=input_value)
+        assert options.keep_intermediate == expected_level
+
+    @pytest.mark.parametrize(
+        "invalid_input, error_type, error_match",
+        [
+            (3, ValueError, "Invalid int for keep_intermediate: 3. Valid integers are 0, 1, 2."),
+            (-1, ValueError, "Invalid int for keep_intermediate: -1. Valid integers are 0, 1, 2."),
+            (
+                "invalid_string",
+                ValueError,
+                "Invalid string for keep_intermediate: invalid_string. Valid strings are 'none',"
+                " 'pipeline', 'pass'.",
+            ),
+            (3.0, TypeError, "Invalid type for keep_intermediate: <class 'float'>."),
+            ([], TypeError, "Invalid type for keep_intermediate: <class 'list'>."),
+        ],
+    )
+    def test_keep_intermediate_invalid_inputs(self, invalid_input, error_type, error_match):
+        """Test that invalid inputs for keep_intermediate raise appropriate errors."""
+        with pytest.raises(error_type, match=error_match):
+            CompileOptions(keep_intermediate=invalid_input)
+
+    def test_options_to_cli_flags_keep_intermediate_none(self):
+        """Test _options_to_cli_flags with KeepIntermediateLevel.NONE."""
+        options = CompileOptions(keep_intermediate=KeepIntermediateLevel.NONE)
+        flags = _options_to_cli_flags(options)
+        assert "--keep-intermediate" not in flags
+        assert "--save-ir-after-each=pass" not in flags
+
+    def test_options_to_cli_flags_keep_intermediate_basic(self):
+        """Test _options_to_cli_flags with KeepIntermediateLevel.PIPELINE."""
+        options = CompileOptions(keep_intermediate=KeepIntermediateLevel.PIPELINE)
+        flags = _options_to_cli_flags(options)
+        assert "--keep-intermediate" in flags
+        assert "--save-ir-after-each=pass" not in flags
+        assert flags.count("--save-ir-after-each=pass") <= 1
+
+    def test_options_to_cli_flags_keep_intermediate_debug(self):
+        """Test _options_to_cli_flags with KeepIntermediateLevel.PASS."""
+        options = CompileOptions(keep_intermediate=KeepIntermediateLevel.PASS)
+        flags = _options_to_cli_flags(options)
+        assert "--keep-intermediate" in flags
+        assert "--save-ir-after-each=pass" in flags
+        assert flags.count("--save-ir-after-each=pass") == 1
+
 
 class TestCompilerWarnings:
     """Test compiler's warning messages."""