[MLIR][python bindings] invalidate ops after PassManager run

[makslevental]

Summary

Apropos of this discussion and this subsequently filed bug, this patch invalidates operations below the root after PassManager run on a root operation.

Suggestions on a better factorization are appreciated.

[makslevental]

cleanup/refactor/add comments

[makslevental]

more comments

[rkayaith]

mlir/lib/Bindings/Python/Pass.cpp

136–137

the while loops look like they could be for loops

mlir/lib/CAPI/IR/IR.cpp

563

I would use a switch to translate the enum, so we don't rely on the encoding being the same

mlir/test/python/pass_manager.py

181–183

could you add some more nesting here to show the traversal works, and also print module to show it's still valid

[makslevental]

incorporate comments

[makslevental]

newline

[rkayaith]

mlir/include/mlir-c/IR.h

606–611

Thinking about this some more, this is missing the options to stop/skip through WalkResult, and the reverse iteration order option. We could maybe add extra functions for some of those later, but we also don't want to end up with too many variants...

Maybe open a separate revision for just the new C api? I think it probably deserves its own dedicated tests anyways.

mlir/lib/Bindings/Python/IRModule.h

273

does this need to be a class? it doesn't seem like it's doing anything

mlir/test/python/pass_manager.py

191–196

these logs aren't actually CHECKd

194

should try printing the nested ones as well

[stellaraccident]

mlir/lib/Bindings/Python/Pass.cpp

119

When Mehdi and I were discussing this (years ago), I believe we had decided that it should be possible to disable such safety features for situations where you take your own responsibility to do the right thing. Can we add an optional argument like invalidate_ops = True? This is a rather heavyweight thing being done, and for performance-sensitive code, I would always opt to just do it the right way vs commingling uses of the API.

123

I thought I had devised a clever way to do this without a walk, but I can't remember now. I think it had something to do with maintaining a global table of visible Py ops by parent. Then you could just zip through that instead of having all of these API calls for what (I would think) is most likely a very sparse set of actually visible Python instances.

But I think there are some corners that make that hard.

[makslevental]

mlir/lib/Bindings/Python/Pass.cpp

119

i definitely support "take your own responsibility" in all possible cases so i'll add that.

123

correct me if i'm wrong but don't you need union-find to make this faster if you don't want to miss any transitive children? but then that union-find is only amortized faster and since this is the first and only place the walk is done, there's no actual amortization.

but maybe i'm misunderstanding what you're describing.

[makslevental]

quick comment incorporation

mlir/lib/Bindings/Python/Pass.cpp

119

added but not sure how to test since if you do indeed pass invalidate_ops=False and then try to print you'll get a segfault...

[stellaraccident]

mlir/lib/Bindings/Python/Pass.cpp

119

That's fine. There are a couple other untestable cases.

123

I'm not sure - I mainly remember thinking about it and believing there was another way. I think the observation was that the python wrappers do not come into existence randomly but as a result of descending down a path of python reachable wrappers. And then you could just keep a flat lookup table of live parents.

Mainly thinking that the common case is to not have dense, live references and wishing we weren't paying an extra traversal.

But probably not worth optimizing now.

[makslevental]

fix overcomplicated impl

[stellaraccident]

Looks good to me modulo comments about the c API.

[makslevental]

mlir/lib/Bindings/Python/Pass.cpp

123

i changed the walk/callback code since @rkayaith comment got me thinking i overcomplicated things. it reduces the number api calls but it's still a traversal. let me know if i've made some grievous error/assumption in writing it this way.

[makslevental]

In D155543#4521343, @stellaraccident wrote:

Looks good to me modulo comments about the c API.

yea i'm gonna add that before landing.

[stellaraccident]

The other thing we could do: define the invalidation to be very course, invalidating every op except for the op the pass is run on (and possibly its parents). This would only require iterating over the live set we already have and flipping the invalid but on everything, then un-invalidating the target op. That would scale with the number of live references vs the size of the IR.

[rkayaith]

mlir/lib/Bindings/Python/Pass.cpp

119

in the test you have, printing the func should work with invalidate_ops=False but fail with True

mlir/include/mlir-c/IR.h

