fix cffi life cycle issue

kalcohol · Feb 9, 2025 · e28b968 · e28b968
1 parent 3dd8d86
commit e28b968
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/axengine/_axe.py b/axengine/_axe.py
@@ -216,32 +216,39 @@ def __init__(
         _inputs= engine_cffi.new(
             "AX_ENGINE_IO_BUFFER_T[{}]".format(self._io[0].nInputSize)
         )
-        self._io[0].pInputs = _inputs
         _outputs = engine_cffi.new(
             "AX_ENGINE_IO_BUFFER_T[{}]".format(self._io[0].nOutputSize)
         )
+        self._io_buffers = (_inputs, _outputs)
+        self._io[0].pInputs = _inputs
         self._io[0].pOutputs = _outputs
+
+        self._io_inputs_pool = []
         for i in range(len(self.get_inputs())):
             max_buf = 0
             for j in range(self._shape_count):
                 max_buf = max(max_buf, self._info[j][0].pInputs[i].nSize)
             self._io[0].pInputs[i].nSize = max_buf
             phy = engine_cffi.new("AX_U64*")
             vir = engine_cffi.new("AX_VOID**")
+            self._io_inputs_pool.append((phy, vir))
             ret = sys_lib.AX_SYS_MemAllocCached(
                 phy, vir, self._io[0].pInputs[i].nSize, self._align, self._cmm_token
             )
             if 0 != ret:
                 raise RuntimeError("Failed to allocate memory for input.")
             self._io[0].pInputs[i].phyAddr = phy[0]
             self._io[0].pInputs[i].pVirAddr = vir[0]
+
+        self._io_outputs_pool = []
         for i in range(len(self.get_outputs())):
             max_buf = 0
             for j in range(self._shape_count):
                 max_buf = max(max_buf, self._info[j][0].pOutputs[i].nSize)
             self._io[0].pOutputs[i].nSize = max_buf
             phy = engine_cffi.new("AX_U64*")
             vir = engine_cffi.new("AX_VOID**")
+            self._io_outputs_pool.append((phy, vir))
             ret = sys_lib.AX_SYS_MemAllocCached(
                 phy, vir, self._io[0].pOutputs[i].nSize, self._align, self._cmm_token
             )