[mlir][Python] Return and accept OpView for all functions.

[stellaraccident]

Summary

• All functions that return an Operation now return an OpView.
• All functions that accept an Operation now accept an _OperationBase, which both Operation and OpView extend and can resolve to the backing Operation.
• Moves user-facing instance methods from Operation -> _OperationBase so that both can have the same API.
• Concretely, this means that if there are custom op classes defined (i.e. in Python), any iteration or creation will return the appropriate instance (i.e. if you get/create an std.addf, you will get an instance of the mlir.dialects.std.AddFOp class, getting full access to any custom API it exposes).
• Refactors all eq methods after realizing the proper way to do this for _OperationBase.

[zhanghb97]

mlir/lib/Bindings/Python/Globals.h

55

Should we also change other const std::string & parameters in the class for the consistency. Is there a trade-off between llvm::StringRef and const std::string &?

[stellaraccident]

mlir/lib/Bindings/Python/Globals.h

55

I should take another look through them: I had placed them originally so as to have fewer copies between the places that needed strings and those that didn't. In general, StringRef is preferred where possible, but in a mixed codebase like pybind11, where std::string is used for many things, it can result in extra copies.

The real answer is to create a type caster for StringRef and SmallVector and then switch everything to them. I had looked into that briefly and it was a bigger change than any one patch and just needs to be done. If you're interested in giving it a try, it will be a while before I get to trying it.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

1043

What kind of extra caching do you have in mind, it seems to me here that we're down do a string hash map lookup?
(i.e. isn't the caching is already done in lookupRawOpViewClass?)

mlir/lib/Bindings/Python/MainModule.cpp

123

Isn't the negative cache preventing someone from having their dynamic loading of dialects? Or is the expectation are we flushing the cache on dialect load?

Looking closer, I think the negative cache would prevent someone from adding a path through append_dialect_search_prefix after the negative cache has been populated.
Can we flush the cache around this kind of manipulation? I'm wary of library behavior changing based on the implicit state of the cache.

[mehdi_amini]

(I even wonder if OpView wouldn't be a better name for the C++ class)

[zhanghb97]

mlir/lib/Bindings/Python/Globals.h

55

Looking forward to the type caster, and I feel that influence scope of the caster will be quite large.

[stellaraccident]

Comments and rebase.

mlir/lib/Bindings/Python/Globals.h

55

Feel free to give it a try, if you have the time. The caster for SmallVector is likely easier than StringRef fwiw.

mlir/lib/Bindings/Python/IRModules.cpp

1043

I was more concerned with the cost of the temporaries and construction overhead. There is likely something clever to keep from doing that but for now, I've removed the TODO.

mlir/lib/Bindings/Python/MainModule.cpp

123

Yep - totally right and thanks for spotting this.

I've reworked this to do proper caching and invalidation on appending to the search path. This caching could get in the way of various (nefarious) monkey patching, but I'm fine with that.

Also added some gil acquires. Since these are re-entrant, better to avoid the races.

mlir/lib/Bindings/Python/Globals.h

public:
   /// implementation classes.
   std::vector<std::string> &getDialectSearchPrefixes() {
     return dialectSearchPrefixes;
   }
   void setDialectSearchPrefixes(std::vector<std::string> newValues) {
     dialectSearchPrefixes.swap(newValues);
   }
 
+  /// Clears positive and negative caches regarding what implementations are
+  /// available. Future lookups will do more expensive existence checks.
+  void clearImportCache();
+
   /// Loads a python module corresponding to the given dialect namespace.
   /// No-ops if the module has already been loaded or is not found. Raises
   /// an error on any evaluation issues.
   /// Note that this returns void because it is expected that the module
   /// contains calls to decorators and helpers that register the salient
   /// entities.
-  void loadDialectModule(const std::string &dialectNamespace);
+  void loadDialectModule(llvm::StringRef dialectNamespace);
 
   /// Decorator for registering a custom Dialect class. The class object must
   /// have a DIALECT_NAMESPACE attribute.
   pybind11::object registerDialectDecorator(pybind11::object pyClass);
 
   /// Adds a concrete implementation dialect class.
   /// Raises an exception if the mapping already exists.
   /// This is intended to be called by implementation code.
   void registerDialectImpl(const std::string &dialectNamespace,
                            pybind11::object pyClass);
 
   /// Adds a concrete implementation operation class.
   /// Raises an exception if the mapping already exists.
   /// This is intended to be called by implementation code.
   void registerOperationImpl(const std::string &operationName,
                              pybind11::object pyClass,
-                             pybind11::object rawClass);
+                             pybind11::object rawOpViewClass);
 
   /// Looks up a registered dialect class by namespace. Note that this may
   /// trigger loading of the defining module and can arbitrarily re-enter.
   llvm::Optional<pybind11::object>
   lookupDialectClass(const std::string &dialectNamespace);
 
+  /// Looks up a registered raw OpView class by operation name. Note that this
+  /// may trigger a load of the dialect, which can arbitrarily re-enter.
+  llvm::Optional<pybind11::object>
+  lookupRawOpViewClass(llvm::StringRef operationName);
+
 private:
   static PyGlobals *instance;
   /// Module name prefixes to search under for dialect implementation modules.
   std::vector<std::string> dialectSearchPrefixes;
-  /// Map of dialect namespace to bool flag indicating whether the module has
-  /// been successfully loaded or resolved to not found.
-  llvm::StringSet<> loadedDialectModules;
   /// Map of dialect namespace to external dialect class object.
   llvm::StringMap<pybind11::object> dialectClassMap;
   /// Map of full operation name to external operation class object.
   llvm::StringMap<pybind11::object> operationClassMap;
   /// Map of operation name to custom subclass that directly initializes
   /// the OpView base class (bypassing the user class constructor).
-  llvm::StringMap<pybind11::object> rawOperationClassMap;
+  llvm::StringMap<pybind11::object> rawOpViewClassMap;
+
+  /// Set of dialect namespaces that we have attempted to import implementation
+  /// modules for.
+  llvm::StringSet<> loadedDialectModulesCache;
+  /// Cache of operation name to custom OpView subclass that directly
+  /// initializes the OpView base class (or an undefined object for negative
+  /// lookup). This is maintained on loopup as a shadow of rawOpViewClassMap
+  /// in order for repeat lookups of the OpView classes to only incur the cost
+  /// of one hashtable lookup.
+  llvm::StringMap<pybind11::object> rawOpViewClassMapCache;
 };
 
 } // namespace python
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_GLOBALS_H

