[mlir][python] Need based additions to work with unregistered ops and do simple transforms.

[stellaraccident]

Summary

• A few fixes to make unregistered dialects more useful for this API:
  • Add 'register_all_dialects=' (default True) Context constructor arg. This bypasses the built-in registration of all dialects (which we really need to retire anyway but this is at least a step).
  • Only attempts to create custom OpViews for registered ops. Previously, if you happened to have generated OpView code on your path but had not registered the dialect with the context, it would be loaded anyway, which is, in general, not correct (i.e. unregistered ops can be anything and the contract for generated code assumes they are specifically what is defined).
• Moves some methods that were only defined on Operation to _OperationBase which makes them shared with OpViews. Afaict, these were just bugs and were put in the wrong place.
• Add Value.replace_all_uses_with.

mlir/include/mlir-c/IR.h

 MLIR_CAPI_EXPORTED bool mlirValueEqual(MlirValue value1, MlirValue value2);
 
 /// Returns 1 if the value is a block argument, 0 otherwise.
 MLIR_CAPI_EXPORTED bool mlirValueIsABlockArgument(MlirValue value);
 
 /// Returns 1 if the value is an operation result, 0 otherwise.
 MLIR_CAPI_EXPORTED bool mlirValueIsAOpResult(MlirValue value);
 
+/// Replaces all uses of a value with a new value.
+MLIR_CAPI_EXPORTED void mlirValueReplaceAllUsesWith(MlirValue value,
+                                                    MlirValue newValue);
+
+/// Replaces uses of a value with a new value if a predicate returns true.
+/// The predicate takes a user operation and the index into its operands of the
+/// use.
+MLIR_CAPI_EXPORTED void mlirValueReplaceUsesWithIf(
+    MlirValue value, MlirValue newValue,
+    bool (*shouldReplace)(MlirOperation operation, intptr_t operandIndex,
+                          void *userData),
+    void *userData);
+
 /// Returns the block in which this value is defined as an argument. Asserts if
 /// the value is not a block argument.
 MLIR_CAPI_EXPORTED MlirBlock mlirBlockArgumentGetOwner(MlirValue value);
 
 /// Returns the position of the value in the argument list of its block.
 MLIR_CAPI_EXPORTED intptr_t mlirBlockArgumentGetArgNumber(MlirValue value);
 
 /// Sets the type of the block argument to the given type.

mlir/lib/Bindings/Python/IRCore.cpp

 
 py::object PyMlirContext::createFromCapsule(py::object capsule) {
   MlirContext rawContext = mlirPythonCapsuleToContext(capsule.ptr());
   if (mlirContextIsNull(rawContext))
     throw py::error_already_set();
   return forContext(rawContext).releaseObject();
 }
 
