diff --git a/lighthouse/execution/__init__.py b/lighthouse/execution/__init__.py index 49d73978..6769d027 100644 --- a/lighthouse/execution/__init__.py +++ b/lighthouse/execution/__init__.py @@ -1,6 +1,7 @@ -from .memory_manager import MemoryManager, GPUMemoryManager +from .memory_manager import MemoryManager, GPUMemoryManager, ExternalMemoryManager __all__ = [ + "ExternalMemoryManager", "GPUMemoryManager", "MemoryManager", ] diff --git a/lighthouse/execution/memory_manager.py b/lighthouse/execution/memory_manager.py index 8491eac6..98726054 100644 --- a/lighthouse/execution/memory_manager.py +++ b/lighthouse/execution/memory_manager.py @@ -43,6 +43,29 @@ def emit_memory_management_funcs(payload_module: ir.Module, **kwargs): pass +@dataclass +class ExternalMemoryManager(MemoryManager, abc.ABC): + """Abstract base class for pass-through memory management. + + It does not manage memory itself, but provides a uniform interface for + converting buffers to memref descriptors. + """ + + @abc.abstractmethod + def get_memrefs(self, inputs): + """Convert inputs to memref descriptors.""" + pass + + def alloc(self, name: str = None, **kwargs) -> ctypes.Structure: + raise NotImplementedError("Allocation not supported.") + + def get(self, name: str) -> ctypes.Structure: + raise NotImplementedError("Lookup not supported.") + + def deallocate_all(self): + raise NotImplementedError("Deallocation not supported.") + + @dataclass class DeviceMemoryManager(MemoryManager, abc.ABC): """Abstract base class for handling memory on accelerator devices.""" diff --git a/lighthouse/execution/runner.py b/lighthouse/execution/runner.py index ba907b3e..0eb64d84 100644 --- a/lighthouse/execution/runner.py +++ b/lighthouse/execution/runner.py @@ -21,7 +21,7 @@ from lighthouse.schedule import schedule_boilerplate from lighthouse.utils.memref import to_packed_args from lighthouse.utils.mlir import get_mlir_library_path -from .memory_manager import GPUMemoryManager, MemoryManager +from .memory_manager import GPUMemoryManager, ExternalMemoryManager, MemoryManager class RunnerCallable(typing.Protocol): @@ -60,7 +60,9 @@ def __init__( if c_runner_lib not in shared_libs: shared_libs.append(c_runner_lib) self.lib_dir = get_mlir_library_path() - self.shared_libs = self._find_shared_libs(shared_libs) + shared_libs = self._find_shared_libs(shared_libs) + # Remove duplicates, the same library cannot be loaded multiple times. + self.shared_libs = list(dict.fromkeys(shared_libs)) self.opt_level = opt_level self.engine = self._get_engine() @@ -100,7 +102,7 @@ def _numpy_to_memref_manager(self, inputs): def _execute_kernel( self, - host_input_buffers: list[np.ndarray], + host_input_buffers: list, payload_function_name: str = "", argument_access_callback: Optional[ Callable[[list[ctypes.Structure], ExecutionEngine, MemoryManager], None] @@ -129,11 +131,18 @@ def _execute_kernel( raise ValueError("host_input_buffers must be provided") if self.mem_manager_cls is None: + if any(not isinstance(buf, np.ndarray) for buf in host_input_buffers): + raise ValueError( + "host_input_buffers must be numpy arrays when no mem_manager_cls is provided" + ) mem_manager = None allocator = partial(self._numpy_to_memref_manager, host_input_buffers) elif self.mem_manager_cls is GPUMemoryManager: mem_manager = self.mem_manager_cls(self.engine) allocator = partial(mem_manager.clone_host_buffers, host_input_buffers) + elif issubclass(self.mem_manager_cls, ExternalMemoryManager): + mem_manager = self.mem_manager_cls(self.engine) + allocator = partial(mem_manager.get_memrefs, host_input_buffers) else: raise ValueError( f"Unsupported mem_manager_cls type: {self.mem_manager_cls}" @@ -176,7 +185,7 @@ def _prepare_args(inputs): def benchmark( self, - host_input_buffers: list[np.ndarray], + host_input_buffers: list, argument_access_callback: RunnerCallable = None, nruns: int = 100, nwarmup: int = 10, @@ -195,7 +204,7 @@ def benchmark( def execute( self, payload_function_name: str, - host_input_buffers: list[np.ndarray], + host_input_buffers: list, argument_access_callback: RunnerCallable = None, ) -> None: """ diff --git a/lighthouse/ingress/torch/compile.py b/lighthouse/ingress/torch/compile.py index 7dba58ae..7be471f6 100644 --- a/lighthouse/ingress/torch/compile.py +++ b/lighthouse/ingress/torch/compile.py @@ -6,10 +6,11 @@ from lighthouse import utils as lh_utils from lighthouse.ingress.torch import import_from_model +from lighthouse.execution.runner import Runner +from lighthouse.execution import ExternalMemoryManager from mlir import ir from mlir.dialects import bufferization from mlir.dialects import func -from mlir.execution_engine import ExecutionEngine import torch @@ -35,6 +36,16 @@ class BufferMetadata: device: torch.device +@dataclass +class TorchMemoryManager(ExternalMemoryManager): + """Pass-through memory manager for PyTorch tensors.""" + + @contextlib.contextmanager + def get_memrefs(self, inputs): + """Convert PyTorch tensors to memref descriptors.""" + yield [lh_utils.torch.to_memref(t) for t in inputs] + + class JITFunction: """ Wrapper around JIT-compiled MLIR function. @@ -57,9 +68,10 @@ def __init__( entry_func: str = "main", n_outputs: int | None = None, ): - self.eng = ExecutionEngine(module, opt_level=3, shared_libs=shared_libs) - self.eng.initialize() - self.fn = self.eng.lookup(entry_func) + self.runner = Runner( + module, mem_manager_cls=TorchMemoryManager, shared_libs=shared_libs + ) + self.entry_func = entry_func self.results = results self.n_outputs = n_outputs if n_outputs is not None else len(results) @@ -86,14 +98,7 @@ def __call__( # input data followed by output storage buffers. all_tensors = [arg.detach() for arg in args] all_tensors.extend(outs) - # Keep the ctypes chain alive for the duration of the MLIR call. - # to_packed_args() returns only the packed void* array; the intermediate - # ctypes structures (memref descriptors and pointers) would be freed - # immediately otherwise, causing a use-after-free in the MLIR function. - _memrefs = [lh_utils.torch.to_memref(t) for t in all_tensors] - _ctype_ptrs = [lh_utils.memref.to_ctype(m) for m in _memrefs] - mlir_args = lh_utils.memref.get_packed_arg(_ctype_ptrs) - self.fn(mlir_args) + self.runner.execute(self.entry_func, all_tensors) # Return only the outputs corresponding to the FX graph's actual return # values. torch-mlir may prepend extra results for in-place state