AMReX Initialize/Finalize as Context Manager

[ax3l]

We cannot rely on the last del/reference drop to an object to really release and free the object in Python:
https://stackoverflow.com/a/17644536/2719194

import amrex

amrex.initialize([])
mfab = amrex.MultiFab( ... )
del mfab  # freed in CPython, but not guaranteed by Python per se
amrex.finalize()

This code can lead to Python holding on to mfab over the lifetime of amrex.finalize(). When using memory in an AMReX area, their actual free after finalize() will be UB and usually segfault.

Refs.:

• https://stackoverflow.com/a/17644536/2719194
• https://doc.pypy.org/en/latest/gc_info.html

[WeiqunZhang]

The Arenas in amrex is fine in this. That's why BArena never dies.

[ax3l]

Yes, but tests that run on pure device MultiFabs should be tearing down cleanly as well before finalize().

import amrex

amrex.initialize([])
mfab = amrex.MultiFab(
    boxarr,
    distmap,
    num_components,
    num_ghost,
    amrex.MFInfo().set_arena(amrex.The_Device_Arena()),
)
del mfab  # freed in CPython, but not guaranteed by Python per se
amrex.finalize()

[WeiqunZhang]

The device arena is also fine.

Arena*
The_Device_Arena ()
{   
    if        (the_device_arena) {
        return the_device_arena;
    } else {
        return The_Null_Arena();
    }
}

[WeiqunZhang]

The_Null_Arena also never dies. The device memory has been freed at that point. As long as the memory in MultiFab is not being used, it will be fine.

[ax3l]

Hm, I get a double free after finalize on process shutdown when I keep such an mfab alive #30

[ax3l]

Also, I kinda want to be sure that I have a context where all Arena (CPU/GPU) memory is free'd again. In testing, we need to run 100s of small init/finalize tests in the same process quickly after reach other - they should not leak memory after each of of these tests.

[ax3l]

Oh I see what you mean, finalize will free all memory we allocated in arenas? @WeiqunZhang

[WeiqunZhang]

Correct. amrex::Finalize will free all the device memory and then delete the device arena. After that The_Device_Arena() will return a pointer to the Null Arena that is a static variable in a function, and the Null Arena does not do anything when its free function is called. In C/C++ codes, the static Null Arena is guaranteed to outlive a variable on the stack. Even if a MultiFab is a static variable in the main function, as long as it's declared after amrex::Initialize is called, we are still fine because the Null Arena is created in amrex::Initialize and the static objects will be destroyed in reverse orders when the program terminates.

But maybe things are different when python is involved.

[ax3l]

That sounds good. I'll try to build a reproducer for the issue I see and change to not use a fixture that is cached in pytest, when I need a multifab.

Another bandage I'll apply for the pattern above is to call gc.collect() programmatically when amrex.finalize() is called, to make the above example standard compliant without even requiring for this simple case that refs can outlive the memory arena.

[ax3l]

#214 binds FabArray::clear, which we can use in context managers to explicitly free the memory in MultiFabs when we do not need it anymore, independent of its destructor call time :)

[sayerhs]

@ax3l is there a reason why amrex_init is function scope instead of session scope?

[ax3l]

@sayerhs yes, the idea is to clean it out for every test and start fresh, so they are isolated. Example: different parameters in ParmParse and other lurking/static state that is reset in init/finalize cycles.