diff --git a/.github/workflows/build-wheel-linux-x86_64.yaml b/.github/workflows/build-wheel-linux-x86_64.yaml index c505d891ae..5c2d104742 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 qcc - name: Build wheel run: | diff --git a/.github/workflows/build-wheel-macos-arm64.yaml b/.github/workflows/build-wheel-macos-arm64.yaml index 0f018e61d6..3d30d88fab 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 qcc - 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..ff78654b3d 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 qcc - 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..830116af37 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 qcc # Copy files needed for the wheel where they are expected cp /catalyst/runtime-build/lib/*/*/*/*/librtd* /catalyst/runtime-build/lib diff --git a/frontend/catalyst/compiler.py b/frontend/catalyst/compiler.py index e49363c719..94b655a9e4 100644 --- a/frontend/catalyst/compiler.py +++ b/frontend/catalyst/compiler.py @@ -544,7 +544,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): @@ -601,7 +600,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 @@ -630,7 +628,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 @@ -638,12 +636,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..849d57af42 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 qcc .PHONY: test test: diff --git a/mlir/include/Driver/CompilerDriver.h b/mlir/include/Driver/CompilerDriver.h index 19e51e27ba..0e6c0c36cb 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 @@ -103,7 +106,16 @@ 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 { diff --git a/mlir/include/Driver/Pipelines.h b/mlir/include/Driver/Pipelines.h new file mode 100644 index 0000000000..d4fdd0798e --- /dev/null +++ b/mlir/include/Driver/Pipelines.h @@ -0,0 +1,61 @@ +// 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(); + +/// Pipeline descriptor +struct Pipeline { + using Name = std::string; + using PassList = llvm::SmallVector; + using PipelineFunc = void (*)(mlir::OpPassManager &); + Name name; + PassList passes; + PipelineFunc registerFunc = nullptr; + + mlir::LogicalResult addPipeline(mlir::OpPassManager &pm) + { + if (registerFunc) { + registerFunc(pm); + return mlir::success(); + } + else { + return mlir::failure(); + } + } +}; + +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..44ae87d4c4 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 { @@ -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,53 +446,37 @@ 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; - } + std::ifstream file(filename); + if (!file.is_open()) { + return ""; } + std::stringstream buffer; + buffer << file.rdbuf(); + return buffer.str(); +} - if (options.keepIntermediate && options.checkpointStage == "") { - llvm::raw_string_ostream s{outputs["mlir"]}; +LogicalResult runLowering(const CompilerOptions &options, MLIRContext *ctx, ModuleOp moduleOp, + CompilerOutput &output, TimingScope &timing) + +{ + 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"), - outputs["mlir"]); + tmp); } catalyst::utils::Timer timer{}; @@ -492,62 +490,87 @@ 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(); + auto pm = PassManager::on(ctx, PassManager::Nesting::Implicit); + 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))); - - // Run the lowering pipelines - if (failed(pm.run(moduleOp))) { + 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 (clHasIndividulaPass) + 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.name)) { + continue; + } + pm.clear(); + if (!clHasManualPipeline && failed(pipeline.addPipeline(pm))) { + llvm::errs() << "Pipeline creation function not found: " << pipeline.name << "\n"; + return failure(); + } + if (clHasManualPipeline && + failed(parsePassPipeline(joinPasses(pipeline.passes), pm, options.diagnosticStream))) { + return failure(); + } + if (options.dumpPassPipeline) { + pm.dump(); + llvm::errs() << "\n"; + } + 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.name, ".mlir"), tmp); + } + } return success(); } -LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput &output) +LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput &output, + DialectRegistry ®istry) { using timer = catalyst::utils::Timer; - DialectRegistry registry; - static bool initialized = false; - if (!initialized) { - registerAllPasses(); - } - initialized |= true; - registerAllCatalystPasses(); - mhlo::registerAllMhloPasses(); - - registerAllCatalystDialects(registry); - registerLLVMTranslations(registry); MLIRContext ctx(registry); ctx.printOpOnDiagnostic(true); ctx.printStackTraceOnDiagnostic(options.verbosity >= Verbosity::Debug); @@ -559,6 +582,7 @@ LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput & llvm::LLVMContext llvmContext; std::shared_ptr llvmModule; + enum InputType inType = InputType::OTHER; llvm::raw_string_ostream outIRStream(output.outIR); @@ -567,62 +591,91 @@ 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"; - // 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))) { - CO_MSG(options, Verbosity::Urgent, "Failed to lower MLIR module\n"); + 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); + } - output.outIR.clear(); - outIRStream << *op; + if (inType == InputType::OTHER) { + CO_MSG(options, Verbosity::Urgent, "Wrong or unsupported input\n"); + return failure(); + } + parserTiming.stop(); - 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(); - } + TimingScope optTiming = timing.nest("Optimization"); + // 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"; - catalyst::utils::LinesCount::Module(*llvmModule); + enum Action currentAction = options.loweringAction; + if (options.loweringAction == Action::All) + currentAction = Action::OPT; - 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"]); - } + if (inType == InputType::MLIR && currentAction == Action::OPT) { + 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 << *mlirModule; } - 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"; - + optTiming.stop(); + + TimingScope translateTiming = timing.nest("Translate"); + if (options.loweringAction == Action::All) + currentAction = Action::Translate; + if (inType == InputType::MLIR && currentAction == Action::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) { + TimingScope llcTiming = timing.nest("llc"); + if (options.loweringAction == Action::All) + currentAction = Action::LLC; + if (inType == InputType::LLVMIR && currentAction == Action::LLC) { // Set data layout before LLVM passes or the default one is used. std::string targetTriple = llvm::sys::getDefaultTargetTriple(); @@ -645,38 +698,217 @@ 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(); if (failed(timer::timer(compileObjectFile, "compileObjFile", /* add_endl */ true, options, std::move(llvmModule), targetMachine, outfile))) { return failure(); } output.objectFilename = 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); + + Pipeline::PassList passes; + for (auto &pass : passList) { + passes.push_back(pass.trim().str()); + } + + Pipeline pipeline; + pipeline.name = pipelineName.str(); + pipeline.passes = std::move(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); + + std::unique_ptr output(new CompilerOutput()); + assert(output); + 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; + } + + 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..eb376cc951 --- /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 registerAllCatalystPipelines() +{ + registerEnforceRuntimeInvariantsPipeline(); + registerHloLoweringPipeline(); + registerQuantumCompilationPipeline(); + registerBufferizationPipeline(); + registerLLVMDialectLoweringPipeline(); + registerDefaultCatalystPipeline(); +} + +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); +} + +std::vector getDefaultPipeline() +{ + using PipelineFunc = void (*)(mlir::OpPassManager &); + std::vector pipelineFuncs = { + &createEnforceRuntimeInvariantsPipeline, &createHloLoweringPipeline, + &createQuantumCompilationPipeline, &createBufferizationPipeline, + &createLLVMDialectLoweringPipeline}; + + Pipeline::PassList defaultPipelineNames = { + "enforce-runtime-invariants-pipeline", "hlo-lowering-pipeline", + "quantum-compilation-pipeline", "bufferization-pipeline", "llvm-dialect-lowering-pipeline"}; + + std::vector defaultPipelines; + for (size_t i = 0; i < defaultPipelineNames.size(); ++i) { + Pipeline pipeline; + pipeline.name = defaultPipelineNames[i]; + pipeline.passes.push_back(defaultPipelineNames[i]); + pipeline.registerFunc = pipelineFuncs[i]; + defaultPipelines.push_back(std::move(pipeline)); + } + return defaultPipelines; +} + +} // namespace driver +} // namespace catalyst diff --git a/mlir/python/PyCompilerDriver.cpp b/mlir/python/PyCompilerDriver.cpp index cd91f12e23..a97f12246c 100644 --- a/mlir/python/PyCompilerDriver.cpp +++ b/mlir/python/PyCompilerDriver.cpp @@ -53,19 +53,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; }) .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 +71,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..67b214bcd3 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 + qcc ) set(TEST_SUITES diff --git a/mlir/test/qcc/DumpAfterFailure.mlir b/mlir/test/qcc/DumpAfterFailure.mlir new file mode 100644 index 0000000000..567b71d57e --- /dev/null +++ b/mlir/test/qcc/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 qcc --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: \ No newline at end of file diff --git a/mlir/test/qcc/DumpBeforeAfterPass.mlir b/mlir/test/qcc/DumpBeforeAfterPass.mlir new file mode 100644 index 0000000000..70a4f79812 --- /dev/null +++ b/mlir/test/qcc/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: qcc --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: qcc --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: qcc --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: qcc --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: qcc --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"}> ({ \ No newline at end of file diff --git a/mlir/test/qcc/DumpPipeline.mlir b/mlir/test/qcc/DumpPipeline.mlir new file mode 100644 index 0000000000..9702d056b5 --- /dev/null +++ b/mlir/test/qcc/DumpPipeline.mlir @@ -0,0 +1,33 @@ +// 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: qcc %s --dump-catalyst-pipeline --verify-diagnostics 2>&1 | FileCheck %s +// RUN: qcc --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: qcc --tool=opt %s -cse --dump-catalyst-pipeline --verify-diagnostics 2>&1 | FileCheck %s --check-prefix=CHECK-ONE-PASS +// RUN: not qcc --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: Pass Manager with 1 passes + +// CHECK-ONE-PASS: Pass Manager with 1 passes + +// CHECK-FAIL: --catalyst-pipline option can't be used with individual pass options or -pass-pipeline. +// CHECK-FAIL: Compilation failed \ No newline at end of file diff --git a/mlir/test/qcc/DumpTiming.mlir b/mlir/test/qcc/DumpTiming.mlir new file mode 100644 index 0000000000..76f7aa92b4 --- /dev/null +++ b/mlir/test/qcc/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: qcc %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 \ No newline at end of file diff --git a/mlir/tools/CMakeLists.txt b/mlir/tools/CMakeLists.txt index 883fc2086e..5bfb0d9053 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(qcc) \ No newline at end of file diff --git a/mlir/tools/qcc/CMakeLists.txt b/mlir/tools/qcc/CMakeLists.txt new file mode 100644 index 0000000000..1635579257 --- /dev/null +++ b/mlir/tools/qcc/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(qcc qcc.cpp) +target_link_libraries(qcc PRIVATE ${LIBS}) +llvm_update_compile_flags(qcc) +mlir_check_all_link_libraries(qcc) diff --git a/mlir/tools/qcc/qcc.cpp b/mlir/tools/qcc/qcc.cpp new file mode 100644 index 0000000000..8dca3fe0e8 --- /dev/null +++ b/mlir/tools/qcc/qcc.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); } \ No newline at end of file