catalyst
catalyst copied to clipboard
[Runtime] Each `qnode` should manage its own memory.
Right now, all qnodes share a global memory manager. The role of the global memory manager is to deallocate tensors generated during the execution of the JIT compiled context. We can follow the same strategy on a qnode basis. Each qnode can have their own memory manager that will insert tensors at run time into a set. Once the qnode returns some return values, the caller of qnode would take owernship of the values returned by removing them from the set. These new values are placed in the host memory manager. After the caller of the qnode takes ownership of the return values, we deallocate all the memory allocated during the qnode. I think it is possible to implement this fully in the runtime without any changes in the compiler but careful attention must be paid to qnodes returning global constants.