MLIR_CAPI_EXPORTED void mlirOperationMoveAfter(MlirOperation op,
                                                MlirOperation other);
 
 /// Moves the given operation immediately before the other operation in its
 /// parent block. The given operation may be owner by the caller or by its
 /// current block. The other operation must belong to a block. In any case, the
 /// ownership is transferred to the block of the other operation.
 MLIR_CAPI_EXPORTED void mlirOperationMoveBefore(MlirOperation op,
                                                 MlirOperation other);
+
+typedef enum MlirWalkOrder {
+  MlirWalkPreOrder,
+  MlirWalkPostOrder
+} MlirWalkOrder;
+
+typedef void (*MlirOperationWalkCallback)(MlirOperation, void *);
+
+/// Walks operation `op` in `walkOrder` and calls `callback` on that operation.
+/// `*userData` is passed to the callback as well and can be used to tunnel some
+/// some context or other data into the callback.
+MLIR_CAPI_EXPORTED
+void mlirOperationWalk(MlirOperation op, MlirOperationWalkCallback callback,
+                       void *userData, MlirWalkOrder walkOrder);
+
 //===----------------------------------------------------------------------===//
 // Region API.
 //===----------------------------------------------------------------------===//
 
 /// Creates a new empty region and transfers ownership to the caller.
 MLIR_CAPI_EXPORTED MlirRegion mlirRegionCreate(void);
 
 /// Takes a region owned by the caller and destroys it.

mlir/lib/Bindings/Python/IRCore.cpp

 size_t PyMlirContext::clearLiveOperations() {
   for (auto &op : liveOperations)
     op.second.second->setInvalid();
   size_t numInvalidated = liveOperations.size();
   liveOperations.clear();
   return numInvalidated;
 }
 
