a[bugfix] [ir_utils] d-hoc solution to memory leak

[egebeysel]

Ad-Hoc Bugfix against Memory Leak

The default current behaviour in the upstream keeps all (or a big portion) of the torch.Tensors and np.ndarrays (that are kept as memorviews) in memory until the process is killed. An upstream issue is opened, a proper and more well-rounded solution will come at one point.

The source of the problem is the ModuleBuilder object attributes: global_ref_tracker and fx_py_attr_tracker, their class is imported from iree.compile.

RefTracker class that is defined in iree.compiler.extras.fx_importer has these lines in it's track method:

if referrent is not Empty:
            weakref.finalize(referrent, self._ref_finalizer, ref_id)

referrent corresponds to a torch.Tensor (or a memoryview of a np.ndarray) and this line causes it to be kept in memory until the process dies. Normally, the function passed into the finalize is called automatically, sort of as an additional finalization callback, when referrent is garbage collected. Unfortunately, if these lines are present, both the weakref.finalize objects and their corresponding tensors are not released. In the abscence of the above lines however, the tensors and the arrays are garbage collected automatically, except for the last one that used the aot.export path, which can though be garbage collected with gc.collect() in the end.

The memory usage w/o the changes:

The memory usage w/ the changes:

Last drop is with gc.collect() and drops the last pair of tensors and arrays.

• Report this to the upstream issue

[egebeysel]

• a[bugfix] [ir_utils] d-hoc solution to memory leak #3 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[chrsmcgrr]

Looks good. I just remembered we should merge this into a feature branch until it gets merged to upstream main. I will quickly create a branch and you can change the base target.

[chrsmcgrr]

nit: remove

[chrsmcgrr]

Can you add a small description about why this RefTracker is here and what makes it different from the torch.RefTracker implementation?

[chrsmcgrr]

Can you retarget to integrate-iree-turbine-20241203

[chrsmcgrr]

And can you open a PR to update the requirements for attic to point to this branch (integrate-iree-turbine-20241203): https://github.com/RooflineAI/roof-mlir/blob/main/packages/attic/requirements.txt#L12

pip should be able to install this package as a git link