Value propagation produces a large memory footprint

[neNasko1]

To demonstrate the issue one may consult the following example:

import tracemalloc
import spox
from spox.opset.ai.onnx import v20 as op
import numpy as np
import gc

def app():
    a = op.const(np.zeros(10000))
    b = op.const(np.zeros(10000))
    for i in range(N):
        a = op.mul(a, a)
        gc.collect()
    return a

tracemalloc.start()

var = app()

print(tracemalloc.get_traced_memory())

tracemalloc.stop()

Upon execution, the outputs are as follows:

(spox) ➜  spox git:(main) ✗ python3 memtest10.py
(4494527, 4900843)
(spox) ➜  spox git:(main) ✗ python3 memtest100.py
(11821814, 12228285)

This example highlights that the Var-s persist with their values even when they are considered "unreachable" through the stable APIs. This issue stems from the current practice of consolidating all Var information in a singular linked location. Since most information is crucial during the subsequent building process, Node-s maintain references to their Input and Output entities, thus creating the present challenge.

Me and @cbourjau have been talking about ways to circumvent the issue at hand. In order of how hacky they are:

• Have the ownership mechanism be represented by some weakref-s. This would be extremely hacky to pull off, since we only need to garbage collect the computed values when all the paths are from the form User->Var->Node->...->Var->_value
• Have the user visible Var-s be a different object than the Var inside the creation graph. I am still a bit unsure whether this is something feasible as spox seems to rely on the fact that copying does not actually copy, however this invariant may be only needed for the building stage.
• Create a InnerVar which holds all information except the _value which is held onto by every object in the graph and the current Var becomes a wrapper around that attaching a _value attribute. I am not extremely sure if this would yield a breaking change somewhere downstream.

This is the current stage of affairs, can you weight in? @jbachurski @adityagoel4512

[cbourjau]

Thanks for the summary, @neNasko1! Could you please elaborate a bit more on the difference between bullet points two and three?

[neNasko1]

Thanks for the summary, @neNasko1! Could you please elaborate a bit more on the difference between bullet points two and three?

Sorry for not being clear enough, they should be mostly the same in functionality, but in 2) we have 2 Var-s instead of a Var and a InnerVar. I gave 2) to have a complete list of directions, but IMO is strictly worse than 3).

As this is pretty abstract at this point, I will try to have a draft PR implementing 3) this week to enable us to talk about things more concretely.