diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index a0c22c992e..43e1a17a5a 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -60,6 +60,10 @@ argument is at an arbitrary position in the argument list. [(#2836)](https://github.com/PennyLaneAI/catalyst/pull/2836) +* A new `catalyst.debug.compile_mlir` function has been added, allowing standalone MLIR files to be + compiled through the full Catalyst pipeline and returned as a callable Python object. + [(#2832)](https://github.com/PennyLaneAI/catalyst/pull/2832) + * PPRs and PPMs can now be lowered properly into MLIR directly in the non-capture workflow. [(#2816)](https://github.com/PennyLaneAI/catalyst/pull/2816) @@ -242,6 +246,10 @@

Internal changes ⚙️

+* A new compiler pass `mark-entry-points` has been added, which annotates externally-callable + entry functions with the `catalyst.entry_point` attribute. + [(#2899)](https://github.com/PennyLaneAI/catalyst/pull/2899) + * The `dim` argument of the `quantum.pcphase` operation has been changed to a static integer attribute (previously a dynamic float operand). This allows, among other things, the decomposition graph to distinguish pcphase gates with different `dim` values, since they need different decomposition rules. diff --git a/frontend/catalyst/__init__.py b/frontend/catalyst/__init__.py index a3afe7f8c9..dc9e66a9f9 100644 --- a/frontend/catalyst/__init__.py +++ b/frontend/catalyst/__init__.py @@ -68,7 +68,7 @@ sys.modules["mlir_quantum.ir"] = __import__("jaxlib.mlir.ir").mlir.ir sys.modules["mlir_quantum._mlir_libs"] = __import__("jaxlib.mlir._mlir_libs").mlir._mlir_libs -from catalyst import debug, logging, passes +from catalyst import debug, kernel, logging, passes from catalyst.api_extensions import * from catalyst.api_extensions import __all__ as _api_extension_list from catalyst.autograph import * @@ -196,6 +196,7 @@ "CompileOptions", "debug", "draw_graph", + "kernel", "passes", "pipeline", "compile_without_static_conditionals", diff --git a/frontend/catalyst/api_extensions/__init__.py b/frontend/catalyst/api_extensions/__init__.py index 6c020597c8..5791fc95d1 100644 --- a/frontend/catalyst/api_extensions/__init__.py +++ b/frontend/catalyst/api_extensions/__init__.py @@ -40,6 +40,7 @@ measure, pauli_measure, ) +from catalyst.api_extensions.target import target __all__ = ( "accelerate", @@ -60,4 +61,5 @@ "pauli_measure", "adjoint", "ctrl", + "target", ) diff --git a/frontend/catalyst/api_extensions/target.py b/frontend/catalyst/api_extensions/target.py new file mode 100644 index 0000000000..4434cdf42d --- /dev/null +++ b/frontend/catalyst/api_extensions/target.py @@ -0,0 +1,85 @@ +# 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. + +"""Tag a PennyLane device as a separate cross-compilation target, optionally remote.""" + +from dataclasses import dataclass +from typing import Optional + +_TARGET_ATTR = "_catalyst_target" +_DISPATCH_ATTR = "_catalyst_dispatch" + + +@dataclass(frozen=True) +class Target: + """Cross-compilation target spec attached to a device by :func:`target`. + + Args: + pipeline: Optional name of a lowering pipeline registered with the compiler. + triple: Optional LLVM target triple. Defaults to the host triple. + """ + + pipeline: Optional[str] = None + triple: Optional[str] = None + + +@dataclass(frozen=True) +class RemoteDispatch: + """Remote dispatch spec attached to a device by :func:`target` with an ``address``. + + Args: + address: Executor address, e.g. ``"127.0.0.1:1373"``. + """ + + address: str + + +def target( + device, + *, + pipeline: Optional[str] = None, + triple: Optional[str] = None, + address: Optional[str] = None, +): + """Tag a PennyLane device as a separate cross-compilation target and return it. + + Any QNode wrapping the returned device is kept as a separate compilation unit which carries + ``catalyst.target = {pipeline, triple}`` and is cross-compiled to a standalone object rather + than being inlined into the host module. With ``address`` set, the object is additionally + dispatched to a remote executor (``catalyst.dispatch = {address}``); without it, the object is + statically linked and runs in-process. + + Args: + device: A PennyLane device. + pipeline: Optional lowering-pipeline name registered with the compiler. + triple: Optional LLVM target triple. Defaults to the host triple. + address: Optional executor address; when set, the target is dispatched remotely. + + Returns: + The same device, now tagged with target (and, if ``address`` is given, dispatch) metadata. + """ + setattr(device, _TARGET_ATTR, Target(pipeline=pipeline, triple=triple)) + if address is not None: + setattr(device, _DISPATCH_ATTR, RemoteDispatch(address=address)) + return device + + +def get_target(device) -> Optional[Target]: + """Return the :class:`Target` attached via :func:`target`, or ``None``.""" + return getattr(device, _TARGET_ATTR, None) + + +def get_dispatch(device) -> Optional[RemoteDispatch]: + """Return the :class:`RemoteDispatch` attached via :func:`target` (``address=...``), or ``None``.""" + return getattr(device, _DISPATCH_ATTR, None) diff --git a/frontend/catalyst/compiled_functions.py b/frontend/catalyst/compiled_functions.py index db9c196e6e..31a8b0253c 100644 --- a/frontend/catalyst/compiled_functions.py +++ b/frontend/catalyst/compiled_functions.py @@ -109,15 +109,11 @@ def load_symbols(self): return function, setup, teardown, mem_transfer def __enter__(self): - params_to_setup = [b"jitted-function"] - argc = len(params_to_setup) - array_of_char_ptrs = (ctypes.c_char_p * len(params_to_setup))() - array_of_char_ptrs[:] = params_to_setup - self.setup(ctypes.c_int(argc), array_of_char_ptrs) + wrapper.invoke_setup(self.setup, ["jitted-function"]) return self def __exit__(self, _type, _value, _traceback): - self.teardown() + wrapper.invoke_teardown(self.teardown) class CompiledFunction: diff --git a/frontend/catalyst/compiler.py b/frontend/catalyst/compiler.py index c906b09428..a3307328b7 100644 --- a/frontend/catalyst/compiler.py +++ b/frontend/catalyst/compiler.py @@ -176,11 +176,23 @@ def get_default_flags(options): "-lrt_decoder", ] + rt_remote_so = "librt_remote" + file_extension + if os.path.isfile(os.path.join(rt_lib_path, rt_remote_so)): + default_flags.append("-lrt_remote") + # If OQD runtime capi is built, link to it as well # TODO: This is not ideal and should be replaced when the compiler is device aware if os.path.isfile(os.path.join(rt_lib_path, "librt_OQD_capi" + file_extension)): default_flags.append("-lrt_OQD_capi") + for artifact_path in options.runtime_artifacts: + dir_name = os.path.dirname(artifact_path) + default_flags += [ + f"-Wl,-rpath,{dir_name}", + f"-L{dir_name}", + artifact_path, + ] + return default_flags @staticmethod @@ -210,9 +222,9 @@ def _available_compilers(fallback_compilers): yield compiler @staticmethod - def _attempt_link(compiler, flags, infile, outfile, options): + def _attempt_link(compiler, flags, infile, outfile, options, extra_objects=None): try: - command = [compiler] + flags + [infile, "-o", outfile] + command = [compiler] + flags + [infile] + list(extra_objects or []) + ["-o", outfile] run_writing_command(command, options) return True except subprocess.CalledProcessError as e: @@ -239,7 +251,9 @@ def get_output_filename(infile): @staticmethod @debug_logger - def run(infile, outfile=None, flags=None, fallback_compilers=None, options=None): + def run( + infile, outfile=None, flags=None, fallback_compilers=None, options=None, extra_objects=None + ): """ Link the infile against the necessary libraries and produce the outfile. @@ -249,6 +263,8 @@ def run(infile, outfile=None, flags=None, fallback_compilers=None, options=None) flags (Optional[List[str]]): flags to be passed down to the compiler fallback_compilers (Optional[List[str]]): name of executables to be looked for in PATH compile_options (Optional[CompileOptions]): generic compilation options. + extra_objects (Optional[List[str]]): additional object files to statically link + (e.g. locally cross-compiled catalyst.target kernels). Raises: EnvironmentError: The exception is raised when no compiler succeeded. """ @@ -261,7 +277,9 @@ def run(infile, outfile=None, flags=None, fallback_compilers=None, options=None) if fallback_compilers is None: fallback_compilers = LinkerDriver._default_fallback_compilers for compiler in LinkerDriver._available_compilers(fallback_compilers): - success = LinkerDriver._attempt_link(compiler, flags, infile, outfile, options) + success = LinkerDriver._attempt_link( + compiler, flags, infile, outfile, options, extra_objects + ) if options.verbose: print("Shared object linking successful", file=options.logfile) if success: @@ -498,7 +516,16 @@ def run_from_ir(self, ir: str, module_name: str, workspace: Directory): out_IR = None if self.options.link: - output = LinkerDriver.run(output_object_name, options=self.options) + # Locally cross-compiled catalyst.target kernels are emitted as separate objects; the + # driver lists them in `{module}.objects` for the linker to statically include. + extra_objects = [] + objects_manifest = os.path.join(str(workspace), f"{module_name}.objects") + if os.path.exists(objects_manifest): + with open(objects_manifest, "r", encoding="utf-8") as f: + extra_objects = [line.strip() for line in f if line.strip()] + output = LinkerDriver.run( + output_object_name, options=self.options, extra_objects=extra_objects + ) output = str(pathlib.Path(output).absolute()) else: output = None @@ -706,41 +733,26 @@ def get_output_of(self, pipeline, workspace) -> Optional[str]: Returns (Optional[str]): output IR """ - 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: + workspace = str(workspace) + filenames = [f for f in os.listdir(workspace) if f.endswith(".mlir") or f.endswith(".ll")] + if not filenames: 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 file_content + + 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 in filenames if pipeline in os.path.splitext(f)[0]] + 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) + + with open(os.path.join(workspace, selected_file[0]), "r", encoding="utf-8") as file: + return file.read() diff --git a/frontend/catalyst/debug/__init__.py b/frontend/catalyst/debug/__init__.py index e70e96027f..798254cba4 100644 --- a/frontend/catalyst/debug/__init__.py +++ b/frontend/catalyst/debug/__init__.py @@ -17,6 +17,7 @@ from catalyst.debug.callback import callback from catalyst.debug.compiler_functions import ( compile_executable, + compile_mlir, get_cmain, get_compilation_stage, get_compilation_stages_groups, @@ -27,6 +28,7 @@ __all__ = ( "callback", + "compile_mlir", "print", "print_memref", "get_compilation_stage", diff --git a/frontend/catalyst/debug/compiler_functions.py b/frontend/catalyst/debug/compiler_functions.py index f5094532a5..c85919b4e9 100644 --- a/frontend/catalyst/debug/compiler_functions.py +++ b/frontend/catalyst/debug/compiler_functions.py @@ -18,16 +18,24 @@ import logging import os +import pathlib import platform import re import shutil import subprocess +from jax._src.lib.mlir import ir + import catalyst -from catalyst.compiler import LinkerDriver +from catalyst.compiled_functions import CompiledFunction +from catalyst.compiler import Compiler, LinkerDriver from catalyst.logging import debug_logger +from catalyst.pipelines import CompileOptions from catalyst.tracing.contexts import EvaluationContext from catalyst.tracing.type_signatures import filter_static_args, promote_arguments +from catalyst.utils.filesystem import WorkspaceManager +from catalyst.utils.runtime_artifacts import collect_runtime_artifacts_from_text +from catalyst.utils.types import convert_numpy_dtype_to_mlir logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -55,16 +63,17 @@ def get_compilation_stage(fn, stage): .. note:: In order to use this function, ``keep_intermediate=True`` must be - set in the :func:`~.qjit` decorator of the input function. + set in the :func:`~.qjit` decorator of the input function, or + passed to :func:`~.compile_mlir`. Args: - fn (QJIT): a qjit-decorated function + fn (QJIT | CompiledMLIR): a qjit-decorated function or a :func:`~.compile_mlir` result stage (str): string corresponding with the name of the stage to be printed Returns: str: output ir from the target compiler stage - .. seealso:: :doc:`/dev/debugging`, :func:`~.replace_ir`. + .. seealso:: :doc:`/dev/debugging`, :func:`~.replace_ir`, :func:`~.compile_mlir`. **Example** @@ -92,8 +101,11 @@ def func(x: float): """ EvaluationContext.check_is_not_tracing("C interface cannot be generated from tracing context.") - if not isinstance(fn, catalyst.QJIT): - raise TypeError(f"First argument needs to be a 'QJIT' object, got a {type(fn)}.") + if not isinstance(fn, (catalyst.QJIT, CompiledMLIR)): + raise TypeError( + f"First argument needs to be a 'QJIT' object or 'CompiledMLIR' instance," + f" got a {type(fn)}." + ) return fn.compiler.get_output_of(stage, fn.workspace) @@ -307,3 +319,96 @@ def f(x): ) return output_file + + +class CompiledMLIR: + """Thin wrapper around a compiled shared object returned by :func:`~.compile_mlir`. + + Call it like a regular Python function to execute the compiled code. + Pass it to :func:`~.get_compilation_stage` to inspect intermediate IR. + """ + + def __init__(self, compiled_function, compiler, workspace): + self.compiled_function = compiled_function + self.compiler = compiler + self.workspace = workspace + + def __call__(self, *args): + return self.compiled_function(*args) + + +@debug_logger +def compile_mlir(mlir_source, *, func_name, result_types, **options): + r"""Compile a standalone MLIR module and return a callable. + + Reads an MLIR module from a file path or a raw string, runs it through the full + Catalyst compilation pipeline (``catalyst-cli``), and returns a :class:`CompiledMLIR` + object that can be called like a regular Python function. + + Note that the MLIR module structure should adhere to the Catalyst runtime conventions. + Therefore the module must contain ``func.func @setup()`` and ``func.func @teardown()`` lifecycle + hooks that call ``quantum.init`` / ``quantum.finalize`` (required by the Catalyst runtime + even when no quantum device is used). + + .. warning:: + + Input argument shapes and dtypes are not validated before compilation. Passing arguments + that do not match the compiled function's signature could produce incorrect results + or undefined behavior. + + Args: + mlir_source (str | pathlib.Path): Path to a ``.mlir`` file, or the raw MLIR text. + func_name (str): Name of the entry-point function (without ``@``), e.g. ``"jit_my_func"``. + The function must be annotated with ``{llvm.emit_c_interface}``. + result_types (Sequence[jax.ShapeDtypeStruct]): Output tensor types, in order. + Each element must have ``.shape`` (tuple of ints) and ``.dtype`` (numpy dtype). + **options: Any :class:`~catalyst.CompileOptions` keyword argument + + Returns: + CompiledMLIR: A callable that accepts the same arguments as the compiled function + + **Example** + + .. code-block:: python + + import numpy as np + import jax + from catalyst.debug import compile_mlir, get_compilation_stage + + fn = compile_mlir( + "my_module.mlir", + func_name="jit_my_func", + result_types=[jax.ShapeDtypeStruct((), np.float64)], + keep_intermediate=True, + ) + + result = fn(np.float64(3.14)) + + print(get_compilation_stage(fn, "QuantumCompilationStage")) + """ + if isinstance(mlir_source, pathlib.Path): + ir_text = mlir_source.read_text(encoding="utf-8") + else: + ir_text = mlir_source + + compile_options = CompileOptions(**options) + + if not compile_options.runtime_artifacts: + collect_runtime_artifacts_from_text(ir_text, compile_options) + + keep = compile_options.keep_intermediate + preferred_dir = os.getcwd() if keep else None + workspace = WorkspaceManager.get_or_create_workspace(func_name, preferred_dir) + + compiler = Compiler(options=compile_options) + shared_object, _ = compiler.run_from_ir(ir_text, func_name, workspace) + + ctx = ir.Context() + with ctx, ir.Location.unknown(): + restype = [ + ir.RankedTensorType.get(list(s.shape), convert_numpy_dtype_to_mlir(s.dtype)) + for s in result_types + ] + + compiled_fn = CompiledFunction(shared_object, func_name, restype, None, compile_options) + return CompiledMLIR(compiled_fn, compiler, workspace) diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index df3b80cb49..2c94ad7b26 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -60,6 +60,7 @@ # pylint: disable=ungrouped-imports from catalyst.jax_extras.patches import mock_attributes from catalyst.utils.patching import Patcher +from catalyst.utils.runtime_artifacts import record_runtime_artifact with Patcher( ( @@ -76,6 +77,7 @@ AssertionOp, CallbackCallOp, CallbackOp, + CustomCallOp, PrintOp, SymbolicArrayOp, ) @@ -349,6 +351,8 @@ class Folding(Enum): quantum_kernel_p.multiple_results = True decomprule_p = core.Primitive("decomposition_rule") decomprule_p.multiple_results = True +runtime_call_p = Primitive("runtime_call") +runtime_call_p.multiple_results = True def decomposition_rule(func=None, *, is_qreg=True, num_params=0, pauli_word=None, op_type=None): @@ -546,6 +550,42 @@ def _python_callback_lowering( return retval +# +# runtime_call +# +@runtime_call_p.def_abstract_eval +def _runtime_call_abstract_eval(*avals, kernel_descriptor): + """Infer output shapes/dtypes from the kernel descriptor.""" + return tuple( + core.ShapedArray(shape, np.dtype(dtype_str)) + for shape, dtype_str in kernel_descriptor.output_spec + ) + + +@runtime_call_p.def_impl +def _runtime_call_impl(*args, **kwargs): # pragma: no cover + raise NotImplementedError() + + +def _runtime_call_lowering(jax_ctx: mlir.LoweringRuleContext, *args, kernel_descriptor): + """Lower runtime_call to catalyst.custom_call and record the artifact on the enclosing module. + The artifact path is written as a catalyst.runtime_artifacts ArrayAttr on calling module. + """ + results_ty = list(convert_shaped_arrays_to_tensors(jax_ctx.avals_out)) + + if kernel_descriptor.remote: + call_op = CustomCallOp(results_ty, list(args), kernel_descriptor.name) + address = kernel_descriptor.remote_address or "" + call_op.operation.attributes["backend_config"] = ir.DictAttr.get( + {"dispatch": ir.StringAttr.get(address)} + ) + return call_op.results + + call_op = CustomCallOp(results_ty, list(args), kernel_descriptor.name) + record_runtime_artifact(jax_ctx.module_context.module.operation, kernel_descriptor.artifact) + return call_op.results + + # # print # @@ -3129,6 +3169,7 @@ def _symbolic_array_lowering(ctx, *, shape, dtype): (print_p, _print_lowering), (assert_p, _assert_lowering), (python_callback_p, _python_callback_lowering), + (runtime_call_p, _runtime_call_lowering), (value_and_grad_p, _value_and_grad_lowering), (set_state_p, _set_state_lowering), (set_basis_state_p, _set_basis_state_lowering), diff --git a/frontend/catalyst/jax_primitives_utils.py b/frontend/catalyst/jax_primitives_utils.py index 1365b45371..d282c3b1a8 100644 --- a/frontend/catalyst/jax_primitives_utils.py +++ b/frontend/catalyst/jax_primitives_utils.py @@ -14,6 +14,7 @@ """This module contains some helper functions for translating JAX primitives to MLIR.""" import copy +import dataclasses import functools import pennylane as qp @@ -26,6 +27,7 @@ from mlir_quantum.dialects.catalyst import LaunchKernelOp from pennylane.transforms.core import BoundTransform +from catalyst.api_extensions.target import get_dispatch, get_target from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval from catalyst.passes import PassPlugin @@ -244,8 +246,16 @@ def lower_qnode_to_funcop(ctx, callable_, call_jaxpr, pipelines): assert isinstance(callable_, qp.QNode), "This function expects qnodes" name = "module_" + callable_.__name__ + target = get_target(callable_.device) + dispatch = get_dispatch(callable_.device) # pylint: disable-next=no-member with NestedModule(ctx, name) as module, ir.InsertionPoint(module.regions[0].blocks[0]) as ip: + if target is not None: + fields = {k: v for k, v in dataclasses.asdict(target).items() if v is not None} + module.operation.attributes["catalyst.target"] = get_mlir_attribute_from_pyval(fields) + if dispatch is not None: + fields = {k: v for k, v in dataclasses.asdict(dispatch).items() if v is not None} + module.operation.attributes["catalyst.dispatch"] = get_mlir_attribute_from_pyval(fields) transform_module_lowering(ctx, pipelines) ctx.module_context.ip = ip func_op = get_or_create_funcop(ctx, callable_, call_jaxpr, pipelines) diff --git a/frontend/catalyst/jit.py b/frontend/catalyst/jit.py index 8d137a8268..18e1e7f6d5 100644 --- a/frontend/catalyst/jit.py +++ b/frontend/catalyst/jit.py @@ -59,6 +59,7 @@ from catalyst.utils.filesystem import WorkspaceManager from catalyst.utils.gen_mlir import inject_functions from catalyst.utils.patching import Patcher +from catalyst.utils.runtime_artifacts import collect_runtime_artifacts logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -867,6 +868,8 @@ def generate_ir(self): self.jaxpr, self.__name__, get_arg_names(self.jaxpr.in_avals, self.original_function) ) + collect_runtime_artifacts(mlir_module, self.compile_options) + # Inject Runtime Library-specific functions (e.g. setup/teardown). inject_functions(mlir_module, ctx, self.compile_options.seed) diff --git a/frontend/catalyst/kernel.py b/frontend/catalyst/kernel.py new file mode 100644 index 0000000000..1868639bed --- /dev/null +++ b/frontend/catalyst/kernel.py @@ -0,0 +1,177 @@ +# 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. + +"""Declarations and calls for pre-compiled external kernels callable from @qjit programs.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Optional + +import jax.numpy as jnp + + +@dataclass(frozen=True) +class KernelDescriptor: + """Describes the ABI of an external kernel callable via :func:`runtime_call`. + + One of two shapes, distinguished by ``artifact``: + * **local kernel** — ``artifact`` is a path to the shared library; ``remote_address`` is None. + * **remote symbol** — ``artifact`` is None (the symbol lives on the executor); + ``remote_address`` is the bound executor (``"host:port"``) or None to inherit the program's + single executor. + + Attributes: + name: C symbol name. For a local kernel it is exported by ``artifact``; for a + remote kernel it is a symbol already loaded on the executor. + artifact: Absolute path to the shared library (.so / .dylib), or ``None`` for a + remote symbol (which lives on the executor, not in a local artifact). + output_spec: Tuple of (shape_tuple, dtype_str) for each output tensor. + remote_address: For a remote kernel, the executor address it targets, e.g. + ``"ADDR:PORT"``; ``None`` means "inherit the program's single + executor". Must be ``None`` for a local kernel. + """ + + name: str + artifact: Optional[str] # absolute path to .so / .dylib, or None for a remote symbol + output_spec: tuple # ((shape_tuple, dtype_str), ...) + remote_address: Optional[str] = None # bound executor for a remote symbol; None => inherit + + @property + def remote(self) -> bool: + """True iff this is an executor-side symbol (no local artifact).""" + return self.artifact is None + + def __post_init__(self): + # `remote_address` only makes sense for a remote symbol; a local kernel has no executor. + if self.remote_address is not None and self.artifact is not None: + raise ValueError( + "kernel: `remote_address` is only valid for a remote kernel (artifact=None)" + ) + + +def _to_hashable(spec): + """Convert a ShapeDtypeStruct or tuple of them to a hashable spec tuple.""" + specs = (spec,) if hasattr(spec, "shape") else spec + entries = [] + for s in specs: + shape = tuple(s.shape) + if any(d is None for d in shape): + raise ValueError( + f"kernel: dynamic shapes unsupported; got shape {shape}. " + "All dimensions must be static integers." + ) + entries.append((shape, str(s.dtype))) + return tuple(entries) + + +def declare( + name: str, artifact: Optional[str] = None, outputs=None, *, remote=False +) -> KernelDescriptor: + """Declare an external kernel for use with :func:`kernel.runtime_call`. + + Args: + name: C symbol name. For a local kernel, exported by ``artifact``; for a remote + kernel, a symbol already loaded on the executor (e.g. + ``"fpga_trampoline_a_setup"``). + artifact: Path to the shared library for a local kernel (resolved relative to + ``os.getcwd()`` if not absolute; must exist at declare time). Omit for a + remote kernel. + outputs: :class:`jax.ShapeDtypeStruct` or tuple of them describing each output tensor. + JAX needs these at trace time to infer what the call returns. + remote: Mark the symbol as executor-side. Pass the **remote device** + (``target(..., address=...)``) to bind this call explicitly to that + executor's address. ``True`` inheriting the program's single remote executor. + + Returns: + A :class:`KernelDescriptor` + """ + output_spec = _to_hashable(outputs) if outputs is not None else () + + if remote: + address = None + if remote is not True: + from catalyst.api_extensions.target import get_dispatch + + dispatch = get_dispatch(remote) + if dispatch is None: + raise ValueError( + "kernel.declare(remote=): the device has no remote dispatch; create it " + "with target(..., address=...) first, or pass remote=True to inherit " + "the program's single executor." + ) + address = dispatch.address + # Remote symbol: no local artifact; `remote` is derived from artifact=None. + return KernelDescriptor( + name=name, artifact=None, output_spec=output_spec, remote_address=address + ) + + if artifact is None: + raise ValueError( + "kernel.declare: a local kernel requires `artifact`; pass remote=True for an " + "executor-side symbol" + ) + artifact = os.path.abspath(artifact) + if not os.path.isfile(artifact): + raise FileNotFoundError(f"kernel.declare: artifact not found: {artifact!r}") + + return KernelDescriptor(name=name, artifact=artifact, output_spec=output_spec) + + +def define(builder, *, name: Optional[str] = None, outputs): + """Build a kernel with ``builder`` and declare it, as a single decorator. + + Args: + builder: A backend-specific object implementing ``build(kernel_fn, *, name) -> path``, + where ``path`` points to a shared library exporting ``name`` with the + :func:`runtime_call` ABI. + name: Symbol the artifact must export. Defaults to ``kernel_fn.__name__``; + passed to both ``builder.build`` and :func:`declare`. + outputs: :class:`jax.ShapeDtypeStruct` or tuple of them, forwarded to :func:`declare`. + + Returns: + KernelDescriptor: the declared kernel. + """ + + def wrap(kernel_fn): + sym = name or getattr(kernel_fn, "__name__", None) + artifact = builder.build(kernel_fn, name=sym) + return declare(sym, artifact=str(artifact), outputs=outputs) + + return wrap + + +def runtime_call(kernel_descriptor, *args): + """Call an external kernel from inside ``@qjit`` — local or remote. + + If ``kernel_descriptor.remote`` is set, the symbol is dispatched on the program's remote + executor (rewritten to a ``remote.call``); otherwise it calls the local artifact. + + Args: + kernel_descriptor: A :class:`KernelDescriptor` returned by :func:`declare` / :func:`define`. + *args: Input tensors. JAX infers shapes and dtypes at trace time. + + Returns: + tuple: Output tensors. For a single output, unpack explicitly:: + + (result,) = kernel.runtime_call(my_kernel, data) + + Raises: + NotImplementedError: On use under ``jax.grad`` or ``jax.vmap``. + """ + from catalyst.jax_primitives import runtime_call_p # pylint: disable=import-outside-toplevel + + jax_args = [jnp.asarray(a) for a in args] + return tuple(runtime_call_p.bind(*jax_args, kernel_descriptor=kernel_descriptor)) diff --git a/frontend/catalyst/pipelines.py b/frontend/catalyst/pipelines.py index 241b37a42d..fded1038f9 100644 --- a/frontend/catalyst/pipelines.py +++ b/frontend/catalyst/pipelines.py @@ -170,6 +170,7 @@ class CompileOptions: dialect_plugins: Optional[Set[Path]] = None capture: bool | Literal["global"] = "global" skip_preprocess: bool = False + runtime_artifacts: tuple = () def __post_init__(self): # Convert keep_intermediate to Enum diff --git a/frontend/catalyst/utils/gen_mlir.py b/frontend/catalyst/utils/gen_mlir.py index a90b2a762d..8888e2d59f 100644 --- a/frontend/catalyst/utils/gen_mlir.py +++ b/frontend/catalyst/utils/gen_mlir.py @@ -62,8 +62,9 @@ def inject_functions(module, ctx, seed): """ This function appends functions to the input module. """ - # Add C interface for the quantum function. - module.body.operations[0].attributes["llvm.emit_c_interface"] = ir.UnitAttr.get(context=ctx) + # Add the C interface for the quantum function so it is callable across the C ABI. + mlir_qfunc = module.body.operations[0] + mlir_qfunc.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get(context=ctx) setup_module = gen_setup(ctx, seed) setup_func = setup_module.body.operations[0] @@ -73,5 +74,4 @@ def inject_functions(module, ctx, seed): teardown_func = teardown_module.body.operations[0] module.body.append(teardown_func) - mlir_qfunc = module.body.operations[0] assert isinstance(mlir_qfunc, FuncOp) diff --git a/frontend/catalyst/utils/libcustom_calls.cpp b/frontend/catalyst/utils/libcustom_calls.cpp index bc02ffc670..933d4e4ea6 100644 --- a/frontend/catalyst/utils/libcustom_calls.cpp +++ b/frontend/catalyst/utils/libcustom_calls.cpp @@ -29,6 +29,7 @@ struct EncodedMemref { int64_t rank; void *data_aligned; int8_t dtype; + int64_t *sizes; }; #define DEFINE_LAPACK_FUNC(FUNC_NAME, DATA_SIZE, OUT_SIZE, KERNEL) \ diff --git a/frontend/catalyst/utils/runtime_artifacts.py b/frontend/catalyst/utils/runtime_artifacts.py new file mode 100644 index 0000000000..ddd6e5b9db --- /dev/null +++ b/frontend/catalyst/utils/runtime_artifacts.py @@ -0,0 +1,80 @@ +# 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. + +"""Helpers for the ``catalyst.runtime_artifacts`` module attribute. + +The attribute records shared-library paths that compiled programs need at link time. The write +side is used by JAX primitive lowering (when emitting a ``catalyst.custom_call`` that targets an +external symbol); the read side is used by the compiler driver to populate the link command. +""" + +import re + +from jax._src.lib.mlir import ir + +RUNTIME_ARTIFACTS_ATTR = "catalyst.runtime_artifacts" + + +def record_runtime_artifact(module_op, artifact_path): + """Append `artifact_path` to the module's `catalyst.runtime_artifacts` attr.""" + attrs = module_op.attributes + existing = ( + [ir.StringAttr(a).value for a in attrs[RUNTIME_ARTIFACTS_ATTR]] + if RUNTIME_ARTIFACTS_ATTR in attrs + else [] + ) + if artifact_path in existing: + return + existing.append(artifact_path) + attrs[RUNTIME_ARTIFACTS_ATTR] = ir.ArrayAttr.get([ir.StringAttr.get(p) for p in existing]) + + +def collect_runtime_artifacts(mlir_module, compile_options): + """Walk all nested modules and collect artifact paths into compile_options. + + Looks for catalyst.runtime_artifacts ArrayAttr on all nested modules and aggregates them so + that the linker receives the full set of artifacts. + """ + seen = set() + + def _walk(op): + attrs = op.attributes + if RUNTIME_ARTIFACTS_ATTR in attrs: + for string_attr in attrs[RUNTIME_ARTIFACTS_ATTR]: + seen.add(ir.StringAttr(string_attr).value) + for region in op.regions: + for block in region: + for child_op in block: + _walk(child_op) + + _walk(mlir_module.operation) + compile_options.runtime_artifacts = tuple(seen) + + +_RUNTIME_ARTIFACTS_TEXT_RE = re.compile(rf"{re.escape(RUNTIME_ARTIFACTS_ATTR)}\s*=\s*\[([^\]]*)\]") +_RUNTIME_ARTIFACTS_PATH_RE = re.compile(r'"([^"]*)"') + + +def collect_runtime_artifacts_from_text(ir_text, compile_options): + """Text-mode sibling of `collect_runtime_artifacts`. + + Used by entry points that only have raw MLIR text (e.g. `compile_mlir`), where parsing into + an `ir.Module` would require registering every Catalyst/StableHLO/etc. dialect on a fresh + context. Scans the text for all `catalyst.runtime_artifacts = [...]` occurrences (handling + nested modules) and populates `compile_options.runtime_artifacts`. + """ + seen = set() + for match in _RUNTIME_ARTIFACTS_TEXT_RE.finditer(ir_text): + seen.update(_RUNTIME_ARTIFACTS_PATH_RE.findall(match.group(1))) + compile_options.runtime_artifacts = tuple(seen) diff --git a/frontend/catalyst/utils/wrapper.cpp b/frontend/catalyst/utils/wrapper.cpp index f9e5c29b8f..38548e2e33 100644 --- a/frontend/catalyst/utils/wrapper.cpp +++ b/frontend/catalyst/utils/wrapper.cpp @@ -13,8 +13,12 @@ // limitations under the License. #include +#include +#include #include "nanobind/nanobind.h" +#include "nanobind/stl/string.h" +#include "nanobind/stl/vector.h" // TODO: Periodically check and increment version. // https://endoflife.date/numpy @@ -230,6 +234,36 @@ nb::list wrap(nb::object func, nb::tuple py_args, nb::object result_desc, nb::ob return returns; } +template static FnPtr extract_fn_ptr(nb::object fn_obj) +{ + auto ctypes = nb::module_::import_("ctypes"); + return *reinterpret_cast(nb::cast(ctypes.attr("addressof")(fn_obj))); +} + +void invoke_setup(nb::object setup_fn, std::vector argv) +{ + using setup_fn_t = void (*)(int, char **); + setup_fn_t fn = extract_fn_ptr(setup_fn); + + std::vector argv_c; + argv_c.reserve(argv.size()); + for (auto &s : argv) { + argv_c.push_back(s.data()); + } + + nb::gil_scoped_release lock; + fn(static_cast(argv.size()), argv_c.data()); +} + +void invoke_teardown(nb::object teardown_fn) +{ + using teardown_fn_t = void (*)(); + teardown_fn_t fn = extract_fn_ptr(teardown_fn); + + nb::gil_scoped_release lock; + fn(); +} + NB_MODULE(wrapper, m) { m.doc() = "wrapper module"; @@ -237,6 +271,12 @@ NB_MODULE(wrapper, m) // See https://nanobind.readthedocs.io/en/latest/functions.html#none-arguments m.def("wrap", &wrap, "A wrapper function.", nb::arg("func"), nb::arg("py_args"), nb::arg("result_desc").none(), nb::arg("transfer"), nb::arg("numpy_arrays")); + + m.def("invoke_setup", &invoke_setup, "Call the JIT'd setup(argc, argv)", nb::arg("setup_fn"), + nb::arg("argv")); + + m.def("invoke_teardown", &invoke_teardown, "Call the JIT'd teardown()", nb::arg("teardown_fn")); + int retval = _import_array(); bool success = retval >= 0; if (!success) { diff --git a/frontend/test/lit/test_target.py b/frontend/test/lit/test_target.py new file mode 100644 index 0000000000..e13ecaffe7 --- /dev/null +++ b/frontend/test/lit/test_target.py @@ -0,0 +1,60 @@ +# 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: %PYTHON %s | FileCheck %s +"""catalyst.target attaches catalyst.target to the nested module.""" + +import pennylane as qp + +import catalyst +from catalyst import qjit, target + + +def test_full_target(): + """All target fields should appear in the emitted DictAttr.""" + + dev = target( + qp.device("null.qubit", wires=1), + pipeline="my-pipeline", + triple="my-triple", + ) + + @qjit(target="mlir") + @qp.qnode(dev) + def circuit(): + qp.Hadamard(0) + return qp.state() + + # CHECK: module @module_circuit attributes {catalyst.target = {pipeline = "my-pipeline", triple = "my-triple"}} + print(circuit.mlir) + + +test_full_target() + + +def test_single_field(): + """Only one field is set; the others must be absent from the dict.""" + + dev = target(qp.device("null.qubit", wires=1), triple="my-triple") + + @qjit(target="mlir") + @qp.qnode(dev) + def circuit_min(): + return qp.state() + + # CHECK: module @module_circuit_min attributes {catalyst.target = {triple = "my-triple"}} + print(circuit_min.mlir) + + +test_single_field() diff --git a/frontend/test/pytest/python_interface/transforms/qecl/test_xdsl_convert_quantum_to_qecl.py b/frontend/test/pytest/python_interface/transforms/qecl/test_xdsl_convert_quantum_to_qecl.py index 5bae4900d7..acbd814523 100644 --- a/frontend/test/pytest/python_interface/transforms/qecl/test_xdsl_convert_quantum_to_qecl.py +++ b/frontend/test/pytest/python_interface/transforms/qecl/test_xdsl_convert_quantum_to_qecl.py @@ -714,6 +714,29 @@ def test_no_apply_t_subroutine_for_clifford_circuit( """ run_filecheck(program, quantum_to_qecl_pipeline_k_1) + def test_no_apply_t_subroutine_for_clifford_circuit( + self, run_filecheck, quantum_to_qecl_pipeline_k_1 + ): + """Test that the `apply_T` magic-state subroutine is NOT emitted for a Clifford-only + circuit (one with no T gates). + + """ + program = """ + builtin.module @module_circuit { + func.func @test_func() attributes {quantum.node} { + // CHECK: qecl.hadamard + %0 = "test.op"() : () -> !qecl.codeblock<1> + %1 = builtin.unrealized_conversion_cast %0 : !qecl.codeblock<1> to !quantum.bit + %2 = quantum.custom "Hadamard"() %1 : !quantum.bit + %3 = "test.op"(%2) : (!quantum.bit) -> !quantum.bit // To prevent DCE + return + } + // CHECK-NOT: @apply_T + // CHECK-NOT: qecl.fabricate + } + """ + run_filecheck(program, quantum_to_qecl_pipeline_k_1) + def test_gate_cnot_k_1(self, run_filecheck, quantum_to_qecl_pipeline_k_1): """Test that CNOT gates (`quantum.custom "CNOT"() ops) are converted to their corresponding `qecl.cnot` ops for k = 1. diff --git a/frontend/test/pytest/runtime_call/libxor_ref.c b/frontend/test/pytest/runtime_call/libxor_ref.c new file mode 100644 index 0000000000..09d420deee --- /dev/null +++ b/frontend/test/pytest/runtime_call/libxor_ref.c @@ -0,0 +1,37 @@ +// 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 +#include + +struct EncodedMemref { + int64_t rank; + void *data_aligned; + int8_t dtype; + int64_t *sizes; +}; + +void xor_reduce(void **args, void **results) +{ + struct EncodedMemref *in = (struct EncodedMemref *)args[0]; + struct EncodedMemref *out = (struct EncodedMemref *)results[0]; + + int8_t *in_data = (int8_t *)in->data_aligned; + int32_t *out_data = (int32_t *)out->data_aligned; + int64_t n = in->sizes[0]; + int32_t acc = 0; + for (int64_t i = 0; i < n; ++i) + acc ^= (int32_t)in_data[i]; + out_data[0] = acc; +} diff --git a/frontend/test/pytest/runtime_call/test_runtime_call.py b/frontend/test/pytest/runtime_call/test_runtime_call.py new file mode 100644 index 0000000000..0da2ef4b81 --- /dev/null +++ b/frontend/test/pytest/runtime_call/test_runtime_call.py @@ -0,0 +1,118 @@ +# 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. + +"""Tests for kernel.declare and kernel.runtime_call.""" + +# pylint: disable=redefined-outer-name + +import itertools +import os +import platform +import subprocess +import tempfile + +import jax.numpy as jnp +import pytest +from jax import ShapeDtypeStruct + +from catalyst import kernel, qjit +from catalyst.kernel import KernelDescriptor + +_C_SOURCE = os.path.join(os.path.dirname(__file__), "libxor_ref.c") +_EXT = ".so" if platform.system() != "Darwin" else ".dylib" + + +@pytest.fixture(scope="module") +def libxor_ref(): + """Compile libxor_ref.c to a shared library; return its absolute path.""" + if not os.path.isfile(_C_SOURCE): + pytest.skip(f"Reference C source not found: {_C_SOURCE}") + + with tempfile.TemporaryDirectory() as tmp_dir: + lib_path = os.path.join(tmp_dir, f"libxor_ref{_EXT}") + subprocess.run( + ["cc", "-shared", "-fPIC", "-o", lib_path, _C_SOURCE], + capture_output=True, + text=True, + check=True, + ) + yield lib_path + + +@pytest.fixture(scope="module") +def xor_reduce(libxor_ref): + """Declare the xor_reduce kernel against the compiled shared library.""" + return kernel.declare( + "xor_reduce", + artifact=libxor_ref, + outputs=ShapeDtypeStruct((1,), jnp.int32), + ) + + +class TestKernelDeclare: + """Unit tests for kernel.declare error paths and descriptor properties.""" + + def test_basic(self, xor_reduce, libxor_ref): + """KernelDescriptor fields are set correctly.""" + assert isinstance(xor_reduce, KernelDescriptor) + assert xor_reduce.name == "xor_reduce" + assert xor_reduce.artifact == libxor_ref + assert xor_reduce.output_spec == (((1,), "int32"),) + + def test_missing_artifact(self): + """declare raises FileNotFoundError for a non-existent artifact.""" + with pytest.raises(FileNotFoundError, match="artifact not found"): + kernel.declare( + "xor_reduce", + artifact="/blah/libxor.so", + outputs=ShapeDtypeStruct((1,), jnp.int32), + ) + + def test_dynamic_shape_rejected(self, libxor_ref): + """declare raises ValueError when output shapes contain None dimensions.""" + with pytest.raises(ValueError, match="dynamic shapes unsupported"): + kernel.declare( + "xor_reduce", + artifact=libxor_ref, + outputs=ShapeDtypeStruct((None,), jnp.int32), + ) + + def test_descriptor_is_hashable(self, xor_reduce): + """KernelDescriptor can be used as a dict key.""" + assert {xor_reduce: 1}[xor_reduce] == 1 + + def test_multiple_outputs(self, libxor_ref): + """declare accepts a tuple of ShapeDtypeStructs for multiple outputs.""" + desc = kernel.declare( + "xor_reduce", + artifact=libxor_ref, + outputs=(ShapeDtypeStruct((1,), jnp.int32), ShapeDtypeStruct((3,), jnp.float32)), + ) + assert len(desc.output_spec) == 2 + + +class TestRuntimeCallIntegration: + """End-to-end tests: compile, link, and execute a kernel via runtime_call.""" + + def test_xor_truth_table(self, xor_reduce): + """runtime_call produces correct XOR reduction for all 3-bit inputs.""" + + @qjit + def circuit(x): + (result,) = kernel.runtime_call(xor_reduce, x) + return result + + for bits in itertools.product([0, 1], repeat=3): + x = jnp.array(bits, dtype=jnp.int8) + assert int(circuit(x)[0]) == bits[0] ^ bits[1] ^ bits[2] diff --git a/frontend/test/pytest/test_compiler.py b/frontend/test/pytest/test_compiler.py index c624f8ccb1..4e509133a7 100644 --- a/frontend/test/pytest/test_compiler.py +++ b/frontend/test/pytest/test_compiler.py @@ -154,6 +154,41 @@ def test_options_to_cli_flags_keep_intermediate_pass(self): assert "--save-ir-after-each=pass" in flags +class TestDefaultFlags: + """Unit tests for optional libraries discovered in LinkerDriver.get_default_flags.""" + + @staticmethod + def _lib_name(stem): + """Build the platform-specific shared library file name for a library stem.""" + ext = ".so" if platform.system() == "Linux" else ".dylib" + return stem + ext + + def _patch_isfile(self, monkeypatch, overrides): + """Patch isfile to force given library file names present/absent.""" + real_isfile = os.path.isfile + overrides = {os.path.basename(name): value for name, value in overrides.items()} + + def fake_isfile(p): + name = os.path.basename(p) + if name in overrides: + return overrides[name] + return real_isfile(p) + + monkeypatch.setattr("catalyst.compiler.os.path.isfile", fake_isfile) + + def test_rt_remote_linked_when_present(self, monkeypatch): + """-lrt_remote is added when librt_remote exists in the runtime lib dir.""" + self._patch_isfile(monkeypatch, overrides={self._lib_name("librt_remote"): True}) + flags = LinkerDriver.get_default_flags(CompileOptions()) + assert "-lrt_remote" in flags + + def test_rt_remote_not_linked_when_absent(self, monkeypatch): + """-lrt_remote is omitted when librt_remote is missing.""" + self._patch_isfile(monkeypatch, overrides={self._lib_name("librt_remote"): False}) + flags = LinkerDriver.get_default_flags(CompileOptions()) + assert "-lrt_remote" not in flags + + class TestCompilerWarnings: """Test compiler's warning messages.""" diff --git a/frontend/test/pytest/test_cross_compile_targets.py b/frontend/test/pytest/test_cross_compile_targets.py new file mode 100644 index 0000000000..36e1bcbf32 --- /dev/null +++ b/frontend/test/pytest/test_cross_compile_targets.py @@ -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. + +"""Integration test for the target cross-compilation + remote-dispatch chain.""" + +import os + +import numpy as np +import pennylane as qp + +from catalyst import qjit, target +from catalyst.compiler import CompileOptions +from catalyst.debug import get_compilation_stage +from catalyst.jit import QJIT + + +def test_remote_dispatch_full_chain(): + """A ``target(..., address=...)`` QNode is cross-compiled to its own object and its host call is + rewritten for remote dispatch.""" + dev = target(qp.device("null.qubit", wires=2), address="ADDR:PORT") + + @qp.qnode(dev) + def circuit(): + qp.Hadamard(0) + qp.CNOT([0, 1]) + return qp.state() + + jitted = QJIT(circuit, CompileOptions(keep_intermediate=True, link=False)) + try: + target_object = os.path.join(str(jitted.workspace), "module_circuit", "module_circuit.o") + assert os.path.exists(target_object) + + ir = get_compilation_stage(jitted, "CrossCompileTargets") + assert "remote.open" in ir + assert "remote.send_binary" in ir + assert "remote.launch" in ir + # Teardown close is handled by the runtime at process exit, not emitted as an op. + assert "remote.close" not in ir + assert "ADDR:PORT" in ir + assert target_object in ir + assert "module @module_circuit" not in ir + finally: + jitted.workspace.cleanup() + + +def test_local_target_compiles_links_and_runs(): + """A ``target``-tagged (non-remote) QNode is cross-compiled to its own object, statically linked, + and executed in-process — its host launch_kernel is flattened to a flat call into the object.""" + # No ``address`` and no explicit triple -> host triple -> local static-link path. + dev = target(qp.device("lightning.qubit", wires=2)) + + @qjit + @qp.qnode(dev) + def circuit(x: float): + qp.RX(x, wires=0) + qp.CNOT(wires=[0, 1]) + return qp.expval(qp.PauliZ(0)) + + assert np.isclose(circuit(0.0), 1.0) + assert np.isclose(circuit(np.pi), -1.0) diff --git a/frontend/test/pytest/test_debug.py b/frontend/test/pytest/test_debug.py index b7d7f0624c..6d197b7003 100644 --- a/frontend/test/pytest/test_debug.py +++ b/frontend/test/pytest/test_debug.py @@ -16,6 +16,7 @@ import subprocess import textwrap +import jax import jax.numpy as jnp import numpy as np import pennylane as qp @@ -25,10 +26,34 @@ from catalyst import debug, qjit, value_and_grad from catalyst.compiler import _options_to_cli_flags, to_llvmir, to_mlir_opt -from catalyst.debug import compile_executable, get_cmain, get_compilation_stage, replace_ir +from catalyst.debug import ( + compile_executable, + compile_mlir, + get_cmain, + get_compilation_stage, + replace_ir, +) from catalyst.pipelines import CompileOptions from catalyst.utils.exceptions import CompileError +_IDENTITY_MLIR = """ +module @identity { + func.func public @jit_identity(%arg0: tensor) -> tensor attributes {llvm.emit_c_interface} { + return %arg0 : tensor + } + func.func @setup() { + quantum.init + return + } + func.func @teardown() { + quantum.finalize + return + } +} +""" + +_IDENTITY_RESULT_TYPES = [jax.ShapeDtypeStruct((), np.float64)] + class TestDebugPrint: """Test suite for the runtime print functionality.""" @@ -632,5 +657,100 @@ def test_catalyst_error(self): to_mlir_opt(stdin=mlir) +class TestCompileMLIR: + """Tests for compile_mlir tocompile and run a standalone MLIR module.""" + + def test_compile_from_string(self): + """compile_mlir accepts a raw MLIR string and returns a callable.""" + fn = compile_mlir( + _IDENTITY_MLIR, func_name="jit_identity", result_types=_IDENTITY_RESULT_TYPES + ) + result = fn(np.float64(1.5)) + np.testing.assert_allclose(result[0], 1.5) + + def test_get_compilation_stage(self): + """get_compilation_stage works on a CompiledMLIR object.""" + fn = compile_mlir( + _IDENTITY_MLIR, + func_name="jit_identity", + result_types=_IDENTITY_RESULT_TYPES, + keep_intermediate=True, + ) + ir = get_compilation_stage(fn, "HLOLoweringStage") + assert "@jit_identity" in ir + fn.workspace.cleanup() + + def test_no_outputs(self): + """compile_mlir handles functions with no return values.""" + no_return_mlir = """ + module @noop { + func.func public @jit_noop() attributes {llvm.emit_c_interface} { + return + } + func.func @setup() { quantum.init return } + func.func @teardown() { quantum.finalize return } + } + """ + fn = compile_mlir(no_return_mlir, func_name="jit_noop", result_types=()) + fn() + + def test_ranked_tensor(self): + """compile_mlir handles ranked tensor inputs and outputs.""" + ranked_mlir = """ + module @ranked { + func.func public @jit_ranked(%arg0: tensor<3xf64>) -> tensor<3xf64> + attributes {llvm.emit_c_interface} { + return %arg0 : tensor<3xf64> + } + func.func @setup() { quantum.init return } + func.func @teardown() { quantum.finalize return } + } + """ + fn = compile_mlir( + ranked_mlir, + func_name="jit_ranked", + result_types=[jax.ShapeDtypeStruct((3,), np.float64)], + ) + result = fn(np.array([1.0, 2.0, 3.0])) + np.testing.assert_allclose(result[0], [1.0, 2.0, 3.0]) + + def test_multiple_inputs_outputs(self): + """compile_mlir handles multiple inputs and outputs (returns them swapped).""" + multi_mlir = """ + module @multi { + func.func public @jit_multi(%arg0: tensor, %arg1: tensor) + -> (tensor, tensor) attributes {llvm.emit_c_interface} { + return %arg1, %arg0 : tensor, tensor + } + func.func @setup() { quantum.init return } + func.func @teardown() { quantum.finalize return } + } + """ + fn = compile_mlir( + multi_mlir, + func_name="jit_multi", + result_types=[ + jax.ShapeDtypeStruct((), np.float64), + jax.ShapeDtypeStruct((), np.float64), + ], + ) + result = fn(np.float64(1.0), np.float64(2.0)) + np.testing.assert_allclose(result[0], 2.0) + np.testing.assert_allclose(result[1], 1.0) + + def test_compile_from_pathlib(self, tmp_path): + """compile_mlir accepts a pathlib.Path without needing a pre-existing fixture file.""" + mlir_file = tmp_path / "identity.mlir" + mlir_file.write_text(_IDENTITY_MLIR, encoding="utf-8") + fn = compile_mlir(mlir_file, func_name="jit_identity", result_types=_IDENTITY_RESULT_TYPES) + result = fn(np.float64(4.0)) + np.testing.assert_allclose(result[0], 4.0) + + def test_bad_ir_raises(self): + """compile_mlir raises on invalid MLIR.""" + with pytest.raises(Exception): + compile_mlir("this is not valid mlir", func_name="jit_foo") + + if __name__ == "__main__": pytest.main(["-x", __file__]) diff --git a/frontend/test/pytest/test_kernel.py b/frontend/test/pytest/test_kernel.py new file mode 100644 index 0000000000..bb70f2d795 --- /dev/null +++ b/frontend/test/pytest/test_kernel.py @@ -0,0 +1,60 @@ +# 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. + +"""Unit tests for the ``catalyst.kernel`` declaration API (declare/define).""" + +import os + +import jax +import jax.numpy as jnp +import pytest + +from catalyst import kernel +from catalyst.kernel import KernelDescriptor + + +def test_declare(tmp_path): + """declare resolves the artifact to an absolute path and records the output spec.""" + artifact = tmp_path / "lib.so" + artifact.write_bytes(b"") + desc = kernel.declare("sym", str(artifact), jax.ShapeDtypeStruct((3,), jnp.int32)) + assert isinstance(desc, KernelDescriptor) + assert desc.name == "sym" + assert desc.artifact == os.path.abspath(str(artifact)) + assert desc.output_spec == (((3,), "int32"),) + + +def test_declare_missing_artifact(): + """A non-existent artifact is rejected at declare time.""" + with pytest.raises(FileNotFoundError): + kernel.declare("sym", "/no/such/lib.so", jax.ShapeDtypeStruct((1,), jnp.float32)) + + +def test_define(tmp_path): + """define builds via the backend builder and declares, inferring the symbol name.""" + artifact = tmp_path / "k.so" + artifact.write_bytes(b"") + + class Builder: + def build(self, kernel_fn, *, name): # pylint: disable=unused-argument + assert name == "my_kernel" + return str(artifact) + + @kernel.define(Builder(), outputs=jax.ShapeDtypeStruct((1,), jnp.int32)) + def my_kernel(): # pragma: no cover + pass + + assert isinstance(my_kernel, KernelDescriptor) + assert my_kernel.name == "my_kernel" + assert my_kernel.artifact == str(artifact) diff --git a/frontend/test/pytest/test_wrapper_exceptions.py b/frontend/test/pytest/test_wrapper_exceptions.py new file mode 100644 index 0000000000..ea9f5549ca --- /dev/null +++ b/frontend/test/pytest/test_wrapper_exceptions.py @@ -0,0 +1,57 @@ +# 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. +""" +Test that a RuntimeError thrown from the JIT'd setup function be caught by the wrapper. +""" + +import pytest +from jax.interpreters.mlir import ir + +from catalyst import qjit +from catalyst.utils import gen_mlir + + +@pytest.fixture +def gen_setup_with_failing_assert(monkeypatch): + """ + Patch the gen_setup function to return an MLIR module with a setup function that + executes a catalyst.assert false after quantum.init. + """ + + def patched_gen_setup(ctx, seed): + seed_attr = f" {{seed = {seed} : i32}}" if seed is not None else "" + txt = f""" +func.func @setup() -> () {{ + "quantum.init"(){seed_attr} : () -> () + %f = arith.constant false + "catalyst.assert"(%f) <{{error = "triggered from setup"}}> : (i1) -> () + return +}} +""" + return ir.Module.parse(txt, ctx) + + monkeypatch.setattr(gen_mlir, "gen_setup", patched_gen_setup) + + +def test_setup_failure_surfaces_as_runtime_error( + gen_setup_with_failing_assert, +): # pylint: disable=redefined-outer-name,unused-argument + """Test that a RuntimeError thrown from the JIT'd setup function be caught by the wrapper.""" + + @qjit + def circuit(): + return 0 + + with pytest.raises(RuntimeError, match="triggered from setup"): + circuit() 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/CMakeLists.txt b/mlir/include/CMakeLists.txt index a94e4af0f9..628ccbde97 100644 --- a/mlir/include/CMakeLists.txt +++ b/mlir/include/CMakeLists.txt @@ -10,5 +10,6 @@ add_subdirectory(QecLogical) add_subdirectory(QecPhysical) add_subdirectory(Quantum) add_subdirectory(QRef) +add_subdirectory(Remote) add_subdirectory(RTIO) add_subdirectory(Test) diff --git a/mlir/include/Catalyst/IR/CatalystOps.td b/mlir/include/Catalyst/IR/CatalystOps.td index b0d075bfa4..881515f38b 100644 --- a/mlir/include/Catalyst/IR/CatalystOps.td +++ b/mlir/include/Catalyst/IR/CatalystOps.td @@ -110,7 +110,8 @@ def CustomCallOp: Catalyst_Op<"custom_call", let arguments = (ins Variadic:$inputs, StrAttr:$call_target_name, - OptionalAttr: $number_original_arg + OptionalAttr: $number_original_arg, + OptionalAttr: $backend_config ); let results = (outs Variadic); 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 9c652e7732..86413136de 100644 --- a/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h +++ b/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h @@ -147,6 +147,7 @@ const PipelineList pipelineList{ "memref-to-llvm-tbaa", "finalize-memref-to-llvm{use-generic-functions}", "convert-index-to-llvm", + "convert-remote-to-llvm", "convert-catalyst-to-llvm", // TODO: remove this once PBC has its own pipeline "convert-pbc-to-llvm", @@ -167,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(), @@ -178,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(), @@ -188,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 " + @@ -210,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/RegisterAllPasses.h b/mlir/include/RegisterAllPasses.h index 791fe4a784..717d776a7a 100644 --- a/mlir/include/RegisterAllPasses.h +++ b/mlir/include/RegisterAllPasses.h @@ -28,6 +28,8 @@ #include "Test/Transforms/Passes.h" #include "hlo-extensions/Transforms/Passes.h" +#include "Remote/Transforms/Passes.h" + namespace catalyst { inline void registerAllPasses() @@ -43,6 +45,7 @@ inline void registerAllPasses() qecp::registerQecPhysicalPasses(); qref::registerQRefPasses(); quantum::registerQuantumPasses(); + remote::registerRemotePasses(); rtio::registerRTIOPasses(); test::registerTestPasses(); } diff --git a/mlir/include/Remote/CMakeLists.txt b/mlir/include/Remote/CMakeLists.txt new file mode 100644 index 0000000000..9f57627c32 --- /dev/null +++ b/mlir/include/Remote/CMakeLists.txt @@ -0,0 +1,2 @@ +add_subdirectory(IR) +add_subdirectory(Transforms) diff --git a/mlir/include/Remote/IR/CMakeLists.txt b/mlir/include/Remote/IR/CMakeLists.txt new file mode 100644 index 0000000000..ffd2338c1a --- /dev/null +++ b/mlir/include/Remote/IR/CMakeLists.txt @@ -0,0 +1,3 @@ +add_mlir_dialect(RemoteOps remote) +add_mlir_doc(RemoteDialect RemoteDialect Remote/ -gen-dialect-doc) +add_mlir_doc(RemoteOps RemoteOps Remote/ -gen-op-doc) diff --git a/mlir/include/Remote/IR/RemoteDialect.h b/mlir/include/Remote/IR/RemoteDialect.h new file mode 100644 index 0000000000..84b1b8953a --- /dev/null +++ b/mlir/include/Remote/IR/RemoteDialect.h @@ -0,0 +1,21 @@ +// 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. + +#pragma once + +#include "mlir/Bytecode/BytecodeOpInterface.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/OpDefinition.h" + +#include "Remote/IR/RemoteOpsDialect.h.inc" diff --git a/mlir/include/Remote/IR/RemoteDialect.td b/mlir/include/Remote/IR/RemoteDialect.td new file mode 100644 index 0000000000..e45dbe206b --- /dev/null +++ b/mlir/include/Remote/IR/RemoteDialect.td @@ -0,0 +1,48 @@ +// 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. + +#ifndef REMOTE_DIALECT +#define REMOTE_DIALECT + +include "mlir/IR/OpBase.td" +include "mlir/IR/DialectBase.td" + +//===----------------------------------------------------------------------===// +// Remote Dialect Definition +//===----------------------------------------------------------------------===// + +def Remote_Dialect : Dialect { + let summary = "Operations for dispatching kernels and calling symbols in a remote executor."; + let description = [{ + The `remote` dialect models a small set of operations that describe how a host program + interacts with a remote executor service: + + * Opening a session with a remote endpoint + * Shipping a cross-compiled kernel object file or assets to that endpoint + * Dispatching a previously-shipped qnode kernel + * Invoking an arbitrary symbol in a remote shared library + }]; + + let name = "remote"; + let cppNamespace = "::catalyst::remote"; +} + +//===----------------------------------------------------------------------===// +// Remote Operation Base +//===----------------------------------------------------------------------===// + +class Remote_Op traits = []> : + Op; + +#endif // REMOTE_DIALECT diff --git a/mlir/include/Remote/IR/RemoteOps.h b/mlir/include/Remote/IR/RemoteOps.h new file mode 100644 index 0000000000..69dff11735 --- /dev/null +++ b/mlir/include/Remote/IR/RemoteOps.h @@ -0,0 +1,25 @@ +// 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. + +#pragma once + +#include "mlir/Bytecode/BytecodeOpInterface.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" + +#include "Remote/IR/RemoteDialect.h" + +#define GET_OP_CLASSES +#include "Remote/IR/RemoteOps.h.inc" diff --git a/mlir/include/Remote/IR/RemoteOps.td b/mlir/include/Remote/IR/RemoteOps.td new file mode 100644 index 0000000000..d3cdf66fb8 --- /dev/null +++ b/mlir/include/Remote/IR/RemoteOps.td @@ -0,0 +1,120 @@ +// 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. + +#ifndef REMOTE_OPS +#define REMOTE_OPS + +include "mlir/IR/OpBase.td" +include "mlir/IR/BuiltinAttributes.td" +include "Remote/IR/RemoteDialect.td" + +//===----------------------------------------------------------------------===// +// remote.open +//===----------------------------------------------------------------------===// + +def Remote_OpenOp : Remote_Op<"open"> { + let summary = "Open a session with a remote executor."; + let description = [{ + Establish a connection to the remote executor reachable at `address` + Subsequent remote operations targeting the same `address` reuse this session. + }]; + + let arguments = (ins StrAttr:$address); + + let assemblyFormat = [{ + `(` $address `)` attr-dict + }]; +} + +//===----------------------------------------------------------------------===// +// remote.send_binary +//===----------------------------------------------------------------------===// + +def Remote_SendBinaryOp : Remote_Op<"send_binary"> { + let summary = "Ship a compiled kernel object file to the remote executor."; + let description = [{ + Transfer the compiled binary at `binary_path` to the remote executor at `address`. + The remote loads the file and exposes all of its symbols. + `format` selects the wire encoding which is handled by the runtime. + }]; + + let arguments = (ins + StrAttr:$address, + StrAttr:$binary_path, + DefaultValuedAttr:$format + ); + + let assemblyFormat = [{ + `(` $address `,` $binary_path `)` attr-dict + }]; +} + +//===----------------------------------------------------------------------===// +// remote.launch +//===----------------------------------------------------------------------===// + +def Remote_LaunchOp : Remote_Op<"launch"> { + let summary = "Dispatch a remote-compiled kernel and receive its results."; + let description = [{ + Invoke the kernel previously registered via `remote.send_binary`. + The `inputs` operand list mirrors the original host-side `func.call` + to the kernel; `results` mirrors the kernel's return values. + }]; + + let arguments = (ins + Variadic:$inputs, + StrAttr:$address, + StrAttr:$kernel_callee, + StrAttr:$object + ); + + let results = (outs Variadic:$results); + + let assemblyFormat = [{ + `(` $kernel_callee `,` $address `,` $object `)` `(` $inputs `)` attr-dict + `:` functional-type($inputs, $results) + }]; +} + +//===----------------------------------------------------------------------===// +// remote.call +//===----------------------------------------------------------------------===// + +def Remote_CallOp : Remote_Op<"call"> { + let summary = "Invoke a symbol in a previously-loaded shared library."; + let description = [{ + Invoke `symbol` in a shared library already loaded on the remote executor. + + The operand list packs input buffers followed by output buffers. + `num_input_args`, when present, marks the boundary: operands `[0, num_input_args)` + are inputs and the remainder are output buffers to be filled by the callee. + Without `num_input_args` all operands are treated as inputs. + }]; + + let arguments = (ins + Variadic:$inputs, + StrAttr:$address, + StrAttr:$symbol, + OptionalAttr:$num_input_args + ); + + let results = (outs Variadic:$results); + + let assemblyFormat = [{ + `(` $symbol `,` $address `)` `(` $inputs `)` attr-dict + `:` functional-type($inputs, $results) + }]; +} + +#endif // REMOTE_OPS diff --git a/mlir/include/Remote/Transforms/CMakeLists.txt b/mlir/include/Remote/Transforms/CMakeLists.txt new file mode 100644 index 0000000000..590e0a582d --- /dev/null +++ b/mlir/include/Remote/Transforms/CMakeLists.txt @@ -0,0 +1,4 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls -name Remote) +add_public_tablegen_target(MLIRRemotePassIncGen) +add_mlir_doc(Passes RemotePasses ./ -gen-pass-doc) diff --git a/mlir/include/Remote/Transforms/Passes.h b/mlir/include/Remote/Transforms/Passes.h new file mode 100644 index 0000000000..ef7fc110bd --- /dev/null +++ b/mlir/include/Remote/Transforms/Passes.h @@ -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. + +#pragma once + +#include "mlir/Pass/Pass.h" + +#include "Remote/IR/RemoteDialect.h" + +namespace catalyst { +namespace remote { + +#define GEN_PASS_DECL +#define GEN_PASS_REGISTRATION +#include "Remote/Transforms/Passes.h.inc" + +} // namespace remote +} // namespace catalyst diff --git a/mlir/include/Remote/Transforms/Passes.td b/mlir/include/Remote/Transforms/Passes.td new file mode 100644 index 0000000000..a0b0fb6ad0 --- /dev/null +++ b/mlir/include/Remote/Transforms/Passes.td @@ -0,0 +1,127 @@ +// 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. + +#ifndef REMOTE_PASSES +#define REMOTE_PASSES + +include "mlir/Pass/PassBase.td" + +def CrossCompileRemoteKernelsPass : Pass<"cross-compile-remote-kernels", "mlir::ModuleOp"> { + let summary = "Cross-compile each `qnode` to a `.o` and emit `remote.*` ops at the host call sites."; + let description = [{ + For every `func.func` in the host module carrying the `qnode` attribute, this pass: + 1. Clones the func into a standalone module. + 2. Sanitizes it for lowering to LLVM IR and emits the `.o` to `/.o`. + 3. Marks every `func.call` targeting the qnode with `catalyst.remote_kernel_path` + holding the absolute path of the produced `.o`. + }]; + + let dependentDialects = [ + "mlir::LLVM::LLVMDialect", + "mlir::func::FuncDialect", + "catalyst::CatalystDialect", + "catalyst::remote::RemoteDialect", + ]; + + 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=*/"target", + /*CLI arg name=*/"target", + /*type=*/"std::string", + /*default=*/"\"x86_64\"", + /*description=*/ + "LLVM target triple used for object emission." + >, + Option< + /*C++ var name=*/"cpu", + /*CLI arg name=*/"cpu", + /*type=*/"std::string", + /*default=*/"\"generic\"", + /*description=*/ + "LLVM CPU model fed to `createTargetMachine`. Defaults to " + "`generic`, which emits baseline code for the triple." + >, + Option< + /*C++ var name=*/"features", + /*CLI arg name=*/"features", + /*type=*/"std::string", + /*default=*/"\"\"", + /*description=*/ + "Comma-separated subtarget features fed to `createTargetMachine` " + "(each token must be `+feat` or `-feat`, e.g. `+crc,+aes,+sha2`). " + "Empty (default) lets the CPU choose its own feature set." + >, + Option< + /*C++ var name=*/"address", + /*CLI arg name=*/"address", + /*type=*/"std::string", + /*default=*/"\"\"", + /*description=*/ + "Executor's TCP `host:port` address for remote dispatch." + > + ]; +} + +def ConvertRemoteToLLVMPass : Pass<"convert-remote-to-llvm", "mlir::ModuleOp"> { + let summary = "Lower the `remote` dialect to direct calls into the Catalyst remote runtime."; + let description = [{ + Replaces each `remote.*` op with the corresponding `__catalyst__remote__*` + runtime entry point: + + * `remote.open` -> `__catalyst__remote__open(addr)` + * `remote.send_binary` -> `__catalyst__remote__send_binary(addr, path, format)` + * `remote.launch` -> `__catalyst__remote__launch(addr, "_catalyst_pyface_", ...)` + * `remote.call` -> `__catalyst__remote__call_wrapper(addr, sym, ...)` + `__catalyst__remote__free_result` + }]; + + let dependentDialects = [ + "mlir::LLVM::LLVMDialect", + "catalyst::remote::RemoteDialect", + ]; +} + +def DispatchRemoteTargetsPass : Pass<"dispatch-remote-targets", "mlir::ModuleOp"> { + let summary = "Ship cross-compiled catalyst.target modules to a remote executor."; + let description = [{ + For every `builtin.module` carrying a `catalyst.dispatch` attribute and the + `catalyst.object_file` path, this pass will perform the following: + + 1. Injects `remote.open` into `setup()` (once per unique address) and `remote.send_binary` + into `setup()` (once per module). + 2. Replaces every host-side `func.call` to the module's entry functions with a + `remote.launch` op carrying the executor address and the entry callee. + 3. Erases the bodyless external declarations and the nested module from the host. + + Session teardown is handled by the runtime, which closes every open session when the + process exits, so no explicit close op is emitted. + + The pass is a no-op when no `catalyst.dispatch` modules are present. + }]; + + let dependentDialects = [ + "catalyst::remote::RemoteDialect", + "catalyst::CatalystDialect", + "mlir::func::FuncDialect" + ]; +} + +#endif // REMOTE_PASSES diff --git a/mlir/lib/CMakeLists.txt b/mlir/lib/CMakeLists.txt index b4ff29b1a5..394b772188 100644 --- a/mlir/lib/CMakeLists.txt +++ b/mlir/lib/CMakeLists.txt @@ -12,5 +12,6 @@ add_subdirectory(QecLogical) add_subdirectory(QecPhysical) add_subdirectory(QRef) add_subdirectory(Quantum) +add_subdirectory(Remote) add_subdirectory(RTIO) add_subdirectory(Test) diff --git a/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp index 1e4a1ce3a2..7eafcf1719 100644 --- a/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp @@ -182,9 +182,12 @@ struct CustomCallOpInterface int32_t numArguments = static_cast(customCallOp.getNumOperands()); IntegerAttr numArgumentsAttr = rewriter.getI32IntegerAttr(numArguments); - // Create an updated custom call operation - CustomCallOp::create(rewriter, op->getLoc(), TypeRange{}, bufferArgs, - customCallOp.getCallTargetName(), numArgumentsAttr); + // Create an updated custom call operation. `backend_config` is a declared attribute, so it + // must be passed to the builder explicitly; discardable attrs are carried separately below. + 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()); bufferization::replaceOpWithBufferizedValues(rewriter, op, bufferResults); @@ -335,6 +338,95 @@ struct CallbackCallOpInterface } }; +/// Bufferization of catalyst.launch_kernel. The callee reads its operands and returns each result +/// in its own freshly allocated buffer. Bufferization therefore converts the operands to buffers +/// and the tensor results to memref results (return-by-value): no operand is written in place and +/// no result aliases an operand. +struct LaunchKernelOpInterface + : public bufferization::BufferizableOpInterface::ExternalModel { + // Each result is returned in a buffer the op allocates. + bool bufferizesToAllocation(Operation *op, Value value) const { return true; } + + // The callee reads its operands. + bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const + { + return true; + } + + // Operands are not written in place: every result is returned in a separate allocation. + bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const + { + return false; + } + + // Results are fresh allocations, so they alias none of the operands. + bufferization::AliasingValueList + getAliasingValues(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const + { + return {}; + } + + LogicalResult bufferize(Operation *op, RewriterBase &rewriter, + const bufferization::BufferizationOptions &options, + bufferization::BufferizationState &state) const + { + auto launchOp = cast(op); + + SmallVector bufferOperands; + for (Value operand : launchOp.getInputs()) { + Value buffer; + if (isa(operand.getType())) { + FailureOr opBuffer = getBuffer(rewriter, operand, options, state); + if (failed(opBuffer)) { + return failure(); + } + buffer = *opBuffer; + } + else { + buffer = operand; + } + + // The callee receives each operand as a contiguous block, addressed through the + // memref's aligned pointer with its strides/offset ignored. Copy any + // non-identity-layout operand into a fresh contiguous buffer first so the callee sees + // the intended elements. + if (auto memrefTy = dyn_cast(buffer.getType())) { + if (!memrefTy.getLayout().isIdentity()) { + MemRefType contiguousTy = + MemRefType::get(memrefTy.getShape(), memrefTy.getElementType()); + auto alloc = memref::AllocOp::create(rewriter, op->getLoc(), contiguousTy); + memref::CopyOp::create(rewriter, op->getLoc(), buffer, alloc.getResult()); + buffer = alloc.getResult(); + } + } + bufferOperands.push_back(buffer); + } + + SmallVector memrefResultTypes; + for (Type resultType : launchOp.getResultTypes()) { + if (auto tensorType = dyn_cast(resultType)) { + memrefResultTypes.push_back( + bufferization::getMemRefTypeWithStaticIdentityLayout(tensorType)); + } + else { + memrefResultTypes.push_back(resultType); + } + } + + auto newLaunchOp = LaunchKernelOp::create( + rewriter, op->getLoc(), memrefResultTypes, launchOp.getCallee(), bufferOperands, + launchOp.getArgAttrsAttr(), launchOp.getResAttrsAttr()); + // Carry over any discardable attributes, since create() only sets the declared ones. + newLaunchOp->setDiscardableAttrs(launchOp->getDiscardableAttrDictionary()); + bufferization::replaceOpWithBufferizedValues(rewriter, op, newLaunchOp.getResults()); + return success(); + } +}; + /// Bufferization of catalyst.symbolic_array. This op is a placeholder with no /// buffer semantics and is expected to be consumed before bufferization. If it /// reaches this stage, emit an informative diagnostic instead of the generic @@ -379,6 +471,7 @@ void catalyst::registerBufferizableOpInterfaceExternalModels(DialectRegistry &re PrintOp::attachInterface(*ctx); CallbackOp::attachInterface(*ctx); CallbackCallOp::attachInterface(*ctx); + LaunchKernelOp::attachInterface(*ctx); SymbolicArrayOp::attachInterface(*ctx); }); } 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..5fbae87797 --- /dev/null +++ b/mlir/lib/Catalyst/Transforms/CrossCompileTargets.cpp @@ -0,0 +1,429 @@ +// 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-remote-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 remote 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-remote) target must use the host triple '") + << llvm::sys::getDefaultTargetTriple() + << "'; use remote() 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: + // - remote (catalyst.dispatch): leave the module intact for dispatch-remote-targets. + // - local: statically link — flatten its host calls and record the object for the linker. + SmallVector localObjectFiles; + for (auto nested : targetMods) { + 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.getName().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/Catalyst/Transforms/InlineNestedModules.cpp b/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp index a3a0986cfc..c11bf9e36c 100644 --- a/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp +++ b/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp @@ -61,6 +61,10 @@ * ``` * 5. Cleanup: remove the catalyst.fully_qualified_name attribute * + * Exception: modules carrying `catalyst.target` are NOT inlined, renamed, or flattened. They are + * compiled to standalone objects elsewhere (cross-compile-targets), so their host launch_kernel is + * left in place for a later consumer. + * */ #include @@ -206,8 +210,7 @@ LogicalResult RenameFunctionsPattern::matchAndRewrite(Operation *child, { bool isSymbolTable = child->hasTrait(); bool hasBeenRenamed = child->hasAttr(hasBeenRenamedAttrName); - // TODO: isQnode - bool mustRename = isSymbolTable && !hasBeenRenamed; + bool mustRename = isSymbolTable && !hasBeenRenamed && !child->hasAttr("catalyst.target"); if (!mustRename) { return failure(); } @@ -289,8 +292,8 @@ struct InlineNestedModule : public RewritePattern { LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override { bool isSymbolTable = op->hasTrait(); - // TODO: isQnode - bool mustInline = isSymbolTable; + // Modules annotated with catalyst.target are compiled separately; skip inlining. + bool mustInline = isSymbolTable && !op->hasAttr("catalyst.target"); if (!mustInline) { return failure(); } @@ -395,13 +398,14 @@ struct NestedToFlatCallPattern : public OpRewritePatternfind(op.getCallee()) != _map->end(); - if (!found) { + // Dispatch-module entries are excluded from the map (their modules are left intact for + // dispatch-remote-targets), so their launch_kernel is not matched here and stays as-is. + auto found = _map->find(op.getCallee()); + if (found == _map->end()) { return failure(); } - auto newSymbolRefAttr = _map->find(op.getCallee())->getSecond(); - rewriter.replaceOpWithNewOp(op, newSymbolRefAttr, op.getResultTypes(), + rewriter.replaceOpWithNewOp(op, found->getSecond(), op.getResultTypes(), op.getOperands()); return success(); } @@ -417,7 +421,8 @@ struct CleanupPattern : public RewritePattern { { bool hasQualifiedName = op->hasAttr(fullyQualifiedNameAttr); bool hasQNodeAttr = op->hasAttr(quantumNodeAttr); - if (!hasQualifiedName && !hasQNodeAttr) { + bool hasBeenRenamed = op->hasAttr(hasBeenRenamedAttrName); + if (!hasQualifiedName && !hasQNodeAttr && !hasBeenRenamed) { return failure(); } @@ -429,6 +434,9 @@ struct CleanupPattern : public RewritePattern { if (hasQualifiedName) { op->removeAttr(fullyQualifiedNameAttr); } + if (hasBeenRenamed) { + op->removeAttr(hasBeenRenamedAttrName); + } }); return success(); @@ -513,22 +521,20 @@ struct InlineNestedSymbolTablePass : PassWrapper old_to_new; - for (auto ®ion : symbolTable->getRegions()) { - for (auto &block : region.getBlocks()) { - for (auto &op : block) { - if (!isa(op)) - continue; - auto hasQualifiedName = op.hasAttr(fullyQualifiedNameAttr); - if (!hasQualifiedName) - continue; - - SymbolRefAttr old = op.getAttrOfType(fullyQualifiedNameAttr); - auto symbol = cast(op); - SymbolRefAttr _new = SymbolRefAttr::get(symbol); - old_to_new.insert({old, _new}); - } - } - } + symbolTable->walk([&](Operation *nested) { + if (nested == symbolTable) + return WalkResult::advance(); + // Any module still nested at this point is a catalyst.target module (ordinary ones were + // inlined above): leave it and its contents entirely untouched. + if (isa(nested)) + return WalkResult::skip(); + if (!isa(nested) || !nested->hasAttr(fullyQualifiedNameAttr)) + return WalkResult::advance(); + SymbolRefAttr old = nested->getAttrOfType(fullyQualifiedNameAttr); + SymbolRefAttr _new = SymbolRefAttr::get(cast(nested)); + old_to_new.insert({old, _new}); + return WalkResult::advance(); + }); RewritePatternSet nestedToFlat(context); nestedToFlat.add( diff --git a/mlir/lib/Catalyst/Transforms/catalyst_to_llvm.cpp b/mlir/lib/Catalyst/Transforms/catalyst_to_llvm.cpp index 50fba06a95..2271c24558 100644 --- a/mlir/lib/Catalyst/Transforms/catalyst_to_llvm.cpp +++ b/mlir/lib/Catalyst/Transforms/catalyst_to_llvm.cpp @@ -279,18 +279,19 @@ struct AssertionOpPattern : public OpConversionPattern { // { // i64 rank, // void* data, -// i8 type_encoding +// i8 type_encoding, +// int64_t* sizes // } Value EncodeDataMemRef(Location loc, PatternRewriter &rewriter, MemRefType memrefType, Type llvmMemrefType, Value memrefLlvm) { auto ctx = rewriter.getContext(); - // Encoded memref type: !llvm.struct<(i64, ptr, i8)>. + // Encoded memref type: !llvm.struct<(i64, ptr, i8, ptr)>. Type i8 = rewriter.getI8Type(); Type i64 = rewriter.getI64Type(); Type ptr = LLVM::LLVMPointerType::get(ctx); - auto type = LLVM::LLVMStructType::getLiteral(ctx, {i64, ptr, i8}); + auto type = LLVM::LLVMStructType::getLiteral(ctx, {i64, ptr, i8, ptr}); std::optional elementDtype = encodeNumericType(memrefType.getElementType()); @@ -316,6 +317,19 @@ Value EncodeDataMemRef(Location loc, PatternRewriter &rewriter, MemRefType memre // Dtype memref = LLVM::InsertValueOp::create(rewriter, loc, memref, dtype, SmallVector{2}); + // Sizes + int64_t rankInt = memrefType.getRank(); + Type i64ArrType = LLVM::LLVMArrayType::get(i64, rankInt); + Value sizesAlloca = getStaticAlloca(loc, rewriter, i64ArrType, 1); + for (int64_t i = 0; i < rankInt; ++i) { + Value dimSize = desc.size(rewriter, loc, i); + Value gep = LLVM::GEPOp::create(rewriter, loc, ptr, i64ArrType, sizesAlloca, + SmallVector{0, static_cast(i)}); + LLVM::StoreOp::create(rewriter, loc, dimSize, gep); + } + memref = + LLVM::InsertValueOp::create(rewriter, loc, memref, sizesAlloca, SmallVector{3}); + return memref; } diff --git a/mlir/lib/Driver/CMakeLists.txt b/mlir/lib/Driver/CMakeLists.txt index c371e2e8e0..c41dc7036c 100644 --- a/mlir/lib/Driver/CMakeLists.txt +++ b/mlir/lib/Driver/CMakeLists.txt @@ -76,6 +76,8 @@ set(LIBS ion-transforms MLIRRTIO rtio-transforms + MLIRRemote + remote-transforms MLIRCatalystTest ${ENZYME_LIB} fmt::fmt diff --git a/mlir/lib/Driver/CompilerDriver.cpp b/mlir/lib/Driver/CompilerDriver.cpp index f3aac8efc9..ab64af3198 100644 --- a/mlir/lib/Driver/CompilerDriver.cpp +++ b/mlir/lib/Driver/CompilerDriver.cpp @@ -89,6 +89,8 @@ #include "RegisterAllPasses.h" +#include "Remote/IR/RemoteDialect.h" + using namespace mlir; using namespace catalyst; using namespace catalyst::driver; @@ -181,6 +183,7 @@ void registerAllCatalystDialects(DialectRegistry ®istry) registry.insert(); registry.insert(); registry.insert(); + registry.insert(); registry.insert(); registry.insert(); registry.insert(); @@ -480,6 +483,39 @@ 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-remote-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) { + for (mlir::Attribute pathAttr : objFiles) { + if (auto s = mlir::dyn_cast(pathAttr)) { + manifest << s.getValue() << "\n"; + } + } + } + } + } } return success(); } diff --git a/mlir/lib/Remote/CMakeLists.txt b/mlir/lib/Remote/CMakeLists.txt new file mode 100644 index 0000000000..9f57627c32 --- /dev/null +++ b/mlir/lib/Remote/CMakeLists.txt @@ -0,0 +1,2 @@ +add_subdirectory(IR) +add_subdirectory(Transforms) diff --git a/mlir/lib/Remote/IR/CMakeLists.txt b/mlir/lib/Remote/IR/CMakeLists.txt new file mode 100644 index 0000000000..3456fba154 --- /dev/null +++ b/mlir/lib/Remote/IR/CMakeLists.txt @@ -0,0 +1,10 @@ +add_mlir_library(MLIRRemote + RemoteDialect.cpp + RemoteOps.cpp + + ADDITIONAL_HEADER_DIRS + ${PROJECT_SOURCE_DIR}/include/Remote + + DEPENDS + MLIRRemoteOpsIncGen +) diff --git a/mlir/lib/Remote/IR/RemoteDialect.cpp b/mlir/lib/Remote/IR/RemoteDialect.cpp new file mode 100644 index 0000000000..907aa766c7 --- /dev/null +++ b/mlir/lib/Remote/IR/RemoteDialect.cpp @@ -0,0 +1,34 @@ +// 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 "Remote/IR/RemoteDialect.h" + +#include "Remote/IR/RemoteOps.h" + +using namespace mlir; +using namespace catalyst::remote; + +//===----------------------------------------------------------------------===// +// Remote Dialect +//===----------------------------------------------------------------------===// + +#include "Remote/IR/RemoteOpsDialect.cpp.inc" + +void catalyst::remote::RemoteDialect::initialize() +{ + addOperations< +#define GET_OP_LIST +#include "Remote/IR/RemoteOps.cpp.inc" + >(); +} diff --git a/mlir/lib/Remote/IR/RemoteOps.cpp b/mlir/lib/Remote/IR/RemoteOps.cpp new file mode 100644 index 0000000000..54d68707d4 --- /dev/null +++ b/mlir/lib/Remote/IR/RemoteOps.cpp @@ -0,0 +1,30 @@ +// 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 "Remote/IR/RemoteOps.h" + +#include "mlir/IR/Builders.h" +#include "mlir/IR/ImplicitLocOpBuilder.h" +#include "mlir/IR/OpImplementation.h" +#include "mlir/IR/PatternMatch.h" + +using namespace mlir; +using namespace catalyst::remote; + +//===----------------------------------------------------------------------===// +// Remote Operations +//===----------------------------------------------------------------------===// + +#define GET_OP_CLASSES +#include "Remote/IR/RemoteOps.cpp.inc" diff --git a/mlir/lib/Remote/Transforms/CMakeLists.txt b/mlir/lib/Remote/Transforms/CMakeLists.txt new file mode 100644 index 0000000000..c879fbd6e0 --- /dev/null +++ b/mlir/lib/Remote/Transforms/CMakeLists.txt @@ -0,0 +1,34 @@ +set(LIBRARY_NAME remote-transforms) + +set(LLVM_LINK_COMPONENTS + AllTargetsAsmParsers + AllTargetsCodeGens +) + +file(GLOB SRC + CrossCompileRemoteKernels.cpp + DispatchRemoteTargets.cpp + RemoteToLLVM.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} + MLIRRemote + MLIRCatalyst +) + +set(DEPENDS + MLIRRemotePassIncGen +) + +add_mlir_library(${LIBRARY_NAME} STATIC ${SRC} LINK_LIBS PRIVATE ${LIBS} DEPENDS ${DEPENDS}) +target_compile_features(${LIBRARY_NAME} PUBLIC cxx_std_20) +target_include_directories(${LIBRARY_NAME} PUBLIC + . + ${PROJECT_SOURCE_DIR}/include + ${CMAKE_BINARY_DIR}/include) diff --git a/mlir/lib/Remote/Transforms/CrossCompileRemoteKernels.cpp b/mlir/lib/Remote/Transforms/CrossCompileRemoteKernels.cpp new file mode 100644 index 0000000000..3161cb31ba --- /dev/null +++ b/mlir/lib/Remote/Transforms/CrossCompileRemoteKernels.cpp @@ -0,0 +1,315 @@ +// 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/StringExtras.h" +#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/Triple.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/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Target/LLVMIR/Export.h" + +#include "Catalyst/IR/CatalystOps.h" +#include "Driver/DefaultPipelines/DefaultPipelines.h" + +#include "Remote/IR/RemoteOps.h" +#include "Remote/Transforms/Passes.h" + +using namespace mlir; + +namespace catalyst { +namespace remote { + +#define GEN_PASS_DEF_CROSSCOMPILEREMOTEKERNELSPASS +#include "Remote/Transforms/Passes.h.inc" + +namespace { + +/// Sanitize the qnode so it can be compiled into a standalone module: +/// 1. Drop `llvm.linkage` and `qnode`. +/// 2. Promote to public visibility. +/// 3. Add `llvm.emit_c_interface` so it's callable from C as +/// `_mlir_ciface_` (and `_catalyst_pyface_` is provided by +/// downstream wrapping). +void sanitizeQNode(func::FuncOp fn) +{ + OpBuilder builder(fn.getContext()); + fn.setVisibility(SymbolTable::Visibility::Public); + fn->setAttr("llvm.emit_c_interface", builder.getUnitAttr()); + fn->setAttr("catalyst.remote_kernel", builder.getUnitAttr()); + fn->removeAttr("llvm.linkage"); + fn->removeAttr("qnode"); +} + +struct CrossCompileRemoteKernelsPass + : impl::CrossCompileRemoteKernelsPassBase { + using CrossCompileRemoteKernelsPassBase::CrossCompileRemoteKernelsPassBase; + + void runOnOperation() final + { + ModuleOp host = getOperation(); + + SmallVector qnodes = + llvm::filter_to_vector(host.getOps(), [&](func::FuncOp fn) { + return fn->hasAttr("qnode") && !fn->hasAttr("catalyst.remote_kernel"); + }); + + SmallVector calls; + host.walk([&](catalyst::CustomCallOp call) { + if (call.getCallTargetName() == "remote_call") { + calls.push_back(call); + } + }); + + if (qnodes.empty() && calls.empty()) { + return; + } + + if (!qnodes.empty()) { + if (workspace.empty()) { + host.emitError("Missing `workspace` option for remote kernel cross-compilation"); + return signalPassFailure(); + } + + llvm::InitializeAllTargetInfos(); + llvm::InitializeAllTargets(); + llvm::InitializeAllTargetMCs(); + llvm::InitializeAllAsmParsers(); + llvm::InitializeAllAsmPrinters(); + } + + injectRemoteOpenIntoSetup(host); + + for (auto qnode : qnodes) { + if (failed(compileQNode(host, qnode))) { + return signalPassFailure(); + } + } + + rewriteLegacyLibCalls(calls); + } + + /// Rewrite legacy `catalyst.custom_call fn("remote_call")` ops into + /// typed `remote.call` ops carrying the executor address. + void rewriteLegacyLibCalls(ArrayRef libCalls) + { + MLIRContext *ctx = &getContext(); + auto addressAttr = StringAttr::get(ctx, address); + for (catalyst::CustomCallOp call : libCalls) { + auto symAttr = call->getAttrOfType("catalyst.remote_symbol"); + if (!symAttr) { + call->emitOpError("legacy remote_call missing `catalyst.remote_symbol` attribute"); + return signalPassFailure(); + } + OpBuilder b(call); + IntegerAttr numInputAttr = nullptr; + if (auto n = call.getNumberOriginalArg()) { + numInputAttr = b.getI32IntegerAttr(*n); + } + auto remoteCall = + remote::CallOp::create(b, call.getLoc(), call.getResultTypes(), call.getOperands(), + /*address=*/addressAttr, /*symbol=*/symAttr, + /*num_input_args=*/numInputAttr); + call.replaceAllUsesWith(remoteCall.getResults()); + call.erase(); + } + } + + /// Insert `remote.open` into the host's `setup` function exactly once. + void injectRemoteOpenIntoSetup(ModuleOp host) + { + auto setupFn = host.lookupSymbol("setup"); + if (!setupFn || setupFn.getBody().empty()) { + return; + } + Block &setupBody = setupFn.getBody().front(); + Operation *terminator = setupBody.getTerminator(); + if (!terminator) { + return; + } + OpBuilder b(terminator); + remote::OpenOp::create(b, setupFn.getLoc(), + /*address=*/StringAttr::get(&getContext(), address)); + } + + /// Insert `remote.send_binary` into the host's `setup` function for the + /// kernel we just produced. + void injectRemoteSendBinaryIntoSetup(ModuleOp host, StringAttr addressAttr, StringAttr pathAttr) + { + auto setupFn = host.lookupSymbol("setup"); + if (!setupFn || setupFn.getBody().empty()) { + return; + } + Block &setupBody = setupFn.getBody().front(); + Operation *terminator = setupBody.getTerminator(); + if (!terminator) { + return; + } + OpBuilder b(terminator); + remote::SendBinaryOp::create(b, setupFn.getLoc(), + /*address=*/addressAttr, + /*binary_path=*/pathAttr); + } + + /// Clone the qnode into a fresh module, sanitize it, run the LLVM-dialect + /// lowering stage, then translate to LLVM IR. + std::unique_ptr loweringQNode(func::FuncOp qnode, llvm::LLVMContext &llvmCtx) + { + MLIRContext *ctx = &getContext(); + StringRef qnodeName = qnode.getName(); + + OpBuilder builder(ctx); + auto clone = ModuleOp::create(builder.getUnknownLoc(), qnodeName); + auto qnodeClone = cast(qnode->clone()); + clone.getBody()->push_back(qnodeClone); + sanitizeQNode(qnodeClone); + + PassManager nested(ctx); + std::string passList = llvm::join(catalyst::driver::getLLVMDialectLoweringStage(), ","); + if (failed(parsePassPipeline(passList, nested))) { + qnode.emitError("Failed to parse LLVM-dialect lowering pipeline"); + clone->erase(); + return nullptr; + } + if (failed(nested.run(clone))) { + qnode.emitError("Lowering failed on cloned qnode"); + clone->erase(); + return nullptr; + } + + std::unique_ptr llvmModule = + translateModuleToLLVMIR(clone, llvmCtx, /*name=*/qnodeName); + if (!llvmModule) { + qnode.emitError("Failed to translate lowered qnode to LLVM IR"); + clone->erase(); + return nullptr; + } + clone->erase(); + return llvmModule; + } + + /// Emit the LLVM module to `/.o`. + std::string emitObjectFile(std::unique_ptr &&llvmModule, StringRef qnodeName, + StringRef targetTriple, StringRef cpuModel, StringRef featureList) + { + llvm::Triple parsedTriple{targetTriple}; + std::string err; + const llvm::Target *llvmTarget = llvm::TargetRegistry::lookupTarget(parsedTriple, err); + if (!llvmTarget) { + llvm::errs() << "Target triple '" << targetTriple + << "' not registered in this LLVM build: " << err << "\n"; + return ""; + } + llvm::TargetOptions opt; + std::unique_ptr targetMachine(llvmTarget->createTargetMachine( + parsedTriple, cpuModel, featureList, opt, llvm::Reloc::Model::PIC_)); + if (!targetMachine) { + llvm::errs() << "Could not create TargetMachine for triple '" << targetTriple + << "' cpu='" << cpuModel << "' features='" << featureList << "'\n"; + return ""; + } + + targetMachine->setOptLevel(llvm::CodeGenOptLevel::Aggressive); + llvmModule->setDataLayout(targetMachine->createDataLayout()); + llvmModule->setTargetTriple(parsedTriple); + + llvm::SmallString<128> p(workspace); + llvm::sys::path::append(p, qnodeName.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 for the requested target" + << "\n"; + return ""; + } + codegenPM.run(*llvmModule); + dest.flush(); + + return objPath; + } + + /// Compile a qnode to a `.o`, then replace every host-side `func.call` to + /// it with a `remote.launch` op and inject the matching `remote.send_binary` + /// into setup. + LogicalResult compileQNode(ModuleOp host, func::FuncOp qnode) + { + MLIRContext *ctx = &getContext(); + StringRef qnodeName = qnode.getName(); + + llvm::LLVMContext llvmCtx; + auto llvmModule = loweringQNode(qnode, llvmCtx); + if (!llvmModule) { + return failure(); + } + + std::string objPath = + emitObjectFile(std::move(llvmModule), qnodeName, target, cpu, features); + if (objPath.empty()) { + return failure(); + } + + SmallVector callsToReplace; + if (auto uses = SymbolTable::getSymbolUses(qnode.getNameAttr(), host)) { + for (const SymbolTable::SymbolUse &use : *uses) { + if (auto call = dyn_cast(use.getUser())) { + callsToReplace.push_back(call); + } + } + } + + auto pathAttr = StringAttr::get(ctx, objPath); + auto addressAttr = StringAttr::get(ctx, address); + auto calleeAttr = StringAttr::get(ctx, qnodeName); + + for (func::CallOp call : callsToReplace) { + OpBuilder builder(call); + auto launch = + remote::LaunchOp::create(builder, call.getLoc(), call.getResultTypes(), + call.getOperands(), addressAttr, calleeAttr, pathAttr); + call.replaceAllUsesWith(launch.getResults()); + call.erase(); + } + + injectRemoteSendBinaryIntoSetup(host, addressAttr, pathAttr); + + qnode.erase(); + return success(); + } +}; + +} // namespace + +} // namespace remote +} // namespace catalyst diff --git a/mlir/lib/Remote/Transforms/DispatchRemoteTargets.cpp b/mlir/lib/Remote/Transforms/DispatchRemoteTargets.cpp new file mode 100644 index 0000000000..52e8c086b5 --- /dev/null +++ b/mlir/lib/Remote/Transforms/DispatchRemoteTargets.cpp @@ -0,0 +1,267 @@ +// 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 "Remote/IR/RemoteOps.h" +#include "Remote/Transforms/Passes.h" + +using namespace mlir; + +namespace catalyst { +namespace remote { + +#define GEN_PASS_DEF_DISPATCHREMOTETARGETSPASS +#include "Remote/Transforms/Passes.h.inc" + +namespace { + +// Ships cross-compiled `catalyst.target` modules to a remote executor using the `remote` dialect. +// +// This pass should be run 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 `remote.open` into `setup()` (once per unique address) and `remote.send_binary` +// into `setup()` (once per module; the object holds every entry). +// 2. Rewrites every host-side `func.call` to an entry function into a `remote.launch` op carrying +// the executor address and the entry callee. Its lowering resolves `_catalyst_pyface_` +// in the shipped object. +// 3. Erases the bodyless external declarations and the nested module from the host. +// +// Session teardown is handled by the runtime, which closes every open session when the process +// exits, so no explicit close op is emitted (matching `cross-compile-remote-kernels`). +struct DispatchRemoteTargetsPass : impl::DispatchRemoteTargetsPassBase { + using DispatchRemoteTargetsPassBase::DispatchRemoteTargetsPassBase; + + 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 (remoteDispatchOf(call)) { + libCalls.push_back(call); + } + }); + + if (targetMods.empty() && libCalls.empty()) { + return; + } + + // Modules sharing an executor get a single remote.open. + llvm::SmallSet openedAddresses; + StringAttr executorAddress; + + for (auto nested : targetMods) { + if (failed(dispatchViaOrcRemote(host, nested, openedAddresses, executorAddress))) { + return signalPassFailure(); + } + nested.erase(); + } + + const size_t numQnodeExecutors = openedAddresses.size(); + + if (!libCalls.empty()) { + for (catalyst::CustomCallOp call : libCalls) { + StringAttr dispatch = remoteDispatchOf(call); + if ((!dispatch || dispatch.getValue().empty()) && numQnodeExecutors > 1) { + call.emitOpError("ambiguous remote executor"); + return signalPassFailure(); + } + StringAttr addrAttr = libCallAddress(call, executorAddress); + if (!addrAttr) { + call.emitOpError("remote runtime_call has no executor address"); + return signalPassFailure(); + } + if (openedAddresses.insert(addrAttr.str()).second) { + injectRemoteOpenIntoSetup(host, addrAttr); + } + } + if (failed(rewriteRemoteLibCalls(libCalls, executorAddress))) { + return signalPassFailure(); + } + } + } + + static StringAttr remoteDispatchOf(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 = remoteDispatchOf(call)) { + if (!dispatch.getValue().empty()) { + return dispatch; + } + } + return fallbackAddress; + } + + LogicalResult rewriteRemoteLibCalls(ArrayRef libCalls, + StringAttr fallbackAddress) + { + MLIRContext *ctx = &getContext(); + for (catalyst::CustomCallOp call : libCalls) { + StringAttr addressAttr = libCallAddress(call, fallbackAddress); + if (!addressAttr) { + call.emitOpError("remote runtime_call 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 remoteCall = + remote::CallOp::create(b, call.getLoc(), call.getResultTypes(), call.getOperands(), + /*address=*/addressAttr, /*symbol=*/symAttr, + /*num_input_args=*/numInputAttr); + call.replaceAllUsesWith(remoteCall.getResults()); + call.erase(); + } + return success(); + } + + // Insert `remote.open` before the terminator of `setup`. No-op if setup is absent or bodyless. + void injectRemoteOpenIntoSetup(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); + remote::OpenOp::create(b, setupFn.getLoc(), addressAttr); + } + + // Insert `remote.send_binary` before the terminator of `setup`. No-op if setup is absent or + // bodyless. Always called after `injectRemoteOpenIntoSetup` so the session is opened first. + void injectRemoteSendBinaryIntoSetup(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); + remote::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 func.call into a `remote.launch`. + LogicalResult dispatchViaOrcRemote(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("remote 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("remote 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 remote lib calls reuse it. + executorAddress = addressAttr; + + // Inject remote.open (setup) once per unique address. + if (!openedAddresses.count(moduleAddress)) { + openedAddresses.insert(moduleAddress); + injectRemoteOpenIntoSetup(host, addressAttr); + } + + // Ship the object once per module: a single `catalyst.object_file` holds every + // entry function, and remote.launch resolves individual symbols within it. + injectRemoteSendBinaryIntoSetup(host, addressAttr, pathAttr); + + // Rewrite each host-side launch_kernel targeting this module into a remote.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) { + // remote.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-remote-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("remote 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, same value shipped by send_binary) keys the + // per-kernel JITDylib on the executor so the entry resolves in its own object. + auto launch = remote::LaunchOp::create( + b, launchKernel.getLoc(), launchKernel.getResultTypes(), launchKernel.getOperands(), + addressAttr, calleeAttr, pathAttr); + launchKernel.replaceAllUsesWith(launch.getResults()); + launchKernel.erase(); + } + return success(); + } +}; + +} // namespace + +} // namespace remote +} // namespace catalyst diff --git a/mlir/lib/Remote/Transforms/RemoteToLLVM.cpp b/mlir/lib/Remote/Transforms/RemoteToLLVM.cpp new file mode 100644 index 0000000000..89d2ec5029 --- /dev/null +++ b/mlir/lib/Remote/Transforms/RemoteToLLVM.cpp @@ -0,0 +1,485 @@ +// 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/Support/Path.h" +#include "mlir/Conversion/LLVMCommon/ConversionTarget.h" +#include "mlir/Conversion/LLVMCommon/MemRefBuilder.h" +#include "mlir/Conversion/LLVMCommon/TypeConverter.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "Catalyst/Utils/EnsureFunctionDeclaration.h" +#include "Catalyst/Utils/StaticAllocas.h" + +#include "Remote/IR/RemoteOps.h" +#include "Remote/Transforms/Passes.h" + +using namespace mlir; + +namespace catalyst { +namespace remote { + +#define GEN_PASS_DEF_CONVERTREMOTETOLLVMPASS +#include "Remote/Transforms/Passes.h.inc" + +namespace { + +//===----------------------------------------------------------------------===// +// Local conversion helpers +//===----------------------------------------------------------------------===// + +Value getGlobalString(Location loc, OpBuilder &rewriter, StringRef key, StringRef value, + ModuleOp mod) +{ + auto type = LLVM::LLVMArrayType::get(IntegerType::get(rewriter.getContext(), 8), value.size()); + LLVM::GlobalOp glb = mod.lookupSymbol(key); + if (!glb) { + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(mod.getBody()); + glb = LLVM::GlobalOp::create(rewriter, loc, type, true, LLVM::Linkage::Internal, key, + rewriter.getStringAttr(value)); + } + return LLVM::GEPOp::create(rewriter, loc, LLVM::LLVMPointerType::get(rewriter.getContext()), + type, LLVM::AddressOfOp::create(rewriter, loc, glb), + ArrayRef{0, 0}, LLVM::GEPNoWrapFlags::inbounds); +} + +// Get or create a `internal constant !llvm.array` global. +// And return a `!llvm.ptr` to its first element. +// +// llvm.mlir.global internal constant @(dense<[v0, v1, ...]> : tensor) +// : !llvm.array +// +// Call site: +// +// %addr = llvm.mlir.addressof @ : !llvm.ptr +// %ptr = llvm.getelementptr inbounds %addr[0, 0] +// : (!llvm.ptr) -> !llvm.ptr, !llvm.array +Value getGlobalI64Array(Location loc, OpBuilder &rewriter, StringRef key, ArrayRef values, + ModuleOp mod) +{ + Type ptrTy = LLVM::LLVMPointerType::get(rewriter.getContext()); + // Skip and return a null pointer if the array is empty. + if (values.empty()) { + return LLVM::ZeroOp::create(rewriter, loc, ptrTy); + } + Type i64Ty = rewriter.getI64Type(); + auto arrTy = LLVM::LLVMArrayType::get(i64Ty, values.size()); + LLVM::GlobalOp glb = mod.lookupSymbol(key); + // Create a new global if it doesn't exist. + if (!glb) { + auto tensorTy = RankedTensorType::get({static_cast(values.size())}, i64Ty); + auto valuesAttr = DenseElementsAttr::get(tensorTy, values); + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(mod.getBody()); + glb = LLVM::GlobalOp::create(rewriter, loc, arrTy, /*isConstant=*/true, + LLVM::Linkage::Internal, key, valuesAttr); + } + return LLVM::GEPOp::create(rewriter, loc, ptrTy, arrTy, + LLVM::AddressOfOp::create(rewriter, loc, glb), + ArrayRef{0, 0}, LLVM::GEPNoWrapFlags::inbounds); +} + +// Allocate a stack buffer of `!llvm.array`. +// +// Outputs: +// +// %slot = llvm.alloca ... : !llvm.array +// %a0 = llvm.undef : !llvm.array +// %a1 = llvm.insertvalue %ptr0, %a0[0] : !llvm.array +// %a2 = llvm.insertvalue %ptr1, %a1[1] : !llvm.array +// ... +// llvm.store %aN, %slot : !llvm.array, !llvm.ptr +// +Value buildStackPtrArray(Location loc, RewriterBase &rewriter, ArrayRef ptrs) +{ + Type ptrTy = LLVM::LLVMPointerType::get(rewriter.getContext()); + if (ptrs.empty()) { + return LLVM::ZeroOp::create(rewriter, loc, ptrTy); + } + auto arrTy = LLVM::LLVMArrayType::get(ptrTy, ptrs.size()); + Value alloca = getStaticAlloca(loc, rewriter, arrTy, 1); + Value arr = LLVM::UndefOp::create(rewriter, loc, arrTy); + for (auto [i, p] : llvm::enumerate(ptrs)) { + arr = LLVM::InsertValueOp::create(rewriter, loc, arr, p, SmallVector{(int64_t)i}); + } + LLVM::StoreOp::create(rewriter, loc, arr, alloca); + return alloca; +} + +// Calculate the byte size of one memref element. +// For complex, the size is 2 * sizeof(T). 1 for real, 1 for imaginary. +int64_t memrefElemSizeBytes(MemRefType ty) +{ + Type elem = ty.getElementType(); + if (auto cplx = dyn_cast(elem)) { + return 2 * ((cplx.getElementType().getIntOrFloatBitWidth() + 7) / 8); + } + return (elem.getIntOrFloatBitWidth() + 7) / 8; +} + +/// Byte size of a primitive element, used for the byte-buffer marshalling path. +int64_t primitiveByteSize(Type ty) +{ + if (auto i = dyn_cast(ty)) { + return (i.getWidth() + 7) / 8; + } + if (auto f = dyn_cast(ty)) { + return (f.getWidth() + 7) / 8; + } + if (isa(ty)) { + return 8; + } + if (auto c = dyn_cast(ty)) { + int64_t inner = primitiveByteSize(c.getElementType()); + return inner < 0 ? -1 : 2 * inner; + } + return -1; +} + +//===----------------------------------------------------------------------===// +// remote.open -> __catalyst__remote__open(addr) +//===----------------------------------------------------------------------===// + +struct OpenOpLowering : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(remote::OpenOp op, OpAdaptor /*adaptor*/, + ConversionPatternRewriter &rewriter) const override + { + Location loc = op.getLoc(); + MLIRContext *ctx = rewriter.getContext(); + Type ptrTy = LLVM::LLVMPointerType::get(ctx); + Type i64Ty = rewriter.getI64Type(); + ModuleOp mod = op->getParentOfType(); + + Type openSig = LLVM::LLVMFunctionType::get(i64Ty, {ptrTy}); + LLVM::LLVMFuncOp openFn = catalyst::ensureFunctionDeclaration( + rewriter, op, "__catalyst__remote__open", openSig); + + Value addrPtr = + getGlobalString(loc, rewriter, "remote_setup_addr", op.getAddress().str() + '\0', mod); + + LLVM::CallOp::create(rewriter, loc, openFn, ValueRange{addrPtr}); + rewriter.eraseOp(op); + return success(); + } +}; + +//===----------------------------------------------------------------------===// +// remote.send_binary -> __catalyst__remote__send_binary(addr, path, format) +//===----------------------------------------------------------------------===// + +struct SendBinaryOpLowering : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(remote::SendBinaryOp op, OpAdaptor /*adaptor*/, + ConversionPatternRewriter &rewriter) const override + { + Location loc = op.getLoc(); + MLIRContext *ctx = rewriter.getContext(); + Type ptrTy = LLVM::LLVMPointerType::get(ctx); + Type i32Ty = rewriter.getI32Type(); + Type i64Ty = rewriter.getI64Type(); + ModuleOp mod = op->getParentOfType(); + + Type sendBinSig = LLVM::LLVMFunctionType::get(i64Ty, {ptrTy, ptrTy, i32Ty}); + LLVM::LLVMFuncOp sendBinFn = catalyst::ensureFunctionDeclaration( + rewriter, op, "__catalyst__remote__send_binary", sendBinSig); + + std::string tag = llvm::sys::path::stem(op.getBinaryPath()).str(); + Value addrPtr = + getGlobalString(loc, rewriter, "remote_addr_" + tag, op.getAddress().str() + '\0', mod); + Value pathPtr = getGlobalString(loc, rewriter, "remote_path_" + tag, + op.getBinaryPath().str() + '\0', mod); + Value formatTag = + LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32IntegerAttr(op.getFormat())); + + LLVM::CallOp::create(rewriter, loc, sendBinFn, ValueRange{addrPtr, pathPtr, formatTag}); + rewriter.eraseOp(op); + return success(); + } +}; + +//===----------------------------------------------------------------------===// +// remote.launch -> __catalyst__remote__launch(addr, sym, +// num_in, in_descs, in_ranks, in_sizes, +// num_out, out_descs, out_ranks, out_sizes) +//===----------------------------------------------------------------------===// + +struct LaunchOpLowering : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(remote::LaunchOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override + { + Location loc = op.getLoc(); + MLIRContext *ctx = rewriter.getContext(); + Type ptrTy = LLVM::LLVMPointerType::get(ctx); + Type i64Ty = rewriter.getI64Type(); + Type voidTy = LLVM::LLVMVoidType::get(ctx); + ModuleOp mod = op->getParentOfType(); + + // parameters: + // - addr: the address of the remote executor + // - symbol: the symbol to invoke + // - num_inputs: the number of input arguments + // - input_descs: the input descriptor array + // - input_ranks: the input rank array + // - input_sizes: the input size array + // - num_outputs: the number of output arguments + // - output_descs: the output descriptor array + // - output_ranks: the output rank array + // Signature: (addr, symbol, object, num_in, in_descs, in_ranks, in_sizes, + // num_out, out_descs, out_ranks, out_sizes). `object` keys the per-kernel + // JITDylib on the executor so same-named entries in different objects don't + // collide in one symbol namespace. + Type launchSig = LLVM::LLVMFunctionType::get( + voidTy, {ptrTy, ptrTy, ptrTy, i64Ty, ptrTy, ptrTy, ptrTy, i64Ty, ptrTy, ptrTy, ptrTy}); + LLVM::LLVMFuncOp launchFn = catalyst::ensureFunctionDeclaration( + rewriter, op, "__catalyst__remote__launch", launchSig); + + std::string callee = op.getKernelCallee().str(); + Value addrPtr = getGlobalString(loc, rewriter, "remote_addr_" + callee, + op.getAddress().str() + '\0', mod); + + std::string symbolName = "_catalyst_pyface_" + callee; + Value symbolPtr = + getGlobalString(loc, rewriter, "remote_sym_" + callee, symbolName + '\0', mod); + + Value objectPtr = getGlobalString(loc, rewriter, "remote_obj_" + callee, + op.getObject().str() + '\0', mod); + + SmallVector inputDescPtrs; + SmallVector inputRanks, inputElemSizes; + for (auto [origInput, llvmInput] : llvm::zip(op.getInputs(), adaptor.getInputs())) { + auto memrefTy = cast(origInput.getType()); + inputRanks.push_back(memrefTy.getRank()); + inputElemSizes.push_back(memrefElemSizeBytes(memrefTy)); + Value alloca = getStaticAlloca(loc, rewriter, llvmInput.getType(), 1); + LLVM::StoreOp::create(rewriter, loc, llvmInput, alloca); + inputDescPtrs.push_back(alloca); + } + + SmallVector outputDescPtrs; + SmallVector outputRanks, outputElemSizes; + for (Type resultTy : op.getResultTypes()) { + auto memrefTy = cast(resultTy); + outputRanks.push_back(memrefTy.getRank()); + outputElemSizes.push_back(memrefElemSizeBytes(memrefTy)); + Type llvmDescTy = getTypeConverter()->convertType(resultTy); + Value alloca = getStaticAlloca(loc, rewriter, llvmDescTy, 1); + outputDescPtrs.push_back(alloca); + } + + Value inputDescsArr = buildStackPtrArray(loc, rewriter, inputDescPtrs); + Value outputDescsArr = buildStackPtrArray(loc, rewriter, outputDescPtrs); + + Value inputRanksArr = + getGlobalI64Array(loc, rewriter, "remote_in_ranks_" + callee, inputRanks, mod); + Value inputSizesArr = + getGlobalI64Array(loc, rewriter, "remote_in_sizes_" + callee, inputElemSizes, mod); + Value outputRanksArr = + getGlobalI64Array(loc, rewriter, "remote_out_ranks_" + callee, outputRanks, mod); + Value outputSizesArr = + getGlobalI64Array(loc, rewriter, "remote_out_sizes_" + callee, outputElemSizes, mod); + + Value numInputs = LLVM::ConstantOp::create( + rewriter, loc, rewriter.getI64IntegerAttr(inputDescPtrs.size())); + Value numOutputs = LLVM::ConstantOp::create( + rewriter, loc, rewriter.getI64IntegerAttr(outputDescPtrs.size())); + + LLVM::CallOp::create(rewriter, loc, launchFn, + ValueRange{addrPtr, symbolPtr, objectPtr, numInputs, inputDescsArr, + inputRanksArr, inputSizesArr, numOutputs, outputDescsArr, + outputRanksArr, outputSizesArr}); + + SmallVector results; + for (auto [descPtr, resultTy] : llvm::zip(outputDescPtrs, op.getResultTypes())) { + Type llvmDescTy = getTypeConverter()->convertType(resultTy); + Value loaded = LLVM::LoadOp::create(rewriter, loc, llvmDescTy, descPtr); + results.push_back(loaded); + } + rewriter.replaceOp(op, results); + return success(); + } +}; + +//===----------------------------------------------------------------------===// +// remote.call -> __catalyst__remote__call_wrapper(addr, sym, +// args_buf, args_size, +// &out_buf, &out_size) +//===----------------------------------------------------------------------===// + +struct CallOpLowering : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + /// Static-shape memref byte count, or -1 on unsupported types. + static int64_t memrefBufferBytes(Type ty, Operation *op) + { + auto memrefTy = dyn_cast(ty); + if (!memrefTy) { + op->emitOpError("remote.call requires memref-typed operands; got ") << ty; + return -1; + } + if (!memrefTy.hasStaticShape()) { + op->emitOpError("remote.call requires static-shape memref args; got ") << memrefTy; + return -1; + } + int64_t elemSz = primitiveByteSize(memrefTy.getElementType()); + if (elemSz < 0) { + op->emitOpError("unsupported memref element type for remote.call: ") + << memrefTy.getElementType(); + return -1; + } + return memrefTy.getNumElements() * elemSz; + } + + LogicalResult matchAndRewrite(remote::CallOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override + { + Location loc = op.getLoc(); + MLIRContext *ctx = rewriter.getContext(); + Type ptrTy = LLVM::LLVMPointerType::get(ctx); + Type i32Ty = rewriter.getI32Type(); + Type i64Ty = rewriter.getI64Type(); + Type voidTy = LLVM::LLVMVoidType::get(ctx); + Type i8Ty = rewriter.getI8Type(); + ModuleOp mod = op->getParentOfType(); + + // parameters: + // - addr: the address of the remote executor + // - symbol: the symbol to invoke + // - args_buf: the input buffer + // - args_size: the size of the input buffer + // - out_buf: the output buffer + // - out_size: the size of the output buffer + Type callSig = + LLVM::LLVMFunctionType::get(i32Ty, {ptrTy, ptrTy, ptrTy, i64Ty, ptrTy, ptrTy}); + LLVM::LLVMFuncOp callFn = catalyst::ensureFunctionDeclaration( + rewriter, op, "__catalyst__remote__call_wrapper", callSig); + Type freeSig = LLVM::LLVMFunctionType::get(voidTy, {ptrTy}); + LLVM::LLVMFuncOp freeFn = catalyst::ensureFunctionDeclaration( + rewriter, op, "__catalyst__remote__free_result", freeSig); + + unsigned numInputs = + op.getNumInputArgs().value_or(static_cast(op.getInputs().size())); + if (numInputs > op.getInputs().size()) { + return op->emitOpError("num_input_args exceeds operand count"); + } + + SmallVector inputOffsets; + int64_t totalInputBytes = 0; + for (unsigned i = 0; i < numInputs; ++i) { + int64_t argSize = memrefBufferBytes(op.getInputs()[i].getType(), op); + if (argSize < 0) { + return failure(); + } + inputOffsets.push_back(totalInputBytes); + totalInputBytes += argSize; + } + for (unsigned i = numInputs; i < op.getInputs().size(); ++i) { + if (memrefBufferBytes(op.getInputs()[i].getType(), op) < 0) { + return failure(); + } + } + + Type bufTy = LLVM::LLVMArrayType::get(i8Ty, totalInputBytes > 0 ? totalInputBytes : 1); + std::string sym = op.getSymbol().str(); + + Value addrPtr = getGlobalString(loc, rewriter, "remote_lib_addr_" + sym, + op.getAddress().str() + '\0', mod); + Value symPtr = getGlobalString(loc, rewriter, "remote_lib_sym_" + sym, sym + '\0', mod); + + Value argsBuf = getStaticAlloca(loc, rewriter, bufTy, 1); + for (unsigned i = 0; i < numInputs; ++i) { + auto memrefTy = cast(op.getInputs()[i].getType()); + int64_t numBytes = + memrefTy.getNumElements() * primitiveByteSize(memrefTy.getElementType()); + Value src = MemRefDescriptor(adaptor.getInputs()[i]).alignedPtr(rewriter, loc); + Value slot = LLVM::GEPOp::create( + rewriter, loc, ptrTy, bufTy, argsBuf, + ArrayRef{0, static_cast(inputOffsets[i])}, + LLVM::GEPNoWrapFlags::inbounds); + Value sizeVal = + LLVM::ConstantOp::create(rewriter, loc, rewriter.getI64IntegerAttr(numBytes)); + LLVM::MemcpyOp::create(rewriter, loc, slot, src, sizeVal, /*isVolatile=*/false); + } + Value argsSize = + LLVM::ConstantOp::create(rewriter, loc, rewriter.getI64IntegerAttr(totalInputBytes)); + + Value outBufSlot = getStaticAlloca(loc, rewriter, ptrTy, 1); + Value outSizeSlot = getStaticAlloca(loc, rewriter, i64Ty, 1); + + LLVM::CallOp::create( + rewriter, loc, callFn, + ValueRange{addrPtr, symPtr, argsBuf, argsSize, outBufSlot, outSizeSlot}); + + Value outBuf = LLVM::LoadOp::create(rewriter, loc, ptrTy, outBufSlot); + int64_t outOffset = 0; + for (unsigned i = numInputs; i < op.getInputs().size(); ++i) { + auto memrefTy = cast(op.getInputs()[i].getType()); + int64_t numBytes = + memrefTy.getNumElements() * primitiveByteSize(memrefTy.getElementType()); + Value destPtr = MemRefDescriptor(adaptor.getInputs()[i]).alignedPtr(rewriter, loc); + Value src = LLVM::GEPOp::create(rewriter, loc, ptrTy, i8Ty, outBuf, + ArrayRef{static_cast(outOffset)}, + LLVM::GEPNoWrapFlags::inbounds); + Value sizeVal = + LLVM::ConstantOp::create(rewriter, loc, rewriter.getI64IntegerAttr(numBytes)); + LLVM::MemcpyOp::create(rewriter, loc, destPtr, src, sizeVal, /*isVolatile=*/false); + outOffset += numBytes; + } + + LLVM::CallOp::create(rewriter, loc, freeFn, ValueRange{outBuf}); + rewriter.eraseOp(op); + return success(); + } +}; + +//===----------------------------------------------------------------------===// +// Pass +//===----------------------------------------------------------------------===// + +struct ConvertRemoteToLLVMPass : impl::ConvertRemoteToLLVMPassBase { + using ConvertRemoteToLLVMPassBase::ConvertRemoteToLLVMPassBase; + + void runOnOperation() final + { + MLIRContext *ctx = &getContext(); + ModuleOp mod = getOperation(); + + LLVMTypeConverter typeConverter(ctx); + + RewritePatternSet patterns(ctx); + patterns.add( + typeConverter, ctx); + + LLVMConversionTarget target(*ctx); + target.addLegalOp(); + target.addIllegalDialect(); + + if (failed(applyPartialConversion(mod, target, std::move(patterns)))) { + signalPassFailure(); + } + } +}; + +} // namespace + +} // namespace remote +} // namespace catalyst diff --git a/mlir/lib/hlo-extensions/Transforms/HloCustomCallPatterns.cpp b/mlir/lib/hlo-extensions/Transforms/HloCustomCallPatterns.cpp index 40d501954d..e921423b48 100644 --- a/mlir/lib/hlo-extensions/Transforms/HloCustomCallPatterns.cpp +++ b/mlir/lib/hlo-extensions/Transforms/HloCustomCallPatterns.cpp @@ -134,7 +134,8 @@ struct HloCustomCallOpRewritePattern : public mlir::OpRewritePattern(op, resultsType, newOperands, callTargetAttr, - nullptr); + /*number_original_arg=*/nullptr, + /*backend_config=*/nullptr); return success(); } diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt index 8806b8fce6..5e7cb66950 100644 --- a/mlir/test/CMakeLists.txt +++ b/mlir/test/CMakeLists.txt @@ -31,6 +31,7 @@ set(TEST_SUITES PauliFrame cli PBC + Remote RTIO QecLogical QecPhysical diff --git a/mlir/test/Catalyst/BufferizationTest.mlir b/mlir/test/Catalyst/BufferizationTest.mlir index 0748b0b64a..86119f494c 100644 --- a/mlir/test/Catalyst/BufferizationTest.mlir +++ b/mlir/test/Catalyst/BufferizationTest.mlir @@ -100,6 +100,17 @@ func.func @custom_call_copy(%arg0: tensor<2x3xf64>) -> tensor<2x2xf64> { // ----- +// `backend_config` attribute survives bufferization. +// CHECK-LABEL: func.func @custom_call_backend_config +// CHECK: catalyst.custom_call fn("lapack_dgesdd") (%{{.*}}, %{{.*}}) {backend_config = {foo = "bar"}, number_original_arg = 1 : i32} : +// CHECK-SAME: (memref<3x3xf64>, memref<3x3xf64>) -> () +func.func @custom_call_backend_config(%arg0: tensor<3x3xf64>) -> tensor<3x3xf64> { + %0 = catalyst.custom_call fn("lapack_dgesdd") (%arg0) {backend_config = {foo = "bar"}} : (tensor<3x3xf64>) -> (tensor<3x3xf64>) + return %0 : tensor<3x3xf64> +} + +// ----- + // CHECK-LABEL: @test0 module @test0 { // CHECK: catalyst.callback @callback_1(memref, memref) diff --git a/mlir/test/Catalyst/ConversionTest.mlir b/mlir/test/Catalyst/ConversionTest.mlir index 6bbc384651..664f311342 100644 --- a/mlir/test/Catalyst/ConversionTest.mlir +++ b/mlir/test/Catalyst/ConversionTest.mlir @@ -103,17 +103,23 @@ func.func @custom_call(%arg0: memref<3x3xf64>, %arg1: memref<3x3xf64>) -> () { // CHECK: [[alloca0:%.+]] = llvm.alloca [[c1]] x !llvm.array<1 x ptr> // CHECK: [[c1:%.+]] = llvm.mlir.constant(1 : i64) - // CHECK: [[alloca1:%.+]] = llvm.alloca [[c1]] x !llvm.struct<(i64, ptr, i8)> + // CHECK: [[alloca1:%.+]] = llvm.alloca [[c1]] x !llvm.struct<(i64, ptr, i8, ptr)> + + // CHECK: [[c1:%.+]] = llvm.mlir.constant(1 : i64) + // CHECK: [[sizes0:%.+]] = llvm.alloca [[c1]] x !llvm.array<2 x i64> // CHECK: [[c1:%.+]] = llvm.mlir.constant(1 : i64) // CHECK: [[alloca2:%.+]] = llvm.alloca [[c1]] x !llvm.array<1 x ptr> // CHECK: [[c1:%.+]] = llvm.mlir.constant(1 : i64) - // CHECK: [[alloca3:%.+]] = llvm.alloca [[c1]] x !llvm.struct<(i64, ptr, i8)> + // CHECK: [[alloca3:%.+]] = llvm.alloca [[c1]] x !llvm.struct<(i64, ptr, i8, ptr)> + + // CHECK: [[c1:%.+]] = llvm.mlir.constant(1 : i64) + // CHECK: [[sizes1:%.+]] = llvm.alloca [[c1]] x !llvm.array<2 x i64> // CHECK: llvm.call @lapack_dgesdd([[alloca2]], [[alloca0]]) catalyst.custom_call fn("lapack_dgesdd") (%arg0, %arg1) {number_original_arg = 1 : i32} : (memref<3x3xf64>, memref<3x3xf64>) -> () - return + 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/LaunchKernelBufferization.mlir b/mlir/test/Catalyst/LaunchKernelBufferization.mlir new file mode 100644 index 0000000000..1c8f7e4dee --- /dev/null +++ b/mlir/test/Catalyst/LaunchKernelBufferization.mlir @@ -0,0 +1,60 @@ +// 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 --split-input-file \ +// RUN: --pass-pipeline="builtin.module( \ +// RUN: one-shot-bufferize{unknown-type-conversion=identity-layout-map} \ +// RUN: )" %s | FileCheck %s + +// launch_kernel must survive bufferization with its results intact: tensor operands become memref +// buffers and tensor results become memref results (no out-param marshalling), matching +// remote.launch's contract. + +// CHECK-LABEL: func.func @host +// CHECK: %[[ARG:.*]] = bufferization.to_buffer %{{.*}} : tensor<4xf64> to memref<4xf64> +// CHECK: %[[RES:.*]] = catalyst.launch_kernel @target::@entry(%[[ARG]]) : (memref<4xf64>) -> memref<4xf64> +// CHECK: bufferization.to_tensor %[[RES]] +func.func @host(%arg0: tensor<4xf64>) -> tensor<4xf64> { + %0 = catalyst.launch_kernel @target::@entry(%arg0) : (tensor<4xf64>) -> tensor<4xf64> + return %0 : tensor<4xf64> +} + +module @target { + func.func public @entry(%arg0: tensor<4xf64>) -> tensor<4xf64> { + return %arg0 : tensor<4xf64> + } +} + +// ----- + +// A non-identity-layout (strided/offset) operand must be copied into a fresh contiguous buffer +// before the call: the remote/kernel ABI sends a contiguous block from the aligned pointer and +// ignores strides/offset, so passing the subview directly would ship the wrong bytes. + +// CHECK-LABEL: func.func @host_strided +// CHECK: %[[SUB:.*]] = memref.subview %{{.*}} : memref<8xf64> to memref<4xf64, strided<[1], offset: 2>> +// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<4xf64> +// CHECK: memref.copy %[[SUB]], %[[ALLOC]] +// CHECK: catalyst.launch_kernel @target::@entry(%[[ALLOC]]) : (memref<4xf64>) -> memref<4xf64> +func.func @host_strided(%arg0: tensor<8xf64>) -> tensor<4xf64> { + %slice = tensor.extract_slice %arg0[2] [4] [1] : tensor<8xf64> to tensor<4xf64> + %0 = catalyst.launch_kernel @target::@entry(%slice) : (tensor<4xf64>) -> tensor<4xf64> + return %0 : tensor<4xf64> +} + +module @target { + func.func public @entry(%arg0: tensor<4xf64>) -> tensor<4xf64> { + return %arg0 : tensor<4xf64> + } +} diff --git a/mlir/test/Catalyst/NestedModuleTarget.mlir b/mlir/test/Catalyst/NestedModuleTarget.mlir new file mode 100644 index 0000000000..8990c25dff --- /dev/null +++ b/mlir/test/Catalyst/NestedModuleTarget.mlir @@ -0,0 +1,142 @@ +// 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. + +// Verify that modules annotated with catalyst.target are not inlined, renamed, or flattened: their +// launch_kernel survives intact for a later consumer (e.g. dispatch-remote-targets). + +// RUN: quantum-opt --inline-nested-module --split-input-file %s | FileCheck %s + + +// CHECK-LABEL: module @host +module @host { + // No root declaration is emitted; the launch_kernel is kept (not flattened to a func.call). + // CHECK-NOT: func.func private @ghz + // CHECK: func.func public @jit_main + func.func public @jit_main() attributes {llvm.emit_c_interface} { + // CHECK: catalyst.launch_kernel @module_accel::@ghz() : () -> () + catalyst.launch_kernel @module_accel::@ghz() : () -> () + func.return + } + + // CHECK: module @module_accel attributes {catalyst.target + // CHECK-NOT: catalyst.unique_names + // CHECK: func.func @ghz() + module @module_accel attributes {catalyst.target = {backend = "accel"}} { + func.func @ghz() { + func.return + } + } + + func.func @setup() { func.return } + func.func @teardown() { func.return } +} + +// ----- + +// A mix: the non-target module is inlined+flattened (call @work_0); the catalyst.target module +// keeps its nested launch_kernel. +// CHECK-LABEL: module @mixed +// CHECK: func.func public @jit_run +// CHECK: call @work_0() +// CHECK: catalyst.launch_kernel @module_accel2::@kernel() +// CHECK: func.func @work_0() +// CHECK-NOT: module @module_cpu +// CHECK: module @module_accel2 attributes {catalyst.target +// CHECK-NOT: catalyst.unique_names +// CHECK: func.func @kernel() +module @mixed { + + func.func public @jit_run() attributes {llvm.emit_c_interface} { + catalyst.launch_kernel @module_cpu::@work() : () -> () + catalyst.launch_kernel @module_accel2::@kernel() : () -> () + func.return + } + + module @module_cpu { + func.func @work() { + func.return + } + } + + module @module_accel2 attributes {catalyst.target = {backend = "accel"}} { + func.func @kernel() { + func.return + } + } + + func.func @setup() { func.return } + func.func @teardown() { func.return } +} + +// ----- + +// Two target modules with the same entry name: since neither is inlined or renamed, both keep the +// name @kernel in their own module (no collision — separate symbol tables), and both launch_kernels +// survive. +// CHECK-LABEL: module @duplicate_target_names +// CHECK: func.func public @jit_run +// CHECK: catalyst.launch_kernel @module_accel_a::@kernel() +// CHECK: catalyst.launch_kernel @module_accel_b::@kernel() +// CHECK-DAG: module @module_accel_a attributes {catalyst.target +// CHECK-DAG: module @module_accel_b attributes {catalyst.target +// CHECK-NOT: catalyst.unique_names +module @duplicate_target_names { + + func.func public @jit_run() attributes {llvm.emit_c_interface} { + catalyst.launch_kernel @module_accel_a::@kernel() : () -> () + catalyst.launch_kernel @module_accel_b::@kernel() : () -> () + func.return + } + + module @module_accel_a attributes {catalyst.target = {backend = "accel"}} { + func.func @kernel() attributes {attr = "value"} { + func.return + } + } + + module @module_accel_b attributes {catalyst.target = {backend = "accel"}} { + func.func @kernel() { + func.return + } + } +} + +// ----- + +// Dispatch modules (catalyst.dispatch) are not inlined, not flattened, AND not renamed: +// dispatch-remote-targets consumes the launch_kernel directly and resolves the entry in the +// object's own JITDylib on the executor, so the kernel keeps its original name and no root +// declaration is emitted. +// CHECK-LABEL: module @host_dispatch +// CHECK-NOT: func.func private @ghz +// CHECK: func.func public @jit_main +// CHECK: catalyst.launch_kernel @module_remote::@ghz() : () -> () +// CHECK: module @module_remote attributes {catalyst.dispatch +// CHECK: func.func @ghz() +module @host_dispatch { + + func.func public @jit_main() attributes {llvm.emit_c_interface} { + catalyst.launch_kernel @module_remote::@ghz() : () -> () + func.return + } + + module @module_remote attributes {catalyst.target = {backend = "accel"}, catalyst.dispatch = {address = "h:1"}} { + func.func @ghz() { + func.return + } + } + + func.func @setup() { func.return } + func.func @teardown() { func.return } +} diff --git a/mlir/test/Catalyst/StaticAllocaTest.mlir b/mlir/test/Catalyst/StaticAllocaTest.mlir index 5c62250782..d88f252560 100644 --- a/mlir/test/Catalyst/StaticAllocaTest.mlir +++ b/mlir/test/Catalyst/StaticAllocaTest.mlir @@ -41,11 +41,15 @@ module @static_alloca_custom_call { // CHECK: [[one:%.+]] = llvm.mlir.constant(1 : i64) // CHECK: llvm.alloca [[one]] x !llvm.array<1 x ptr> // CHECK: [[one:%.+]] = llvm.mlir.constant(1 : i64) - // CHECK: llvm.alloca [[one]] x !llvm.struct<(i64, ptr, i8)> + // CHECK: llvm.alloca [[one]] x !llvm.struct<(i64, ptr, i8, ptr)> + // CHECK: [[one:%.+]] = llvm.mlir.constant(1 : i64) + // CHECK: llvm.alloca [[one]] x !llvm.array<2 x i64> // CHECK: [[one:%.+]] = llvm.mlir.constant(1 : i64) // CHECK: llvm.alloca [[one]] x !llvm.array<1 x ptr> // CHECK: [[one:%.+]] = llvm.mlir.constant(1 : i64) - // CHECK: llvm.alloca [[one]] x !llvm.struct<(i64, ptr, i8)> + // CHECK: llvm.alloca [[one]] x !llvm.struct<(i64, ptr, i8, ptr)> + // CHECK: [[one:%.+]] = llvm.mlir.constant(1 : i64) + // CHECK: llvm.alloca [[one]] x !llvm.array<2 x i64> // CHECK: ^bb1: cf.br ^bb1 ^bb1: diff --git a/mlir/test/Remote/ConvertRemoteToLLVM.mlir b/mlir/test/Remote/ConvertRemoteToLLVM.mlir new file mode 100644 index 0000000000..5a3f2dbfd2 --- /dev/null +++ b/mlir/test/Remote/ConvertRemoteToLLVM.mlir @@ -0,0 +1,56 @@ +// 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 --convert-remote-to-llvm --split-input-file | FileCheck %s + +// CHECK: llvm.func @__catalyst__remote__open(!llvm.ptr) -> i64 +// CHECK-LABEL: func.func @open +func.func @open() { + // CHECK: llvm.call @__catalyst__remote__open + remote.open("127.0.0.1:9000") + return +} + +// ----- + +// CHECK: llvm.func @__catalyst__remote__send_binary(!llvm.ptr, !llvm.ptr, i32) -> i64 +// CHECK-LABEL: func.func @send_binary +func.func @send_binary() { + // CHECK: llvm.call @__catalyst__remote__send_binary + remote.send_binary("127.0.0.1:9000", "/tmp/qnode_0.o") + return +} + +// ----- + +// CHECK: llvm.func @__catalyst__remote__launch +// CHECK-LABEL: func.func @launch +func.func @launch(%arg0: memref) -> memref { + // CHECK: llvm.call @__catalyst__remote__launch + %0 = remote.launch("qnode_0", "127.0.0.1:9000", "/tmp/qnode_0.o") (%arg0) : (memref) -> memref + return %0 : memref +} + +// ----- + +// CHECK-DAG: llvm.func @__catalyst__remote__call_wrapper +// CHECK-DAG: llvm.func @__catalyst__remote__free_result +// CHECK-LABEL: func.func @call +func.func @call(%arg0: memref<4xf64>, %arg1: memref<4xf64>) { + // CHECK: llvm.call @__catalyst__remote__call_wrapper + // CHECK: llvm.call @__catalyst__remote__free_result + remote.call("foo", "127.0.0.1:9000") (%arg0, %arg1) + {num_input_args = 1 : i32} : (memref<4xf64>, memref<4xf64>) -> () + return +} diff --git a/mlir/test/Remote/DispatchRemoteLibCalls.mlir b/mlir/test/Remote/DispatchRemoteLibCalls.mlir new file mode 100644 index 0000000000..e6d645193a --- /dev/null +++ b/mlir/test/Remote/DispatchRemoteLibCalls.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-remote-targets --verify-diagnostics | FileCheck %s + +// dispatch-remote-targets also rewrites standalone remote library calls which are +// catalyst.custom_call ops whose backend_config carries a `dispatch` entry into +// remote.call ops invoking a symbol already loaded on the executor. + +// ----- + +// When backend_config.dispatch names the executor explicitly, it is rewritten to a remote.call on +// that address, and a session is opened once in setup(). +// CHECK-LABEL: func.func @setup() +// CHECK: remote.open("ADDR:PORT") +// CHECK: return +// CHECK-LABEL: func.func @main +// CHECK: remote.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: remote.open("ADDR:PORT") +// CHECK-LABEL: func.func @main +// CHECK: remote.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 remote 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/Remote/DispatchRemoteTargetModules.mlir b/mlir/test/Remote/DispatchRemoteTargetModules.mlir new file mode 100644 index 0000000000..1b6b6f882a --- /dev/null +++ b/mlir/test/Remote/DispatchRemoteTargetModules.mlir @@ -0,0 +1,70 @@ +// 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-remote-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: remote.open("ADDR:PORT") +// CHECK: remote.send_binary("ADDR:PORT", "/tmp/target_compute.o") +// CHECK-NOT: remote.send_binary +// CHECK: return + +// teardown() is left untouched: the runtime closes sessions at process exit. +// CHECK-LABEL: func.func @teardown() +// CHECK-NOT: remote. +// CHECK: return + +// Each host-side launch_kernel is rewritten to its own remote.launch, preserving the call's +// operand/result types (a typed call exercises the memref-result path). +// CHECK-LABEL: func.func public @jit_main +// CHECK: remote.launch("noop", "ADDR:PORT", "/tmp/target_compute.o") () : () -> () +// CHECK: remote.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 remote-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> + } + } +} diff --git a/mlir/test/Remote/RemoteOps.mlir b/mlir/test/Remote/RemoteOps.mlir new file mode 100644 index 0000000000..141d85e30d --- /dev/null +++ b/mlir/test/Remote/RemoteOps.mlir @@ -0,0 +1,53 @@ +// 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 | FileCheck %s + +// CHECK-LABEL: func.func @open +// CHECK: remote.open("127.0.0.1:9000") +func.func @open() { + remote.open("127.0.0.1:9000") + return +} + +// ----- + +// CHECK-LABEL: func.func @send_binary +// CHECK: remote.send_binary("127.0.0.1:9000", "/tmp/qnode_0.o") +func.func @send_binary() { + remote.send_binary("127.0.0.1:9000", "/tmp/qnode_0.o") + return +} + +// ----- + +// CHECK-LABEL: func.func @launch +// CHECK: remote.launch("qnode_0", "127.0.0.1:9000", "/tmp/qnode.o") (%{{.*}}) : +// CHECK-SAME: (memref) -> memref +func.func @launch(%arg0: memref) -> memref { + %0 = remote.launch("qnode_0", "127.0.0.1:9000", "/tmp/qnode.o") (%arg0) : (memref) -> memref + return %0 : memref +} + +// ----- + +// CHECK-LABEL: func.func @call +// CHECK: remote.call("foo", "127.0.0.1:9000") +// CHECK-SAME: num_input_args = 1 : i32 +// CHECK-SAME: (memref<4xf64>, memref<4xf64>) -> () +func.func @call(%arg0: memref<4xf64>, %arg1: memref<4xf64>) { + remote.call("foo", "127.0.0.1:9000") (%arg0, %arg1) + {num_input_args = 1 : i32} : (memref<4xf64>, memref<4xf64>) -> () + return +} diff --git a/mlir/tools/catalyst-cli/CMakeLists.txt b/mlir/tools/catalyst-cli/CMakeLists.txt index a1de2381d3..ec3fa54340 100644 --- a/mlir/tools/catalyst-cli/CMakeLists.txt +++ b/mlir/tools/catalyst-cli/CMakeLists.txt @@ -47,6 +47,8 @@ set(LIBS ion-transforms MLIRRTIO rtio-transforms + MLIRRemote + remote-transforms MLIRQecLogical MLIRQecPhysical MLIRCatalystTest diff --git a/mlir/tools/quantum-lsp-server/CMakeLists.txt b/mlir/tools/quantum-lsp-server/CMakeLists.txt index 5a6bc26be6..cc89753535 100644 --- a/mlir/tools/quantum-lsp-server/CMakeLists.txt +++ b/mlir/tools/quantum-lsp-server/CMakeLists.txt @@ -16,6 +16,7 @@ set(LIBS MLIRPauliFrame MLIRIon MLIRRTIO + MLIRRemote MLIRQecLogical MLIRQecPhysical ) diff --git a/mlir/tools/quantum-lsp-server/quantum-lsp-server.cpp b/mlir/tools/quantum-lsp-server/quantum-lsp-server.cpp index 9cc18cbcc2..4ffedb60d6 100644 --- a/mlir/tools/quantum-lsp-server/quantum-lsp-server.cpp +++ b/mlir/tools/quantum-lsp-server/quantum-lsp-server.cpp @@ -30,6 +30,8 @@ #include "Quantum/IR/QuantumDialect.h" #include "RTIO/IR/RTIODialect.h" +#include "Remote/IR/RemoteDialect.h" + int main(int argc, char **argv) { mlir::DialectRegistry registry; @@ -44,6 +46,7 @@ int main(int argc, char **argv) registry.insert(); registry.insert(); registry.insert(); + registry.insert(); registry.insert(); registry.insert(); diff --git a/mlir/tools/quantum-opt/CMakeLists.txt b/mlir/tools/quantum-opt/CMakeLists.txt index ef16ce934d..bbfdee58e9 100644 --- a/mlir/tools/quantum-opt/CMakeLists.txt +++ b/mlir/tools/quantum-opt/CMakeLists.txt @@ -33,6 +33,8 @@ set(LIBS ion-transforms MLIRRTIO rtio-transforms + MLIRRemote + remote-transforms MLIRQecLogical MLIRQecPhysical MLIRCatalystTest diff --git a/mlir/tools/quantum-opt/quantum-opt.cpp b/mlir/tools/quantum-opt/quantum-opt.cpp index 5b9a279fd8..0d7d625e57 100644 --- a/mlir/tools/quantum-opt/quantum-opt.cpp +++ b/mlir/tools/quantum-opt/quantum-opt.cpp @@ -50,6 +50,8 @@ #include "RegisterAllPasses.h" +#include "Remote/IR/RemoteDialect.h" + namespace test { void registerTestDialect(mlir::DialectRegistry &); } // namespace test @@ -77,6 +79,7 @@ int main(int argc, char **argv) registry.insert(); registry.insert(); registry.insert(); + registry.insert(); registry.insert(); registry.insert(); registry.insert(); diff --git a/runtime/Makefile b/runtime/Makefile index 84c9ba328e..b7595ff6bb 100644 --- a/runtime/Makefile +++ b/runtime/Makefile @@ -13,9 +13,11 @@ CODE_COVERAGE ?= OFF BUILD_TYPE ?= RelWithDebInfo ENABLE_OPENQASM ?= ON ENABLE_OQD ?= OFF +ENABLE_REMOTE ?= OFF ENABLE_ASAN ?= OFF STRICT_WARNINGS ?= ON LLVM_DIR ?= $(MK_DIR)/../mlir/llvm-project/ +LLVM_BUILD_DIR ?= $(LLVM_DIR)/build PLATFORM := $(shell uname -s) @@ -53,6 +55,10 @@ ifeq ($(ENABLE_OQD), ON) TEST_TARGETS += runner_tests_oqd endif +ifeq ($(ENABLE_REMOTE), ON) + BUILD_TARGETS += rt_remote +endif + .PHONY: help help: @echo "Please use \`make ' where is one of" @@ -73,11 +79,13 @@ configure: -DCMAKE_LIBRARY_OUTPUT_DIRECTORY=$(RT_BUILD_DIR)/lib \ -DCMAKE_C_COMPILER=$(C_COMPILER) \ -DMLIR_INCLUDE_DIRS=$(LLVM_DIR)/mlir/include \ + -DLLVM_DIR=$(LLVM_BUILD_DIR)/lib/cmake/llvm \ -DCMAKE_CXX_COMPILER=$(CXX_COMPILER) \ -DCMAKE_C_COMPILER_LAUNCHER=$(COMPILER_LAUNCHER) \ -DCMAKE_CXX_COMPILER_LAUNCHER=$(COMPILER_LAUNCHER) \ -DENABLE_OPENQASM=$(ENABLE_OPENQASM) \ -DENABLE_OQD=$(ENABLE_OQD) \ + -DENABLE_REMOTE=$(ENABLE_REMOTE) \ -DENABLE_CODE_COVERAGE=$(CODE_COVERAGE) \ -DPython_EXECUTABLE=$(PYTHON) \ -DENABLE_ADDRESS_SANITIZER=$(ENABLE_ASAN) \ diff --git a/runtime/include/RemoteCAPI.h b/runtime/include/RemoteCAPI.h new file mode 100644 index 0000000000..562f07fd71 --- /dev/null +++ b/runtime/include/RemoteCAPI.h @@ -0,0 +1,46 @@ +// 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. + +#pragma once +#ifndef REMOTECAPI_H +#define REMOTECAPI_H + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// Remote Runtime API. This is expected to be called by the host program to establish a remote +// session. Including open a session, send the binary to the remote, launch the kernel on the remote +// and close the session. +int __catalyst__remote__open(const char *addr); +int __catalyst__remote__send_binary(const char *addr, const char *path, uint32_t format); +void __catalyst__remote__launch(const char *addr, const char *entry_symbol, const char *object, + size_t num_inputs, void *const *input_descs, + const size_t *input_ranks, const size_t *input_elem_sizes, + size_t num_outputs, void *const *output_descs, + const size_t *output_ranks, const size_t *output_elem_sizes); +int __catalyst__remote__call_wrapper(const char *addr, const char *symbol, const char *args_buf, + size_t args_size, void **out_buf, size_t *out_size); +void __catalyst__remote__free_result(void *buf); + +int __catalyst__remote__close(); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif diff --git a/runtime/include/RuntimeCAPI.h b/runtime/include/RuntimeCAPI.h index 0ee58432f0..c70bab6182 100644 --- a/runtime/include/RuntimeCAPI.h +++ b/runtime/include/RuntimeCAPI.h @@ -120,6 +120,9 @@ RESULT *__catalyst__mbqc__measure_in_basis(QUBIT *, uint32_t, double, int32_t); // Async runtime error void __catalyst__host__rt__unrecoverable_error(); +// Allocate a host buffer of `size` bytes and register it with the runtime's memory manager +void *__catalyst__rt__alloc_managed(size_t size); + #ifdef __cplusplus } // extern "C" #endif diff --git a/runtime/lib/CMakeLists.txt b/runtime/lib/CMakeLists.txt index eb7a77b63d..99eb651bbc 100644 --- a/runtime/lib/CMakeLists.txt +++ b/runtime/lib/CMakeLists.txt @@ -68,3 +68,7 @@ add_subdirectory(QEC) if(ENABLE_OQD) add_subdirectory(OQDcapi) endif() + +if(ENABLE_REMOTE) +add_subdirectory(remote) +endif() diff --git a/runtime/lib/capi/ExecutionContext.hpp b/runtime/lib/capi/ExecutionContext.hpp index 02c7400afb..d83da45182 100644 --- a/runtime/lib/capi/ExecutionContext.hpp +++ b/runtime/lib/capi/ExecutionContext.hpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -25,6 +26,7 @@ #include #include +#include "DataView.hpp" #include "Exception.hpp" #include "QuantumDevice.hpp" @@ -81,15 +83,37 @@ class SharedLibraryManager final { { #ifdef __APPLE__ auto rtld_flags = RTLD_LAZY; + constexpr const char *dl_ext = ".dylib"; #else // Closing the dynamic library of Lightning simulators with dlclose() where OpenMP // directives (in Lightning simulators) are in use would raise memory segfaults. // Note that we use RTLD_NODELETE as a workaround to fix the issue. auto rtld_flags = RTLD_LAZY | RTLD_NODELETE; + constexpr const char *dl_ext = ".so"; #endif _handler = dlopen(filename.c_str(), rtld_flags); - RT_FAIL_IF(!_handler, dlerror()); + if (_handler) { + return; + } + + // dlerror() is destructive: it returns the pending error and clears the state. Capture it + // into a local before any further use. + const char *primary_dlerror = dlerror(); + std::string primary_error = primary_dlerror ? primary_dlerror : "unknown dlopen error"; + auto basename = std::filesystem::path(filename).filename().string(); + + // rewrite the extension if it differs from the platform one. + if (!basename.ends_with(dl_ext)) { + auto dot = basename.find_last_of('.'); + if (dot != std::string::npos) { + basename.resize(dot); + } + basename += dl_ext; + } + + _handler = dlopen(basename.c_str(), rtld_flags); + RT_FAIL_IF(!_handler, ("dlopen failed to load " + filename + ": " + primary_error).c_str()); } ~SharedLibraryManager() diff --git a/runtime/lib/capi/RuntimeCAPI.cpp b/runtime/lib/capi/RuntimeCAPI.cpp index ec9863686b..c01141c667 100644 --- a/runtime/lib/capi/RuntimeCAPI.cpp +++ b/runtime/lib/capi/RuntimeCAPI.cpp @@ -170,6 +170,8 @@ void *_mlir_memref_to_llvm_alloc(size_t size) return ptr; } +void *__catalyst__rt__alloc_managed(size_t size) { return _mlir_memref_to_llvm_alloc(size); } + void *_mlir_memref_to_llvm_aligned_alloc(size_t alignment, size_t size) { void *ptr = aligned_alloc(alignment, size); diff --git a/runtime/lib/remote/CMakeLists.txt b/runtime/lib/remote/CMakeLists.txt new file mode 100644 index 0000000000..1030e54f26 --- /dev/null +++ b/runtime/lib/remote/CMakeLists.txt @@ -0,0 +1,63 @@ +############################################### +# library catalyst_remote_session + rt_remote # +############################################### + +find_package(LLVM CONFIG REQUIRED) + +llvm_map_components_to_libnames(_remote_llvm_libs + core support orcjit jitlink object passes + ${LLVM_TARGETS_TO_BUILD} +) + +# catalyst_remote_session +add_library(catalyst_remote_session SHARED RemoteSession.cpp) + +target_include_directories(catalyst_remote_session + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} +) + +target_include_directories(catalyst_remote_session SYSTEM PRIVATE + ${LLVM_INCLUDE_DIRS} +) + +target_link_libraries(catalyst_remote_session PRIVATE + rt_capi + ${_remote_llvm_libs} + pthread + ${CMAKE_DL_LIBS} +) + +target_compile_options(catalyst_remote_session PRIVATE -fno-rtti) + +set_property(TARGET catalyst_remote_session PROPERTY POSITION_INDEPENDENT_CODE ON) + +if(NOT APPLE) + set_property(TARGET catalyst_remote_session APPEND PROPERTY BUILD_RPATH $ORIGIN) +else() + set_property(TARGET catalyst_remote_session APPEND PROPERTY BUILD_RPATH @loader_path) +endif() + +# rt_remote + +add_library(rt_remote SHARED RemoteRuntime.cpp) + +target_include_directories(rt_remote + PRIVATE + ${runtime_includes} + ${CMAKE_CURRENT_SOURCE_DIR} +) + +target_link_libraries(rt_remote PRIVATE + catalyst_remote_session + rt_capi + pthread + ${CMAKE_DL_LIBS} +) + +set_property(TARGET rt_remote PROPERTY POSITION_INDEPENDENT_CODE ON) + +if(NOT APPLE) + set_property(TARGET rt_remote APPEND PROPERTY BUILD_RPATH $ORIGIN) +else() + set_property(TARGET rt_remote APPEND PROPERTY BUILD_RPATH @loader_path) +endif() diff --git a/runtime/lib/remote/RemoteRuntime.cpp b/runtime/lib/remote/RemoteRuntime.cpp new file mode 100644 index 0000000000..a4d3b7a307 --- /dev/null +++ b/runtime/lib/remote/RemoteRuntime.cpp @@ -0,0 +1,270 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include + +#include "Exception.hpp" +#include "RemoteCAPI.h" +#include "RemoteSession.hpp" + +namespace { + +struct RemoteEntry { + catalyst::remote::RemoteSession *session = nullptr; // The remote session handle. + std::set loaded_paths; // The paths of the binaries that are going to be loaded + // into the remote session. + std::mutex mu; // The mutex to protect the loaded paths. + + ~RemoteEntry() + { + if (session) { + catalyst::remote::close(session); + session = nullptr; + } + } +}; + +std::mutex g_map_mu; + +// Each address has its own RemoteEntry, so we can dispatch the object file to different remote +// sessions. +std::map> remote_sessions; + +// DEBUG logs +bool remote_verbose() +{ + static const bool v = []() { + const char *e = std::getenv("CATALYST_REMOTE_VERBOSE"); + return e && *e && *e != '0'; + }(); + return v; +} + +// Look up or create the entry for `addr`. +RemoteEntry *find_or_create_entry(const char *addr, bool create_if_missing) +{ + if (!addr || !*addr) { + return nullptr; + } + std::lock_guard lock(g_map_mu); + auto it = remote_sessions.find(addr); + if (it != remote_sessions.end()) { + return it->second.get(); + } + if (!create_if_missing) { + return nullptr; + } + auto inserted = remote_sessions.emplace(std::string(addr), std::make_unique()); + return inserted.first->second.get(); +} + +} // namespace + +extern "C" { + +int __catalyst__remote__open(const char *addr) +{ + if (!addr || !*addr) { + RT_FAIL("Empty address"); + } + RemoteEntry *entry = find_or_create_entry(addr, /*create_if_missing=*/true); + std::string err_msg; + { + std::lock_guard lock(entry->mu); + if (entry->session) { + return 0; // idempotent per addr + } + if (remote_verbose()) { + std::fprintf(stderr, "[remote] open(addr=%s)\n", addr); + } + entry->session = catalyst::remote::open(addr); + if (entry->session) { + if (remote_verbose()) { + std::fprintf(stderr, "[remote] open(%s) OK\n", addr); + } + return 0; + } + err_msg = "Could not connect to catalyst-executor at "; + err_msg += addr; + err_msg += ": "; + err_msg += catalyst::remote::last_error(); + } + { + std::lock_guard mapLock(g_map_mu); + remote_sessions.erase(addr); + } + RT_FAIL(err_msg.c_str()); +} + +int __catalyst__remote__send_binary(const char *addr, const char *path, uint32_t format) +{ + RemoteEntry *entry = find_or_create_entry(addr, /*create_if_missing=*/false); + if (!entry) { + RT_FAIL("No session found, call __catalyst__remote__open first."); + } + std::lock_guard lock(entry->mu); + if (!entry->session) { + std::string msg = "__catalyst__remote__send_binary("; + msg += addr; + msg += "): session is closed."; + RT_FAIL(msg.c_str()); + } + if (!path || !*path) { + return 0; + } + std::string key(path); + if (!entry->loaded_paths.insert(key).second) { + return 0; + } + if (remote_verbose()) { + std::fprintf(stderr, "[remote] send_binary(addr=%s, path=%s, format=%u)\n", addr, path, + format); + } + + int rc = 0; + switch (format) { + case 0: + rc = catalyst::remote::load_object_path(entry->session, path); + break; + case 1: + rc = catalyst::remote::load_asset_path(entry->session, path); + break; + default: { + std::string msg = "unknown binary format tag "; + msg += std::to_string(format); + entry->loaded_paths.erase(key); + RT_FAIL(msg.c_str()); + } + } + + if (rc != 0) { + std::string msg = catalyst::remote::last_error(); + entry->loaded_paths.erase(key); + RT_FAIL(msg.c_str()); + } + return 0; +} + +/** + * @brief Generic ORC wrapper-function call by symbol name. Returns 0 on success, -1 on error. + * + * @param addr The address of the remote session. + * @param symbol The symbol of the function to call. + * @param args_buf The buffer of the arguments. + * @param args_size The size of the arguments. + * @param out_buf The buffer of the result. + * @param out_size The size of the result. + * @return int 0 on success, -1 on error. + */ +int __catalyst__remote__call_wrapper(const char *addr, const char *symbol, const char *args_buf, + size_t args_size, void **out_buf, size_t *out_size) +{ + if (out_buf) { + *out_buf = nullptr; + } + if (out_size) { + *out_size = 0; + } + RemoteEntry *entry = find_or_create_entry(addr, /*create_if_missing=*/false); + if (!entry) { + RT_FAIL("No session found, call __catalyst__remote__open first."); + } + std::lock_guard lock(entry->mu); + if (!entry->session) { + RT_FAIL("Session is closed"); + } + if (!symbol || !*symbol) { + RT_FAIL("Empty symbol passed to __catalyst__remote__call_wrapper"); + } + if (remote_verbose()) { + std::fprintf(stderr, "[remote] call_wrapper(addr=%s, sym=%s, in_size=%zu)\n", addr, symbol, + args_size); + } + char *buf = nullptr; + size_t n = 0; + int rc = + catalyst::remote::call_wrapper_raw(entry->session, symbol, args_buf, args_size, &buf, &n); + if (rc != 0) { + RT_FAIL(catalyst::remote::last_error()); + } + if (out_buf) { + *out_buf = buf; + } + else { + std::free(buf); // caller didn't want the bytes back + } + if (out_size) { + *out_size = n; + } + return 0; +} + +void __catalyst__remote__free_result(void *buf) { std::free(buf); } + +int __catalyst__remote__close() +{ + std::lock_guard mapLock(g_map_mu); + for (auto &[addr, entry] : remote_sessions) { + std::lock_guard lock(entry->mu); + if (entry->session) { + catalyst::remote::close(entry->session); + entry->session = nullptr; + entry->loaded_paths.clear(); + } + } + remote_sessions.clear(); + return 0; +} + +void __catalyst__remote__launch(const char *addr, const char *entry_symbol, const char *object, + size_t num_inputs, void *const *input_descs, + const size_t *input_ranks, const size_t *input_elem_sizes, + size_t num_outputs, void *const *output_descs, + const size_t *output_ranks, const size_t *output_elem_sizes) +{ + RemoteEntry *entry = find_or_create_entry(addr, /*create_if_missing=*/false); + if (!entry) { + RT_FAIL("Can't find opened session"); + } + + std::lock_guard lock(entry->mu); + if (remote_verbose()) { + std::fprintf(stderr, + "[remote] launch(addr=%s, symbol=%s, object=%s, n_in=%zu, n_out=%zu)\n", addr, + entry_symbol, object ? object : "", num_inputs, num_outputs); + } + if (!entry->session) { + RT_FAIL("Session is closed"); + } + // Resolve the entry within the JITDylib of its own object, so same-named entries shipped from + // different kernel objects don't collide in one symbol namespace. + uint64_t entry_addr = catalyst::remote::lookup(entry->session, entry_symbol, object); + if (!entry_addr) { + RT_FAIL(catalyst::remote::last_error()); + } + if (catalyst::remote::invoke_kernel(entry->session, entry_addr, num_inputs, input_descs, + input_ranks, input_elem_sizes, num_outputs, output_descs, + output_ranks, output_elem_sizes) != 0) { + RT_FAIL(catalyst::remote::last_error()); + } +} + +} // extern "C" diff --git a/runtime/lib/remote/RemoteSession.cpp b/runtime/lib/remote/RemoteSession.cpp new file mode 100644 index 0000000000..c0e5b6e62e --- /dev/null +++ b/runtime/lib/remote/RemoteSession.cpp @@ -0,0 +1,889 @@ +// 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 "RemoteSession.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "llvm/ADT/StringRef.h" +#include "llvm/ExecutionEngine/Orc/Core.h" +#include "llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h" +#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h" +#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" +#include "llvm/ExecutionEngine/Orc/Mangling.h" +#include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h" +#include "llvm/ExecutionEngine/Orc/MemoryAccess.h" +#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h" +#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h" +#include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h" +#include "llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h" +#include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h" +#include "llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h" +#include "llvm/ExecutionEngine/Orc/SimpleRemoteMemoryMapper.h" +#include "llvm/ExecutionEngine/Orc/TaskDispatch.h" +#include "llvm/Support/Error.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/TargetSelect.h" + +#include +#include + +using namespace llvm; +using namespace llvm::orc; + +// Public CAPI from librt_capi (declared in runtime/include/RuntimeCAPI.h). +// Allocates a host buffer and registers it with CTX->getMemoryManager(). +extern "C" void *__catalyst__rt__alloc_managed(size_t size); + +namespace { + +// The connection behaviours below are mirrored from `llvm/tools/llvm-jitlink/llvm-jitlink.cpp`, +// and the Session class is mirrored from the official LLVM JIT tutorial: +// `https://llvm.org/docs/tutorial/BuildingAJIT1.html`. + +// Memref descriptor layout: allocated*, aligned*, offset, sizes[], strides[]. +// Offsets of the fixed-size prefix are independent of rank. +// The information can be obtained from `mlir/ExecutionEngine/CRunnerUtils.h`. +constexpr size_t kAllocatedOff = 0; +constexpr size_t kAlignedOff = sizeof(void *); +constexpr size_t kOffsetOff = sizeof(void *) * 2; +constexpr size_t kShapeOff = sizeof(void *) * 2 + sizeof(size_t); + +void initialize_targets() +{ + static const bool inited = []() { + InitializeAllTargets(); + InitializeAllTargetMCs(); + InitializeAllAsmPrinters(); + return true; + }(); + (void)inited; +} + +// A SimpleRemoteEPC transport that wraps the stock FDSimpleRemoteEPCTransport and adds a +// shutdown() in disconnect(). +// +// Teardown of a socket-backed session otherwise deadlocks: SimpleRemoteEPC::disconnect() (called +// from ExecutionSession::endSession()) does `T->disconnect()` then waits for the transport's +// listenLoop to observe end-of-stream and call handleDisconnect. FDSimpleRemoteEPCTransport's +// disconnect() only close()s the socket fd — but on Linux close() neither wakes a read() already +// in flight on the listener thread nor sends a FIN while that read holds a reference, so the +// listener never returns and the wait blocks forever (and the peer never sees EOF either). +// +// shutdown(SHUT_RDWR) forces the in-flight read() to return and sends FIN to the peer. Doing it +// here — inside disconnect(), which runs only after endSession() has finished its executor +// messaging (e.g. deallocating JIT'd memory) — avoids tearing the socket down prematurely. +class ShutdownFDTransport : public SimpleRemoteEPCTransport { + public: + static Expected> Create(SimpleRemoteEPCTransportClient &C, + int InFD, int OutFD) + { + auto Inner = FDSimpleRemoteEPCTransport::Create(C, InFD, OutFD); + if (!Inner) { + return Inner.takeError(); + } + return std::make_unique(std::move(*Inner), InFD, OutFD); + } + + ShutdownFDTransport(std::unique_ptr inner, int inFD, int outFD) + : Inner(std::move(inner)), InFD(inFD), OutFD(outFD) + { + } + + Error start() override { return Inner->start(); } + + Error sendMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo, ExecutorAddr TagAddr, + ArrayRef ArgBytes) override + { + return Inner->sendMessage(OpC, SeqNo, TagAddr, ArgBytes); + } + + void disconnect() override + { + ::shutdown(InFD, SHUT_RDWR); + if (OutFD != InFD) { + ::shutdown(OutFD, SHUT_RDWR); + } + Inner->disconnect(); + } + + private: + std::unique_ptr Inner; + int InFD, OutFD; +}; + +// For avoiding the error message being overwritten by subsequent errors in async jobs. +// We use thread_local to store the error message. +thread_local std::string g_last_error; +void set_error(const std::string &msg) { g_last_error = msg; } +void clear_error() { g_last_error.clear(); } + +void check(Error E, const Twine &what) +{ + if (E) { + throw std::runtime_error((what + ": " + toString(std::move(E))).str()); + } +} + +// unwrap LLVM Expected to C++ exception +template T unwrap(Expected v, const Twine &what) +{ + check(v.takeError(), what); + return std::move(*v); +} + +// TCP connect (mirrored from llvm-jitlink) +std::string OutOfProcessExecutorConnect; + +Error createTCPSocketError(Twine Details) +{ + return make_error("Failed to connect TCP socket '" + + Twine(OutOfProcessExecutorConnect) + "': " + Details, + inconvertibleErrorCode()); +} + +// Seconds to wait for the catalyst-executor ORC bootstrap handshake before giving up. Overridable +// via CATALYST_REMOTE_CONNECT_TIMEOUT; defaults to 10s. A value of 0 disables the timeout (the +// handshake then blocks indefinitely, matching the upstream llvm-jitlink behaviour). +unsigned remoteConnectTimeoutSeconds() +{ + if (const char *env = std::getenv("CATALYST_REMOTE_CONNECT_TIMEOUT")) { + char *end = nullptr; + unsigned long secs = std::strtoul(env, &end, 10); + if (end != env) { + return static_cast(secs); + } + } + return 10; +} + +Expected connectTCPSocket(std::string Host, std::string PortStr) +{ + addrinfo *AI; + addrinfo Hints{}; + Hints.ai_family = AF_INET; + Hints.ai_socktype = SOCK_STREAM; + Hints.ai_flags = AI_NUMERICSERV; + + if (int EC = getaddrinfo(Host.c_str(), PortStr.c_str(), &Hints, &AI)) { + return createTCPSocketError("Address resolution failed (" + StringRef(gai_strerror(EC)) + + ")"); + } + + // Cycle through the returned addrinfo structures and connect to the first + // reachable endpoint. + int SockFD; + addrinfo *Server; + for (Server = AI; Server != nullptr; Server = Server->ai_next) { + // socket might fail, e.g. if the address family is not supported. Skip to + // the next addrinfo structure in such a case. + if ((SockFD = socket(AI->ai_family, AI->ai_socktype, AI->ai_protocol)) < 0) + continue; + + // If connect returns null, we exit the loop with a working socket. + if (connect(SockFD, Server->ai_addr, Server->ai_addrlen) == 0) + break; + + close(SockFD); + } + freeaddrinfo(AI); + + // If we reached the end of the loop without connecting to a valid endpoint, + // dump the last error that was logged in socket() or connect(). + if (Server == nullptr) { + return createTCPSocketError(std::strerror(errno)); + } + + return SockFD; +} + +// Slab-based JIT-link memory manager: reserve a 1 GB slab on the remote once +Expected> +createSimpleRemoteMemoryManager(SimpleRemoteEPC &SREPC) +{ + SimpleRemoteMemoryMapper::SymbolAddrs SAs; + if (auto Err = SREPC.getBootstrapSymbols( + {{SAs.Instance, rt::SimpleExecutorMemoryManagerInstanceName}, + {SAs.Reserve, rt::SimpleExecutorMemoryManagerReserveWrapperName}, + {SAs.Initialize, rt::SimpleExecutorMemoryManagerInitializeWrapperName}, + {SAs.Deinitialize, rt::SimpleExecutorMemoryManagerDeinitializeWrapperName}, + {SAs.Release, rt::SimpleExecutorMemoryManagerReleaseWrapperName}})) { + return std::move(Err); + } + // 1 GB for object's sections (e.g .text, .rodata, ...) + // It will be released once the Session is destroyed. + // TODO: we might want to make this configurable. + size_t SlabSize = 1024 * 1024 * 1024; + return MapperJITLinkMemoryManager::CreateWithMapper(SlabSize, SREPC, + SAs); +} + +Expected> getFile(const Twine &filename) +{ + auto F = MemoryBuffer::getFile(filename); + if (F) { + return std::move(*F); + } + return createFileError(filename, F.getError()); +} + +} // namespace + +namespace catalyst::remote { + +struct RemoteSession { + std::unique_ptr ES; + + DataLayout DL; + + MangleAndInterner Mangle; + ObjectLinkingLayer ObjectLayer; + + // Shared namespace: the process-symbol generator (QIR runtime, libc) plus library-call objects. + JITDylib &MainJD; + // One JITDylib per shipped kernel object, keyed by its object-file path. Each links against + // MainJD for its external deps, but keeps its own entry symbol isolated — so two objects can + // export the same `_catalyst_pyface_` without colliding. + std::map KernelJDs; + + ExecutorAddr alloc_fn{0}; + ExecutorAddr free_fn{0}; + ExecutorAddr invoke_fn{0}; + ExecutorAddr store_asset_fn{0}; + + RemoteSession(std::unique_ptr es, DataLayout dl) + : ES(std::move(es)), DL(std::move(dl)), Mangle(*this->ES, this->DL), ObjectLayer(*this->ES), + MainJD(this->ES->createBareJITDylib("
")) + { + MainJD.addGenerator( + cantFail(EPCDynamicLibrarySearchGenerator::GetForTargetProcess(*this->ES))); + } + + ~RemoteSession() + { + if (auto Err = ES->endSession()) { + ES->reportError(std::move(Err)); + } + } + + static Expected> Create(StringRef remote_addr) + { + initialize_targets(); + + OutOfProcessExecutorConnect = remote_addr.str(); + auto [Host, PortStr] = remote_addr.split(':'); + if (Host.empty()) { + return createTCPSocketError("Host name for -" + OutOfProcessExecutorConnect + + " can not be empty"); + } + if (PortStr.empty()) { + return createTCPSocketError("Port number in -" + OutOfProcessExecutorConnect + + " can not be empty"); + } + + auto SockFD = connectTCPSocket(Host.str(), PortStr.str()); + if (!SockFD) { + return SockFD.takeError(); + } + + auto setup = SimpleRemoteEPC::Setup(); + setup.CreateMemoryManager = createSimpleRemoteMemoryManager; + // The ORC bootstrap handshake inside SimpleRemoteEPC::Create is an unbounded blocking read + // on the socket. If the peer is not a live catalyst-executor (e.g. a stale port-forward or + // a dead SSH tunnel accepted the connection), it would hang forever. A watchdog shuts the + // socket down after the timeout, which forces the blocked read to fail and Create to return + // an error instead of hanging. + const int sockFd = *SockFD; + const unsigned timeoutSecs = remoteConnectTimeoutSeconds(); + std::mutex mtx; + std::condition_variable cv; + bool handshakeDone = false; + bool timedOut = false; + std::thread watchdog; + if (timeoutSecs > 0) { + watchdog = std::thread([&] { + std::unique_lock lock(mtx); + if (!cv.wait_for(lock, std::chrono::seconds(timeoutSecs), + [&] { return handshakeDone; })) { + timedOut = true; + ::shutdown(sockFd, SHUT_RDWR); + } + }); + } + + // ShutdownFDTransport instead of FDSimpleRemoteEPCTransport so teardown shutdown()s the + // socket and doesn't deadlock — see the class comment above. + auto EPC = SimpleRemoteEPC::Create( + std::make_unique(std::nullopt), std::move(setup), + sockFd, sockFd); + + if (watchdog.joinable()) { + { + std::lock_guard lock(mtx); + handshakeDone = true; + } + cv.notify_all(); + watchdog.join(); + } + + if (timedOut) { + // Consume any error the forced socket shutdown produced before reporting the timeout. + if (!EPC) { + consumeError(EPC.takeError()); + } + return createTCPSocketError("handshake with catalyst-executor timed out after " + + Twine(timeoutSecs) + + "s (is a catalyst-executor actually listening there?). " + "Override with CATALYST_REMOTE_CONNECT_TIMEOUT"); + } + if (!EPC) { + return EPC.takeError(); + } + + JITTargetMachineBuilder JTMB((*EPC)->getTargetTriple()); + auto DL = JTMB.getDefaultDataLayoutForTarget(); + if (!DL) { + return DL.takeError(); + } + + auto ES = std::make_unique(std::move(*EPC)); + return std::make_unique(std::move(ES), std::move(*DL)); + } + + // Load a kernel object into its own JITDylib (keyed by `path`), linked against MainJD for deps. + Error addObjectFile(StringRef path, std::unique_ptr Buf) + { + JITDylib &jd = ES->createBareJITDylib(("kernel:" + path).str()); + jd.addToLinkOrder(MainJD); + KernelJDs[path.str()] = &jd; + return ObjectLayer.add(jd, std::move(Buf)); + } + + // Resolve `Name` in the JITDylib of the object identified by `path` (kernel entry points). + // Falls back to MainJD if the object is unknown. + ExecutorAddr lookupSym(StringRef path, StringRef Name) + { + auto it = KernelJDs.find(path.str()); + JITDylib *jd = (it != KernelJDs.end()) ? it->second : &MainJD; + auto Sym = unwrap(ES->lookup({jd}, Mangle(Name.str())), "lookup(" + Name + ")"); + return Sym.getAddress(); + } + + // Resolve `Name` in the shared process namespace (library-call symbols). + ExecutorAddr lookupSym(StringRef Name) + { + auto Sym = unwrap(ES->lookup({&MainJD}, Mangle(Name.str())), "lookup(" + Name + ")"); + return Sym.getAddress(); + } + + ExecutorProcessControl &getEPC() { return ES->getExecutorProcessControl(); } +}; + +// --------------------------------------------------------------------------- +// Memref Marshalling Helpers +// --------------------------------------------------------------------------- + +namespace { + +ExecutorAddr remote_alloc(RemoteSession *s, size_t size) +{ + ExecutorAddr ret; + auto &epc = s->getEPC(); + std::string error_prefix = "alloc(" + std::to_string(size) + ")"; + check(epc.callSPSWrapper(s->alloc_fn, ret, + static_cast(size)), + error_prefix); + if (!ret) { + throw std::runtime_error(error_prefix + " out of memory"); + } + return ret; +} + +void remote_free(RemoteSession *s, ExecutorAddr addr) +{ + auto &epc = s->getEPC(); + if (auto err = epc.callSPSWrapper(s->free_fn, addr)) { + consumeError(std::move(err)); + } +} + +void remote_write(RemoteSession *s, ExecutorAddr addr, const void *data, size_t size) +{ + auto &epc = s->getEPC(); + tpctypes::BufferWrite w{addr, ArrayRef(static_cast(data), size)}; + check(epc.getMemoryAccess().writeBuffers({w}), "write"); +} + +void remote_read(RemoteSession *s, ExecutorAddr addr, void *data, size_t size) +{ + ExecutorAddrRange r(addr, addr + size); + auto out = unwrap(s->getEPC().getMemoryAccess().readBuffers({r}), "read"); + if (out.empty()) { + throw std::runtime_error("read: empty read"); + } + if (out[0].size() != size) { + throw std::runtime_error("read: size mismatch"); + } + std::memcpy(data, out[0].data(), size); +} + +void remote_invoke(RemoteSession *s, ExecutorAddr entry, ArrayRef arg_addrs) +{ + auto &epc = s->getEPC(); + check(epc.callSPSWrapper)>(s->invoke_fn, + entry, arg_addrs), + "remote_invoke"); +} + +// Memref descriptors are layout-equivalent to: +// { void *allocated; void *aligned; size_t offset; size_t sizes[rank]; size_t strides[rank]; } +size_t memref_desc_size(size_t rank) +{ + return sizeof(void *) // void *allocated + + sizeof(void *) // void *aligned + + sizeof(size_t) // size_t offset + + sizeof(size_t) * rank // size_t sizes[rank] + + sizeof(size_t) * rank // size_t strides[rank] + ; +} + +size_t memref_data_size(const char *desc, size_t rank, size_t elem_size) +{ + if (rank == 0) { + return elem_size; + } + const size_t *shape = reinterpret_cast(desc + kShapeOff); + return std::accumulate(shape, shape + rank, elem_size, std::multiplies()); +} + +class RemoteAllocator { + private: + RemoteSession *sess; + std::vector addrs; + + public: + explicit RemoteAllocator(RemoteSession *s) : sess(s) {} + ~RemoteAllocator() + { + for (ExecutorAddr a : addrs) { + remote_free(sess, a); + } + } + RemoteAllocator(const RemoteAllocator &) = delete; + RemoteAllocator &operator=(const RemoteAllocator &) = delete; + + ExecutorAddr alloc(size_t size) + { + ExecutorAddr addr = remote_alloc(sess, size); + addrs.push_back(addr); + return addr; + } +}; + +} // namespace + +// --------------------------------------------------------------------------- +// RemoteSession's Exported APIs +// --------------------------------------------------------------------------- + +/** + * @brief Open a session to the remote device. + * + * @param remote_addr the remote address + * @return RemoteSession * the session object + */ +RemoteSession *open(const char *remote_addr) +{ + clear_error(); + try { + auto s = unwrap(RemoteSession::Create(remote_addr), "open(" + Twine(remote_addr) + ")"); + check(s->getEPC().getBootstrapSymbols({{s->alloc_fn, "catalyst_remote_alloc"}, + {s->free_fn, "catalyst_remote_free"}, + {s->invoke_fn, "catalyst_remote_invoke"}, + {s->store_asset_fn, "catalyst_remote_store_asset"}}), + "getBootstrapSymbols"); + return s.release(); + } + catch (const std::exception &e) { + set_error(e.what()); + return nullptr; + } +} + +/** + * @brief Close the session to the remote device. + * The remote session's destructor will handle the cleanup of the session. + * + * @param s the session object + */ +void close(RemoteSession *s) { delete s; } + +/** + * @brief Load an object file (cross-compiled for the remote arch) into the remote JIT. + * + * @param s the session object + * @param path the path to the object file + * @return int 0 on success, non-zero on error + */ +int load_object_path(RemoteSession *s, const char *path) +{ + clear_error(); + try { + auto buf = unwrap(getFile(path), "getFile(" + Twine(path) + ")"); + check(s->addObjectFile(path, std::move(buf)), "addObjectFile"); + return 0; + } + catch (const std::exception &e) { + set_error(e.what()); + return -1; + } +} + +/** + * @brief Load an asset file into the remote JIT. + * + * @param s the session object + * @param path the path to the asset file + * @return int 0 on success, -1 on error + */ +int load_asset_path(RemoteSession *s, const char *path) +{ + clear_error(); + try { + auto buf = unwrap(getFile(path), "getFile(" + Twine(path) + ")"); + ArrayRef bytes(buf->getBufferStart(), buf->getBufferSize()); + + int32_t rc = 0; + check(s->getEPC().callSPSWrapper, shared::SPSString)>( + s->store_asset_fn, rc, bytes, std::string(path)), + "store_asset"); + if (rc != 0) { + throw std::runtime_error("Got non-zero status from store_asset(" + std::string(path) + + "): " + std::to_string(rc)); + } + return 0; + } + catch (const std::exception &e) { + set_error(e.what()); + return -1; + } +} + +/** + * @brief Generic raw ORC wrapper-function call. Looks `sym` up on the executor, invokes its wrapper + * with `(args_buf, args_size)`, and copies the resulting byte buffer into a `out_buf` with the size + * of `out_size`. The caller is responsible for `free()` the result buffer. + * + * @param s The session object. + * @param sym The symbol of the function to call. + * @param args_buf The buffer of the arguments. + * @param args_size The size of the arguments. + * @param out_buf The buffer of the result. + * @param out_size The size of the result. + * @return int 0 on success, -1 on error. + */ +int call_wrapper_raw(RemoteSession *s, const char *sym, const char *args_buf, size_t args_size, + char **out_buf, size_t *out_size) +{ + clear_error(); + *out_buf = nullptr; + *out_size = 0; + try { + ExecutorAddr fn = s->lookupSym(sym); + if (!fn) { + throw std::runtime_error(std::string("symbol not found: ") + sym); + } + auto result = s->getEPC().callWrapper(fn, ArrayRef(args_buf, args_size)); + size_t n = result.size(); + char *buf = nullptr; + if (n > 0) { + buf = static_cast(std::malloc(n)); + if (!buf) { + throw std::runtime_error("malloc failed for wrapper result"); + } + std::memcpy(buf, result.data(), n); + } + *out_buf = buf; + *out_size = n; + return 0; + } + catch (const std::exception &e) { + set_error(e.what()); + return -1; + } +} + +/** + * @brief Lookup the address of a symbol in the remote device. + * + * @param s the session object + * @param name the name of the symbol + * @return uint64_t the address of the symbol + */ +uint64_t lookup(RemoteSession *s, const char *name, const char *object) +{ + clear_error(); + try { + if (object && *object) { + return s->lookupSym(object, name).getValue(); + } + return s->lookupSym(name).getValue(); + } + catch (const std::exception &e) { + set_error(e.what()); + return 0; + } +} + +/** + * @brief Run the kernel as a main function (take argv as arguments, argc is the length of argv). + * + * @param s the session object + * @param entry the entry function address + * @param argv the command line arguments + * @return int32_t the exit code + */ +int32_t run_as_main(RemoteSession *s, uint64_t entry_addr, int argc, const char *const *argv) +{ + clear_error(); + try { + std::vector args; + args.reserve(argc); + for (int i = 0; i < argc; ++i) { + args.emplace_back(argv[i]); + } + return unwrap(s->getEPC().runAsMain(ExecutorAddr(entry_addr), args), "run_as_main"); + } + catch (const std::exception &e) { + set_error(e.what()); + return -1; + } +} + +/** + * @brief Push one host memref to the remote: + * 1. allocates the data buffer on the remote + * 2. allocates the descriptor on the remote (which has a pointer to the data buffer) + * 3. returns the descriptor's remote addr. + * 4. If `copy_data` is true, the data will be copied to the remote. + * It's used for input memrefs like arguments. + * In the case of output memrefs, the data will be copied back to the host, + * so we don't need to copy the data to the remote. + * + * @param s the session object + * @param alloc the remote allocator + * @param host_desc the host memref descriptor + * @param rank the rank of the memref + * @param elem_size the element size of the memref + * @param copy_data whether to copy the data to the remote + * @return ExecutorAddr the remote address of the memref descriptor + */ +ExecutorAddr push_memref(RemoteSession *s, RemoteAllocator &alloc, void *host_desc, size_t rank, + size_t elem_size, bool copy_data) +{ + char *desc_host = static_cast(host_desc); + size_t desc_size = memref_desc_size(rank); + size_t data_size = memref_data_size(desc_host, rank, elem_size); + + // memref descriptor layout: + // ┌──────────────────────┐ ┌──────┐ + // │memref .allocated┼┬─►│buffer│ + // │descriptor .aligned ─┼┘ └──────┘ + // │ .offset │ + // │ .shape │ + // │ .strides │ + // └──────────────────────┘ + + std::vector desc(desc_size); + std::memcpy(desc.data(), desc_host, desc_size); + + ExecutorAddr data_remote = ExecutorAddr(0); + if (data_size > 0) { + data_remote = alloc.alloc(data_size); + if (copy_data) { + void *aligned_host = *reinterpret_cast(desc_host + kAlignedOff); + if (aligned_host) { + remote_write(s, data_remote, aligned_host, data_size); + } + } + } + std::memcpy(desc.data() + kAllocatedOff, &data_remote, sizeof(uintptr_t)); + std::memcpy(desc.data() + kAlignedOff, &data_remote, sizeof(uintptr_t)); + std::memset(desc.data() + kOffsetOff, 0, sizeof(int64_t)); + + ExecutorAddr desc_remote = alloc.alloc(desc_size); + remote_write(s, desc_remote, desc.data(), desc.size()); + return desc_remote; +} + +/** + * @brief Pull a remote memref descriptor + its data back into the host descriptor. + * + * @param s the session object + * @param remote_desc the remote address of the memref descriptor + * @param host_desc the host memref descriptor + * @param rank the rank of the memref + * @param elem_size the element size of the memref + */ +void pull_memref(RemoteSession *s, ExecutorAddr remote_desc, void *host_desc, size_t rank, + size_t elem_size) +{ + size_t desc_size = memref_desc_size(rank); + std::vector desc(desc_size); + remote_read(s, remote_desc, desc.data(), desc.size()); + + uintptr_t aligned_remote; + std::memcpy(&aligned_remote, desc.data() + kAlignedOff, sizeof(uintptr_t)); + + size_t data_size = memref_data_size(desc.data(), rank, elem_size); + size_t alloc_size = std::max(data_size, 1); + void *aligned_host = __catalyst__rt__alloc_managed(alloc_size); + if (data_size && aligned_remote) { + remote_read(s, ExecutorAddr(aligned_remote), aligned_host, data_size); + } + uintptr_t aligned_addr = reinterpret_cast(aligned_host); + std::memcpy(desc.data() + kAllocatedOff, &aligned_addr, sizeof(uintptr_t)); + std::memcpy(desc.data() + kAlignedOff, &aligned_addr, sizeof(uintptr_t)); + std::memcpy(host_desc, desc.data(), desc_size); +} + +/** + * @brief Invoke a remote kernel. + * + * @param s the session object + * @param entry_addr the address of the kernel entry function + * @param num_inputs the number of input memrefs + * @param input_descs the input memref descriptors + * @param input_ranks the ranks of the input memrefs + * @param input_elem_sizes the element sizes of the input memrefs + * @param num_outputs the number of output memrefs + * @param output_descs the output memref descriptors + * @param output_ranks the ranks of the output memrefs + * @param output_elem_sizes the element sizes of the output memrefs + * @return int 0 on success, non-zero on error + */ +int invoke_kernel(RemoteSession *s, uint64_t entry_addr, size_t num_inputs, + void *const *input_descs, const size_t *input_ranks, + const size_t *input_elem_sizes, size_t num_outputs, void *const *output_descs, + const size_t *output_ranks, const size_t *output_elem_sizes) +{ + clear_error(); + RemoteAllocator allocator(s); + try { + // The remote executor's catalyst_remote_invoke calls the entry as Catalyst's pyface ABI: + // `void(rv*, av*)`. + + // Layout (av): + // av (argument) is a struct whose Nth field is a pointer to the Nth input memref + // descriptor. So av = [N x uintptr_t] (array of remote descriptor addresses). + // + // av ──►┌───────────┐ ┌──────────────────────┐ ┌──────┐ + // │slot0 (ptr)┼────►│memref .allocated┼┬─►│buffer│ + // │slot1 │ │descriptor .aligned ─┼┘ └──────┘ + // │slot2 │ │ .offset │ + // │ ... │ │ .shape │ + // │ │ │ .strides │ + // └───────────┘ └──────────────────────┘ + + // Layout (rv): + // rv is a struct whose Nth field is the Nth output memref descriptor (not a pointer) + // + // rv ──►┌─────┬─────┬─────┬─────┬─────────┐ + // │desc0│desc1│desc2│desc3│ ... │ + // └─────┴─────┴─────┴─────┴─────────┘ + // Each slot maps to a output memref descriptor + + // Step 1. Push the input memref descriptors (with data) to the remote. + std::vector input_remote_descs(num_inputs); + for (size_t i = 0; i < num_inputs; ++i) { + input_remote_descs[i] = push_memref(s, allocator, input_descs[i], input_ranks[i], + input_elem_sizes[i], /*copy_data=*/true); + } + + // Step 2. Allocate a remote buffer holding the input memref descriptors' pointers. + ExecutorAddr av_remote = ExecutorAddr(0); + if (num_inputs > 0) { + av_remote = allocator.alloc(sizeof(uintptr_t) * num_inputs); + std::vector av_buf(num_inputs); + for (size_t i = 0; i < num_inputs; ++i) { + av_buf[i] = input_remote_descs[i].getValue(); + } + remote_write(s, av_remote, av_buf.data(), sizeof(uintptr_t) * num_inputs); + } + + // Step 3. Allocate a remote buffer for kernel to write the output memref descriptors. + std::vector output_offsets(num_outputs); + size_t rv_total = 0; + for (size_t i = 0; i < num_outputs; ++i) { + output_offsets[i] = rv_total; + rv_total += memref_desc_size(output_ranks[i]); + } + ExecutorAddr rv_remote = ExecutorAddr(0); + if (rv_total > 0) { + rv_remote = allocator.alloc(rv_total); + } + + // Step 4. Invoke the kernel remotely. + std::vector arg_addrs = {rv_remote, av_remote}; + remote_invoke(s, ExecutorAddr(entry_addr), arg_addrs); + + // Step 5. Pull each output descriptor back from rv buffer. + if (rv_total > 0) { + std::vector rv_buf(rv_total); + remote_read(s, rv_remote, rv_buf.data(), rv_total); + for (size_t i = 0; i < num_outputs; ++i) { + size_t desc_size = memref_desc_size(output_ranks[i]); + size_t elem_size = output_elem_sizes[i]; + char *desc = rv_buf.data() + output_offsets[i]; + + uintptr_t aligned_remote; + std::memcpy(&aligned_remote, desc + kAlignedOff, sizeof(uintptr_t)); + + size_t data_size = memref_data_size(desc, output_ranks[i], elem_size); + size_t alloc_size = std::max(data_size, 1); + void *aligned_host = __catalyst__rt__alloc_managed(alloc_size); + if (data_size && aligned_remote) { + remote_read(s, ExecutorAddr(aligned_remote), aligned_host, data_size); + } + uintptr_t aligned_addr = reinterpret_cast(aligned_host); + std::memcpy(desc + kAllocatedOff, &aligned_addr, sizeof(uintptr_t)); + std::memcpy(desc + kAlignedOff, &aligned_addr, sizeof(uintptr_t)); + std::memcpy(output_descs[i], desc, desc_size); + } + } + return 0; + } + catch (const std::exception &e) { + set_error(e.what()); + return -1; + } +} + +const char *last_error() { return g_last_error.c_str(); } + +} // namespace catalyst::remote diff --git a/runtime/lib/remote/RemoteSession.hpp b/runtime/lib/remote/RemoteSession.hpp new file mode 100644 index 0000000000..2afebc487b --- /dev/null +++ b/runtime/lib/remote/RemoteSession.hpp @@ -0,0 +1,57 @@ +// 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. + +#pragma once + +#include +#include + +namespace catalyst::remote { + +// Opaque session handle. Created by open(), released by close(). +struct RemoteSession; + +// Open a TCP session to a `host:port` executor. Returns nullptr on error. +RemoteSession *open(const char *remote_addr); + +void close(RemoteSession *s); + +// Load an object file into the remote JIT. Returns 0 on success, -1 on error. +int load_object_path(RemoteSession *s, const char *path); + +// Load an asset file into the remote JIT. Returns 0 on success, -1 on error. +int load_asset_path(RemoteSession *s, const char *path); + +// Generic ORC wrapper-function call by symbol name. Returns 0 on success, -1 on error. +int call_wrapper_raw(RemoteSession *s, const char *sym, const char *args_buf, size_t args_size, + char **out_buf, size_t *out_size); + +// Look up a symbol address on the remote. When `object` is non-null/non-empty, the lookup is +// scoped to that object's JITDylib (kernel entries); otherwise it resolves against the shared +// process namespace (library calls). Returns 0 on error. +uint64_t lookup(RemoteSession *s, const char *name, const char *object = nullptr); + +// Run a remote function as `main(argc, argv)`. +int32_t run_as_main(RemoteSession *s, uint64_t entry_addr, int argc, const char *const *argv); + +// Invoke a remote kernel. Returns 0 on success, -1 on error. +int invoke_kernel(RemoteSession *s, uint64_t entry_addr, size_t num_inputs, + void *const *input_descs, const size_t *input_ranks, + const size_t *input_elem_sizes, size_t num_outputs, void *const *output_descs, + const size_t *output_ranks, const size_t *output_elem_sizes); + +// Last error message. +const char *last_error(); + +} // namespace catalyst::remote diff --git a/runtime/tests/CMakeLists.txt b/runtime/tests/CMakeLists.txt index bd9f9ba4d0..40d2b3ad71 100644 --- a/runtime/tests/CMakeLists.txt +++ b/runtime/tests/CMakeLists.txt @@ -57,6 +57,7 @@ target_sources(runner_tests_qir_runtime PRIVATE Test_DataView.cpp Test_NullQubit.cpp Test_ResourceTracker.cpp + Test_SharedLibraryManager.cpp ) # For tests we do require libpython in order to embed a Python interpreter. @@ -67,6 +68,19 @@ target_link_libraries(runner_tests_qir_runtime PRIVATE rtd_null_qubit ) +# Test_SharedLibraryManager dlopens +if(NOT APPLE) + set_property(TARGET runner_tests_qir_runtime APPEND PROPERTY BUILD_RPATH $ORIGIN) + set_property(TARGET runner_tests_qir_runtime APPEND PROPERTY BUILD_RPATH $ORIGIN/../lib) +else() + set_property(TARGET runner_tests_qir_runtime APPEND PROPERTY BUILD_RPATH @loader_path) + set_property(TARGET runner_tests_qir_runtime APPEND PROPERTY BUILD_RPATH @loader_path/../lib) +endif() + +target_compile_definitions(runner_tests_qir_runtime PRIVATE + RTD_NULL_QUBIT_LIB="$" +) + catch_discover_tests(runner_tests_qir_runtime) if(ENABLE_OPENQASM) diff --git a/runtime/tests/Test_SharedLibraryManager.cpp b/runtime/tests/Test_SharedLibraryManager.cpp new file mode 100644 index 0000000000..530cb68294 --- /dev/null +++ b/runtime/tests/Test_SharedLibraryManager.cpp @@ -0,0 +1,69 @@ +// 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 + +#include "catch2/catch_test_macros.hpp" +#include "catch2/matchers/catch_matchers_string.hpp" + +#include "Exception.hpp" +#include "ExecutionContext.hpp" + +using namespace Catalyst::Runtime; +using Catch::Matchers::ContainsSubstring; + +#ifdef __APPLE__ +constexpr const char *kPlatformExt = ".dylib"; +constexpr const char *kWrongExt = ".so"; +#else +constexpr const char *kPlatformExt = ".so"; +constexpr const char *kWrongExt = ".dylib"; +#endif + +static const std::string kNullQubitAbs = RTD_NULL_QUBIT_LIB; +static const std::string kNullQubitBasename = std::string("librtd_null_qubit") + kPlatformExt; + +TEST_CASE("SharedLibraryManager loads from absolute path", "[shared_lib]") +{ + REQUIRE_NOTHROW(SharedLibraryManager{kNullQubitAbs}); +} + +TEST_CASE("SharedLibraryManager rewrites wrong extension", "[shared_lib]") +{ + std::string mangled = kNullQubitAbs; + auto dot = mangled.find_last_of('.'); + REQUIRE(dot != std::string::npos); + mangled.resize(dot); + mangled += kWrongExt; + + REQUIRE_NOTHROW(SharedLibraryManager{mangled}); +} + +TEST_CASE("SharedLibraryManager falls back when directory is stale", "[shared_lib]") +{ + std::string mangled = "/this/path/does/not/exist/" + kNullQubitBasename; + REQUIRE_NOTHROW(SharedLibraryManager{mangled}); +} + +TEST_CASE("SharedLibraryManager throws with original filename on hard failure", "[shared_lib]") +{ + const std::string bogus = "lib_not_real.so"; + REQUIRE_THROWS_WITH(SharedLibraryManager{bogus}, ContainsSubstring(bogus)); +} + +TEST_CASE("SharedLibraryManager resolves a known symbol", "[shared_lib]") +{ + SharedLibraryManager mgr{kNullQubitAbs}; + REQUIRE(mgr.getSymbol("NullQubitFactory") != nullptr); +}