-PyMlirContext *PyMlirContext::createNewContextForInit() {
+PyMlirContext *
+PyMlirContext::createNewContextForInit(bool registerAllDialects) {
   MlirContext context = mlirContextCreate();
+  // TODO: Remove implicit registration of dialects.
+  if (registerAllDialects) {
     mlirRegisterAllDialects(context);
+  }
   return new PyMlirContext(context);
 }
 
 PyMlirContextRef PyMlirContext::forContext(MlirContext context) {
   py::gil_scoped_acquire acquire;
   auto &liveContexts = getLiveContexts();
   auto it = liveContexts.find(context.ptr);
   if (it == liveContexts.end()) {
if (ip)
       ip->insert(*created.get());
   }
 
   return created->createOpView();
 }
 
 py::object PyOperation::createOpView() {
   checkValid();
-  MlirIdentifier ident = mlirOperationGetName(get());
+  MlirOperation cOp = get();
+  MlirTypeID registeredTypeID = mlirOperationGetTypeID(cOp);
+  // For registered ops, lookup an appropriate, generated op view.
+  if (!mlirTypeIDIsNull(registeredTypeID)) {
+    MlirIdentifier ident = mlirOperationGetName(cOp);
     MlirStringRef identStr = mlirIdentifierStr(ident);
     auto opViewClass = PyGlobals::get().lookupRawOpViewClass(
         StringRef(identStr.data, identStr.length));
     if (opViewClass)
       return (*opViewClass)(getRef().getObject());
+  }
   return py::cast(PyOpView(getRef().getObject()));
 }
 
 void PyOperation::erase() {
   checkValid();
   // TODO: Fix memory hazards when erasing a tree of operations for which a deep
   // Python reference to a child operation is live. All children should also
   // have their `valid` bit set to false.
py::class_<PyDiagnosticHandler>(m, "DiagnosticHandler", py::module_local())
       .def_property_readonly("had_error", &PyDiagnosticHandler::getHadError)
       .def("__enter__", &PyDiagnosticHandler::contextEnter)
       .def("__exit__", &PyDiagnosticHandler::contextExit);
 
   //----------------------------------------------------------------------------
   // Mapping of MlirContext.
   //----------------------------------------------------------------------------
   py::class_<PyMlirContext>(m, "Context", py::module_local())
-      .def(py::init<>(&PyMlirContext::createNewContextForInit))
+      .def(py::init<>(&PyMlirContext::createNewContextForInit),
+           py::arg("register_all_dialects") = true)
       .def_static("_get_live_count", &PyMlirContext::getLiveCount)
       .def("_get_context_again",
            [](PyMlirContext &self) {
              PyMlirContextRef ref = PyMlirContext::forContext(self.get());
              return ref.releaseObject();
            })
       .def("_get_live_operation_count", &PyMlirContext::getLiveOperationCount)
       .def("_get_live_module_count", &PyMlirContext::getLiveModuleCount)
py::class_<PyOperationBase>(m, "_OperationBase", py::module_local())
            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,
            py::arg("assume_verified") = false, kOperationGetAsmDocstring)
       .def(
-          "verify",
+          "detach_from_parent",
           [](PyOperationBase &self) {
-            return mlirOperationVerify(self.getOperation());
+            PyOperation &operation = self.getOperation();
+            operation.checkValid();
+            if (!operation.isAttached())
+              throw py::value_error("Detached operation has no parent.");
+
+            operation.detachFromParent();
+            return operation.createOpView();
           },
-          "Verify the operation and return true if it passes, false if it "
-          "fails.")
+          "Detaches the operation from its parent block.")
+      .def("erase", [](PyOperationBase &self) { self.getOperation().erase(); })
       .def("move_after", &PyOperationBase::moveAfter, py::arg("other"),
            "Puts self immediately after the other operation in its parent "
            "block.")
       .def("move_before", &PyOperationBase::moveBefore, py::arg("other"),
            "Puts self immediately before the other operation in its parent "
            "block.")
-      .def(
-          "detach_from_parent",
+      .def_property_readonly("name",
                              [](PyOperationBase &self) {
                                PyOperation &operation = self.getOperation();
                                operation.checkValid();
-            if (!operation.isAttached())
-              throw py::value_error("Detached operation has no parent.");
-
-            operation.detachFromParent();
-            return operation.createOpView();
+                               MlirStringRef name = mlirIdentifierStr(
+                                   mlirOperationGetName(operation.get()));
+                               return py::str(name.data, name.length);
+                             })
+      .def_property_readonly("parent",
+                             [](PyOperationBase &self) -> py::object {
+                               PyOperation &operation = self.getOperation();
+                               operation.checkValid();
+                               auto parent = operation.getParentOperation();
+                               if (parent)
+                                 return parent->getObject();
+                               return py::none();
+                             })
+      .def(
+          "verify",
+          [](PyOperationBase &self) {
+            return mlirOperationVerify(self.getOperation());
           },
-          "Detaches the operation from its parent block.");
+          "Verify the operation and return true if it passes, false if it "
+          "fails.");
 
   py::class_<PyOperation, PyOperationBase>(m, "Operation", py::module_local())
       .def_static("create", &PyOperation::create, py::arg("name"),
                   py::arg("results") = py::none(),
                   py::arg("operands") = 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("parent",
-                             [](PyOperation &self) -> py::object {
-                               auto parent = self.getParentOperation();
-                               if (parent)
-                                 return parent->getObject();
-                               return py::none();
-                             })
-      .def("erase", &PyOperation::erase)
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyOperation::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyOperation::createFromCapsule)
-      .def_property_readonly("name",
-                             [](PyOperation &self) {
-                               self.checkValid();
-                               MlirOperation operation = self.get();
-                               MlirStringRef name = mlirIdentifierStr(
-                                   mlirOperationGetName(operation));
-                               return py::str(name.data, name.length);
-                             })
       .def_property_readonly(
           "context",
           [](PyOperation &self) {
             self.checkValid();
             return self.getContext().getObject();
           },
           "Context that owns the Operation")
       .def_property_readonly("opview", &PyOperation::createOpView);
