[hrx] Retain executable per kernel handle to fix run_chain use-after-free (#3545)

Author: Abhishek-Varma

python/utils/hostruntime/hrxruntime/_bindings.py

@@ -347,6 +347,7 @@ def decl(fn, restype, argtypes):
             _status_t,
             [_handle, ctypes.c_char_p, ctypes.POINTER(ctypes.c_uint32)],
         )
+        self.hrx_executable_retain = decl("hrx_executable_retain", None, [_handle])
         self.hrx_executable_release = decl("hrx_executable_release", None, [_handle])
 
         # Dispatch / sync

python/utils/hostruntime/hrxruntime/context.py

@@ -279,6 +279,10 @@ def lookup_export(self, exe, name: str) -> int:
         )
         return ordv.value
 
+    def retain_executable(self, exe):
+        if exe:
+            lib.hrx_executable_retain(exe)
+
     def release_executable(self, exe):
         if exe:
             lib.hrx_executable_release(exe)

python/utils/hostruntime/hrxruntime/hostruntime.py

@@ -106,13 +106,31 @@ class HRXKernelHandle(KernelHandle):
     """Handle for a loaded HRX executable (one XADX export)."""
 
     def __init__(
-        self, executable, export_ordinal, kernel_name, xclbin_path, insts_path
+        self, executable, export_ordinal, kernel_name, xclbin_path, insts_path, ctx=None
     ):
         self.executable = executable
         self.export_ordinal = export_ordinal
         self.kernel_name = kernel_name
         self.xclbin_path = xclbin_path
         self.insts_path = insts_path
+        # Own an independent libhrx reference to the executable. The executable
+        # cache holds only a single reference and drops it on LRU eviction; a
+        # live handle (e.g. every step of a batched run_chain, kept in the
+        # sequence callable for the whole dispatch) must not be left dangling
+        # when an unrelated load evicts its cache entry. Balanced in __del__.
+        self._ctx = ctx
+        if ctx is not None and executable:
+            ctx.retain_executable(executable)
+
+    def __del__(self):
+        ctx = getattr(self, "_ctx", None)
+        exe = getattr(self, "executable", None)
+        if ctx is not None and exe:
+            try:
+                ctx.release_executable(exe)
+            except Exception:
+                pass
+            self.executable = None
 
 
 class HRXKernelResult(KernelResult):
@@ -216,7 +234,9 @@ def load(self, npu_kernel, **kwargs) -> HRXKernelHandle:
         xclbin_path, insts_path, kernel_name = self._resolve_kernel(npu_kernel)
         exe, ordv = self._build_executable(xclbin_path, insts_path, kernel_name)
         self._executables.append(exe)