+void PyMlirContext::setOperationInvalid(MlirOperation op) {
+  if (liveOperations.contains(op.ptr))
+    liveOperations[op.ptr].second->setInvalid();
+}
+
 size_t PyMlirContext::getLiveModuleCount() { return liveModules.size(); }
 
 pybind11::object PyMlirContext::contextEnter() {
   return PyThreadContextEntry::pushContext(*this);
 }
 
 void PyMlirContext::contextExit(const pybind11::object &excType,
                                 const pybind11::object &excVal,

mlir/lib/Bindings/Python/IRModule.h

public:
   size_t getLiveOperationCount();
 
   /// Clears the live operations map, returning the number of entries which were
   /// invalidated. To be used as a safety mechanism so that API end-users can't
   /// corrupt by holding references they shouldn't have accessed in the first
   /// place.
   size_t clearLiveOperations();
 
+  /// Sets an operation invalid. This is useful for when some non-bindings
+  /// code destroys the operation and the bindings need to made aware. For
+  /// example, in the case when pass manager is run.
+  void setOperationInvalid(MlirOperation op);
+
   /// Gets the count of live modules associated with this context.
   /// Used for testing.
   size_t getLiveModuleCount();
 
   /// Enter and exit the context manager.
   pybind11::object contextEnter();
   void contextExit(const pybind11::object &excType,
                    const pybind11::object &excVal,
private:
 
   bool emitErrorDiagnostics = false;
 
   MlirContext context;
   friend class PyModule;
   friend class PyOperation;
 };
 
+/// Bundles together a MlirOperationWalkCallback and PyMlirContext,
+/// the latter is to be passed as *userData to the MlirOperationWalkCallback
+/// for callbacks that need to manipulate the context.
+class PyOperationWalkerCallback {
+public:
+  PyOperationWalkerCallback(PyMlirContext *context,
+                            MlirOperationWalkCallback callback)
+      : context(context), callback(callback) {}
+
+  PyMlirContext *getContext() { return context; }
+
+  MlirOperationWalkCallback getCallback() { return callback; }
+
+private:
+  PyMlirContext *context;
+  MlirOperationWalkCallback callback;
+};
+
 /// Used in function arguments when None should resolve to the current context
 /// manager set instance.
 class DefaultingPyMlirContext
     : public Defaulting<DefaultingPyMlirContext, PyMlirContext> {
 public:
   using Defaulting::Defaulting;
   static constexpr const char kTypeDescription[] = "mlir.ir.Context";
   static PyMlirContext &resolve();

mlir/lib/Bindings/Python/Pass.cpp

py::class_<PyPassManager>(m, "PassManager", py::module_local())
             if (mlirLogicalResultIsFailure(status))
               throw py::value_error(std::string(errorMsg.join()));
           },
           py::arg("pipeline"),
           "Add textual pipeline elements to the pass manager. Throws a "
           "ValueError if the pipeline can't be parsed.")
       .def(
           "run",
           [](PyPassManager &passManager, PyOperationBase &op) {
+            // Mark all ops below the op that the passmanager will be rooted at
+            // as invalid.
+            auto *context = op.getOperation().getContext().get();
+            MlirOperationWalkCallback invalidatingCallback =
+                [](MlirOperation op, void *userData) {
+                  PyOperationWalkerCallback *walker =
+                      static_cast<PyOperationWalkerCallback *>(userData);
+                  walker->getContext()->setOperationInvalid(op);
+                };
+            PyOperationWalkerCallback operationWalker(context,
+                                                      invalidatingCallback);
+            auto numRegions =
+                mlirOperationGetNumRegions(op.getOperation().get());
+            for (int i = 0; i < numRegions; ++i) {
+              MlirRegion region =
+                  mlirOperationGetRegion(op.getOperation().get(), i);
+              for (MlirBlock block = mlirRegionGetFirstBlock(region);
+                   !mlirBlockIsNull(block);
+                   block = mlirBlockGetNextInRegion(block))
+                for (MlirOperation childOp = mlirBlockGetFirstOperation(block);
+                     !mlirOperationIsNull(childOp);
+                     childOp = mlirOperationGetNextInBlock(childOp))
+                  mlirOperationWalk(childOp, operationWalker.getCallback(),
+                                    &operationWalker, MlirWalkPostOrder);
+            }
+            // Actually run the pass manager.
             PyMlirContext::ErrorCapture errors(op.getOperation().getContext());
             MlirLogicalResult status = mlirPassManagerRunOnOp(
                 passManager.get(), op.getOperation().get());
             if (mlirLogicalResultIsFailure(status))
               throw MLIRError("Failure while executing pass pipeline",
                               errors.take());
           },
           py::arg("operation"),

mlir/lib/CAPI/IR/IR.cpp

 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/Verifier.h"
+#include "mlir/IR/Visitors.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Parser/Parser.h"
 
 #include <cstddef>
 #include <optional>
 
 using namespace mlir;
 

 void mlirOperationMoveAfter(MlirOperation op, MlirOperation other) {
   return unwrap(op)->moveAfter(unwrap(other));
 }
 
 void mlirOperationMoveBefore(MlirOperation op, MlirOperation other) {
   return unwrap(op)->moveBefore(unwrap(other));
 }
 
+void mlirOperationWalk(MlirOperation op, MlirOperationWalkCallback callback,
+                       void *userData, MlirWalkOrder walkOrder) {
+  switch (walkOrder) {
+
+  case MlirWalkPreOrder:
+    unwrap(op)->walk<mlir::WalkOrder::PreOrder>(
+        [callback, userData](Operation *op) { callback(wrap(op), userData); });
+    break;
+  case MlirWalkPostOrder:
+    unwrap(op)->walk<mlir::WalkOrder::PostOrder>(
+        [callback, userData](Operation *op) { callback(wrap(op), userData); });
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Region API.
 //===----------------------------------------------------------------------===//
 
 MlirRegion mlirRegionCreate() { return wrap(new Region); }
 
 bool mlirRegionEqual(MlirRegion region, MlirRegion other) {
   return unwrap(region) == unwrap(other);

mlir/test/python/pass_manager.py

 # RUN: %PYTHON %s 2>&1 | FileCheck %s
 
 import gc, sys
 from mlir.ir import *
 from mlir.passmanager import *
 from mlir.dialects.func import FuncOp
+from mlir.dialects.builtin import ModuleOp
+
 
 # Log everything to stderr and flush so that we have a unified stream to match
 # errors/info emitted by MLIR to stderr.
 def log(*args):
     print(*args, file=sys.stderr)
     sys.stderr.flush()
 
 
with Context():
         assert '"mlir.passmanager.PassManager._CAPIPtr"' in repr(pm_capsule)
         pm._testing_release()
         pm1 = PassManager._CAPICreate(pm_capsule)
         assert pm1 is not None  # And does not crash.
 
 
 run(testCapsule)
 
+
 # CHECK-LABEL: TEST: testConstruct
 @run
 def testConstruct():
     with Context():
         # CHECK: pm1: 'any()'
         # CHECK: pm2: 'builtin.module()'
         pm1 = PassManager()
         pm2 = PassManager("builtin.module")
with Context():
         # A registered pass should parse successfully.
         pm = PassManager.parse("builtin.module(func.func(print-op-stats{json=false}))")
         # CHECK: Roundtrip: builtin.module(func.func(print-op-stats{json=false}))
         log("Roundtrip: ", pm)
 
 
 run(testParseSuccess)
 
+
 # Verify successful round-trip.
 # CHECK-LABEL: TEST: testParseSpacedPipeline
 def testParseSpacedPipeline():
     with Context():
         # A registered pass should parse successfully even if has extras spaces for readability
         pm = PassManager.parse(
             """builtin.module(
         func.func( print-op-stats{ json=false } )
     )"""
         )
         # CHECK: Roundtrip: builtin.module(func.func(print-op-stats{json=false}))
         log("Roundtrip: ", pm)
 
 
 run(testParseSpacedPipeline)
 
+
 # Verify failure on unregistered pass.
 # CHECK-LABEL: TEST: testParseFail
 def testParseFail():
     with Context():
         try:
             pm = PassManager.parse("any(unknown-pass)")
         except ValueError as e:
             #      CHECK: ValueError exception: MLIR Textual PassPipeline Parser:1:1: error:
             # CHECK-SAME: 'unknown-pass' does not refer to a registered pass or pass pipeline
             #      CHECK: unknown-pass
             #      CHECK: ^
             log("ValueError exception:", e)
         else:
             log("Exception not produced")
 
 
 run(testParseFail)
 
+
 # Check that adding to a pass manager works
 # CHECK-LABEL: TEST: testAdd
 @run
 def testAdd():
     pm = PassManager("any", Context())
     # CHECK: pm: 'any()'
     log(f"pm: '{pm}'")
     # CHECK: pm: 'any(cse)'
with Context():
         pm.run(func)
 
 
 # CHECK: Operations encountered:
 # CHECK: func.func      , 1
 # CHECK: func.return        , 1
 run(testRunPipeline)
 
+
 # CHECK-LABEL: TEST: testRunPipelineError
 @run
 def testRunPipelineError():
     with Context() as ctx:
         ctx.allow_unregistered_dialects = True
         op = Operation.parse('"test.op"() : () -> ()')
         pm = PassManager.parse("any(cse)")
         try:
             pm.run(op)
         except MLIRError as e:
             # CHECK: Exception: <
             # CHECK:   Failure while executing pass pipeline:
             # CHECK:   error: "-":1:1: 'test.op' op trying to schedule a pass on an unregistered operation
             # CHECK:    note: "-":1:1: see current operation: "test.op"() : () -> ()
             # CHECK: >
-            print(f"Exception: <{e}>")
+            log(f"Exception: <{e}>")
+
+
+# CHECK-LABEL: TEST: testPostPassOpInvalidation
+@run
+def testPostPassOpInvalidation():
+    with Context() as ctx:
+        module = ModuleOp.parse(
+            """
+          module {
+            arith.constant 10
+            func.func @foo() {
+              arith.constant 10
+              return
+            }
+          }
+        """
+        )
+        const_op = module.body.operations[0]
+        log(const_op)
+        PassManager.parse("builtin.module(canonicalize)").run(module)
+        try:
+            log(const_op)
+        except RuntimeError as e:
+            log(e)
+        # CHECK:       module {
+        # CHECK-LABEL:   func.func @foo() {
+        # CHECK:           return
+        # CHECK:         }
+        # CHECK:       }
+        log(module)