Small improvements in catalyst cli docs, option help, and intermediat…

…e file naming (#1405) **Context:** This PR fixes some small bugs or lack of documentation for catalyst-cli gathered from @dime10's feedbacks. - The `--help` flag seems to dump a lot of unrelated options which makes it difficult to navigate through catalyst options. - The possible stages for the --checkpoint-stage option are not mentioned in the documentation - When using `--checkpoint-stage`, `--save-ir-after-each=pipeline` no longer works. - The output from --save-ir-after-each=pass produces one output for each function when dealing with a function pass which results in a large number of outputs and one has to find the function of interest randomly. **Description of the Change:** - `--help` option now prints all the catalyst-cli specific options first before jumping into mlir-opt options - Added more details to documentation - Fixed the bug for `save-ir-after-each` and `save-ir-after-each` coexisting together. - Output from `save-ir-after-each` now appends the name of the function to the file name making it easier to identify the desired output. **Benefits:** easier experience for catalyst-cli user **Possible Drawbacks:** **Related GitHub Issues:** --------- Co-authored-by: Joey Carter <[email protected]> Co-authored-by: erick-xanadu <[email protected]> Co-authored-by: David Ittah <[email protected]>
PennyLaneAI · Jan 7, 2025 · 07ffe5b · 07ffe5b
1 parent b29046b
commit 07ffe5b
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 23 deletions.
diff --git a/doc/catalyst-cli/catalyst-cli.rst b/doc/catalyst-cli/catalyst-cli.rst
@@ -32,11 +32,11 @@ each stage individually. For example:
 
 .. note::
 
-    The Catalyst CLI tool is currently only available when Catalyst is built from source, and is not
-    included when installing Catalyst via pip or from wheels.
+    If Catalyst is built from source, the ``catalyst-cli`` executable will be located in 
+    the ``mlir/build/bin/`` directory relative to the root of your Catalyst source directory.
 
-    After building Catalyst, the ``catalyst-cli`` executable will be available in the
-    ``mlir/build/bin/`` directory.
+    If Catalyst is installed via pip or from wheels, the executable will be located 
+    in the ``catalyst/bin/`` directory relative to the environment’s installation directory.
 
 Usage
 -----
@@ -98,6 +98,23 @@ intermediate files are saved.
 Keep intermediate files after each pipeline in the compilation. By default, no intermediate files
 are saved. Using ``--keep-intermediate`` is equivalent to using ``--save-ir-after-each=pipeline``.
 
+``--{passname}``
+"""""""""""""""
+
+Enable a specific pass. For example, to enable the ``remove-chained-self-inverse`` pass, use
+``--remove-chained-self-inverse``.
+
+Catalyst's main ``mlir`` stage is split up into a sequence of pass pipelines that can also be run
+individually via this option. In that case, the name of the pipeline is substituted for the pass
+name. Currently, the following pipelines are available:
+``enforce-runtime-invariants-pipeline``,
+``hlo_lowering-pipeline``,
+``quantum-compilation-pipeline``,
+``bufferization-pipeline``,
+``llvm-dialect-lowring-pipeline``, and finally 
+``default-catalyst-pipeline`` which encompasses all the above as the default pipeline used by the
+Catalyst CLI tool if no pass option is specified.
+
 ``--catalyst-pipeline=<pipeline1(pass1[;pass2[;...]])[,pipeline2(...)]>``
 """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
 
@@ -113,7 +130,7 @@ applies the pass ``inline-nested-module``, we would specify this pipeline config
 
 .. code-block::
 
-    --catalyst-pipeline=pipe1(split-multiple-tapes;apply-transform-sequence),pipe2(inline-nested-module)
+    --catalyst-pipeline="pipe1(split-multiple-tapes;apply-transform-sequence),pipe2(inline-nested-module)"
 
 ``--workspace=<path>``
 """"""""""""""""""""""
@@ -138,7 +155,14 @@ Enable asynchronous QNodes.
 """""""""""""""""""""""""""""""""""
 
 Define a *checkpoint stage*, used to indicate that the compiler should start only after reaching the
-given pass.
+given stage. The stages that are currently available are:
+
+* MLIR: ``mlir`` (start with first MLIR stage), ``{pipeline}`` such as any of the built-in pipeline
+  names described under the ``--{passname}`` option, OR any custom pipeline names if the
+  ``--catalyst-pipeline={pipeline(...),...}`` option is used.
+* LLVM: ``llvm_ir`` (start with first LLVM stage), ``CoroOpt``, ``O2Opt``, ``Enzyme``.
+  Note that ``CoroOpt`` (Coroutine lowering), ``O2Opt`` (O2 optimization), and ``Enzyme``
+  (automatic differentiation) passes are only run conditionally as needed.
 
 ``--dump-catalyst-pipeline[=<true|false>]``
 """""""""""""""""""""""""""""""""""""""""""

diff --git a/doc/releases/changelog-0.10.0.md b/doc/releases/changelog-0.10.0.md
@@ -167,6 +167,10 @@
   array arguments of the function, in particular when non-64bit datatypes are used.
   [(#1338)](https://github.com/PennyLaneAI/catalyst/pull/1338)
 
+* Fixed a bug in catalyst cli where using `checkpoint-stage` would cause `save-ir-after-each`
+  to not work properly.
+  [(#1405)](https://github.com/PennyLaneAI/catalyst/pull/1405)
+
 <h3>Internal changes ⚙️</h3>
 
 * Catalyst no longer depends on or pins the `scipy` package. Instead, OpenBLAS is sourced directly
@@ -272,6 +276,10 @@
   for transformation passes.
   [(#1368)](https://github.com/PennyLaneAI/catalyst/pull/1368)
 
+* Added more details to catalyst-cli documentation specifiying available options for
+  checkpoint-stage and default pipelines
+  [(#1405)](https://github.com/PennyLaneAI/catalyst/pull/1405)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):

diff --git a/mlir/lib/Driver/CompilerDriver.cpp b/mlir/lib/Driver/CompilerDriver.cpp
@@ -490,7 +490,11 @@ LogicalResult preparePassManager(PassManager &pm, const CompilerOptions &options
             std::string tmp;
             llvm::raw_string_ostream s{tmp};
             s << *op;
-            dumpToFile(options, output.nextPipelineDumpFilename(pipelineName.str()), tmp);
+            std::string fileName = pipelineName.str();
+            if (auto funcOp = dyn_cast<mlir::func::FuncOp>(op)) {
+                fileName += "_" + funcOp.getName().str();
+            }
+            dumpToFile(options, output.nextPipelineDumpFilename(fileName), tmp);
         }
     };
 
@@ -551,7 +555,7 @@ LogicalResult runPipeline(PassManager &pm, const CompilerOptions &options, Compi
         llvm::errs() << "Failed to run pipeline: " << pipeline.getName() << "\n";
         return failure();
     }
-    if (options.keepIntermediate && options.checkpointStage.empty()) {
+    if (options.keepIntermediate && (options.checkpointStage.empty() || output.isCheckpointFound)) {
         std::string tmp;
         llvm::raw_string_ostream s{tmp};
         s << moduleOp;
@@ -564,7 +568,7 @@ LogicalResult runLowering(const CompilerOptions &options, MLIRContext *ctx, Modu
                           CompilerOutput &output, TimingScope &timing)
 
 {
-    if (options.keepIntermediate && options.checkpointStage.empty()) {
+    if (options.keepIntermediate && (options.checkpointStage.empty() || output.isCheckpointFound)) {
         std::string tmp;
         llvm::raw_string_ostream s{tmp};
         s << moduleOp;
@@ -861,26 +865,30 @@ int QuantumDriverMainFromCL(int argc, char **argv)
     // ---------
     // Any modifications made to the command-line interface should be documented in
     // doc/catalyst-cli/catalyst-cli.rst
-    cl::opt<std::string> WorkspaceDir("workspace", cl::desc("Workspace directory"), cl::init("."));
+    cl::OptionCategory CatalystCat("Catalyst-cli Options", "");
+    cl::opt<std::string> WorkspaceDir("workspace", cl::desc("Workspace directory"), cl::init("."),
+                                      cl::cat(CatalystCat));
     cl::opt<std::string> ModuleName("module-name", cl::desc("Module name"),
-                                    cl::init("catalyst_module"));
+                                    cl::init("catalyst_module"), cl::cat(CatalystCat));
 
     cl::opt<enum SaveTemps> SaveAfterEach(
         "save-ir-after-each", cl::desc("Keep intermediate files after each pass or pipeline"),
         cl::values(clEnumValN(SaveTemps::AfterPass, "pass", "Save IR after each pass")),
         cl::values(clEnumValN(SaveTemps::AfterPipeline, "pipeline", "Save IR after each pipeline")),
-        cl::init(SaveTemps::None));
+        cl::init(SaveTemps::None), cl::cat(CatalystCat));
     cl::opt<bool> KeepIntermediate(
         "keep-intermediate", cl::desc("Keep intermediate files"), cl::init(false),
-        cl::callback([&](const bool &) { SaveAfterEach.setValue(SaveTemps::AfterPipeline); }));
+        cl::callback([&](const bool &) { SaveAfterEach.setValue(SaveTemps::AfterPipeline); }),
+        cl::cat(CatalystCat));
     cl::opt<bool> AsyncQNodes("async-qnodes", cl::desc("Enable asynchronous QNodes"),
-                              cl::init(false));
-    cl::opt<bool> Verbose("verbose", cl::desc("Set verbose"), cl::init(false));
-    cl::list<std::string> CatalystPipeline("catalyst-pipeline",
-                                           cl::desc("Catalyst Compiler pass pipelines"),
-                                           cl::ZeroOrMore, cl::CommaSeparated);
+                              cl::init(false), cl::cat(CatalystCat));
+    cl::opt<bool> Verbose("verbose", cl::desc("Set verbose"), cl::init(false),
+                          cl::cat(CatalystCat));
+    cl::list<std::string> CatalystPipeline(
+        "catalyst-pipeline", cl::desc("Catalyst Compiler pass pipelines"), cl::ZeroOrMore,
+        cl::CommaSeparated, cl::cat(CatalystCat));
     cl::opt<std::string> CheckpointStage("checkpoint-stage", cl::desc("Checkpoint stage"),
-                                         cl::init(""));
+                                         cl::init(""), cl::cat(CatalystCat));
     cl::opt<enum Action> LoweringAction(
         "tool", cl::desc("Select the tool to isolate"),
         cl::values(clEnumValN(Action::OPT, "opt", "run quantum-opt on the MLIR input")),
@@ -889,9 +897,10 @@ int QuantumDriverMainFromCL(int argc, char **argv)
         cl::values(clEnumValN(Action::LLC, "llc", "run llc on the llvm IR input")),
         cl::values(clEnumValN(Action::All, "all",
                               "run quantum-opt, mlir-translate, and llc on the MLIR input")),
-        cl::init(Action::All));
-    cl::opt<bool> DumpPassPipeline(
-        "dump-catalyst-pipeline", cl::desc("Print the pipeline that will be run"), cl::init(false));
+        cl::init(Action::All), cl::cat(CatalystCat));
+    cl::opt<bool> DumpPassPipeline("dump-catalyst-pipeline",
+                                   cl::desc("Print the pipeline that will be run"), cl::init(false),
+                                   cl::cat(CatalystCat));
 
     // Create dialect registry
     DialectRegistry registry;
@@ -904,8 +913,13 @@ int QuantumDriverMainFromCL(int argc, char **argv)
 
     // Register and parse command line options.
     std::string inputFilename, outputFilename;
+    std::string helpStr = "Catalyst Command Line Interface options. \n"
+                          "Below, there is a complete list of options for the Catalyst CLI tool"
+                          "In the first section, you can find the options that are used to"
+                          "configure the Catalyst compiler. Next, you can find the options"
+                          "specific to the mlir-opt tool.\n";
     std::tie(inputFilename, outputFilename) =
-        registerAndParseCLIOptions(argc, argv, "quantum compiler", registry);
+        registerAndParseCLIOptions(argc, argv, helpStr, registry);
     llvm::InitLLVM y(argc, argv);
     MlirOptMainConfig config = MlirOptMainConfig::createFromCLOptions();