[mlir][python] Running passes can invalidate objects and lead to crashes

[rkayaith]

If C++ code deletes an operation the PyOperation::valid flag isn't updated, which can lead to crashes:

from mlir import ir
from mlir.dialects import builtin
from mlir.passmanager import PassManager

with ir.Context():
  module = builtin.ModuleOp.parse("""
    module {
      arith.constant 10
    }
  """)
  const_op = module.body.operations[0]
  print(const_op)
  PassManager.parse('builtin.module(canonicalize)').run(module)
  print(const_op)   # Segmentation fault (core dumped)

[llvmbot]

@llvm/issue-subscribers-mlir-python

[rkayaith]

Not sure the best way to avoid this, invalidate all ops in the LiveOperationMap conservatively? Walk the IR after the pass and see which operations still exist?

[ftynse]

I'd suggest walking the IR rooted at the pass manager root operation and invalidating all of those ops. Other ops are fine. We cannot guarantee that operations pointers to which are still live are actually the operation of the same kind, or even operations to start with, the pointer address could have been reused by the allocator (this does happen quite often and cannot be caught by asan).

[makslevental]

diff submitted here.