diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 35bb08012e..128ab1ff0f 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -15,6 +15,10 @@ * A new runtime transport layer for remote/local executors is introduced. [(#3043)](https://github.com/PennyLaneAI/catalyst/pull/3043) +* Catalyst can now cross-compile target nested modules to standalone object files and + either statically link them into the host program or ship them to an executor for dispatch. + [(#3033)](https://github.com/PennyLaneAI/catalyst/pull/3033) + * A `BufferizableOpInterface` implementation is now added for `catalyst.launch_kernel` operation and it is now bufferizable. [(#3024)](https://github.com/PennyLaneAI/catalyst/pull/3024) diff --git a/mlir/Makefile b/mlir/Makefile index d0e643ee6b..4a719ef15a 100644 --- a/mlir/Makefile +++ b/mlir/Makefile @@ -76,7 +76,7 @@ llvm: cmake -G Ninja -S llvm-project/llvm -B $(LLVM_BUILD_DIR) \ -DCMAKE_BUILD_TYPE=$(BUILD_TYPE) \ -DLLVM_BUILD_EXAMPLES=OFF \ - -DLLVM_TARGETS_TO_BUILD=${LLVM_TARGETS_TO_BUILD} \ + -DLLVM_TARGETS_TO_BUILD="${LLVM_TARGETS_TO_BUILD}" \ -DLLVM_ENABLE_PROJECTS="$(LLVM_PROJECTS)" \ -DLLVM_ENABLE_ASSERTIONS=ON \ -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ diff --git a/mlir/include/Catalyst/Transforms/Passes.td b/mlir/include/Catalyst/Transforms/Passes.td index c5f87f0d7c..0d2a6af902 100644 --- a/mlir/include/Catalyst/Transforms/Passes.td +++ b/mlir/include/Catalyst/Transforms/Passes.td @@ -143,6 +143,43 @@ def ApplyTransformSequencePass : Pass<"apply-transform-sequence"> { let summary = "Apply the passes scheduled with the transform dialect."; } +def CrossCompileTargetsPass : Pass<"cross-compile-targets", "mlir::ModuleOp"> { + let summary = "Cross-compile catalyst.target nested modules to standalone `.o` files."; + let description = [{ + For every `builtin.module` carrying a `catalyst.target` attribute, this pass: + + 1. Extracts the module body into a standalone root module. + 2. Runs the default lowering pipeline on it (bufferization + LLVM-dialect lowering). + 3. Translates to LLVM IR and emits an object file. + 4. Records the emitted object path on the module as a `catalyst.object_file` string + attribute and reduces the module to external declarations of its entry functions. + + The target triple is taken from the optional `triple` key on the `catalyst.target` + dict, falling back to the host triple. + }]; + + let options = [ + Option< + /*C++ var name=*/"workspace", + /*CLI arg name=*/"workspace", + /*type=*/"std::string", + /*default=*/"\"\"", + /*description=*/ + "Filesystem directory to write cross-compiled `.o` files into." + >, + Option< + /*C++ var name=*/"dumpIntermediate", + /*CLI arg name=*/"dump-intermediate", + /*type=*/"bool", + /*default=*/"false", + /*description=*/ + "Also write each target module's extracted MLIR and translated LLVM IR " + "(`extracted.mlir`, `.ll`) into its workspace subdirectory. Set by the " + "driver from `keep_intermediate`." + >, + ]; +} + def InlineNestedModulePass : Pass<"inline-nested-module"> { let summary = "Inline nested modules with qnode attribute."; diff --git a/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h b/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h index bea597e7fa..de896557cd 100644 --- a/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h +++ b/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h @@ -168,7 +168,7 @@ const PipelineList pipelineList{ "register-inactive-callback"}}}; // clang-format on -PipelineNames getPipelineNames() +inline PipelineNames getPipelineNames() { static std::vector names = std::accumulate(driver::pipelineList.begin(), driver::pipelineList.end(), @@ -179,7 +179,7 @@ PipelineNames getPipelineNames() return names; } -PassNames getQuantumCompilationStage(bool disableAssertion = true) +inline PassNames getQuantumCompilationStage(bool disableAssertion = true) { PassNames ret; std::copy_if(pipelineList[0].passNames.begin(), pipelineList[0].passNames.end(), @@ -189,11 +189,11 @@ PassNames getQuantumCompilationStage(bool disableAssertion = true) return ret; } -PassNames getHLOLoweringStage() { return pipelineList[1].passNames; } +inline PassNames getHLOLoweringStage() { return pipelineList[1].passNames; } -PassNames getGradientLoweringStage() { return pipelineList[2].passNames; } +inline PassNames getGradientLoweringStage() { return pipelineList[2].passNames; } -PassNames getBufferizationStage(bool asyncQNodes = false) +inline PassNames getBufferizationStage(bool asyncQNodes = false) { const std::string bufferizationOptions = std::string("{bufferize-function-boundaries ") + "allow-return-allocs-from-loops " + @@ -211,7 +211,7 @@ PassNames getBufferizationStage(bool asyncQNodes = false) return ret; } -PassNames getLLVMDialectLoweringStage(bool asyncQNodes = false) +inline PassNames getLLVMDialectLoweringStage(bool asyncQNodes = false) { PassNames ret; std::copy_if(pipelineList[4].passNames.begin(), pipelineList[4].passNames.end(), diff --git a/mlir/include/Executor/Transforms/Passes.td b/mlir/include/Executor/Transforms/Passes.td index 2d90f45cd2..d317823eac 100644 --- a/mlir/include/Executor/Transforms/Passes.td +++ b/mlir/include/Executor/Transforms/Passes.td @@ -36,4 +36,34 @@ def ConvertExecutorToLLVMPass : Pass<"convert-executor-to-llvm", "mlir::ModuleOp ]; } +def DispatchExecutorTargetsPass : Pass<"dispatch-executor-targets", "mlir::ModuleOp"> { + let summary = "Ship cross-compiled `catalyst.target` modules to an executor."; + let description = [{ + For every nested `builtin.module` carrying a `catalyst.dispatch` attribute and a + `catalyst.object_file` path, the pass: + + 1. Injects `executor.open` into `setup()` (once per unique address) and `executor.send_binary` + into `setup()` (once per module). + 2. Rewrites every host-side `catalyst.launch_kernel` targeting that module into an + `executor.launch` carrying the executor address, the entry callee, and the object-file + path. + 3. Erases the nested module from the host after its `catalyst.launch_kernel`s are rewritten. + + A `catalyst.custom_call` carrying a `dispatch` entry in its `backend_config` is rewritten into + an `executor.call` on the resolved executor address. + + Session teardown is handled by the runtime, which closes every open session at process exit, + so no explicit close op is emitted. + + The pass is a no-op when no `catalyst.dispatch` modules and no dispatch `catalyst.custom_call` + ops are present. + }]; + + let dependentDialects = [ + "catalyst::executor::ExecutorDialect", + "catalyst::CatalystDialect", + "mlir::func::FuncDialect" + ]; +} + #endif // EXECUTOR_PASSES diff --git a/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp index 79e6517116..b2296b347d 100644 --- a/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp @@ -183,10 +183,9 @@ struct CustomCallOpInterface IntegerAttr numArgumentsAttr = rewriter.getI32IntegerAttr(numArguments); // Create an updated custom call operation - auto newCustomCallOp = - CustomCallOp::create(rewriter, op->getLoc(), TypeRange{}, bufferArgs, - customCallOp.getCallTargetName(), numArgumentsAttr, - customCallOp.getBackendConfigAttr()); + auto newCustomCallOp = CustomCallOp::create( + rewriter, op->getLoc(), TypeRange{}, bufferArgs, customCallOp.getCallTargetName(), + numArgumentsAttr, customCallOp.getBackendConfigAttr()); newCustomCallOp->setDiscardableAttrs(customCallOp->getDiscardableAttrDictionary()); size_t startIndex = bufferArgs.size() - customCallOp.getNumResults(); SmallVector bufferResults(bufferArgs.begin() + startIndex, bufferArgs.end()); diff --git a/mlir/lib/Catalyst/Transforms/CMakeLists.txt b/mlir/lib/Catalyst/Transforms/CMakeLists.txt index bb794d3b4c..562a8f959f 100644 --- a/mlir/lib/Catalyst/Transforms/CMakeLists.txt +++ b/mlir/lib/Catalyst/Transforms/CMakeLists.txt @@ -7,6 +7,7 @@ file(GLOB SRC BufferDeallocation.cpp BufferizableOpInterfaceImpl.cpp catalyst_to_llvm.cpp + CrossCompileTargets.cpp DetectQNodes.cpp DetensorizeFunctionBoundaryPass.cpp DetensorizeSCFPass.cpp @@ -29,11 +30,22 @@ file(GLOB SRC TBAATagsPass.cpp ) +# Required by the cross-compile-targets pass: object emission for arbitrary target +# triples (AllTargets*) and MLIR->LLVM-IR translation registration +# (registerLLVMDialectTranslation / registerBuiltinDialectTranslation, +# translateModuleToLLVMIR). +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(translation_libs GLOBAL PROPERTY MLIR_TRANSLATION_LIBS) set(LIBS ${dialect_libs} ${conversion_libs} + ${translation_libs} catalyst-analysis ) diff --git a/mlir/lib/Catalyst/Transforms/CrossCompileTargets.cpp b/mlir/lib/Catalyst/Transforms/CrossCompileTargets.cpp new file mode 100644 index 0000000000..cc04a584d6 --- /dev/null +++ b/mlir/lib/Catalyst/Transforms/CrossCompileTargets.cpp @@ -0,0 +1,437 @@ +// Copyright 2026 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 "llvm/ADT/SmallSet.h" +#include "llvm/ADT/StringExtras.h" // llvm::join +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/IR/Module.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/Path.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/Target/TargetMachine.h" +#include "llvm/Target/TargetOptions.h" +#include "llvm/TargetParser/Host.h" +#include "llvm/TargetParser/Triple.h" +#include "mlir/Dialect/DLTI/DLTI.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/InitAllDialects.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Pass/PassRegistry.h" +#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h" +#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" +#include "mlir/Target/LLVMIR/Export.h" +#include "mlir/Target/LLVMIR/Import.h" // translateDataLayout + +#include "Catalyst/IR/CatalystDialect.h" +#include "Catalyst/IR/CatalystOps.h" +#include "Driver/DefaultPipelines/DefaultPipelines.h" +#include "Gradient/IR/GradientDialect.h" +#include "Ion/IR/IonDialect.h" +#include "MBQC/IR/MBQCDialect.h" +#include "Mitigation/IR/MitigationDialect.h" +#include "PBC/IR/PBCDialect.h" +#include "PauliFrame/IR/PauliFrameDialect.h" +#include "QRef/IR/QRefDialect.h" +#include "QecLogical/IR/QecLogicalDialect.h" +#include "QecPhysical/IR/QecPhysicalDialect.h" +#include "Quantum/IR/QuantumDialect.h" +#include "RTIO/IR/RTIODialect.h" + +using namespace mlir; + +namespace catalyst { + +#define GEN_PASS_DECL_CROSSCOMPILETARGETSPASS +#define GEN_PASS_DEF_CROSSCOMPILETARGETSPASS +#include "Catalyst/Transforms/Passes.h.inc" + +namespace { + +// Make `fn` externally callable through the C ABI +void exposeEntryViaCInterface(func::FuncOp fn) +{ + OpBuilder builder(fn.getContext()); + fn.setVisibility(SymbolTable::Visibility::Public); + fn->setAttr("llvm.emit_c_interface", builder.getUnitAttr()); + fn->removeAttr("llvm.linkage"); +} + +// The default lowering applied to a catalyst.target module: bufferization + +// LLVM-dialect lowering, reusing the same stage definitions as the host pipeline. +std::vector defaultLoweringPassList() +{ + auto buf = driver::getBufferizationStage(); + auto llvmPasses = driver::getLLVMDialectLoweringStage(); + std::vector passes; + passes.reserve(buf.size() + llvmPasses.size()); + passes.insert(passes.end(), buf.begin(), buf.end()); + for (const auto &passName : llvmPasses) { + if (passName != "convert-executor-to-llvm") { + passes.push_back(passName); + } + } + return passes; +} + +// Lower a *local* (non-dispatch) target module. Its cross-compiled object is statically linked into +// the final binary, so each host-side launch_kernel into it is rewritten to a flat func.call +// against an external declaration of the entry (resolved at link time by the object's native +// symbol). The object path is recorded on the root module for the linker, and the now-empty module +// is erased. +LogicalResult lowerLocalTargetCalls(ModuleOp host, ModuleOp nested, + SmallVectorImpl &objectFiles) +{ + MLIRContext *ctx = host.getContext(); + + // A statically-linked object must be built for the host architecture. A different triple can + // only be reached via executor dispatch. + if (auto targetAttr = nested->getAttrOfType("catalyst.target")) { + if (auto triple = targetAttr.getAs("triple")) { + if (!triple.getValue().empty() && + triple.getValue() != llvm::sys::getDefaultTargetTriple()) { + nested.emitError("local (non-executor) target must use the host triple '") + << llvm::sys::getDefaultTargetTriple() + << "'; use executor() to dispatch a cross-compiled target"; + return failure(); + } + } + } + + StringRef moduleName = nested.getSymName().value_or(""); + SmallVector launches; + host.walk([&](catalyst::LaunchKernelOp launchKernel) { + if (launchKernel.getCalleeModuleName().getValue() == moduleName) { + launches.push_back(launchKernel); + } + }); + for (catalyst::LaunchKernelOp launchKernel : launches) { + StringRef entry = launchKernel.getCalleeName().getValue(); + // One external declaration per entry, resolved against the linked object's native symbol. + auto decl = host.lookupSymbol(entry); + if (!decl) { + OpBuilder rootBuilder(host.getBody(), host.getBody()->begin()); + auto fnTy = FunctionType::get(ctx, launchKernel.getOperandTypes(), + launchKernel.getResultTypes()); + decl = func::FuncOp::create(rootBuilder, launchKernel.getLoc(), entry, fnTy); + decl.setPrivate(); + } + OpBuilder callBuilder(launchKernel); + auto call = func::CallOp::create(callBuilder, launchKernel.getLoc(), decl, + launchKernel.getOperands()); + launchKernel.replaceAllUsesWith(call.getResults()); + launchKernel.erase(); + } + + if (auto objPath = nested->getAttrOfType("catalyst.object_file")) { + objectFiles.push_back(objPath); + } + nested.erase(); + return success(); +} + +struct CrossCompileTargetsPass : impl::CrossCompileTargetsPassBase { + using CrossCompileTargetsPassBase::CrossCompileTargetsPassBase; + + void getDependentDialects(DialectRegistry ®istry) const override + { + mlir::registerLLVMDialectTranslation(registry); + mlir::registerBuiltinDialectTranslation(registry); + mlir::registerAllDialects(registry); + registry.insert(); + } + + void runOnOperation() final + { + ModuleOp host = getOperation(); + + SmallVector targetMods; + for (auto &op : host.getBody()->getOperations()) { + if (auto mod = dyn_cast(&op)) { + if (mod->hasAttr("catalyst.target")) { + targetMods.push_back(mod); + } + } + } + + if (targetMods.empty()) { + return; + } + + if (workspace.empty()) { + host.emitError("Missing `workspace` option for target cross-compilation"); + return signalPassFailure(); + } + + llvm::InitializeAllTargetInfos(); + llvm::InitializeAllTargets(); + llvm::InitializeAllTargetMCs(); + llvm::InitializeAllAsmParsers(); + llvm::InitializeAllAsmPrinters(); + + // Compile each target module to an object, then dispose of it per delivery mode: + // - executor (catalyst.dispatch): leave the module intact for dispatch-executor-targets. + // - local: statically link — flatten its host calls and record the object for the linker. + SmallVector localObjectFiles; + // Track seen target names to reject duplicates before compiling + llvm::SmallSet seenNames; + for (auto nested : targetMods) { + StringRef name = nested.getSymName().value_or("unnamed"); + if (!seenNames.insert(name).second) { + nested.emitError("duplicate catalyst.target module name '") + << name << "'; each target module must have a unique symbol name"; + return signalPassFailure(); + } + FailureOr objPath = compileTargetModule(nested); + if (failed(objPath)) { + return signalPassFailure(); + } + nested->setAttr("catalyst.object_file", StringAttr::get(&getContext(), *objPath)); + if (!nested->hasAttr("catalyst.dispatch")) { + if (failed(lowerLocalTargetCalls(host, nested, localObjectFiles))) { + return signalPassFailure(); + } + } + } + + // Record the objects to statically link, for the driver to hand to the linker. + if (!localObjectFiles.empty()) { + host->setAttr("catalyst.object_files", ArrayAttr::get(&getContext(), localObjectFiles)); + } + } + + // Returns the kernel-specific subdirectory {workspace}/{name}/, creating it if needed. + std::string makeKernelDir(StringRef name) + { + llvm::SmallString<128> dir(workspace); + llvm::sys::path::append(dir, name); + (void)llvm::sys::fs::create_directories(dir); + return std::string(dir.str()); + } + + // Write `op`/`mod` to {dir}/{filename} (used only when dump-intermediate is set). + void dumpMLIR(mlir::Operation *op, StringRef dir, StringRef filename) + { + llvm::SmallString<128> path(dir); + llvm::sys::path::append(path, filename); + std::error_code ec; + llvm::raw_fd_ostream os(path, ec, llvm::sys::fs::OF_None); + if (!ec) { + op->print(os); + } + } + + void dumpLLVMIR(llvm::Module &mod, StringRef dir, StringRef filename) + { + llvm::SmallString<128> path(dir); + llvm::sys::path::append(path, filename); + std::error_code ec; + llvm::raw_fd_ostream os(path, ec, llvm::sys::fs::OF_None); + if (!ec) { + mod.print(os, nullptr); + } + } + + /** + * @brief Build a TargetMachine for the given triple (generic CPU, no extra features). + * + * @return The TargetMachine, or nullptr (with a diagnostic on stderr) if the + * triple is not available in this LLVM build. + */ + std::unique_ptr createTargetMachine(StringRef triple) + { + llvm::Triple parsedTriple{triple}; + std::string err; + const llvm::Target *llvmTarget = llvm::TargetRegistry::lookupTarget(parsedTriple, err); + if (!llvmTarget) { + llvm::errs() << "Target triple '" << triple + << "' not registered in this LLVM build: " << err << "\n"; + return nullptr; + } + llvm::TargetOptions opt; + std::unique_ptr targetMachine(llvmTarget->createTargetMachine( + parsedTriple, /*cpu=*/"generic", /*features=*/"", opt, llvm::Reloc::Model::PIC_)); + if (!targetMachine) { + llvm::errs() << "Could not create TargetMachine for triple '" << triple << "'\n"; + return nullptr; + } + targetMachine->setOptLevel(llvm::CodeGenOptLevel::Aggressive); + return targetMachine; + } + + /** + * @brief Emit a `.o` file from an LLVM module using a prepared TargetMachine. + * + * The module's triple and data layout were stamped on the source MLIR module before + * lowering and carried into `llvmModule` by translateModuleToLLVMIR, so they already + * match `targetMachine` — no need to re-set them here. + * + * @param llvmModule The LLVM module to emit. + * @param name The name used as the object file's basename. + * @param dir The directory to write the object file into. + * @param targetMachine The target machine to emit for. + * @return std::string The path to the emitted object file, or "" on failure. + */ + std::string emitObjectFile(std::unique_ptr &&llvmModule, StringRef name, + StringRef dir, llvm::TargetMachine &targetMachine) + { + llvm::SmallString<128> p(dir); + llvm::sys::path::append(p, name.str() + ".o"); + std::string objPath = std::string(p.str()); + + std::error_code errCode; + llvm::raw_fd_ostream dest(objPath, errCode, llvm::sys::fs::OF_None); + if (errCode) { + llvm::errs() << "Cannot open " << objPath << " for writing: " << errCode.message() + << "\n"; + return ""; + } + llvm::legacy::PassManager codegenPM; + if (targetMachine.addPassesToEmitFile(codegenPM, dest, nullptr, + llvm::CodeGenFileType::ObjectFile)) { + llvm::errs() << "TargetMachine cannot emit an object file\n"; + return ""; + } + codegenPM.run(*llvmModule); + dest.flush(); + return objPath; + } + + // Cross-compile a catalyst.target module to a `.o` for its triple and return the path. + FailureOr compileTargetModule(ModuleOp nested) + { + MLIRContext *ctx = &getContext(); + StringRef name = nested.getSymName().value_or("unnamed"); + + auto targetAttr = nested->getAttrOfType("catalyst.target"); + if (!targetAttr) { + nested.emitError("catalyst.target module missing catalyst.target dict attr"); + return failure(); + } + + // Triple selection: the optional 'triple' key on the catalyst.target dict + // wins; otherwise the host triple. + std::string moduleTarget; + if (auto tripleAttr = targetAttr.getAs("triple")) { + moduleTarget = tripleAttr.getValue().str(); + } + if (moduleTarget.empty()) { + moduleTarget = llvm::sys::getDefaultTargetTriple(); + } + + std::string kernelDir = makeKernelDir(name); + + std::unique_ptr targetMachine = createTargetMachine(moduleTarget); + if (!targetMachine) { + nested.emitError("failed to create target machine for triple '" + moduleTarget + + "' (target module: " + name.str() + ")"); + return failure(); + } + llvm::DataLayout dataLayout = targetMachine->createDataLayout(); + + // Clone the target module into an unparented root module: leaves `nested` intact in the + // host (its host launch_kernel is consumed later by local flattening or dispatch) and gives + // the sub-pipeline / translateModuleToLLVMIR a top-level module to operate on. + OpBuilder builder(ctx); + mlir::OwningOpRef standalone(cast(nested->clone())); + + Operation *moduleOp = standalone->getOperation(); + moduleOp->setAttr(LLVM::LLVMDialect::getTargetTripleAttrName(), + builder.getStringAttr(moduleTarget)); + moduleOp->setAttr(LLVM::LLVMDialect::getDataLayoutAttrName(), + builder.getStringAttr(dataLayout.getStringRepresentation())); + moduleOp->setAttr(DLTIDialect::kDataLayoutAttrName, + mlir::translateDataLayout(dataLayout, ctx)); + + // The entry points are exactly the functions the host calls into this module, named by the + // surviving launch_kernel call edges. Expose those through the C ABI; privatize the rest so + // they can be internalized / DCE'd and don't leak as exported symbols. + StringRef moduleName = nested.getSymName().value_or(""); + llvm::SmallSet entries; + if (auto host = nested->getParentOfType()) { + host.walk([&](catalyst::LaunchKernelOp launchKernel) { + if (launchKernel.getCalleeModuleName().getValue() == moduleName) { + entries.insert(launchKernel.getCalleeName().getValue()); + } + }); + } + for (auto fn : standalone->getOps()) { + if (entries.contains(fn.getName())) { + exposeEntryViaCInterface(fn); + } + else { + fn.setPrivate(); + } + } + + if (dumpIntermediate) { + dumpMLIR(*standalone, kernelDir, "extracted.mlir"); + } + + // Lower the extracted module. A 'pipeline' key on catalyst.target selects a named + // target-lowering pipeline resolved against the host pipeline registry. + // Without it, fall back to the default bufferization + LLVM-dialect lowering. + std::string pipelineSpec; + if (auto pipelineAttr = targetAttr.getAs("pipeline")) { + pipelineSpec = pipelineAttr.getValue().str(); + } + if (pipelineSpec.empty()) { + pipelineSpec = llvm::join(defaultLoweringPassList(), ","); + } + PassManager subPM(ctx); + if (failed(parsePassPipeline(pipelineSpec, subPM))) { + nested.emitError("failed to build the target-lowering pipeline '" + pipelineSpec + "'"); + return failure(); + } + if (failed(subPM.run(*standalone))) { + nested.emitError("failed to lower target module to LLVM dialect: " + name.str()); + return failure(); + } + + // Translate to LLVM IR and emit the object file. + llvm::LLVMContext llvmCtx; + std::unique_ptr llvmModule = + translateModuleToLLVMIR(*standalone, llvmCtx, name); + if (!llvmModule) { + nested.emitError("failed to translate target module to LLVM IR: " + name.str()); + return failure(); + } + if (dumpIntermediate) { + dumpLLVMIR(*llvmModule, kernelDir, name.str() + ".ll"); + } + std::string objPath = + emitObjectFile(std::move(llvmModule), name, kernelDir, *targetMachine); + if (objPath.empty()) { + nested.emitError("failed to emit object file for target module: " + name.str()); + return failure(); + } + return objPath; + } +}; + +} // namespace + +} // namespace catalyst diff --git a/mlir/lib/Driver/CompilerDriver.cpp b/mlir/lib/Driver/CompilerDriver.cpp index f781e5bfb5..6439371099 100644 --- a/mlir/lib/Driver/CompilerDriver.cpp +++ b/mlir/lib/Driver/CompilerDriver.cpp @@ -483,6 +483,45 @@ llvm::LogicalResult catalyst::driver::runPipeline(PassManager &pm, const Compile return failure(); } catalyst::utils::LinesCount::call(moduleOp); + + // Cross-compile catalyst.target nested modules and dispatch for execution + if (pipeline.getName() == "BufferizationStage" && !options.workspace.empty()) { + Pipeline targetPipeline; + targetPipeline.setName("CrossCompileTargets"); + std::string dumpIntermediate = options.keepIntermediate ? "true" : "false"; + targetPipeline.setPasses({"cross-compile-targets{workspace=\"" + + options.workspace.str() + + "\" dump-intermediate=" + dumpIntermediate + "}", + "dispatch-executor-targets"}); + if (failed(catalyst::utils::Timer<>::timer( + catalyst::driver::runPipeline, targetPipeline.getName(), + /* add_endl */ false, pm, options, output, targetPipeline, + /* clHasManualPipeline */ true, moduleOp))) { + return failure(); + } + catalyst::utils::LinesCount::call(moduleOp); + + // cross-compile-targets records the objects of local (statically-linked) targets on the + // root module. Write them to a manifest the frontend hands to the linker. + if (auto objFiles = moduleOp->getAttrOfType("catalyst.object_files")) { + std::string manifestPath = + options.workspace.str() + "/" + options.moduleName.str() + ".objects"; + std::error_code ec; + llvm::raw_fd_ostream manifest(manifestPath, ec); + if (ec) { + moduleOp->emitError("failed to write object manifest '") + << manifestPath << "': " << ec.message(); + return failure(); + } + else { + for (mlir::Attribute pathAttr : objFiles) { + if (auto s = mlir::dyn_cast(pathAttr)) { + manifest << s.getValue() << "\n"; + } + } + } + } + } } return success(); } diff --git a/mlir/lib/Executor/Transforms/CMakeLists.txt b/mlir/lib/Executor/Transforms/CMakeLists.txt index dc396e2413..e8f76ddbc4 100644 --- a/mlir/lib/Executor/Transforms/CMakeLists.txt +++ b/mlir/lib/Executor/Transforms/CMakeLists.txt @@ -1,14 +1,22 @@ set(LIBRARY_NAME executor-transforms) +set(LLVM_LINK_COMPONENTS + AllTargetsAsmParsers + AllTargetsCodeGens +) + file(GLOB SRC + DispatchExecutorTargets.cpp ExecutorToLLVM.cpp ) get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS) get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS) +get_property(translation_libs GLOBAL PROPERTY MLIR_TRANSLATION_LIBS) set(LIBS ${dialect_libs} ${conversion_libs} + ${translation_libs} MLIRExecutor MLIRCatalyst ) diff --git a/mlir/lib/Executor/Transforms/DispatchExecutorTargets.cpp b/mlir/lib/Executor/Transforms/DispatchExecutorTargets.cpp new file mode 100644 index 0000000000..fdf7539149 --- /dev/null +++ b/mlir/lib/Executor/Transforms/DispatchExecutorTargets.cpp @@ -0,0 +1,272 @@ +// Copyright 2026 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 "llvm/ADT/SmallSet.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +#include "Catalyst/IR/CatalystOps.h" + +#include "Executor/IR/ExecutorOps.h" +#include "Executor/Transforms/Passes.h" + +using namespace mlir; + +namespace catalyst { +namespace executor { + +#define GEN_PASS_DEF_DISPATCHEXECUTORTARGETSPASS +#include "Executor/Transforms/Passes.h.inc" + +namespace { + +// Ships cross-compiled target modules to an executor using the `executor` dialect. +// +// This pass runs after `cross-compile-targets`, which records each module's object file in +// `catalyst.object_file`. For every nested module carrying a `catalyst.dispatch` attribute this +// pass: +// 1. Injects `executor.open` into `setup()` (once per unique address) and `executor.send_binary` +// into `setup()` (once per module; the object holds every entry). +// 2. Rewrites every host-side `catalyst.launch_kernel` targeting the module into an +// `executor.launch` carrying the executor address, the entry callee, and the object-file path. +// 3. Erases the nested module from the host after its `catalyst.launch_kernel`s are rewritten. +// +// A `catalyst.custom_call` carrying a `dispatch` entry in its `backend_config` is rewritten into an +// `executor.call` on the resolved executor address. +// +// Session teardown is handled by the runtime, which closes every open session at process exit, so +// no explicit close op is emitted. +struct DispatchExecutorTargetsPass + : impl::DispatchExecutorTargetsPassBase { + using DispatchExecutorTargetsPassBase::DispatchExecutorTargetsPassBase; + + void runOnOperation() final + { + ModuleOp host = getOperation(); + + SmallVector targetMods; + for (auto &op : host.getBody()->getOperations()) { + if (auto mod = dyn_cast(&op)) { + if (mod->hasAttr("catalyst.dispatch")) { + targetMods.push_back(mod); + } + } + } + + // catalyst.custom_call ops whose backend_config carries a `dispatch` entry + // The call-target name is the executor-side symbol. + SmallVector libCalls; + host.walk([&](catalyst::CustomCallOp call) { + if (executorDispatchOf(call)) { + libCalls.push_back(call); + } + }); + + if (targetMods.empty() && libCalls.empty()) { + return; + } + + // Modules sharing an executor get a single executor.open. + llvm::SmallSet openedAddresses; + StringAttr executorAddress; + + for (auto nested : targetMods) { + if (failed(dispatchTargetModule(host, nested, openedAddresses, executorAddress))) { + return signalPassFailure(); + } + nested.erase(); + } + + const size_t numQnodeExecutors = openedAddresses.size(); + + if (!libCalls.empty()) { + for (catalyst::CustomCallOp call : libCalls) { + StringAttr dispatch = executorDispatchOf(call); + if ((!dispatch || dispatch.getValue().empty()) && numQnodeExecutors > 1) { + call.emitOpError("ambiguous executor"); + return signalPassFailure(); + } + StringAttr addrAttr = libCallAddress(call, executorAddress); + if (!addrAttr) { + call.emitOpError("custom_call dispatch has no executor address"); + return signalPassFailure(); + } + if (openedAddresses.insert(addrAttr.str()).second) { + injectExecutorOpenIntoSetup(host, addrAttr); + } + } + if (failed(rewriteExecutorLibCalls(libCalls, executorAddress))) { + return signalPassFailure(); + } + } + } + + static StringAttr executorDispatchOf(catalyst::CustomCallOp call) + { + if (auto cfg = call.getBackendConfigAttr()) { + return cfg.getAs("dispatch"); + } + return nullptr; + } + + static StringAttr libCallAddress(catalyst::CustomCallOp call, StringAttr fallbackAddress) + { + if (StringAttr dispatch = executorDispatchOf(call)) { + if (!dispatch.getValue().empty()) { + return dispatch; + } + } + return fallbackAddress; + } + + LogicalResult rewriteExecutorLibCalls(ArrayRef libCalls, + StringAttr fallbackAddress) + { + MLIRContext *ctx = &getContext(); + for (catalyst::CustomCallOp call : libCalls) { + StringAttr addressAttr = libCallAddress(call, fallbackAddress); + if (!addressAttr) { + call.emitOpError("custom_call dispatch has no executor address"); + return failure(); + } + auto symAttr = StringAttr::get(ctx, call.getCallTargetName()); + OpBuilder b(call); + IntegerAttr numInputAttr = nullptr; + if (auto n = call.getNumberOriginalArg()) { + numInputAttr = b.getI32IntegerAttr(*n); + } + auto executorCall = executor::CallOp::create( + b, call.getLoc(), call.getResultTypes(), call.getOperands(), + /*address=*/addressAttr, /*symbol=*/symAttr, + /*num_input_args=*/numInputAttr); + call.replaceAllUsesWith(executorCall.getResults()); + call.erase(); + } + return success(); + } + + // Insert `executor.open` before the terminator of `setup`. No-op if setup is absent or + // bodyless. + void injectExecutorOpenIntoSetup(ModuleOp host, StringAttr addressAttr) + { + auto setupFn = host.lookupSymbol("setup"); + if (!setupFn || setupFn.getBody().empty()) { + return; + } + Operation *terminator = setupFn.getBody().front().getTerminator(); + if (!terminator) { + return; + } + OpBuilder b(terminator); + executor::OpenOp::create(b, setupFn.getLoc(), addressAttr); + } + + // Insert `executor.send_binary` before the terminator of `setup`. No-op if setup is absent or + // bodyless. Always called after `injectExecutorOpenIntoSetup` so the session is opened first. + void injectExecutorSendBinaryIntoSetup(ModuleOp host, StringAttr addressAttr, + StringAttr pathAttr) + { + auto setupFn = host.lookupSymbol("setup"); + if (!setupFn || setupFn.getBody().empty()) { + return; + } + Operation *terminator = setupFn.getBody().front().getTerminator(); + if (!terminator) { + return; + } + OpBuilder b(terminator); + executor::SendBinaryOp::create(b, setupFn.getLoc(), addressAttr, pathAttr); + } + + // Open a session (once per address), ship the object recorded in + // `catalyst.object_file`, and rewrite each host-side launch_kernel into an `executor.launch`. + LogicalResult dispatchTargetModule(ModuleOp host, ModuleOp nested, + llvm::SmallSet &openedAddresses, + StringAttr &executorAddress) + { + MLIRContext *ctx = &getContext(); + + // The object path is produced by the cross-compile-targets pass. + auto objPathAttr = nested->getAttrOfType("catalyst.object_file"); + if (!objPathAttr || objPathAttr.getValue().empty()) { + nested.emitError("executor dispatch requires a non-empty 'catalyst.object_file' " + "attribute (run cross-compile-targets first)"); + return failure(); + } + + auto dispatchAttr = nested->getAttrOfType("catalyst.dispatch"); + auto addrAttr = dispatchAttr ? dispatchAttr.getAs("address") : nullptr; + if (!addrAttr || addrAttr.getValue().empty()) { + nested.emitError("executor dispatch requires a non-empty 'address' key in the " + "catalyst.dispatch attribute"); + return failure(); + } + std::string moduleAddress = addrAttr.getValue().str(); + + auto pathAttr = StringAttr::get(ctx, objPathAttr.getValue()); + auto addressAttr = StringAttr::get(ctx, moduleAddress); + // Remember the executor address so standalone executor lib calls reuse it. + executorAddress = addressAttr; + + // Inject executor.open (setup) once per unique address. + if (!openedAddresses.count(moduleAddress)) { + openedAddresses.insert(moduleAddress); + injectExecutorOpenIntoSetup(host, addressAttr); + } + + // Ship the object once per module: a single `catalyst.object_file` holds every + // entry function, and executor.launch resolves individual symbols within it. + injectExecutorSendBinaryIntoSetup(host, addressAttr, pathAttr); + + // Rewrite each host-side launch_kernel targeting this module into an executor.launch. + StringRef moduleName = nested.getSymName().value_or(""); + SmallVector launches; + host.walk([&](catalyst::LaunchKernelOp launchKernel) { + if (launchKernel.getCalleeModuleName().getValue() == moduleName) { + launches.push_back(launchKernel); + } + }); + for (catalyst::LaunchKernelOp launchKernel : launches) { + // executor.launch marshals memref descriptors, so its lowering only accepts + // memref-typed operands and results. Reject anything else here with a clear error + // rather than crashing later in convert-executor-to-llvm. This runs after + // bufferization, so a well-formed entry call is already memref-typed. + auto isMemref = [](Type ty) { return isa(ty); }; + if (!llvm::all_of(launchKernel.getOperandTypes(), isMemref) || + !llvm::all_of(launchKernel.getResultTypes(), isMemref)) { + launchKernel.emitOpError("executor dispatch of '") + << launchKernel.getCalleeName().getValue() + << "' requires memref-typed operands and results"; + return failure(); + } + auto calleeAttr = StringAttr::get(ctx, launchKernel.getCalleeName().getValue()); + OpBuilder b(launchKernel); + // `pathAttr` (the object-file path shipped by send_binary) identifies which object the + // entry resolves in, so executor.launch is resolved against the right binary. + auto launch = executor::LaunchOp::create( + b, launchKernel.getLoc(), launchKernel.getResultTypes(), launchKernel.getOperands(), + addressAttr, calleeAttr, pathAttr); + launchKernel.replaceAllUsesWith(launch.getResults()); + launchKernel.erase(); + } + return success(); + } +}; + +} // namespace + +} // namespace executor +} // namespace catalyst diff --git a/mlir/test/Catalyst/CrossCompileTargetMissingWorkspace.mlir b/mlir/test/Catalyst/CrossCompileTargetMissingWorkspace.mlir new file mode 100644 index 0000000000..9fe0930600 --- /dev/null +++ b/mlir/test/Catalyst/CrossCompileTargetMissingWorkspace.mlir @@ -0,0 +1,29 @@ +// Copyright 2026 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 quantum-opt %s --cross-compile-targets 2>&1 | FileCheck %s + +// CHECK: Missing `workspace` option for target cross-compilation +module @jit_missing_workspace { + + func.func public @jit_main() attributes {llvm.emit_c_interface} { + return + } + + module attributes {catalyst.target = {}} { + func.func public @noop() { + return + } + } +} diff --git a/mlir/test/Catalyst/CrossCompileTargetModules.mlir b/mlir/test/Catalyst/CrossCompileTargetModules.mlir new file mode 100644 index 0000000000..91193d202c --- /dev/null +++ b/mlir/test/Catalyst/CrossCompileTargetModules.mlir @@ -0,0 +1,59 @@ +// Copyright 2026 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: mkdir -p %t +// RUN: quantum-opt %s --cross-compile-targets="workspace=%t" | FileCheck %s + +// A local (non-dispatch) catalyst.target module is compiled to an object and statically linked: the +// object path is recorded on the root as catalyst.object_files, the host launch_kernel is flattened +// to a func.call against an external declaration of the entry, and the module is erased. +// CHECK: module @jit_test_cross_compile_target +// CHECK-SAME: catalyst.object_files = ["{{.*}}.o"] +// CHECK: func.func private @noop() +// CHECK: func.func public @jit_main +// CHECK: call @noop() +// CHECK-NOT: catalyst.launch_kernel +// CHECK-NOT: module @target_compute +// CHECK-NOT: noop_helper + +// dump-intermediate writes the extracted MLIR and translated LLVM IR to the workspace. The entry +// points are derived from the launch_kernel call edges (no separate visibility pass): @noop is the +// callee, so it is exposed through the C ABI (public + llvm.emit_c_interface); @noop_helper is not +// referenced from the host, so it is privatized. +// RUN: rm -rf %t/target_compute +// RUN: quantum-opt %s --cross-compile-targets="workspace=%t dump-intermediate=true" +// RUN: cat %t/target_compute/extracted.mlir | FileCheck %s --check-prefix=EXTRACTED +// RUN: cat %t/target_compute/target_compute.ll | FileCheck %s --check-prefix=LL +// EXTRACTED: func.func @noop() attributes {llvm.emit_c_interface} +// EXTRACTED: func.func private @noop_helper() +// LL: define {{.*}}@noop + +module @jit_test_cross_compile_target { + + func.func public @jit_main() attributes {llvm.emit_c_interface} { + catalyst.launch_kernel @target_compute::@noop() : () -> () + return + } + + // Both functions start public; cross-compile-targets derives the entry set from the host's + // launch_kernel and privatizes the unreferenced helper. + module @target_compute attributes {catalyst.target = {}} { + func.func public @noop() { + return + } + func.func public @noop_helper() { + return + } + } +} diff --git a/mlir/test/Catalyst/CrossCompileTargetNameCollision.mlir b/mlir/test/Catalyst/CrossCompileTargetNameCollision.mlir new file mode 100644 index 0000000000..437675392c --- /dev/null +++ b/mlir/test/Catalyst/CrossCompileTargetNameCollision.mlir @@ -0,0 +1,36 @@ +// Copyright 2026 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: mkdir -p %t +// RUN: not quantum-opt %s --cross-compile-targets="workspace=%t" 2>&1 | FileCheck %s + +// CHECK: duplicate catalyst.target module name 'unnamed' +module @jit_name_collision { + + func.func public @jit_main() attributes {llvm.emit_c_interface} { + return + } + + module attributes {catalyst.target = {}} { + func.func public @noop() { + return + } + } + + module attributes {catalyst.target = {}} { + func.func public @noop() { + return + } + } +} diff --git a/mlir/test/Executor/DispatchExecutorLibCalls.mlir b/mlir/test/Executor/DispatchExecutorLibCalls.mlir new file mode 100644 index 0000000000..f7fef43d6a --- /dev/null +++ b/mlir/test/Executor/DispatchExecutorLibCalls.mlir @@ -0,0 +1,95 @@ +// Copyright 2026 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: quantum-opt %s --split-input-file --dispatch-executor-targets --verify-diagnostics | FileCheck %s + +// dispatch-executor-targets also rewrites standalone executor library calls which are +// catalyst.custom_call ops whose backend_config carries a `dispatch` entry into +// executor.call ops invoking a symbol already loaded on the executor. + +// ----- + +// When backend_config.dispatch names the executor explicitly, it is rewritten to a executor.call on +// that address, and a session is opened once in setup(). +// CHECK-LABEL: func.func @setup() +// CHECK: executor.open("ADDR:PORT") +// CHECK: return +// CHECK-LABEL: func.func @main +// CHECK: executor.call("fpga_trampoline_a_setup", "ADDR:PORT") +// CHECK-SAME: (memref<256xi8>) -> () +// CHECK-NOT: catalyst.custom_call +module @jit_bound { + func.func @setup() { + return + } + func.func @main(%arg0: memref<256xi8>) { + catalyst.custom_call fn("fpga_trampoline_a_setup") (%arg0) {backend_config = {dispatch = "ADDR:PORT"}, number_original_arg = 1 : i32} : (memref<256xi8>) -> () + return + } +} + +// ----- + +// backend_config.dispatch = "" binds to the program's single executor, which +// is supplied by the QNode target module's catalyst.dispatch address. +// CHECK-LABEL: func.func @setup() +// CHECK: executor.open("ADDR:PORT") +// CHECK-LABEL: func.func @main +// CHECK: executor.call("fpga_trampoline_a_teardown", "ADDR:PORT") +// CHECK-NOT: catalyst.custom_call +module @jit_inherit { + func.func @setup() { + return + } + func.func @teardown() { + return + } + func.func @main() { + catalyst.launch_kernel @target::@compute() : () -> () + catalyst.custom_call fn("fpga_trampoline_a_teardown") () {backend_config = {dispatch = ""}} : () -> () + return + } + module @target attributes {catalyst.object_file = "/tmp/target.o", catalyst.dispatch = {address = "ADDR:PORT"}} { + func.func public @compute() { + return + } + } +} + +// ----- + +// When the program targets more than one executor, the custom_call should explicitly bind the +// dispatch otherwise it would be ambiguous. +module @jit_ambiguous { + func.func @setup() { + return + } + func.func @main() { + catalyst.launch_kernel @t1::@c1() : () -> () + catalyst.launch_kernel @t2::@c2() : () -> () + // expected-error @below {{ambiguous executor}} + catalyst.custom_call fn("foo") () {backend_config = {dispatch = ""}} : () -> () + return + } + module @t1 attributes {catalyst.object_file = "/tmp/t1.o", catalyst.dispatch = {address = "host:1"}} { + func.func public @c1() { + return + } + } + module @t2 attributes {catalyst.object_file = "/tmp/t2.o", catalyst.dispatch = {address = "host:2"}} { + func.func public @c2() { + return + } + } +} diff --git a/mlir/test/Executor/DispatchExecutorTargetModules.mlir b/mlir/test/Executor/DispatchExecutorTargetModules.mlir new file mode 100644 index 0000000000..05d0563c4f --- /dev/null +++ b/mlir/test/Executor/DispatchExecutorTargetModules.mlir @@ -0,0 +1,71 @@ +// Copyright 2026 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: quantum-opt %s --dispatch-executor-targets | FileCheck %s + +// This pass consumes the input IR left by cross-compile-targets which carries catalyst.object_file, +// plus a catalyst.dispatch attributes. + +// setup() opens the session once and ships the module's object. +// CHECK-LABEL: func.func @setup() +// CHECK: executor.open("ADDR:PORT") +// CHECK: executor.send_binary("ADDR:PORT", "/tmp/target_compute.o") +// Make sure the binary is shipped exactly once even though the host has two launch_kernels +// CHECK-NOT: executor.send_binary +// CHECK: return + +// teardown() is left untouched: the runtime closes sessions at process exit. +// CHECK-LABEL: func.func @teardown() +// CHECK-NOT: executor. +// CHECK: return + +// Each host-side launch_kernel is rewritten to its own executor.launch, preserving the call's +// operand/result types (a typed call exercises the memref-result path). +// CHECK-LABEL: func.func public @jit_main +// CHECK: executor.launch("noop", "ADDR:PORT", "/tmp/target_compute.o") () : () -> () +// CHECK: executor.launch("compute", "ADDR:PORT", "/tmp/target_compute.o") (%{{.*}}) : (memref<4xf64>) -> memref<4xf64> + +// The launch_kernels and the nested module are gone. +// CHECK-NOT: catalyst.launch_kernel +// CHECK-NOT: module @target_compute + +module @jit_test_dispatch { + + func.func @setup() { + return + } + + func.func @teardown() { + return + } + + // inline-nested-module keeps dispatch-module calls as launch_kernel (no flattening, no host-side + // declarations), so dispatch matches the launch_kernel callee module directly. + func.func public @jit_main(%arg0: memref<4xf64>) -> memref<4xf64> attributes {llvm.emit_c_interface} { + catalyst.launch_kernel @target_compute::@noop() : () -> () + %0 = catalyst.launch_kernel @target_compute::@compute(%arg0) : (memref<4xf64>) -> memref<4xf64> + return %0 : memref<4xf64> + } + + // cross-compile-targets leaves executor-dispatch modules intact (public, with bodies); dispatch + // erases the whole module after rewriting the host launch_kernels. + module @target_compute attributes {catalyst.object_file = "/tmp/target_compute.o", catalyst.dispatch = {address = "ADDR:PORT"}} { + func.func public @noop() { + return + } + func.func public @compute(%arg0: memref<4xf64>) -> memref<4xf64> { + return %arg0 : memref<4xf64> + } + } +}