From 26aab5015d4362c1dc7024a4cdbeb3d2274b2644 Mon Sep 17 00:00:00 2001 From: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com> Date: Thu, 24 Oct 2024 17:35:13 -0400 Subject: [PATCH] New compiler driver interface (#1208) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **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 <79805239+paul0403@users.noreply.github.com> --- .../workflows/build-wheel-linux-x86_64.yaml | 2 +- .../workflows/build-wheel-macos-arm64.yaml | 2 +- .../workflows/build-wheel-macos-x86_64.yaml | 2 +- .../scripts/linux_arm64/rh8/build_catalyst.sh | 2 +- doc/releases/changelog-dev.md | 28 + frontend/catalyst/compiler.py | 43 +- frontend/catalyst/debug/compiler_functions.py | 4 +- frontend/test/pytest/test_compiler.py | 23 +- frontend/test/pytest/test_debug.py | 6 +- mlir/Makefile | 2 +- mlir/include/Driver/CompilerDriver.h | 48 +- mlir/include/Driver/Pipelines.h | 74 +++ mlir/lib/Driver/CMakeLists.txt | 3 + mlir/lib/Driver/CompilerDriver.cpp | 563 +++++++++++++----- mlir/lib/Driver/Pipelines.cpp | 204 +++++++ mlir/python/PyCompilerDriver.cpp | 38 +- mlir/test/CMakeLists.txt | 1 + mlir/test/cli/DumpAfterFailure.mlir | 23 + mlir/test/cli/DumpBeforeAfterPass.mlir | 38 ++ mlir/test/cli/DumpPipeline.mlir | 36 ++ mlir/test/cli/DumpTiming.mlir | 26 + mlir/test/cli/WrongInput.mlir | 26 + mlir/tools/CMakeLists.txt | 1 + mlir/tools/catalyst-cli/CMakeLists.txt | 46 ++ mlir/tools/catalyst-cli/catalyst-cli.cpp | 17 + 25 files changed, 1042 insertions(+), 216 deletions(-) create mode 100644 mlir/include/Driver/Pipelines.h create mode 100644 mlir/lib/Driver/Pipelines.cpp create mode 100644 mlir/test/cli/DumpAfterFailure.mlir create mode 100644 mlir/test/cli/DumpBeforeAfterPass.mlir create mode 100644 mlir/test/cli/DumpPipeline.mlir create mode 100644 mlir/test/cli/DumpTiming.mlir create mode 100644 mlir/test/cli/WrongInput.mlir create mode 100644 mlir/tools/catalyst-cli/CMakeLists.txt create mode 100644 mlir/tools/catalyst-cli/catalyst-cli.cpp diff --git a/.github/workflows/build-wheel-linux-x86_64.yaml b/.github/workflows/build-wheel-linux-x86_64.yaml index c505d891ae..389b5666c3 100644 --- a/.github/workflows/build-wheel-linux-x86_64.yaml +++ b/.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: | diff --git a/.github/workflows/build-wheel-macos-arm64.yaml b/.github/workflows/build-wheel-macos-arm64.yaml index 0f018e61d6..b5d5e674de 100644 --- a/.github/workflows/build-wheel-macos-arm64.yaml +++ b/.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: | diff --git a/.github/workflows/build-wheel-macos-x86_64.yaml b/.github/workflows/build-wheel-macos-x86_64.yaml index e3c161745e..5e73c582aa 100644 --- a/.github/workflows/build-wheel-macos-x86_64.yaml +++ b/.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: | diff --git a/.github/workflows/scripts/linux_arm64/rh8/build_catalyst.sh b/.github/workflows/scripts/linux_arm64/rh8/build_catalyst.sh index d8d0fe1a80..b16b8e0143 100644 --- a/.github/workflows/scripts/linux_arm64/rh8/build_catalyst.sh +++ b/.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 diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 8ddb5b30e2..737c4c3691 100644 --- a/doc/releases/changelog-dev.md +++ b/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)` +

Bug fixes

* Fix a bug in `catalyst.mitigate_with_zne` that would lead diff --git a/frontend/catalyst/compiler.py b/frontend/catalyst/compiler.py index 59a4634ada..5d3d117469 100644 --- a/frontend/catalyst/compiler.py +++ b/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,7 +631,7 @@ 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 @@ -641,12 +639,41 @@ def get_output_of(self, pipeline) -> Optional[str]: 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 diff --git a/frontend/catalyst/debug/compiler_functions.py b/frontend/catalyst/debug/compiler_functions.py index aa8cdb9564..45865ddc08 100644 --- a/frontend/catalyst/debug/compiler_functions.py +++ b/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 diff --git a/frontend/test/pytest/test_compiler.py b/frontend/test/pytest/test_compiler.py index a218aa3d0f..8fcc5deaa4 100644 --- a/frontend/test/pytest/test_compiler.py +++ b/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: diff --git a/frontend/test/pytest/test_debug.py b/frontend/test/pytest/test_debug.py index abf53ffd2f..14f98905ef 100644 --- a/frontend/test/pytest/test_debug.py +++ b/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): diff --git a/mlir/Makefile b/mlir/Makefile index 1d5a126ef6..2cb8de388a 100644 --- a/mlir/Makefile +++ b/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: diff --git a/mlir/include/Driver/CompilerDriver.h b/mlir/include/Driver/CompilerDriver.h index 19e51e27ba..28e3e60fcf 100644 --- a/mlir/include/Driver/CompilerDriver.h +++ b/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; - Name name; - PassList passes; -}; - /// Optional parameters, for which we provide reasonable default values. struct CompilerOptions { /// The textual IR (MLIR or LLVM IR) @@ -58,8 +59,8 @@ 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. @@ -67,10 +68,12 @@ struct CompilerOptions { /// Ordered list of named pipelines to execute, each pipeline is described by a list of MLIR /// passes it includes. std::vector 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 PipelineOutputs; - std::string objectFilename; + typedef std::unordered_map 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 ®istry); + +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 &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; } diff --git a/mlir/include/Driver/Pipelines.h b/mlir/include/Driver/Pipelines.h new file mode 100644 index 0000000000..c2ec1a8ea0 --- /dev/null +++ b/mlir/include/Driver/Pipelines.h @@ -0,0 +1,74 @@ +// Copyright 2024 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. + +#pragma once + +#include "mlir/Pass/Pass.h" + +namespace catalyst { +namespace driver { + +void createEnforceRuntimeInvariantsPipeline(mlir::OpPassManager &pm); +void createHloLoweringPipeline(mlir::OpPassManager &pm); +void createQuantumCompilationPipeline(mlir::OpPassManager &pm); +void createBufferizationPipeline(mlir::OpPassManager &pm); +void createLLVMDialectLoweringPipeline(mlir::OpPassManager &pm); +void createDefaultCatalystPipeline(mlir::OpPassManager &pm); + +void registerEnforceRuntimeInvariantsPipeline(); +void registerHloLoweringPipeline(); +void registerQuantumCompilationPipeline(); +void registerBufferizationPipeline(); +void registerLLVMDialectLoweringPipeline(); +void registerDefaultCatalystPipeline(); +void registerAllCatalystPipelines(); +class Pipeline { + public: + using PipelineFunc = void (*)(mlir::OpPassManager &); + + Pipeline() {} + + void addPass(std::string &pass) { passes.push_back(pass); } + + mlir::LogicalResult addPipeline(mlir::OpPassManager &pm) + { + if (registerFunc) { + registerFunc(pm); + return mlir::success(); + } + return mlir::failure(); + } + + const std::string &getName() const { return name; } + + const llvm::SmallVector &getPasses() const { return passes; } + + PipelineFunc getRegisterFunc() const { return registerFunc; } + + void setRegisterFunc(PipelineFunc func) { registerFunc = func; } + + void setName(const std::string &pipelineName) { name = pipelineName; } + + void setPasses(const llvm::SmallVector &newPasses) { passes = newPasses; } + + private: + std::string name; + llvm::SmallVector passes; + PipelineFunc registerFunc; +}; + +std::vector getDefaultPipeline(); + +} // namespace driver +} // namespace catalyst diff --git a/mlir/lib/Driver/CMakeLists.txt b/mlir/lib/Driver/CMakeLists.txt index 32f2bd76db..14e831f8cf 100644 --- a/mlir/lib/Driver/CMakeLists.txt +++ b/mlir/lib/Driver/CMakeLists.txt @@ -19,10 +19,12 @@ set(LLVM_LINK_COMPONENTS get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS) get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS) get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS) +get_property(translation_libs GLOBAL PROPERTY MLIR_TRANSLATION_LIBS) set(LIBS ${dialect_libs} ${conversion_libs} ${extension_libs} + ${translation_libs} MLIROptLib MLIRCatalyst MLIRCatalystTransforms @@ -42,6 +44,7 @@ set(LIBS add_mlir_library(CatalystCompilerDriver CompilerDriver.cpp CatalystLLVMTarget.cpp + Pipelines.cpp LINK_LIBS PRIVATE ${EXTERNAL_LIB} diff --git a/mlir/lib/Driver/CompilerDriver.cpp b/mlir/lib/Driver/CompilerDriver.cpp index 1f977662a3..8ccba8b742 100644 --- a/mlir/lib/Driver/CompilerDriver.cpp +++ b/mlir/lib/Driver/CompilerDriver.cpp @@ -32,6 +32,7 @@ #include "mlir/InitAllPasses.h" #include "mlir/Parser/Parser.h" #include "mlir/Pass/PassManager.h" +#include "mlir/Support/FileUtilities.h" #include "mlir/Target/LLVMIR/Export.h" #include "stablehlo/dialect/Register.h" #include "llvm/Analysis/CGSCCPassManager.h" @@ -41,8 +42,10 @@ #include "llvm/MC/TargetRegistry.h" #include "llvm/Passes/PassBuilder.h" #include "llvm/Support/FileSystem.h" +#include "llvm/Support/InitLLVM.h" #include "llvm/Support/SourceMgr.h" #include "llvm/Support/TargetSelect.h" +#include "llvm/Support/ToolOutputFile.h" #include "llvm/Target/TargetMachine.h" #include "llvm/Target/TargetOptions.h" #include "llvm/TargetParser/Host.h" @@ -56,6 +59,7 @@ #include "Catalyst/Transforms/Passes.h" #include "Driver/CatalystLLVMTarget.h" #include "Driver/CompilerDriver.h" +#include "Driver/Pipelines.h" #include "Driver/Support.h" #include "Gradient/IR/GradientDialect.h" #include "Gradient/IR/GradientInterfaces.h" @@ -71,6 +75,7 @@ using namespace mlir; using namespace catalyst; using namespace catalyst::driver; +namespace cl = llvm::cl; namespace catalyst::utils { @@ -188,7 +193,7 @@ class LinesCount { namespace { -std::string joinPasses(const Pipeline::PassList &passes) +std::string joinPasses(const llvm::SmallVector &passes) { std::string joined; llvm::raw_string_ostream stream{joined}; @@ -294,16 +299,27 @@ void registerAllCatalystDialects(DialectRegistry ®istry) } } // namespace +// Determines if the compilation stage should be executed if a checkpointStage is given +bool shouldRunStage(const CompilerOptions &options, CompilerOutput &output, + const std::string &stageName) +{ + if (options.checkpointStage.empty()) { + return true; + } + if (!output.isCheckpointFound) { + output.isCheckpointFound = (options.checkpointStage == stageName); + return false; + } + return true; +} + LogicalResult runCoroLLVMPasses(const CompilerOptions &options, std::shared_ptr llvmModule, CompilerOutput &output) { - if (options.checkpointStage != "" && !output.isCheckpointFound) { - output.isCheckpointFound = options.checkpointStage == "CoroOpt"; + if (!shouldRunStage(options, output, "CoroOpt")) { return success(); } - auto &outputs = output.pipelineOutputs; - // Create a pass to lower LLVM coroutines (similar to what happens in O0) llvm::ModulePassManager CoroPM; CoroPM.addPass(llvm::CoroEarlyPass()); @@ -330,10 +346,11 @@ LogicalResult runCoroLLVMPasses(const CompilerOptions &options, CoroPM.run(*llvmModule.get(), MAM); if (options.keepIntermediate) { - llvm::raw_string_ostream rawStringOstream{outputs["CoroOpt"]}; + std::string tmp; + llvm::raw_string_ostream rawStringOstream{tmp}; llvmModule->print(rawStringOstream, nullptr); auto outFile = output.nextPipelineDumpFilename("CoroOpt", ".ll"); - dumpToFile(options, outFile, outputs["CoroOpt"]); + dumpToFile(options, outFile, tmp); } return success(); @@ -345,12 +362,10 @@ LogicalResult runO2LLVMPasses(const CompilerOptions &options, // opt -O2 // As seen here: // https://llvm.org/docs/NewPassManager.html#just-tell-me-how-to-run-the-default-optimization-pipeline-with-the-new-pass-manager - if (options.checkpointStage != "" && !output.isCheckpointFound) { - output.isCheckpointFound = options.checkpointStage == "O2Opt"; + if (!shouldRunStage(options, output, "O2Opt")) { return success(); } - auto &outputs = output.pipelineOutputs; // Create the analysis managers. llvm::LoopAnalysisManager LAM; llvm::FunctionAnalysisManager FAM; @@ -383,10 +398,11 @@ LogicalResult runO2LLVMPasses(const CompilerOptions &options, MPM.run(*llvmModule.get(), MAM); if (options.keepIntermediate) { - llvm::raw_string_ostream rawStringOstream{outputs["O2Opt"]}; + std::string tmp; + llvm::raw_string_ostream rawStringOstream{tmp}; llvmModule->print(rawStringOstream, nullptr); auto outFile = output.nextPipelineDumpFilename("O2Opt", ".ll"); - dumpToFile(options, outFile, outputs["O2Opt"]); + dumpToFile(options, outFile, tmp); } return success(); @@ -395,12 +411,10 @@ LogicalResult runO2LLVMPasses(const CompilerOptions &options, LogicalResult runEnzymePasses(const CompilerOptions &options, std::shared_ptr llvmModule, CompilerOutput &output) { - if (options.checkpointStage != "" && !output.isCheckpointFound) { - output.isCheckpointFound = options.checkpointStage == "Enzyme"; + if (!shouldRunStage(options, output, "Enzyme")) { return success(); } - auto &outputs = output.pipelineOutputs; // Create the new pass manager builder. // Take a look at the PassBuilder constructor parameters for more // customization, e.g. specifying a TargetMachine or various debugging @@ -432,57 +446,31 @@ LogicalResult runEnzymePasses(const CompilerOptions &options, MPM.run(*llvmModule.get(), MAM); if (options.keepIntermediate) { - llvm::raw_string_ostream rawStringOstream{outputs["Enzyme"]}; + std::string tmp; + llvm::raw_string_ostream rawStringOstream{tmp}; llvmModule->print(rawStringOstream, nullptr); auto outFile = output.nextPipelineDumpFilename("Enzyme", ".ll"); - dumpToFile(options, outFile, outputs["Enzyme"]); + dumpToFile(options, outFile, tmp); } return success(); } -LogicalResult runLowering(const CompilerOptions &options, MLIRContext *ctx, ModuleOp moduleOp, - CompilerOutput &output) - +std::string readInputFile(const std::string &filename) { - auto &outputs = output.pipelineOutputs; - auto pm = PassManager::on(ctx, PassManager::Nesting::Implicit); - - // Maps a pass to zero or one pipelines ended by this pass - // Maps a pass to its owning pipeline - std::unordered_map pipelineTailMarkers; - std::unordered_map passPipelineNames; - - // Fill all the pipe-to-pipeline mappings - for (const auto &pipeline : options.pipelinesCfg) { - if (options.checkpointStage != "" && !output.isCheckpointFound) { - output.isCheckpointFound = options.checkpointStage == pipeline.name; - continue; - } - size_t existingPasses = pm.size(); - if (failed(parsePassPipeline(joinPasses(pipeline.passes), pm, options.diagnosticStream))) { - return failure(); - } - if (existingPasses != pm.size()) { - const Pass *pass = nullptr; - for (size_t pn = existingPasses; pn < pm.size(); pn++) { - pass = &(*(pm.begin() + pn)); - passPipelineNames[pass] = pipeline.name; - } - assert(pass != nullptr); - pipelineTailMarkers[pass] = pipeline.name; - } - } - - if (options.keepIntermediate && options.checkpointStage == "") { - llvm::raw_string_ostream s{outputs["mlir"]}; - s << moduleOp; - dumpToFile(options, output.nextPipelineDumpFilename(options.moduleName.str(), ".mlir"), - outputs["mlir"]); + std::ifstream file(filename); + if (!file.is_open()) { + return ""; } + std::stringstream buffer; + buffer << file.rdbuf(); + return buffer.str(); +} - catalyst::utils::Timer timer{}; - +LogicalResult preparePassManager(PassManager &pm, const CompilerOptions &options, + CompilerOutput &output, catalyst::utils::Timer &timer, + TimingScope &timing) +{ auto beforePassCallback = [&](Pass *pass, Operation *op) { if (!timer.is_active()) { timer.start(); @@ -492,62 +480,145 @@ LogicalResult runLowering(const CompilerOptions &options, MLIRContext *ctx, Modu // For each pipeline-terminating pass, print the IR into the corresponding dump file and // into a diagnostic output buffer. Note that one pass can terminate multiple pipelines. auto afterPassCallback = [&](Pass *pass, Operation *op) { - auto res = pipelineTailMarkers.find(pass); - if (res != pipelineTailMarkers.end()) { - timer.dump(res->second, /*add_endl */ false); - catalyst::utils::LinesCount::Operation(op); - } - - if (options.keepIntermediate && res != pipelineTailMarkers.end()) { - auto pipelineName = res->second; - llvm::raw_string_ostream s{outputs[pipelineName]}; + auto pipelineName = pass->getName(); + timer.dump(pipelineName.str(), /*add_endl */ false); + catalyst::utils::LinesCount::Operation(op); + if (options.keepIntermediate >= SaveTemps::AfterPass) { + std::string tmp; + llvm::raw_string_ostream s{tmp}; s << *op; - dumpToFile(options, output.nextPipelineDumpFilename(pipelineName), - outputs[pipelineName]); + dumpToFile(options, output.nextPipelineDumpFilename(pipelineName.str()), tmp); } }; // For each failed pass, print the owner pipeline name into a diagnostic stream. auto afterPassFailedCallback = [&](Pass *pass, Operation *op) { - auto res = passPipelineNames.find(pass); - assert(res != passPipelineNames.end() && "Unexpected pass"); - options.diagnosticStream << "While processing '" << pass->getName() << "' pass " - << "of the '" << res->second << "' pipeline\n"; - llvm::raw_string_ostream s{outputs[res->second]}; + options.diagnosticStream << "While processing '" << pass->getName().str() << "' pass "; + std::string tmp; + llvm::raw_string_ostream s{tmp}; s << *op; if (options.keepIntermediate) { - dumpToFile(options, output.nextPipelineDumpFilename(res->second + "_FAILED"), - outputs[res->second]); + dumpToFile(options, output.nextPipelineDumpFilename(pass->getName().str() + "_FAILED"), + tmp); } }; - // Output pipeline names on failures + MlirOptMainConfig config = MlirOptMainConfig::createFromCLOptions(); + pm.enableVerifier(config.shouldVerifyPasses()); + if (failed(applyPassManagerCLOptions(pm))) + return failure(); + if (failed(config.setupPassPipeline(pm))) + return failure(); + pm.enableTiming(timing); pm.addInstrumentation(std::unique_ptr(new CatalystPassInstrumentation( beforePassCallback, afterPassCallback, afterPassFailedCallback))); + return success(); +} - // Run the lowering pipelines - if (failed(pm.run(moduleOp))) { +LogicalResult configurePipeline(PassManager &pm, const CompilerOptions &options, Pipeline &pipeline, + bool clHasManualPipeline) +{ + pm.clear(); + if (!clHasManualPipeline && failed(pipeline.addPipeline(pm))) { + llvm::errs() << "Pipeline creation function not found: " << pipeline.getName() << "\n"; return failure(); } - + if (clHasManualPipeline && + failed(parsePassPipeline(joinPasses(pipeline.getPasses()), pm, options.diagnosticStream))) { + return failure(); + } + if (options.dumpPassPipeline) { + pm.dump(); + llvm::errs() << "\n"; + } return success(); } -LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput &output) +LogicalResult runLowering(const CompilerOptions &options, MLIRContext *ctx, ModuleOp moduleOp, + CompilerOutput &output, TimingScope &timing) + { - using timer = catalyst::utils::Timer; + if (options.keepIntermediate && options.checkpointStage.empty()) { + std::string tmp; + llvm::raw_string_ostream s{tmp}; + s << moduleOp; + dumpToFile(options, output.nextPipelineDumpFilename(options.moduleName.str(), ".mlir"), + tmp); + } - DialectRegistry registry; - static bool initialized = false; - if (!initialized) { - registerAllPasses(); + catalyst::utils::Timer timer{}; + + auto pm = PassManager::on(ctx, PassManager::Nesting::Implicit); + if (failed(preparePassManager(pm, options, output, timer, timing))) { + llvm::errs() << "Failed to setup pass manager\n"; + return failure(); } - initialized |= true; - registerAllCatalystPasses(); - mhlo::registerAllMhloPasses(); - registerAllCatalystDialects(registry); - registerLLVMTranslations(registry); + bool clHasIndividulaPass = pm.size() > 0; + bool clHasManualPipeline = !options.pipelinesCfg.empty(); + if (clHasIndividulaPass && clHasManualPipeline) { + llvm::errs() << "--catalyst-pipline option can't be used with individual pass options " + "or -pass-pipeline.\n"; + return failure(); + } + + // If individual passes are configured, run them + if (clHasIndividulaPass) { + if (options.dumpPassPipeline) { + pm.dump(); + llvm::errs() << "\n"; + } + return pm.run(moduleOp); + } + + // If pipelines are not cofigured explicitly, use the catalyst default pipeline + std::vector UserPipeline = + clHasManualPipeline ? options.pipelinesCfg : getDefaultPipeline(); + for (auto &pipeline : UserPipeline) { + if (!shouldRunStage(options, output, pipeline.getName()) || + pipeline.getPasses().size() == 0) { + continue; + } + if (failed(configurePipeline(pm, options, pipeline, clHasManualPipeline))) { + llvm::errs() << "Failed to run pipeline: " << pipeline.getName() << "\n"; + return failure(); + } + + if (failed(pm.run(moduleOp))) + return failure(); + + if (options.keepIntermediate && options.checkpointStage.empty()) { + std::string tmp; + llvm::raw_string_ostream s{tmp}; + s << moduleOp; + dumpToFile(options, output.nextPipelineDumpFilename(pipeline.getName(), ".mlir"), tmp); + } + } + return success(); +} + +LogicalResult verifyInputType(const CompilerOptions &options, InputType inType) +{ + if (inType == InputType::OTHER) { + CO_MSG(options, Verbosity::Urgent, "Wrong or unsupported input\n"); + return failure(); + } + if (options.loweringAction == Action::LLC && 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) { + CO_MSG(options, Verbosity::Urgent, "Expected MLIR input but received LLVM IR input.\n"); + return failure(); + } + return success(); +} + +LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput &output, + DialectRegistry ®istry) +{ + using timer = catalyst::utils::Timer; + MLIRContext ctx(registry); ctx.printOpOnDiagnostic(true); ctx.printStackTraceOnDiagnostic(options.verbosity >= Verbosity::Debug); @@ -567,62 +638,90 @@ LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput & sourceMgr->AddNewSourceBuffer(std::move(moduleBuffer), SMLoc()); SourceMgrDiagnosticHandler sourceMgrHandler(*sourceMgr, &ctx, options.diagnosticStream); - OwningOpRef op = + DefaultTimingManager tm; + applyDefaultTimingManagerCLOptions(tm); + TimingScope timing = tm.getRootScope(); + + TimingScope parserTiming = timing.nest("Parser"); + OwningOpRef mlirModule = timer::timer(parseMLIRSource, "parseMLIRSource", /* add_endl */ false, &ctx, *sourceMgr); - catalyst::utils::LinesCount::ModuleOp(*op); - output.isCheckpointFound = options.checkpointStage == "mlir"; + + enum InputType inType = InputType::OTHER; + if (mlirModule) { + inType = InputType::MLIR; + catalyst::utils::LinesCount::ModuleOp(*mlirModule); + output.isCheckpointFound = options.checkpointStage == "mlir"; + } + else { + llvm::SMDiagnostic err; + llvmModule = timer::timer(parseLLVMSource, "parseLLVMSource", false, llvmContext, + options.source, options.moduleName, err); + + if (!llvmModule) { + err.print(options.moduleName.data(), options.diagnosticStream); + CO_MSG(options, Verbosity::Urgent, "Failed to parse module as LLVM or MLIR source\n"); + return failure(); + } + inType = InputType::LLVMIR; + output.isCheckpointFound = options.checkpointStage == "llvm_ir"; + catalyst::utils::LinesCount::Module(*llvmModule); + } + if (failed(verifyInputType(options, inType))) { + return failure(); + } + parserTiming.stop(); // Enzyme always happens after O2Opt. If the checkpoint is O2Opt, enzymeRun must be set to // true so that the enzyme pass can be executed. bool enzymeRun = options.checkpointStage == "O2Opt"; - if (op) { - enzymeRun = containsGradients(*op); - if (failed(runLowering(options, &ctx, *op, output))) { + + bool runAll = (options.loweringAction == Action::All); + bool runOpt = (options.loweringAction == Action::OPT) || runAll; + bool runTranslate = (options.loweringAction == Action::Translate) || runAll; + bool runLLC = (options.loweringAction == Action::LLC) || runAll; + + if (runOpt && (inType == InputType::MLIR)) { + TimingScope optTiming = timing.nest("Optimization"); + // TODO: The enzymeRun flag will not travel correctly in the case where different + // stages of compilation are executed independantly via the catalyst-cli executable. + // Ideally, It should be added to the IR via an attribute. + enzymeRun = containsGradients(*mlirModule); + if (failed(runLowering(options, &ctx, *mlirModule, output, optTiming))) { CO_MSG(options, Verbosity::Urgent, "Failed to lower MLIR module\n"); return failure(); } - output.outIR.clear(); - outIRStream << *op; - - if (options.lowerToLLVM) { - llvmModule = timer::timer(translateModuleToLLVMIR, "translateModuleToLLVMIR", - /* add_endl */ false, *op, llvmContext, "LLVMDialectModule"); - if (!llvmModule) { - CO_MSG(options, Verbosity::Urgent, "Failed to translate LLVM module\n"); - return failure(); - } - - catalyst::utils::LinesCount::Module(*llvmModule); - - if (options.keepIntermediate) { - auto &outputs = output.pipelineOutputs; - llvm::raw_string_ostream rawStringOstream{outputs["llvm_ir"]}; - llvmModule->print(rawStringOstream, nullptr); - auto outFile = output.nextPipelineDumpFilename("llvm_ir", ".ll"); - dumpToFile(options, outFile, outputs["llvm_ir"]); - } - } + outIRStream << *mlirModule; + optTiming.stop(); } - else { - CO_MSG(options, Verbosity::Urgent, - "Failed to parse module as MLIR source, retrying parsing as LLVM source\n"); - llvm::SMDiagnostic err; - llvmModule = timer::timer(parseLLVMSource, "parseLLVMSource", /* add_endl */ false, - llvmContext, options.source, options.moduleName, err); - output.isCheckpointFound = options.checkpointStage == "llvm_ir"; + if (runTranslate && (inType == InputType::MLIR)) { + TimingScope translateTiming = timing.nest("Translate"); + llvmModule = + timer::timer(translateModuleToLLVMIR, "translateModuleToLLVMIR", + /* add_endl */ false, *mlirModule, llvmContext, "LLVMDialectModule"); if (!llvmModule) { - // If both MLIR and LLVM failed to parse, exit. - err.print(options.moduleName.data(), options.diagnosticStream); - CO_MSG(options, Verbosity::Urgent, "Failed to parse module as LLVM source\n"); + CO_MSG(options, Verbosity::Urgent, "Failed to translate LLVM module\n"); return failure(); } + inType = InputType::LLVMIR; catalyst::utils::LinesCount::Module(*llvmModule); + + if (options.keepIntermediate) { + std::string tmp; + llvm::raw_string_ostream rawStringOstream{tmp}; + llvmModule->print(rawStringOstream, nullptr); + auto outFile = output.nextPipelineDumpFilename("llvm_ir", ".ll"); + dumpToFile(options, outFile, tmp); + } + output.outIR.clear(); + outIRStream << *llvmModule; + translateTiming.stop(); } - if (llvmModule) { + if (runLLC && (inType == InputType::LLVMIR)) { + TimingScope llcTiming = timing.nest("llc"); // Set data layout before LLVM passes or the default one is used. std::string targetTriple = llvm::sys::getDefaultTargetTriple(); @@ -645,38 +744,228 @@ LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput & catalyst::utils::LinesCount::Module(*llvmModule.get()); - if (options.asyncQnodes) { - if (failed(timer::timer(runCoroLLVMPasses, "runCoroLLVMPasses", /* add_endl */ false, - options, llvmModule, output))) { - return failure(); - } - catalyst::utils::LinesCount::Module(*llvmModule.get()); + TimingScope coroLLVMPassesTiming = llcTiming.nest("LLVM coroutine passes"); + if (options.asyncQnodes && + failed(timer::timer(runCoroLLVMPasses, "runCoroLLVMPasses", /* add_endl */ false, + options, llvmModule, output))) { + return failure(); } + coroLLVMPassesTiming.stop(); + if (enzymeRun) { + TimingScope o2PassesTiming = llcTiming.nest("LLVM O2 passes"); if (failed(timer::timer(runO2LLVMPasses, "runO2LLVMPasses", /* add_endl */ false, options, llvmModule, output))) { return failure(); } + o2PassesTiming.stop(); catalyst::utils::LinesCount::Module(*llvmModule.get()); + TimingScope enzymePassesTiming = llcTiming.nest("Enzyme passes"); if (failed(timer::timer(runEnzymePasses, "runEnzymePasses", /* add_endl */ false, options, llvmModule, output))) { return failure(); } + enzymePassesTiming.stop(); catalyst::utils::LinesCount::Module(*llvmModule.get()); } + TimingScope outputTiming = llcTiming.nest("compileObject"); output.outIR.clear(); outIRStream << *llvmModule; - // Attempt to infer the name and return type of the module from LLVM IR. This information is - // required when executing a module given as textual IR. - auto outfile = options.getObjectFile(); + bool outputFileSpecified = !output.outputFilename.empty() && output.outputFilename != "-"; + auto outfile = outputFileSpecified ? output.outputFilename : options.getObjectFile(); if (failed(timer::timer(compileObjectFile, "compileObjFile", /* add_endl */ true, options, std::move(llvmModule), targetMachine, outfile))) { return failure(); } - output.objectFilename = outfile; + output.outputFilename = outfile; + outputTiming.stop(); + llcTiming.stop(); } + return success(); } + +std::vector parsePipelines(const cl::list &catalystPipeline) +{ + std::vector allPipelines; + for (const auto &pipelineStr : catalystPipeline) { + llvm::StringRef pipelineRef = llvm::StringRef(pipelineStr).trim(); + + size_t openParenPos = pipelineRef.find('('); + size_t closeParenPos = pipelineRef.find(')', openParenPos); + + if (openParenPos == llvm::StringRef::npos || closeParenPos == llvm::StringRef::npos) { + llvm::errs() << "Error: Invalid pipeline format: " << pipelineStr << "\n"; + continue; + } + + // Extract pipeline name + llvm::StringRef pipelineName = pipelineRef.slice(0, openParenPos).trim(); + llvm::StringRef passesStr = pipelineRef.slice(openParenPos + 1, closeParenPos).trim(); + llvm::SmallVector passList; + passesStr.split(passList, ';', /*MaxSplit=*/-1, /*KeepEmpty=*/false); + + llvm::SmallVector passes; + for (auto &pass : passList) { + passes.push_back(pass.trim().str()); + } + + Pipeline pipeline; + pipeline.setName(pipelineName.str()); + pipeline.setPasses(passes); + allPipelines.push_back(std::move(pipeline)); + } + return allPipelines; +} + +int QuantumDriverMainFromCL(int argc, char **argv) +{ + // Command-line options + cl::opt WorkspaceDir("workspace", cl::desc("Workspace directory"), cl::init(".")); + cl::opt ModuleName("module-name", cl::desc("Module name"), + cl::init("catalyst_module")); + + cl::opt 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::opt KeepIntermediate( + "keep-intermediate", cl::desc("Keep intermediate files"), cl::init(false), + cl::callback([&](const bool &) { SaveAfterEach.setValue(SaveTemps::AfterPipeline); })); + cl::opt AsyncQNodes("async-qnodes", cl::desc("Enable asynchronous QNodes"), + cl::init(false)); + cl::opt Verbose("verbose", cl::desc("Set verbose"), cl::init(false)); + cl::list CatalystPipeline("catalyst-pipeline", + cl::desc("Catalyst Compiler pass pipelines"), + cl::ZeroOrMore, cl::CommaSeparated); + cl::opt CheckpointStage("checkpoint-stage", cl::desc("Checkpoint stage"), + cl::init("")); + cl::opt LoweringAction( + "tool", cl::desc("Select the tool to isolate"), + cl::values(clEnumValN(Action::OPT, "opt", "run quantum-opt on the MLIR input")), + cl::values(clEnumValN(Action::Translate, "translate", + "run mlir-translate on the MLIR LLVM dialect")), + 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 DumpPassPipeline( + "dump-catalyst-pipeline", cl::desc("Print the pipeline that will be run"), cl::init(false)); + + // Create dialect registery + DialectRegistry registry; + registerAllPasses(); + registerAllCatalystPasses(); + registerAllCatalystPipelines(); + mhlo::registerAllMhloPasses(); + registerAllCatalystDialects(registry); + registerLLVMTranslations(registry); + + // Register and parse command line options. + std::string inputFilename, outputFilename; + std::tie(inputFilename, outputFilename) = + registerAndParseCLIOptions(argc, argv, "qunatum compiler", registry); + llvm::InitLLVM y(argc, argv); + MlirOptMainConfig config = MlirOptMainConfig::createFromCLOptions(); + + // Read the input IR file + std::string source = readInputFile(inputFilename); + if (source.empty()) { + llvm::errs() << "Error: Unable to read input file: " << inputFilename << "\n"; + return 1; + } + + std::unique_ptr output(new CompilerOutput()); + assert(output); + output->outputFilename = outputFilename; + llvm::raw_string_ostream errStream{output->diagnosticMessages}; + + CompilerOptions options{.source = source, + .workspace = WorkspaceDir, + .moduleName = ModuleName, + .diagnosticStream = errStream, + .keepIntermediate = SaveAfterEach, + .asyncQnodes = AsyncQNodes, + .verbosity = Verbose ? Verbosity::All : Verbosity::Urgent, + .pipelinesCfg = parsePipelines(CatalystPipeline), + .checkpointStage = CheckpointStage, + .loweringAction = LoweringAction, + .dumpPassPipeline = DumpPassPipeline}; + + mlir::LogicalResult result = QuantumDriverMain(options, *output, registry); + + errStream.flush(); + + if (mlir::failed(result)) { + llvm::errs() << "Compilation failed:\n" << output->diagnosticMessages << "\n"; + return 1; + } + + // If not creating object file, output the IR to the specified file. + if (LoweringAction < Action::LLC) { + std::string errorMessage; + auto outfile = openOutputFile(outputFilename, &errorMessage); + if (!outfile) { + llvm::errs() << errorMessage << "\n"; + return 1; + } + outfile->os() << output->outIR; + outfile->keep(); + } + + return 0; +} + +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 &passPipelines, + const std::string &checkpointStage, + catalyst::driver::CompilerOutput &output) +{ + llvm::raw_string_ostream errStream{output.diagnosticMessages}; + + CompilerOptions options{.source = source, + .workspace = workspace, + .moduleName = moduleName, + .diagnosticStream = errStream, + .keepIntermediate = + keepIntermediate ? SaveTemps::AfterPipeline : SaveTemps::None, + .asyncQnodes = asyncQNodes, + .verbosity = verbose ? Verbosity::All : Verbosity::Urgent, + .pipelinesCfg = passPipelines, + .checkpointStage = checkpointStage, + .loweringAction = lowerToLLVM ? Action::All : Action::OPT, + .dumpPassPipeline = false}; + + DialectRegistry registry; + static bool initialized = false; + if (!initialized) { + registerAllPasses(); + registerAllCatalystPasses(); + registerAllCatalystPipelines(); + } + initialized |= true; + + mhlo::registerAllMhloPasses(); + registerAllCatalystDialects(registry); + registerAsmPrinterCLOptions(); + registerMLIRContextCLOptions(); + registerPassManagerCLOptions(); + registerDefaultTimingManagerCLOptions(); + registerLLVMTranslations(registry); + + mlir::LogicalResult result = QuantumDriverMain(options, output, registry); + + errStream.flush(); + + if (mlir::failed(result)) { + llvm::errs() << "Compilation failed:\n" << output.diagnosticMessages << "\n"; + return 1; + } + return 0; +} diff --git a/mlir/lib/Driver/Pipelines.cpp b/mlir/lib/Driver/Pipelines.cpp new file mode 100644 index 0000000000..7438a51094 --- /dev/null +++ b/mlir/lib/Driver/Pipelines.cpp @@ -0,0 +1,204 @@ +// Copyright 2024 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. + +#include "Driver/Pipelines.h" +#include "Catalyst/Transforms/Passes.h" +#include "Gradient/Transforms/Passes.h" +#include "Mitigation/Transforms/Passes.h" +#include "Quantum/Transforms/Passes.h" +#include "mhlo/transforms/passes.h" +#include "mlir/InitAllDialects.h" +#include "mlir/InitAllPasses.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Transforms/Passes.h" + +using namespace mlir; +namespace catalyst { +namespace driver { + +void createEnforceRuntimeInvariantsPipeline(OpPassManager &pm) +{ + pm.addPass(catalyst::createSplitMultipleTapesPass()); + pm.addPass(catalyst::createApplyTransformSequencePass()); + pm.addPass(catalyst::createInlineNestedModulePass()); +} +void createHloLoweringPipeline(OpPassManager &pm) +{ + pm.addPass(mlir::createCanonicalizerPass()); + pm.addNestedPass(mhlo::createChloLegalizeToHloPass()); + pm.addPass(mlir::mhlo::createStablehloLegalizeToHloPass()); + pm.addNestedPass(mhlo::createLegalizeControlFlowPass()); + pm.addNestedPass(mhlo::createLegalizeHloToLinalgPass()); + pm.addNestedPass(mhlo::createLegalizeToStdPass()); + pm.addNestedPass(mhlo::createLegalizeSortPass()); + pm.addPass(mlir::mhlo::createConvertToSignlessPass()); + pm.addPass(mlir::createCanonicalizerPass()); + pm.addPass(catalyst::createScatterLoweringPass()); + pm.addPass(catalyst::createHloCustomCallLoweringPass()); + pm.addPass(mlir::createCSEPass()); + mlir::LinalgDetensorizePassOptions options; + options.aggressiveMode = true; + pm.addNestedPass(mlir::createLinalgDetensorizePass(options)); + pm.addPass(catalyst::createDetensorizeSCFPass()); + pm.addPass(mlir::createCanonicalizerPass()); +} +void createQuantumCompilationPipeline(OpPassManager &pm) +{ + pm.addPass(catalyst::createAnnotateFunctionPass()); + pm.addPass(catalyst::createMitigationLoweringPass()); + pm.addPass(catalyst::createGradientLoweringPass()); + pm.addPass(catalyst::createAdjointLoweringPass()); + pm.addPass(catalyst::createDisableAssertionPass()); +} +void createBufferizationPipeline(OpPassManager &pm) +{ + mlir::bufferization::OneShotBufferizationOptions options; + options.opFilter.allowDialect(); + pm.addPass(mlir::bufferization::createOneShotBufferizePass(options)); + pm.addPass(mlir::createInlinerPass()); + pm.addPass(catalyst::createGradientPreprocessingPass()); + pm.addPass(catalyst::createGradientBufferizationPass()); + pm.addPass(mlir::createSCFBufferizePass()); + pm.addPass(mlir::createConvertTensorToLinalgPass()); + pm.addPass(mlir::createConvertElementwiseToLinalgPass()); + pm.addPass(mlir::arith::createArithBufferizePass()); + pm.addPass(mlir::bufferization::createEmptyTensorToAllocTensorPass()); + pm.addNestedPass(mlir::bufferization::createBufferizationBufferizePass()); + pm.addNestedPass(mlir::tensor::createTensorBufferizePass()); + pm.addPass(catalyst::createCatalystBufferizationPass()); + pm.addNestedPass(mlir::createLinalgBufferizePass()); + pm.addNestedPass(mlir::tensor::createTensorBufferizePass()); + pm.addPass(catalyst::createQuantumBufferizationPass()); + pm.addPass(mlir::func::createFuncBufferizePass()); + pm.addNestedPass(mlir::bufferization::createFinalizingBufferizePass()); + pm.addPass(mlir::createCanonicalizerPass()); + pm.addPass(catalyst::createGradientPostprocessingPass()); + pm.addNestedPass(mlir::bufferization::createBufferHoistingPass()); + pm.addNestedPass(mlir::bufferization::createBufferLoopHoistingPass()); + pm.addNestedPass(mlir::bufferization::createBufferDeallocationPass()); + pm.addPass(catalyst::createArrayListToMemRefPass()); + pm.addPass(mlir::createBufferizationToMemRefPass()); + pm.addPass(mlir::createCanonicalizerPass()); + pm.addPass(catalyst::createCopyGlobalMemRefPass()); +} +void createLLVMDialectLoweringPipeline(OpPassManager &pm) +{ + pm.addPass(mlir::memref::createExpandReallocPass()); + pm.addPass(catalyst::createGradientConversionPass()); + pm.addPass(catalyst::createMemrefCopyToLinalgCopyPass()); + pm.addNestedPass(mlir::createConvertLinalgToLoopsPass()); + pm.addPass(mlir::createConvertSCFToCFPass()); + pm.addPass(mlir::memref::createExpandStridedMetadataPass()); + pm.addPass(mlir::createLowerAffinePass()); + pm.addPass(mlir::arith::createArithExpandOpsPass()); + pm.addPass(mlir::createConvertComplexToStandardPass()); + pm.addPass(mlir::createConvertComplexToLLVMPass()); + pm.addPass(mlir::createConvertMathToLLVMPass()); + pm.addPass(mlir::createConvertMathToLibmPass()); + pm.addPass(mlir::createArithToLLVMConversionPass()); + pm.addPass(catalyst::createMemrefToLLVMWithTBAAPass()); + FinalizeMemRefToLLVMConversionPassOptions options; + options.useGenericFunctions = true; + pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass(options)); + pm.addPass(mlir::createConvertIndexToLLVMPass()); + pm.addPass(catalyst::createCatalystConversionPass()); + pm.addPass(catalyst::createQuantumConversionPass()); + pm.addPass(catalyst::createAddExceptionHandlingPass()); + pm.addPass(catalyst::createEmitCatalystPyInterfacePass()); + pm.addPass(mlir::createCanonicalizerPass()); + pm.addPass(mlir::createReconcileUnrealizedCastsPass()); + pm.addPass(catalyst::createGEPInboundsPass()); + pm.addPass(catalyst::createRegisterInactiveCallbackPass()); +} + +void createDefaultCatalystPipeline(OpPassManager &pm) +{ + createEnforceRuntimeInvariantsPipeline(pm); + createHloLoweringPipeline(pm); + createQuantumCompilationPipeline(pm); + createBufferizationPipeline(pm); + createLLVMDialectLoweringPipeline(pm); +} + +void registerEnforceRuntimeInvariantsPipeline() +{ + PassPipelineRegistration<>("enforce-runtime-invariants-pipeline", + "Register enforce runtime invariants pipeline as a pass.", + createEnforceRuntimeInvariantsPipeline); +} +void registerHloLoweringPipeline() +{ + PassPipelineRegistration<>("hlo_lowering-pipeline", "Register HLO lowering pipeline as a pass.", + createHloLoweringPipeline); +} +void registerQuantumCompilationPipeline() +{ + PassPipelineRegistration<>("quantum-compilation-pipeline", + "Register quantum compilation pipeline as a pass.", + createQuantumCompilationPipeline); +} +void registerBufferizationPipeline() +{ + PassPipelineRegistration<>("bufferization-pipeline", + "Register bufferization pipeline as a pass.", + createBufferizationPipeline); +} +void registerLLVMDialectLoweringPipeline() +{ + PassPipelineRegistration<>("llvm-dialect-lowring-pipeline", + "Register LLVM dialect lowring pipeline as a pass.", + createLLVMDialectLoweringPipeline); +} + +void registerDefaultCatalystPipeline() +{ + PassPipelineRegistration<>("default-catalyst-pipeline", + "Register full default catalyst pipeline as a pass.", + createDefaultCatalystPipeline); +} + +void registerAllCatalystPipelines() +{ + registerEnforceRuntimeInvariantsPipeline(); + registerHloLoweringPipeline(); + registerQuantumCompilationPipeline(); + registerBufferizationPipeline(); + registerLLVMDialectLoweringPipeline(); + registerDefaultCatalystPipeline(); +} + +std::vector getDefaultPipeline() +{ + using PipelineFunc = void (*)(mlir::OpPassManager &); + std::vector pipelineFuncs = { + &createEnforceRuntimeInvariantsPipeline, &createHloLoweringPipeline, + &createQuantumCompilationPipeline, &createBufferizationPipeline, + &createLLVMDialectLoweringPipeline}; + + llvm::SmallVector defaultPipelineNames = { + "enforce-runtime-invariants-pipeline", "hlo-lowering-pipeline", + "quantum-compilation-pipeline", "bufferization-pipeline", "llvm-dialect-lowering-pipeline"}; + + std::vector defaultPipelines; + defaultPipelines.reserve(defaultPipelineNames.size()); + for (size_t i = 0; i < defaultPipelineNames.size(); ++i) { + defaultPipelines[i].setRegisterFunc(pipelineFuncs[i]); + defaultPipelines[i].setName(defaultPipelineNames[i]); + defaultPipelines[i].addPass(defaultPipelineNames[i]); + } + return defaultPipelines; +} + +} // namespace driver +} // namespace catalyst diff --git a/mlir/python/PyCompilerDriver.cpp b/mlir/python/PyCompilerDriver.cpp index cd91f12e23..7e01850008 100644 --- a/mlir/python/PyCompilerDriver.cpp +++ b/mlir/python/PyCompilerDriver.cpp @@ -38,10 +38,13 @@ std::vector parseCompilerSpec(const py::list &pipelines) auto py_passes = i++; assert(i == t.end()); std::string name = py_name->attr("__str__")().cast(); - Pipeline::PassList passes; + llvm::SmallVector passes; std::transform(py_passes->begin(), py_passes->end(), std::back_inserter(passes), [](py::handle p) { return p.attr("__str__")().cast(); }); - out.push_back(Pipeline({name, passes})); + Pipeline pipeline; + pipeline.setName(name); + pipeline.setPasses(passes); + out.push_back(pipeline); } return out; } @@ -53,19 +56,11 @@ PYBIND11_MODULE(compiler_driver, m) //===--------------------------------------------------------------------===// py::class_ compout_class(m, "CompilerOutput"); compout_class.def(py::init<>()) - .def("get_pipeline_output", - [](const CompilerOutput &co, const std::string &name) -> std::optional { - auto res = co.pipelineOutputs.find(name); - return res != co.pipelineOutputs.end() ? res->second - : std::optional(); - }) .def("get_output_ir", [](const CompilerOutput &co) -> std::string { return co.outIR; }) .def("get_object_filename", - [](const CompilerOutput &co) -> std::string { return co.objectFilename; }) + [](const CompilerOutput &co) -> std::string { return co.outputFilename; }) .def("get_diagnostic_messages", - [](const CompilerOutput &co) -> std::string { return co.diagnosticMessages; }) - .def("get_is_checkpoint_found", - [](const CompilerOutput &co) -> bool { return co.isCheckpointFound; }); + [](const CompilerOutput &co) -> std::string { return co.diagnosticMessages; }); m.def( "run_compiler_driver", @@ -79,22 +74,9 @@ PYBIND11_MODULE(compiler_driver, m) std::unique_ptr output(new CompilerOutput()); assert(output); - llvm::raw_string_ostream errStream{output->diagnosticMessages}; - - CompilerOptions options{.source = source, - .workspace = workspace, - .moduleName = moduleName, - .diagnosticStream = errStream, - .keepIntermediate = keepIntermediate, - .asyncQnodes = asyncQnodes, - .verbosity = verbose ? Verbosity::All : Verbosity::Urgent, - .pipelinesCfg = parseCompilerSpec(pipelines), - .lowerToLLVM = lower_to_llvm, - .checkpointStage = checkpointStage}; - - errStream.flush(); - - if (mlir::failed(QuantumDriverMain(options, *output))) { + if (QuantumDriverMainFromArgs(source, workspace, moduleName, keepIntermediate, + asyncQnodes, verbose, lower_to_llvm, + parseCompilerSpec(pipelines), checkpointStage, *output)) { throw std::runtime_error("Compilation failed:\n" + output->diagnosticMessages); } return output; diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt index f70fc94e2a..2539c3e294 100644 --- a/mlir/test/CMakeLists.txt +++ b/mlir/test/CMakeLists.txt @@ -8,6 +8,7 @@ configure_lit_site_cfg( set(DIALECT_TESTS_DEPEND FileCheck quantum-opt + catalyst-cli ) set(TEST_SUITES diff --git a/mlir/test/cli/DumpAfterFailure.mlir b/mlir/test/cli/DumpAfterFailure.mlir new file mode 100644 index 0000000000..64863a7f7c --- /dev/null +++ b/mlir/test/cli/DumpAfterFailure.mlir @@ -0,0 +1,23 @@ +// Copyright 2024 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. + +// RUN: not catalyst-cli --tool=opt %s --catalyst-pipeline="pipeline(llvm-dialect-lowring-pipeline)" --mlir-print-ir-after-failure --verify-diagnostics 2>&1 | FileCheck %s + +func.func @foo(%arg0: tensor) { + "catalyst.print"(%arg0) : (tensor) -> () + return +} + +// CHECK: IR Dump After CatalystConversionPass Failed +// CHECK: Compilation failed: diff --git a/mlir/test/cli/DumpBeforeAfterPass.mlir b/mlir/test/cli/DumpBeforeAfterPass.mlir new file mode 100644 index 0000000000..abb47813df --- /dev/null +++ b/mlir/test/cli/DumpBeforeAfterPass.mlir @@ -0,0 +1,38 @@ +// Copyright 2024 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. + +// RUN: catalyst-cli --tool=opt %s --catalyst-pipeline="pipe1(split-multiple-tapes;apply-transform-sequence),pipe2(inline-nested-module)" --mlir-print-ir-before-all --verify-diagnostics 2>&1 | FileCheck %s --check-prefix=CHECK-BEFORE +// RUN: catalyst-cli --tool=opt %s --catalyst-pipeline="pipe1(split-multiple-tapes;apply-transform-sequence),pipe2(inline-nested-module)" --mlir-print-ir-after-all --verify-diagnostics 2>&1 | FileCheck %s --check-prefix=CHECK-AFTER +// RUN: catalyst-cli --tool=opt %s --catalyst-pipeline="pipe1(split-multiple-tapes;apply-transform-sequence),pipe2(inline-nested-module)" --mlir-print-ir-before=inline-nested-module --verify-diagnostics 2>&1 | FileCheck %s --check-prefix=CHECK-BEFORE-ONE +// RUN: catalyst-cli --tool=opt %s --catalyst-pipeline="pipe1(split-multiple-tapes;apply-transform-sequence),pipe2(inline-nested-module)" --mlir-print-ir-after=inline-nested-module --verify-diagnostics 2>&1 | FileCheck %s --check-prefix=CHECK-AFTER-ONE +// RUN: catalyst-cli --tool=opt %s --catalyst-pipeline="pipe1(split-multiple-tapes;apply-transform-sequence),pipe2(inline-nested-module)" --mlir-print-op-generic --verify-diagnostics 2>&1 | FileCheck %s --check-prefix=CHECK-GENERIC + +func.func @foo() { + return +} + +// CHECK-BEFORE: IR Dump Before SplitMultipleTapesPass +// CHECK-BEFORE: IR Dump Before ApplyTransformSequencePass +// CHECK-BEFORE: IR Dump Before InlineNestedModulePass + +// CHECK-AFTER: IR Dump After SplitMultipleTapesPass +// CHECK-AFTER: IR Dump After ApplyTransformSequencePass +// CHECK-AFTER: IR Dump After InlineNestedModulePass + +// CHECK-BEFORE-ONE: IR Dump Before InlineNestedModulePass + +// CHECK-AFTER-ONE: IR Dump After InlineNestedModulePass + +// CHECK-GENERIC: "builtin.module"() ({ +// CHECK-GENERIC-NEXT: "func.func"() <{function_type = () -> (), sym_name = "foo"}> ({ diff --git a/mlir/test/cli/DumpPipeline.mlir b/mlir/test/cli/DumpPipeline.mlir new file mode 100644 index 0000000000..16bb2669a6 --- /dev/null +++ b/mlir/test/cli/DumpPipeline.mlir @@ -0,0 +1,36 @@ +// Copyright 2024 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. + +// RUN: catalyst-cli %s --dump-catalyst-pipeline --verify-diagnostics 2>&1 | FileCheck %s +// RUN: catalyst-cli --tool=opt %s --catalyst-pipeline="pipe1(split-multiple-tapes;apply-transform-sequence),pipe2(inline-nested-module)" --dump-catalyst-pipeline --verify-diagnostics 2>&1 | FileCheck %s --check-prefix=CHECK-CUSTOM +// RUN: catalyst-cli --tool=opt %s -cse --dump-catalyst-pipeline --verify-diagnostics 2>&1 | FileCheck %s --check-prefix=CHECK-ONE-PASS +// RUN: not catalyst-cli --tool=opt %s -cse --catalyst-pipeline="pipe1(split-multiple-tapes;apply-transform-sequence),pipe2(inline-nested-module)" --dump-catalyst-pipeline --verify-diagnostics 2>&1 | FileCheck %s --check-prefix=CHECK-FAIL + +func.func @foo() { + return +} + +// CHECK: Pass Manager with +// CHECK: builtin.module + +// CHECK-CUSTOM: Pass Manager with 2 passes +// CHECK-CUSTOM: builtin.module(split-multiple-tapes,apply-transform-sequence) +// CHECK-CUSTOM: Pass Manager with 1 passes +// CHECK-CUSTOM: builtin.module(inline-nested-module{stop-after-step=0}) + +// CHECK-ONE-PASS: Pass Manager with 1 passes +// CHECK-ONE-PASS: builtin.module(cse) + +// CHECK-FAIL: --catalyst-pipline option can't be used with individual pass options or -pass-pipeline. +// CHECK-FAIL: Compilation failed diff --git a/mlir/test/cli/DumpTiming.mlir b/mlir/test/cli/DumpTiming.mlir new file mode 100644 index 0000000000..c84779e626 --- /dev/null +++ b/mlir/test/cli/DumpTiming.mlir @@ -0,0 +1,26 @@ +// Copyright 2024 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. + +// RUN: catalyst-cli %s --mlir-timing --verify-diagnostics 2>&1 | FileCheck %s + +func.func @foo() { + return +} + +// CHECK: Execution time report +// CHECK: Total Execution Time +// CHECK: Parser +// CHECK: Optimization +// CHECK: Translate +// CHECK: llc diff --git a/mlir/test/cli/WrongInput.mlir b/mlir/test/cli/WrongInput.mlir new file mode 100644 index 0000000000..3069ea0a9c --- /dev/null +++ b/mlir/test/cli/WrongInput.mlir @@ -0,0 +1,26 @@ +// Copyright 2024 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. + +// RUN: not catalyst-cli %s --tool=llc --verify-diagnostics 2>&1 | FileCheck %s -check-prefix=CHECK-LLC +// RUN: not catalyst-cli %s --tool=translate --verify-diagnostics 2>&1 | FileCheck %s -check-prefix=CHECK-TRANSLATE + + +func.func @foo() { + return +} + +// CHECK-LLC: Compilation failed: +// CHECK-LLC: Expected LLVM IR input but received MLIR + +// CHECK-TRANSLATE: Failed to translate LLVM module diff --git a/mlir/tools/CMakeLists.txt b/mlir/tools/CMakeLists.txt index 883fc2086e..a483fe1f03 100644 --- a/mlir/tools/CMakeLists.txt +++ b/mlir/tools/CMakeLists.txt @@ -1,2 +1,3 @@ add_subdirectory(quantum-opt) add_subdirectory(quantum-lsp-server) +add_subdirectory(catalyst-cli) diff --git a/mlir/tools/catalyst-cli/CMakeLists.txt b/mlir/tools/catalyst-cli/CMakeLists.txt new file mode 100644 index 0000000000..0eaed84d8e --- /dev/null +++ b/mlir/tools/catalyst-cli/CMakeLists.txt @@ -0,0 +1,46 @@ +find_package(Enzyme REQUIRED CONFIG) +message(STATUS "Using EnzymeConfig.cmake in: ${Enzyme_DIR}") + +# TODO: remove this once Enzyme exports the static library target +find_library(ENZYME_LIB EnzymeStatic-${LLVM_VERSION_MAJOR} PATHS ${Enzyme_DIR}/Enzyme/) +message(STATUS "Found Enzyme lib: ${ENZYME_LIB}") + +include_directories(${ENZYME_SRC_DIR}/enzyme/Enzyme) + +# Experimentally found through removing items +# from llvm/tools/llc/CMakeLists.txt. +# It does make sense that we need the parser to parse MLIR +# and the codegen to codegen. +set(LLVM_LINK_COMPONENTS + AllTargetsAsmParsers + AllTargetsCodeGens + ) + +get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS) +get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS) +get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS) +set(LIBS + ${dialect_libs} + ${conversion_libs} + ${extension_libs} + MLIROptLib + MLIRCatalyst + MLIRCatalystTransforms + MLIRQuantum + quantum-transforms + MLIRGradient + gradient-transforms + MLIRMitigation + mitigation-transforms + MhloRegisterDialects + StablehloRegister + MLIRCatalystTest + ${ALL_MHLO_PASSES} + ${ENZYME_LIB} + CatalystCompilerDriver +) + +add_llvm_executable(catalyst-cli catalyst-cli.cpp) +target_link_libraries(catalyst-cli PRIVATE ${LIBS}) +llvm_update_compile_flags(catalyst-cli) +mlir_check_all_link_libraries(catalyst-cli) diff --git a/mlir/tools/catalyst-cli/catalyst-cli.cpp b/mlir/tools/catalyst-cli/catalyst-cli.cpp new file mode 100644 index 0000000000..43da6ed07b --- /dev/null +++ b/mlir/tools/catalyst-cli/catalyst-cli.cpp @@ -0,0 +1,17 @@ +// Copyright 2024 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. + +#include "Driver/CompilerDriver.h" + +int main(int argc, char **argv) { return QuantumDriverMainFromCL(argc, argv); }