@@ -216,32 +216,39 @@ def __init__(
216
216
_inputs = engine_cffi .new (
217
217
"AX_ENGINE_IO_BUFFER_T[{}]" .format (self ._io [0 ].nInputSize )
218
218
)
219
- self ._io [0 ].pInputs = _inputs
220
219
_outputs = engine_cffi .new (
221
220
"AX_ENGINE_IO_BUFFER_T[{}]" .format (self ._io [0 ].nOutputSize )
222
221
)
222
+ self ._io_buffers = (_inputs , _outputs )
223
+ self ._io [0 ].pInputs = _inputs
223
224
self ._io [0 ].pOutputs = _outputs
225
+
226
+ self ._io_inputs_pool = []
224
227
for i in range (len (self .get_inputs ())):
225
228
max_buf = 0
226
229
for j in range (self ._shape_count ):
227
230
max_buf = max (max_buf , self ._info [j ][0 ].pInputs [i ].nSize )
228
231
self ._io [0 ].pInputs [i ].nSize = max_buf
229
232
phy = engine_cffi .new ("AX_U64*" )
230
233
vir = engine_cffi .new ("AX_VOID**" )
234
+ self ._io_inputs_pool .append ((phy , vir ))
231
235
ret = sys_lib .AX_SYS_MemAllocCached (
232
236
phy , vir , self ._io [0 ].pInputs [i ].nSize , self ._align , self ._cmm_token
233
237
)
234
238
if 0 != ret :
235
239
raise RuntimeError ("Failed to allocate memory for input." )
236
240
self ._io [0 ].pInputs [i ].phyAddr = phy [0 ]
237
241
self ._io [0 ].pInputs [i ].pVirAddr = vir [0 ]
242
+
243
+ self ._io_outputs_pool = []
238
244
for i in range (len (self .get_outputs ())):
239
245
max_buf = 0
240
246
for j in range (self ._shape_count ):
241
247
max_buf = max (max_buf , self ._info [j ][0 ].pOutputs [i ].nSize )
242
248
self ._io [0 ].pOutputs [i ].nSize = max_buf
243
249
phy = engine_cffi .new ("AX_U64*" )
244
250
vir = engine_cffi .new ("AX_VOID**" )
251
+ self ._io_outputs_pool .append ((phy , vir ))
245
252
ret = sys_lib .AX_SYS_MemAllocCached (
246
253
phy , vir , self ._io [0 ].pOutputs [i ].nSize , self ._align , self ._cmm_token
247
254
)
0 commit comments