mlir/lib/Bindings/Python/IRModules.h

public:
   static pybind11::object createFromCapsule(pybind11::object capsule);
 
 private:
   PyModule(PyMlirContextRef contextRef, MlirModule module);
   MlirModule module;
   pybind11::handle handle;
 };
 
+/// Base class for PyOperation and PyOpView which exposes the primary, user
+/// visible methods for manipulating it.
+class PyOperationBase {
+public:
+  virtual ~PyOperationBase() = default;
+  /// Implements the bound 'print' method and helps with others.
+  void print(pybind11::object fileObject, bool binary,
+             llvm::Optional<int64_t> largeElementsLimit, bool enableDebugInfo,
+             bool prettyDebugInfo, bool printGenericOpForm, bool useLocalScope);
+  pybind11::object getAsm(bool binary,
+                          llvm::Optional<int64_t> largeElementsLimit,
+                          bool enableDebugInfo, bool prettyDebugInfo,
+                          bool printGenericOpForm, bool useLocalScope);
+
+  /// Each must provide access to the raw Operation.
+  virtual PyOperation &getOperation() = 0;
+};
+
 /// Wrapper around PyOperation.
 /// Operations exist in either an attached (dependent) or detached (top-level)
 /// state. In the detached state (as on creation), an operation is owned by
 /// the creator and its lifetime extends either until its reference count
 /// drops to zero or it is attached to a parent, at which point its lifetime
 /// is bounded by its top-level parent reference.
 class PyOperation;
 using PyOperationRef = PyObjectRef<PyOperation>;
