[Fix] Don't retain the last tensor in layout-dynamic DLPack fills

python/flydsl/compiler/jit_argument.py

@@ -487,6 +487,8 @@ def ptr_fill(a, s, _open=_open):
         # for this arg, so they share a single ``__dlpack__()`` per launch. The
         # pointer fill opens the live tensor once and hands the dladaptor to the
         # layout fill through a thread-local (thread-safe; never mutates the arg).
+        # The layout fill drops it again: the adaptor keeps the tensor alive, and
+        # the thread-local lives as long as the cached CallState.
         plan = _LayoutPlan(self.shape_dyn_indices, self.stride_dyn_indices, bool(self.use_32bit_stride))
         shared = threading.local()
 
@@ -495,7 +497,7 @@ def ptr_fill(a, s, _open=_open, _shared=shared):
             _shared.dladaptor = ad
             s.value = ad.data_ptr
 
-        body = ["    _ad = _shared.dladaptor"]
+        body = ["    _ad = _shared.dladaptor", "    _shared.dladaptor = None"]
         terms = []
         if plan.shape:
             body.append("    sh = _ad.shape")

tests/unit/test_callstate_dispatch.py

@@ -12,7 +12,9 @@
 """
 
 import ctypes
+import gc
 import struct
+import weakref
 
 import pytest
 
@@ -99,3 +101,20 @@ def func_exe(packed):
         assert layout == _expected_layout_bytes(t)
         assert scalar == ival
         assert auto_stream in (None, 0)  # auto-stream slot stays NULL (default stream)
+
+
+def test_dlpack_dynamic_layout_fill_does_not_retain_tensor():
+    """The layout-dynamic ``DLTensorJitArg`` fills hand one DLPack adaptor from the
+    pointer fill to the layout fill. It must not outlive the launch: the fills live
+    as long as the cached CallState, so a kept adaptor pins each argument's last
+    tensor (and its device memory) until the next launch replaces it."""
+    proto = torch.empty((4, 8), dtype=torch.float32)
+    (dp_ctype, dp_fill), (buf_ctype, pack) = ja.DLTensorJitArg(proto).__c_abi_spec__()
+
+    t = torch.empty((4, 8), dtype=torch.float32)
+    dp_fill(t, dp_ctype(0))
+    pack(t, buf_ctype())
+    ref = weakref.ref(t)
+    del t
+    gc.collect()
+    assert ref() is None