py::class_<PyValue>(m, "Value", py::module_local())
           "owner",
           [](PyValue &self) {
             assert(mlirOperationEqual(self.getParentOperation()->get(),
                                       mlirOpResultGetOwner(self.get())) &&
                    "expected the owner of the value in Python to match that in "
                    "the IR");
             return self.getParentOperation().getObject();
           })
+      .def("replace_all_uses_with",
+           [](PyValue &self, PyValue &newValue) {
+             mlirValueReplaceAllUsesWith(self.get(), newValue.get());
+           })
       .def("__eq__",
            [](PyValue &self, PyValue &other) {
              return self.get().ptr == other.get().ptr;
            })
       .def("__eq__", [](PyValue &self, py::object other) { return false; })
       .def("__hash__",
            [](PyValue &self) {
              return static_cast<size_t>(llvm::hash_value(self.get().ptr));

mlir/lib/Bindings/Python/IRModule.h

public:
   /// For the case of a python __init__ (py::init) method, pybind11 is quite
   /// strict about needing to return a pointer that is not yet associated to
   /// an py::object. Since the forContext() method acts like a pool, possibly
   /// returning a recycled context, it does not satisfy this need. The usual
   /// way in python to accomplish such a thing is to override __new__, but
   /// that is also not supported by pybind11. Instead, we use this entry
   /// point which always constructs a fresh context (which cannot alias an
   /// existing one because it is fresh).
-  static PyMlirContext *createNewContextForInit();
+  static PyMlirContext *createNewContextForInit(bool registerAllDialects);
 
   /// Returns a context reference for the singleton PyMlirContext wrapper for
   /// the given context.
   static PyMlirContextRef forContext(MlirContext context);
   ~PyMlirContext();
 
   /// Accesses the underlying MlirContext.
   MlirContext get() { return context; }

mlir/lib/Bindings/Python/IRTypes.cpp

c.def(
         },
         py::arg("pos"), "Returns the pos-th type in the tuple type.");
     c.def_property_readonly(
         "num_types",
         [](PyTupleType &self) -> intptr_t {
           return mlirTupleTypeGetNumTypes(self);
         },
         "Returns the number of types contained in a tuple.");
+    c.def_property_readonly(
+        "types",
+        [](PyTupleType &self) {
+          py::list result;
+          auto &context = self.getContext();
+          intptr_t n = mlirTupleTypeGetNumTypes(self);
+          for (intptr_t i = 0; i < n; ++i) {
+            result.append(PyType(context, mlirTupleTypeGetType(self, i)));
+          }
+          return result;
+        },
+        "Returns a sequence of the contained types.");
   }
 };
 
 /// Function type.
 class PyFunctionType : public PyConcreteType<PyFunctionType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFunction;
   static constexpr const char *pyClassName = "FunctionType";

mlir/lib/CAPI/IR/IR.cpp

 bool mlirValueIsABlockArgument(MlirValue value) {
   return unwrap(value).isa<BlockArgument>();
 }
 
 bool mlirValueIsAOpResult(MlirValue value) {
   return unwrap(value).isa<OpResult>();
 }
 
