Open christopherbate opened 5 months ago
@llvm/issue-subscribers-mlir-python
Author: Christopher Bate (christopherbate)
This is all true but I just have to say there's basically no way out of this without a major overhaul of the ownership reconciliation system. And maybe not even then because you're essentially dealing with the unsolvable problem of managed memory on the C++ side vs unmanaged memory on the Python side because it's actually C. So we can patch up these issues using more heuristics but it's inevitable that the heuristics become unwieldy (and maybe they already have).
I wish there existed even a tedious but comprehensive solution - I would sit down and write it - but I don't think there is short of rewriting everything in Rust (lol).
The dangling pointers in liveOperations are the biggest problem imo. I think we need to overhaul that, possibly with a different API to scope access... Maybe a context manager that holds the map and guarantees PyOperation uniquing within its schedule.
Alternatively, we could back off of all of this liveness to accounting and just embrace that PyOperations are very ephemeral things vs trying to keep the worlds in sync.
PyOperations are very ephemeral things
It's late so I'm not thinking straight - what does this world look like? Re-creating/materializing a new PyOperation each time someone digs around in a Module? If so then I guess there's a perf hit for anyone doing that a lot? If that's the only thing (and we recover sanity), I'm for it.
Yeah, same as PyAttribute and PyType in that world. I'd be shocked if this caching is paying for itself by reducing an object allocation... You can't go an inch in Python without such an allocation anyway.
We could implement some extra special methods to make the is
check fudge over the difference.
I'd want to test such a change carefully with some downstreams vs just doing a rip and run. I'm trying to remember if there was a more nuanced rationale for the current scheme and not remembering one. I think it was just a premature optimization gone wrong.
I happened to stumble upon another instance of this problem before I reached this thread in my email, so I went ahead and fixed my problem here: https://github.com/llvm/llvm-project/pull/93339/. This will not be sufficient to fix this problem for at least two reasons:
checkValid
before print.Both seem fixable, though the former may require a bit of re-engineering of how modules are handled. So maybe there is a chance to improve the existing solution, potentially with explicit scoping as suggested above. FWIW, my initial "fix" to this was to call _clear_live_operations_inside
.
Yeah, we really need to so this fix. Sorry, I've been out of commission with illness and not staying in top of things. I want to find time next week to put some pressure on this patch.
Background:
The PyMlirContext has a mechanism for creating unique references to
MlirOperation
objects. When aPyOperation
object is created, it is interned into a map (together with an owningpy::object
) into a map calledliveOperations
held by thePyMlirContext
. This allows for various checks (assertions in the C++ code or tests written in Python) to occur to ensure that objects have the appropriate lifetimes. Since MLIR IR has a nested structure, it's possible to create a reference to an Operation in Python and then do something funny (e.g. delete the parent), and then try to use the child op somehow. In C++, printing a freedOperation
will probably crash, but in Python printing aPyOperation
whose underlyingOperation*
has been freed should recover gracefully. For reasons like this (I think), the PyOperation retains avalid
member variable which should be set tofalse
in such situations, and most operations inPyOperation
are guarded by a check to ensurevalid == true
.This seems to work well, but there are several holes in how lifetimes of PyOperations are tracked (and therefore how they are inserted or cleaned up from the
liveOperations
map). These are mostly documented inIRCore.cpp
as TODOs.For example, the following code appears to crash:
The reason is that even though
module
(PyModule
) was dereferenced and cleaned up, it didn't invalidate its children. Therefore, thePyOperation
pointed to byfunc
still has its valid bit set (and is still in theliveOperations
map).Furthermore, even if you don't do anything with
func
, it remaining inliveOperations
is problematic. If you re-use the same Context, as in the below code, the MLIR C++ API and the underlying allocator may return a pointer that is still inliveOperations
(since operations may be freed by MLIR but still exist inliveOperations
, that is a valid state):An assertion will be encountered because of this line: https://github.com/llvm/llvm-project/blame/main/mlir/lib/Bindings/Python/IRCore.cpp#L1164
Here you can see that even if we amend the assertion to include a check if it is
valid
, proceeding to update theliveOperations
map may invalidate a reference still held by the Python user. IMO this is overly cautious because there are too many ways in which theliveOperations
tracking mechanism can loose track of which operations are valid or not.In addition, compiling MLIR with CMAKE_BUILD_TYPE=RelWithDebInfo vs CMAKE_BUILD_TYPE=Debug appears to change the likelihood of encountering the assertion in the second example.