-class PyOperation : public BaseContextObject {
+class PyOperation : public PyOperationBase, public BaseContextObject {
 public:
   ~PyOperation();
+  PyOperation &getOperation() override { return *this; }
+
   /// Returns a PyOperation for the given MlirOperation, optionally associating
   /// it with a parentKeepAlive.
   static PyOperationRef
   forOperation(PyMlirContextRef contextRef, MlirOperation operation,
                pybind11::object parentKeepAlive = pybind11::object());
 
   /// Creates a detached operation. The operation must not be associated with
   /// any existing live operation.
public:
 
   bool isAttached() { return attached; }
   void setAttached() {
     assert(!attached && "operation already attached");
     attached = true;
   }
   void checkValid();
 
-  /// Implements the bound 'print' method and helps with others.
-  void print(pybind11::object fileObject, bool binary,
-             llvm::Optional<int64_t> largeElementsLimit, bool enableDebugInfo,
-             bool prettyDebugInfo, bool printGenericOpForm, bool useLocalScope);
-  pybind11::object getAsm(bool binary,
-                          llvm::Optional<int64_t> largeElementsLimit,
-                          bool enableDebugInfo, bool prettyDebugInfo,
-                          bool printGenericOpForm, bool useLocalScope);
-
   /// Gets the owning block or raises an exception if the operation has no
   /// owning block.
   PyBlock getBlock();
 
   /// Gets the parent operation or raises an exception if the operation has
   /// no parent.
   PyOperationRef getParentOperation();
 
   /// Creates an operation. See corresponding python docstring.
   static pybind11::object
   create(std::string name, llvm::Optional<std::vector<PyValue *>> operands,
          llvm::Optional<std::vector<PyType *>> results,
          llvm::Optional<pybind11::dict> attributes,
          llvm::Optional<std::vector<PyBlock *>> successors, int regions,
          DefaultingPyLocation location, pybind11::object ip);
 
+  /// Creates an OpView suitable for this operation.
+  pybind11::object createOpView();
+
 private:
   PyOperation(PyMlirContextRef contextRef, MlirOperation operation);
   static PyOperationRef createInstance(PyMlirContextRef contextRef,
                                        MlirOperation operation,
                                        pybind11::object parentKeepAlive);
 
   MlirOperation operation;
   pybind11::handle handle;
   // Keeps the parent alive, regardless of whether it is an Operation or
   // Module.
   // TODO: As implemented, this facility is only sufficient for modeling the
   // trivial module parent back-reference. Generalize this to also account for
   // transitions from detached to attached and address TODOs in the
   // ir_operation.py regarding testing corresponding lifetime guarantees.
   pybind11::object parentKeepAlive;
   bool attached = true;
   bool valid = true;
 };
 
 /// A PyOpView is equivalent to the C++ "Op" wrappers: these are the basis for
 /// providing more instance-specific accessors and serve as the base class for
 /// custom ODS-style operation classes. Since this class is subclass on the
 /// python side, it must present an __init__ method that operates in pure
 /// python types.
-class PyOpView {
+class PyOpView : public PyOperationBase {
 public:
-  PyOpView(pybind11::object operation);
+  PyOpView(pybind11::object operationObject);
+  PyOperation &getOperation() override { return operation; }
 
   static pybind11::object createRawSubclass(pybind11::object userClass);
 
   pybind11::object getOperationObject() { return operationObject; }
 
 private:
+  PyOperation &operation;           // For efficient, cast-free access from C++
   pybind11::object operationObject; // Holds the reference.
-  PyOperation *operation;           // For efficient, cast-free access from C++
 };
 
 /// Wrapper around an MlirRegion.
 /// Regions are managed completely by their containing operation. Unlike the
 /// C++ API, the python API does not support detached regions.
 class PyRegion {
 public:
   PyRegion(PyOperationRef parentOperation, MlirRegion region)

 /// reference operation. If the reference operation is null, then appends to
 /// the end of the block.
 class PyInsertionPoint {
 public:
   /// Creates an insertion point positioned after the last operation in the
   /// block, but still inside the block.
   PyInsertionPoint(PyBlock &block);
   /// Creates an insertion point positioned before a reference operation.
-  PyInsertionPoint(PyOperation &beforeOperation);
+  PyInsertionPoint(PyOperationBase &beforeOperationBase);
 
   /// Shortcut to create an insertion point at the beginning of the block.
   static PyInsertionPoint atBlockBegin(PyBlock &block);
   /// Shortcut to create an insertion point before the block terminator.
   static PyInsertionPoint atBlockTerminator(PyBlock &block);
 
   /// Inserts an operation.
-  void insert(PyOperation &operation);
+  void insert(PyOperationBase &operationBase);
 
   /// Enter and exit the context manager.
   pybind11::object contextEnter();
   void contextExit(pybind11::object excType, pybind11::object excVal,
                    pybind11::object excTb);
 
   PyBlock &getBlock() { return block; }
 
 private:
   // Trampoline constructor that avoids null initializing members while
   // looking up parents.
   PyInsertionPoint(PyBlock block, llvm::Optional<PyOperationRef> refOperation)
-      : block(std::move(block)), refOperation(std::move(refOperation)) {}
+      : refOperation(std::move(refOperation)), block(std::move(block)) {}
 
-  PyBlock block;
   llvm::Optional<PyOperationRef> refOperation;
+  PyBlock block;
 };
 
 /// Wrapper around the generic MlirAttribute.
 /// The lifetime of a type is bound by the PyContext that created it.
 class PyAttribute : public BaseContextObject {
 public:
   PyAttribute(PyMlirContextRef contextRef, MlirAttribute attr)
       : BaseContextObject(std::move(contextRef)), attr(attr) {}

mlir/lib/Bindings/Python/IRModules.cpp

py::object dunderNext() {
     parentOperation->checkValid();
     if (mlirOperationIsNull(next)) {
       throw py::stop_iteration();
     }
 
     PyOperationRef returnOperation =
         PyOperation::forOperation(parentOperation->getContext(), next);
     next = mlirOperationGetNextInBlock(next);
-    return returnOperation.releaseObject();
+    return returnOperation->createOpView();
   }
 
   static void bind(py::module &m) {
     py::class_<PyOperationIterator>(m, "OperationIterator")
         .def("__iter__", &PyOperationIterator::dunderIter)
         .def("__next__", &PyOperationIterator::dunderNext);
   }
 
py::object dunderGetItem(intptr_t index) {
     if (index < 0) {
       throw SetPyError(PyExc_IndexError,
                        "attempt to access out of bounds operation");
     }
     MlirOperation childOp = mlirBlockGetFirstOperation(block);
     while (!mlirOperationIsNull(childOp)) {
       if (index == 0) {
         return PyOperation::forOperation(parentOperation->getContext(), childOp)
-            .releaseObject();
+            ->createOpView();
       }
       childOp = mlirOperationGetNextInBlock(childOp);
       index -= 1;
     }
     throw SetPyError(PyExc_IndexError,
                      "attempt to access out of bounds operation");
   }
 

 }
 
 void PyOperation::checkValid() {
   if (!valid) {
     throw SetPyError(PyExc_RuntimeError, "the operation has been invalidated");
   }
 }
 
-void PyOperation::print(py::object fileObject, bool binary,
+void PyOperationBase::print(py::object fileObject, bool binary,
                             llvm::Optional<int64_t> largeElementsLimit,
                             bool enableDebugInfo, bool prettyDebugInfo,
                             bool printGenericOpForm, bool useLocalScope) {
-  checkValid();
+  PyOperation &operation = getOperation();
+  operation.checkValid();
   if (fileObject.is_none())
     fileObject = py::module::import("sys").attr("stdout");
   MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
   if (largeElementsLimit)
     mlirOpPrintingFlagsElideLargeElementsAttrs(flags, *largeElementsLimit);
   if (enableDebugInfo)
     mlirOpPrintingFlagsEnableDebugInfo(flags, /*prettyForm=*/prettyDebugInfo);
   if (printGenericOpForm)
     mlirOpPrintingFlagsPrintGenericOpForm(flags);
 
   PyFileAccumulator accum(fileObject, binary);
   py::gil_scoped_release();
-  mlirOperationPrintWithFlags(get(), flags, accum.getCallback(),
+  mlirOperationPrintWithFlags(operation.get(), flags, accum.getCallback(),
                               accum.getUserData());
   mlirOpPrintingFlagsDestroy(flags);
 }
 
-py::object PyOperation::getAsm(bool binary,
+py::object PyOperationBase::getAsm(bool binary,
                                    llvm::Optional<int64_t> largeElementsLimit,
                                    bool enableDebugInfo, bool prettyDebugInfo,
-                               bool printGenericOpForm, bool useLocalScope) {
+                                   bool printGenericOpForm,
+                                   bool useLocalScope) {
   py::object fileObject;
   if (binary) {
     fileObject = py::module::import("io").attr("BytesIO")();
   } else {
     fileObject = py::module::import("io").attr("StringIO")();
   }
   print(fileObject, /*binary=*/binary,
         /*largeElementsLimit=*/largeElementsLimit,
if (maybeIp.is_none()) {
       ip = PyThreadContextEntry::getDefaultInsertionPoint();
     } else {
       ip = py::cast<PyInsertionPoint *>(maybeIp);
     }
     if (ip)
       ip->insert(*created.get());
   }
 
-  return created.releaseObject();
+  return created->createOpView();
 }
 
-PyOpView::PyOpView(py::object operation)
-    : operationObject(std::move(operation)),
-      operation(py::cast<PyOperation *>(this->operationObject)) {}
+py::object PyOperation::createOpView() {
+  MlirIdentifier ident = mlirOperationGetName(get());
+  MlirStringRef identStr = mlirIdentifierStr(ident);
+  auto opViewClass = PyGlobals::get().lookupRawOpViewClass(
+      llvm::StringRef(identStr.data, identStr.length));
+  if (opViewClass)
+    return (*opViewClass)(getRef().getObject());
+  return py::cast(PyOpView(getRef().getObject()));
+}
+
+PyOpView::PyOpView(py::object operationObject)
+    // Casting through the PyOperationBase base-class and then back to the
+    // Operation lets us accept any PyOperationBase subclass.
+    : operation(py::cast<PyOperationBase &>(operationObject).getOperation()),
+      operationObject(operation.getRef().getObject()) {}
 
 py::object PyOpView::createRawSubclass(py::object userClass) {
   // This is... a little gross. The typical pattern is to have a pure python
   // class that extends OpView like:
   //   class AddFOp(_cext.ir.OpView):
   //     def __init__(self, loc, lhs, rhs):
   //       operation = loc.context.create_operation(
   //           "addf", lhs, rhs, results=[lhs.type])

 }
 
 //------------------------------------------------------------------------------
 // PyInsertionPoint.
 //------------------------------------------------------------------------------
 
 PyInsertionPoint::PyInsertionPoint(PyBlock &block) : block(block) {}
 
-PyInsertionPoint::PyInsertionPoint(PyOperation &beforeOperation)
-    : block(beforeOperation.getBlock()),
-      refOperation(beforeOperation.getRef()) {}
+PyInsertionPoint::PyInsertionPoint(PyOperationBase &beforeOperationBase)
+    : refOperation(beforeOperationBase.getOperation().getRef()),
+      block((*refOperation)->getBlock()) {}
 
-void PyInsertionPoint::insert(PyOperation &operation) {
+void PyInsertionPoint::insert(PyOperationBase &operationBase) {
+  PyOperation &operation = operationBase.getOperation();
   if (operation.isAttached())
     throw SetPyError(PyExc_ValueError,
                      "Attempt to insert operation that is already attached");
   block.getParentOperation()->checkValid();
   MlirOperation beforeOp = {nullptr};
   if (refOperation) {
     // Insert before operation.
     (*refOperation)->checkValid();
py::class_<PyModule>(m, "Module")
                                printAccum.getUserData());
             return printAccum.join();
           },
           kOperationStrDunderDocstring);
 
   //----------------------------------------------------------------------------
   // Mapping of Operation.
   //----------------------------------------------------------------------------
-  py::class_<PyOperation>(m, "Operation")
-      .def_static("create", &PyOperation::create, py::arg("name"),
-                  py::arg("operands") = py::none(),
-                  py::arg("results") = py::none(),
-                  py::arg("attributes") = py::none(),
-                  py::arg("successors") = py::none(), py::arg("regions") = 0,
-                  py::arg("loc") = py::none(), py::arg("ip") = py::none(),
-                  kOperationCreateDocstring)
-      .def_property_readonly(
-          "context",
-          [](PyOperation &self) { return self.getContext().getObject(); },
-          "Context that owns the Operation")
-      .def_property_readonly(
-          "operands",
-          [](PyOperation &self) { return PyOpOperandList(self.getRef()); })
-      .def_property_readonly(
-          "regions",
-          [](PyOperation &self) { return PyRegionList(self.getRef()); })
+  py::class_<PyOperationBase>(m, "_OperationBase")
+      .def("__eq__",
+           [](PyOperationBase &self, PyOperationBase &other) {
+             return &self.getOperation() == &other.getOperation();
+           })
+      .def("__eq__",
+           [](PyOperationBase &self, py::object other) { return false; })
+      .def_property_readonly("operands",
+                             [](PyOperationBase &self) {
+                               return PyOpOperandList(
+                                   self.getOperation().getRef());
+                             })
+      .def_property_readonly("regions",
+                             [](PyOperationBase &self) {
+                               return PyRegionList(
+                                   self.getOperation().getRef());
+                             })
       .def_property_readonly(
           "results",
-          [](PyOperation &self) { return PyOpResultList(self.getRef()); },
+          [](PyOperationBase &self) {
+            return PyOpResultList(self.getOperation().getRef());
+          },
           "Returns the list of Operation results.")
       .def("__iter__",
-           [](PyOperation &self) { return PyRegionIterator(self.getRef()); })
+           [](PyOperationBase &self) {
+             return PyRegionIterator(self.getOperation().getRef());
+           })
       .def(
           "__str__",
-          [](PyOperation &self) {
+          [](PyOperationBase &self) {
             return self.getAsm(/*binary=*/false,
                                /*largeElementsLimit=*/llvm::None,
                                /*enableDebugInfo=*/false,
                                /*prettyDebugInfo=*/false,
                                /*printGenericOpForm=*/false,
                                /*useLocalScope=*/false);
           },
           "Returns the assembly form of the operation.")
-      .def("print", &PyOperation::print,
+      .def("print", &PyOperationBase::print,
            // Careful: Lots of arguments must match up with print method.
            py::arg("file") = py::none(), py::arg("binary") = false,
            py::arg("large_elements_limit") = py::none(),
            py::arg("enable_debug_info") = false,
            py::arg("pretty_debug_info") = false,
            py::arg("print_generic_op_form") = false,
            py::arg("use_local_scope") = false, kOperationPrintDocstring)
-      .def("get_asm", &PyOperation::getAsm,
+      .def("get_asm", &PyOperationBase::getAsm,
            // Careful: Lots of arguments must match up with get_asm method.
            py::arg("binary") = false,
            py::arg("large_elements_limit") = py::none(),
            py::arg("enable_debug_info") = false,
            py::arg("pretty_debug_info") = false,
            py::arg("print_generic_op_form") = false,
            py::arg("use_local_scope") = false, kOperationGetAsmDocstring);
 
-  py::class_<PyOpView>(m, "OpView")
+  py::class_<PyOperation, PyOperationBase>(m, "Operation")
+      .def_static("create", &PyOperation::create, py::arg("name"),
+                  py::arg("operands") = py::none(),
+                  py::arg("results") = py::none(),
+                  py::arg("attributes") = py::none(),
+                  py::arg("successors") = py::none(), py::arg("regions") = 0,
+                  py::arg("loc") = py::none(), py::arg("ip") = py::none(),
+                  kOperationCreateDocstring)
+      .def_property_readonly(
+          "context",
+          [](PyOperation &self) { return self.getContext().getObject(); },
+          "Context that owns the Operation")
+      .def_property_readonly("opview", &PyOperation::createOpView);
+
+  py::class_<PyOpView, PyOperationBase>(m, "OpView")
       .def(py::init<py::object>())
       .def_property_readonly("operation", &PyOpView::getOperationObject)
+      .def_property_readonly(
+          "context",
+          [](PyOpView &self) {
+            return self.getOperation().getContext().getObject();
+          },
+          "Context that owns the Operation")
       .def("__str__",
            [](PyOpView &self) { return py::str(self.getOperationObject()); });
 
   //----------------------------------------------------------------------------
   // Mapping of PyRegion.
   //----------------------------------------------------------------------------
   py::class_<PyRegion>(m, "Region")
       .def_property_readonly(
           "blocks",
           [](PyRegion &self) {
             return PyBlockList(self.getParentOperation(), self.get());
           },
           "Returns a forward-optimized sequence of blocks.")
       .def(
           "__iter__",
           [](PyRegion &self) {
             self.checkValid();
             MlirBlock firstBlock = mlirRegionGetFirstBlock(self.get());
             return PyBlockIterator(self.getParentOperation(), firstBlock);
           },
           "Iterates over blocks in the region.")
-      .def("__eq__", [](PyRegion &self, py::object &other) {
-        try {
-          PyRegion *otherRegion = other.cast<PyRegion *>();
-          return self.get().ptr == otherRegion->get().ptr;
-        } catch (std::exception &e) {
-          return false;
-        }
-      });
+      .def("__eq__",
+           [](PyRegion &self, PyRegion &other) {
+             return self.get().ptr == other.get().ptr;
+           })
+      .def("__eq__", [](PyRegion &self, py::object &other) { return false; });
 
   //----------------------------------------------------------------------------
   // Mapping of PyBlock.
   //----------------------------------------------------------------------------
   py::class_<PyBlock>(m, "Block")
       .def_property_readonly(
           "arguments",
           [](PyBlock &self) {
py::class_<PyBlock>(m, "Block")
             self.checkValid();
             MlirOperation firstOperation =
                 mlirBlockGetFirstOperation(self.get());
             return PyOperationIterator(self.getParentOperation(),
                                        firstOperation);
           },
           "Iterates over operations in the block.")
       .def("__eq__",
-           [](PyBlock &self, py::object &other) {
-             try {
-               PyBlock *otherBlock = other.cast<PyBlock *>();
-               return self.get().ptr == otherBlock->get().ptr;
-             } catch (std::exception &e) {
-               return false;
-             }
+           [](PyBlock &self, PyBlock &other) {
+             return self.get().ptr == other.get().ptr;
            })
+      .def("__eq__", [](PyBlock &self, py::object &other) { return false; })
       .def(
           "__str__",
           [](PyBlock &self) {
             self.checkValid();
             PyPrintAccumulator printAccum;
             mlirBlockPrint(self.get(), printAccum.getCallback(),
                            printAccum.getUserData());
             return printAccum.join();
py::class_<PyInsertionPoint>(m, "InsertionPoint")
           [](py::object & /*class*/) {
             auto *ip = PyThreadContextEntry::getDefaultInsertionPoint();
             if (!ip)
               throw SetPyError(PyExc_ValueError, "No current InsertionPoint");
             return ip;
           },
           "Gets the InsertionPoint bound to the current thread or raises "
           "ValueError if none has been set")
-      .def(py::init<PyOperation &>(), py::arg("beforeOperation"),
+      .def(py::init<PyOperationBase &>(), py::arg("beforeOperation"),
            "Inserts before a referenced operation.")
       .def_static("at_block_begin", &PyInsertionPoint::atBlockBegin,
                   py::arg("block"), "Inserts at the beginning of the block.")
       .def_static("at_block_terminator", &PyInsertionPoint::atBlockTerminator,
                   py::arg("block"), "Inserts before the block terminator.")
       .def("insert", &PyInsertionPoint::insert, py::arg("operation"),
            "Inserts an operation.");
 
py::class_<PyAttribute>(m, "Attribute")
                              })
       .def(
           "get_named",
           [](PyAttribute &self, std::string name) {
             return PyNamedAttribute(self.attr, std::move(name));
           },
           py::keep_alive<0, 1>(), "Binds a name to the attribute")
       .def("__eq__",
-           [](PyAttribute &self, py::object &other) {
-             try {
-               PyAttribute otherAttribute = other.cast<PyAttribute>();
-               return self == otherAttribute;
-             } catch (std::exception &e) {
-               return false;
-             }
-           })
+           [](PyAttribute &self, PyAttribute &other) { return self == other; })
+      .def("__eq__", [](PyAttribute &self, py::object &other) { return false; })
       .def(
           "dump", [](PyAttribute &self) { mlirAttributeDump(self.attr); },
           kDumpDocstring)
       .def(
           "__str__",
           [](PyAttribute &self) {
             PyPrintAccumulator printAccum;
             mlirAttributePrint(self.attr, printAccum.getCallback(),
py::class_<PyType>(m, "Type")
             }
             return PyType(context->getRef(), type);
           },
           py::arg("asm"), py::arg("context") = py::none(),
           kContextParseTypeDocstring)
       .def_property_readonly(
           "context", [](PyType &self) { return self.getContext().getObject(); },
           "Context that owns the Type")
-      .def("__eq__",
-           [](PyType &self, py::object &other) {
-             try {
-               PyType otherType = other.cast<PyType>();
-               return self == otherType;
-             } catch (std::exception &e) {
-               return false;
-             }
-           })
+      .def("__eq__", [](PyType &self, PyType &other) { return self == other; })
+      .def("__eq__", [](PyType &self, py::object &other) { return false; })
       .def(
           "dump", [](PyType &self) { mlirTypeDump(self.type); }, kDumpDocstring)
       .def(
           "__str__",
           [](PyType &self) {
             PyPrintAccumulator printAccum;
             mlirTypePrint(self.type, printAccum.getCallback(),
                           printAccum.getUserData());

mlir/lib/Bindings/Python/MainModule.cpp

 
 PyGlobals::PyGlobals() {
   assert(!instance && "PyGlobals already constructed");
   instance = this;
 }
 
 PyGlobals::~PyGlobals() { instance = nullptr; }
 
-void PyGlobals::loadDialectModule(const std::string &dialectNamespace) {
-  if (loadedDialectModules.contains(dialectNamespace))
+void PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) {
+  py::gil_scoped_acquire();
+  if (loadedDialectModulesCache.contains(dialectNamespace))
     return;
   // Since re-entrancy is possible, make a copy of the search prefixes.
   std::vector<std::string> localSearchPrefixes = dialectSearchPrefixes;
   py::object loaded;
   for (std::string moduleName : localSearchPrefixes) {
     moduleName.push_back('.');
-    moduleName.append(dialectNamespace);
+    moduleName.append(dialectNamespace.data(), dialectNamespace.size());
 
     try {
+      py::gil_scoped_release();
       loaded = py::module::import(moduleName.c_str());
     } catch (py::error_already_set &e) {
       if (e.matches(PyExc_ModuleNotFoundError)) {
         continue;
       } else {
         throw;
       }
     }
     break;
   }
 
   // Note: Iterator cannot be shared from prior to loading, since re-entrancy
   // may have occurred, which may do anything.
-  loadedDialectModules.insert(dialectNamespace);
+  loadedDialectModulesCache.insert(dialectNamespace);
 }
 
 void PyGlobals::registerDialectImpl(const std::string &dialectNamespace,
                                     py::object pyClass) {
+  py::gil_scoped_acquire();
   py::object &found = dialectClassMap[dialectNamespace];
   if (found) {
     throw SetPyError(PyExc_RuntimeError, llvm::Twine("Dialect namespace '") +
                                              dialectNamespace +
                                              "' is already registered.");
   }
   found = std::move(pyClass);
 }
 
 void PyGlobals::registerOperationImpl(const std::string &operationName,
-                                      py::object pyClass, py::object rawClass) {
+                                      py::object pyClass,
+                                      py::object rawOpViewClass) {
+  py::gil_scoped_acquire();
   py::object &found = operationClassMap[operationName];
   if (found) {
     throw SetPyError(PyExc_RuntimeError, llvm::Twine("Operation '") +
                                              operationName +
                                              "' is already registered.");
   }
   found = std::move(pyClass);
-  rawOperationClassMap[operationName] = std::move(rawClass);
+  rawOpViewClassMap[operationName] = std::move(rawOpViewClass);
 }
 
 llvm::Optional<py::object>
 PyGlobals::lookupDialectClass(const std::string &dialectNamespace) {
+  py::gil_scoped_acquire();
   loadDialectModule(dialectNamespace);
   // Fast match against the class map first (common case).
   const auto foundIt = dialectClassMap.find(dialectNamespace);
   if (foundIt != dialectClassMap.end()) {
     if (foundIt->second.is_none())
       return llvm::None;
     assert(foundIt->second && "py::object is defined");
     return foundIt->second;
   }
 
   // Not found and loading did not yield a registration. Negative cache.
   dialectClassMap[dialectNamespace] = py::none();
   return llvm::None;
 }
 
+llvm::Optional<pybind11::object>
+PyGlobals::lookupRawOpViewClass(llvm::StringRef operationName) {
+  {
+    py::gil_scoped_acquire();
+    auto foundIt = rawOpViewClassMapCache.find(operationName);
+    if (foundIt != rawOpViewClassMapCache.end()) {
+      if (foundIt->second.is_none())
+        return llvm::None;
+      assert(foundIt->second && "py::object is defined");
+      return foundIt->second;
+    }
+  }
+
+  // Not found. Load the dialect namespace.
+  auto split = operationName.split('.');
+  llvm::StringRef dialectNamespace = split.first;
+  loadDialectModule(dialectNamespace);
+
+  // Attempt to find from the canonical map and cache.
+  {
+    py::gil_scoped_acquire();
+    auto foundIt = rawOpViewClassMap.find(operationName);
+    if (foundIt != rawOpViewClassMap.end()) {
+      if (foundIt->second.is_none())
+        return llvm::None;
+      assert(foundIt->second && "py::object is defined");
+      // Positive cache.
+      rawOpViewClassMapCache[operationName] = foundIt->second;
+      return foundIt->second;
+    } else {
+      // Negative cache.
+      rawOpViewClassMap[operationName] = py::none();
+      return llvm::None;
+    }
+  }
+}
+
+void PyGlobals::clearImportCache() {
+  py::gil_scoped_acquire();
+  loadedDialectModulesCache.clear();
+  rawOpViewClassMapCache.clear();
+}
+
 // -----------------------------------------------------------------------------
 // Module initialization.
 // -----------------------------------------------------------------------------
 
 PYBIND11_MODULE(_mlir, m) {
   m.doc() = "MLIR Python Native Extension";
 
   py::class_<PyGlobals>(m, "_Globals")
       .def_property("dialect_search_modules",
                     &PyGlobals::getDialectSearchPrefixes,
                     &PyGlobals::setDialectSearchPrefixes)
       .def("append_dialect_search_prefix",
            [](PyGlobals &self, std::string moduleName) {
              self.getDialectSearchPrefixes().push_back(std::move(moduleName));
+             self.clearImportCache();
            })
       .def("_register_dialect_impl", &PyGlobals::registerDialectImpl,
            "Testing hook for directly registering a dialect")
       .def("_register_operation_impl", &PyGlobals::registerOperationImpl,
            "Testing hook for directly registering an operation");
 
   // Aside from making the globals accessible to python, having python manage
   // it is necessary to make sure it is destroyed (and releases its python

mlir/test/Bindings/Python/ir_operation.py

def testOperationPrint():
 
   # Test get_asm with options.
   # CHECK: value = opaque<"", "0xDEADBEEF"> : tensor<4xi32>
   # CHECK: "std.return"(%arg0) : (i32) -> () -:4:7
   module.operation.print(large_elements_limit=2, enable_debug_info=True,
       pretty_debug_info=True, print_generic_op_form=True, use_local_scope=True)
 
 run(testOperationPrint)
+
+
+def testKnownOpView():
+  with Context(), Location.unknown():
+    Context.current.allow_unregistered_dialects = True
+    module = Module.parse(r"""
+      %1 = "custom.f32"() : () -> f32
+      %2 = "custom.f32"() : () -> f32
+      %3 = addf %1, %2 : f32
+    """)
+    print(module)
+
+    # addf should map to a known OpView class in the std dialect.
+    # We know the OpView for it defines an 'lhs' attribute.
+    addf = module.body.operations[2]
+    # CHECK: <mlir.dialects.std._AddFOp object
+    print(repr(addf))
+    # CHECK: "custom.f32"()
+    print(addf.lhs)
+
+    # One of the custom ops should resolve to the default OpView.
+    custom = module.body.operations[0]
+    # CHECK: <_mlir.ir.OpView object
+    print(repr(custom))
+
+    # Check again to make sure negative caching works.
+    custom = module.body.operations[0]
+    # CHECK: <_mlir.ir.OpView object
+    print(repr(custom))
+
+run(testKnownOpView)