+void mlirValueReplaceAllUsesWith(MlirValue value, MlirValue newValue) {
+  unwrap(value).replaceAllUsesWith(unwrap(newValue));
+}
+
+void mlirValueReplaceUsesWithIf(MlirValue value, MlirValue newValue,
+                                bool (*shouldReplace)(MlirOperation operation,
+                                                      intptr_t operandIndex,
+                                                      void *userData),
+                                void *userData) {
+  unwrap(value).replaceUsesWithIf(unwrap(newValue), [&](OpOperand &operand) {
+    return shouldReplace(wrap(operand.getOwner()), operand.getOperandNumber(),
+                         userData);
+  });
+}
+
 MlirBlock mlirBlockArgumentGetOwner(MlirValue value) {
   return wrap(unwrap(value).cast<BlockArgument>().getOwner());
 }
 
 intptr_t mlirBlockArgumentGetArgNumber(MlirValue value) {
   return static_cast<intptr_t>(
       unwrap(value).cast<BlockArgument>().getArgNumber());
 }

mlir/python/mlir/_mlir_libs/_mlir/ir.pyi

class _OperationBase:
     def __eq__(self, arg0: "_OperationBase") -> bool: ...
     @overload
     def __eq__(self, arg0: object) -> bool: ...
     def __hash__(self) -> int: ...
     @property
     def _CAPIPtr(self) -> object: ...
     @property
     def attributes(self) -> "OpAttributeMap": ...
+    def erase(self) -> None: ...
     @property
     def location(self) -> "Location": ...
     @property
+    def name(self) -> str: ...
+    @property
     def operands(self) -> "OpOperandList": ...
     @property
+    def parent(self) -> Optional["_OperationBase"]: ...
+    @property
     def regions(self) -> "RegionSequence": ...
     @property
     def result(self) -> "OpResult": ...
     @property
     def results(self) -> "OpResultList": ...
 
 # TODO: Auto-generated. Audit and fix.
 class AffineExpr:
class ComplexType(Type):
     def get(self, *args, **kwargs) -> Any: ...
     def isinstance(self, *args, **kwargs) -> Any: ...
     @property
     def element_type(self) -> Type: ...
 
 class Context:
     current: ClassVar["Context"] = ...  # read-only
     allow_unregistered_dialects: bool
-    def __init__(self) -> None: ...
+    def __init__(self, register_all_dialects: bool = True) -> None: ...
     def _CAPICreate(self) -> object: ...
     def _get_context_again(self) -> "Context": ...
     @staticmethod
     def _get_live_count() -> int: ...
     def _get_live_module_count(self) -> int: ...
     def _get_live_operation_count(self) -> int: ...
     def attach_diagnostic_handler(self, callback: Callable[["Diagnostic"], bool]) -> "DiagnosticHandler": ...
     def enable_multithreading(self, enable: bool) -> None: ...
class Operation(_OperationBase):
     @staticmethod
     def create(name: str, results: Optional[Sequence["Type"]] = None,
         operands: Optional[Sequence["Value"]] = None,
         attributes: Optional[Dict[str, "Attribute"]] = None,
         successors: Optional[Sequence["Block"]] = None,
         regions: int = 0,
         loc: Optional["Location"] = None,
         ip: Optional["InsertionPoint"] = None) -> "_OperationBase": ...
-    def erase(self) -> None: ...
     @property
     def _CAPIPtr(self) -> object: ...
     @property
     def context(self) -> "Context": ...
     @property
-    def name(self) -> str: ...
-    @property
     def opview(self) -> "OpView": ...
-    @property
-    def parent(self) -> Optional["_OperationBase"]: ...
 
 class OperationIterator:
     def __iter__(self) -> "OperationIterator": ...
     def __next__(self) -> "OpView": ...
 
 class OperationList:
     def __getitem__(self, arg0: int) -> "OpView": ...
     def __iter__(self) -> "OperationIterator": ...