-        return HRXKernelHandle(exe, ordv, kernel_name, xclbin_path, insts_path)
+        return HRXKernelHandle(
+            exe, ordv, kernel_name, xclbin_path, insts_path, ctx=self._ctx
+        )
 
     def _prepare_bindings(self, args):
         """Validate/sync a run's args and return its HRX dispatch bindings.
@@ -453,7 +473,9 @@ def load(self, npu_kernel, **kwargs) -> HRXKernelHandle:
         if key in self._exe_cache:
             self._exe_cache.move_to_end(key)
             exe, ordv = self._exe_cache[key]
-            return HRXKernelHandle(exe, ordv, kernel_name, xclbin_path, insts_path)
+            return HRXKernelHandle(
+                exe, ordv, kernel_name, xclbin_path, insts_path, ctx=self._ctx
+            )
 
         exe, ordv = self._build_executable(xclbin_path, insts_path, kernel_name)
 
@@ -462,7 +484,9 @@ def load(self, npu_kernel, **kwargs) -> HRXKernelHandle:
             self._release_executable(old_exe)
         self._exe_cache[key] = (exe, ordv)
 
-        return HRXKernelHandle(exe, ordv, kernel_name, xclbin_path, insts_path)
+        return HRXKernelHandle(
+            exe, ordv, kernel_name, xclbin_path, insts_path, ctx=self._ctx
+        )
 
     def cleanup(self) -> None:
         """Release cached executables, then any tracked by the base runtime."""

test/python/npu-hrx/test_chain_hrx.py

@@ -39,6 +39,7 @@
 )
 from aie.iron.controlflow import range_
 from aie.utils.npukernel import NPUKernel
+from aie.utils.hostruntime.hrxruntime.hostruntime import CachedHRXRuntime
 
 _TILE = 16
 _SIZE = 1024
@@ -76,6 +77,38 @@ def add_one(input_buf: In, output_buf: Out, *, N: CompileTime[int]):
     return _add_one_design(input_buf, output_buf, N=N)
 
 
+def _add_two_design(input_buf: In, output_buf: Out, N: CompileTime[int]):
+    """Add 2 to every element -- a second, distinct executable (see add_one)."""
+    tile_ty = np.ndarray[(_TILE,), np.dtype[np.int32]]
+    tensor_ty = np.ndarray[(N,), np.dtype[np.int32]]
+
+    of_in = ObjectFifo(tile_ty, name="in")
+    of_out = ObjectFifo(tile_ty, name="out")
+
+    def core_body(of_in, of_out):
+        for _ in range_(N // _TILE):
+            elem_in = of_in.acquire(1)
+            elem_out = of_out.acquire(1)
+            for i in range_(_TILE):
+                elem_out[i] = elem_in[i] + 2
+            of_in.release(1)
+            of_out.release(1)
+
+    worker = Worker(core_body, fn_args=[of_in.cons(), of_out.prod()])
+
+    def sequence(a, b, in_h, out_h):
+        in_h.fill(a)
+        out_h.drain(b, wait=True)
+
+    rt = Runtime(sequence, [tensor_ty, tensor_ty, of_in.prod(), of_out.cons()])
+    return Program(iron.get_current_device(), rt, workers=[worker]).resolve_program()
+
+
+@compileconfig
+def add_two(input_buf: In, output_buf: Out, *, N: CompileTime[int]):
+    return _add_two_design(input_buf, output_buf, N=N)
+
+
 def _hrx_runtime():
     """The default NPU runtime, which is the HRX runtime under NPU_RUNTIME=hrx.
 
@@ -143,3 +176,49 @@ def test_deep_chain(hrx_kernel):
     for k, st in enumerate(stages):
         st.to("cpu")
         np.testing.assert_array_equal(st.numpy(), base + (k + 1))
+
+
+def test_chain_survives_executable_eviction():
+    """A run_chain handle must outlive eviction of its executable cache entry.
+
+    The executable cache holds a single libhrx reference per executable and
+    drops it on LRU eviction. A chain keeps every step's handle live for the
+    whole batched dispatch, so if a later load evicts an earlier step's entry
+    the handle must still own the executable -- otherwise the dispatch touches a
+    freed executable (hrx_stream_dispatch: base_executable == NULL). We force
+    the eviction by shrinking the cache to one entry and loading two distinct
+    executables, then chaining across both. Without the per-handle retain this
+    dispatches a freed executable and fails; with it the chain runs correctly.
+
+    Uses a fresh ``CachedHRXRuntime`` rather than the process-wide
+    ``DefaultNPURuntime`` singleton so the cache starts empty and eviction is
+    deterministic (a shared runtime could already hold entries, making the load
+    a cache hit or evicting an unrelated entry).
+    """
+    rt = CachedHRXRuntime()
+
+    xa, ia = add_one.specialize(N=_SIZE).compile()
+    xb, ib = add_two.specialize(N=_SIZE).compile()
+    ka = NPUKernel(str(xa), str(ia), kernel_name="MLIR_AIE")
+    kb = NPUKernel(str(xb), str(ib), kernel_name="MLIR_AIE")
+
+    rt._cache_size = 1
+    try:
+        h_add1 = rt.load(ka)  # cache: {add_one}
+        h_add2 = rt.load(kb)  # size==1 -> evicts + releases add_one's executable
+        # h_add1 now references an executable the cache no longer keeps alive.
+
+        base = np.arange(1, _SIZE + 1, dtype=np.int32)
+        a1 = iron.tensor(base, dtype=np.int32, device="npu")
+        c1 = iron.zeros(_SIZE, dtype=np.int32, device="npu")
+        a2 = iron.tensor(base, dtype=np.int32, device="npu")
+        c2 = iron.zeros(_SIZE, dtype=np.int32, device="npu")
+
+        rt.run_chain([(h_add1, [a1, c1]), (h_add2, [a2, c2])])
+
+        c1.to("cpu")
+        c2.to("cpu")
+        np.testing.assert_array_equal(c1.numpy(), base + 1)
+        np.testing.assert_array_equal(c2.numpy(), base + 2)
+    finally:
+        rt.cleanup()