Expose MlirOperationClone in Python bindings.

[dominikgrewe]

Summary

Expose MlirOperationClone in Python bindings.

[mehdi_amini]

mlir/lib/Bindings/Python/IRCore.cpp

1233

Except the first line, this is all the same as the end of PyOperation::create(, I wonder if it is worth refactoring in a helper function?

[dominikgrewe]

Factor out insertion point logic.

[ftynse]

mlir/test/python/ir/operation.py

783

Could we check that there are indeed two operations before erasing the original?

[dominikgrewe]

Check in test that op actually got cloned.

[dominikgrewe]

mlir/lib/Bindings/Python/IRCore.cpp

1233

I factored out the insertion logic as that seemed like a natural function to pull out and constituted the majority of the duplicated code. Ptal.

[dominikgrewe]

Upload changes.

mlir/lib/Bindings/Python/IRCore.cpp

py::object PyOperation::createFromCapsule(py::object capsule) {
   MlirOperation rawOperation = mlirPythonCapsuleToOperation(capsule.ptr());
   if (mlirOperationIsNull(rawOperation))
     throw py::error_already_set();
   MlirContext rawCtxt = mlirOperationGetContext(rawOperation);
   return forOperation(PyMlirContext::forContext(rawCtxt), rawOperation)
       .releaseObject();
 }
 
+static void maybeInsertOperation(PyOperationRef &op,
+                                 const py::object &maybeIp) {
+  // InsertPoint active?
+  if (!maybeIp.is(py::cast(false))) {
+    PyInsertionPoint *ip;
+    if (maybeIp.is_none()) {
+      ip = PyThreadContextEntry::getDefaultInsertionPoint();
+    } else {
+      ip = py::cast<PyInsertionPoint *>(maybeIp);
+    }
+    if (ip)
+      ip->insert(*op.get());
+  }
+}
+
 py::object PyOperation::create(
     const std::string &name, llvm::Optional<std::vector<PyType *>> results,
     llvm::Optional<std::vector<PyValue *>> operands,
     llvm::Optional<py::dict> attributes,
     llvm::Optional<std::vector<PyBlock *>> successors, int regions,
     DefaultingPyLocation location, const py::object &maybeIp) {
   llvm::SmallVector<MlirValue, 4> mlirOperands;
   llvm::SmallVector<MlirType, 4> mlirResults;
if (regions) {
     mlirOperationStateAddOwnedRegions(&state, mlirRegions.size(),
                                       mlirRegions.data());
   }
 
   // Construct the operation.
   MlirOperation operation = mlirOperationCreate(&state);
   PyOperationRef created =
       PyOperation::createDetached(location->getContext(), operation);
+  maybeInsertOperation(created, maybeIp);
 
-  // InsertPoint active?
-  if (!maybeIp.is(py::cast(false))) {
-    PyInsertionPoint *ip;
-    if (maybeIp.is_none()) {
-      ip = PyThreadContextEntry::getDefaultInsertionPoint();
-    } else {
-      ip = py::cast<PyInsertionPoint *>(maybeIp);
-    }
-    if (ip)
-      ip->insert(*created.get());
-  }
+  return created->createOpView();
 
-  return created->createOpView();
+py::object PyOperation::clone(const py::object &maybeIp) {
+  MlirOperation clonedOperation = mlirOperationClone(operation);
+  PyOperationRef cloned =
+      PyOperation::createDetached(getContext(), clonedOperation);
+  maybeInsertOperation(cloned, maybeIp);
+
+  return cloned->createOpView();
 }
 
 py::object PyOperation::createOpView() {
   checkValid();
   MlirIdentifier ident = mlirOperationGetName(get());
   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.
   auto &liveOperations = getContext()->liveOperations;
   if (liveOperations.count(operation.ptr))
py::class_<PyOperation, PyOperationBase>(m, "Operation", py::module_local())
       .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("clone", &PyOperation::clone, py::arg("ip") = py::none())
       .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(

mlir/lib/Bindings/Python/IRModule.h

public:
 
   /// Creates an OpView suitable for this operation.
   pybind11::object createOpView();
 
   /// Erases the underlying MlirOperation, removes its pointer from the
   /// parent context's live operations map, and sets the valid bit false.
   void erase();
 
+  /// Clones this operation.
+  pybind11::object clone(const pybind11::object &ip);
+
 private:
   PyOperation(PyMlirContextRef contextRef, MlirOperation operation);
   static PyOperationRef createInstance(PyMlirContextRef contextRef,
                                        MlirOperation operation,
                                        pybind11::object parentKeepAlive);
 
   MlirOperation operation;
   pybind11::handle handle;

mlir/test/python/ir/operation.py

with InsertionPoint(m.body):
 
       # CHECK-NOT: "custom.op1"
       print(m)
 
       # Ensure we can create another operation
       Operation.create("custom.op2")
 
 
+# CHECK-LABEL: TEST: testOperationClone
+@run
+def testOperationClone():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with Location.unknown(ctx):
+    m = Module.create()
+    with InsertionPoint(m.body):
+      op = Operation.create("custom.op1")
+
+      # CHECK: "custom.op1"
+      print(m)
+
+      clone = op.operation.clone()
+      op.operation.erase()
+
+      # CHECK: "custom.op1"
+      print(m)
+
+
 # CHECK-LABEL: TEST: testOperationLoc
 @run
 def testOperationLoc():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with ctx:
     loc = Location.name("loc")
     op = Operation.create("custom.op", loc=loc)