[ingress][torch] Use Runner in PyTorch backend

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",
 ]

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."""

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:
         """

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