Cross-compile target modules and dispatch to executor

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 \

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`, `<name>.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.";
 

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<std::string> 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(),

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

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<Value> bufferResults(bufferArgs.begin() + startIndex, bufferArgs.end());

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
 )