Add mlir python APIs for creating operations, regions and blocks.

[stellaraccident]

Summary

• The API is a bit more verbose than I feel like it needs to be. In a follow-up I'd like to abbreviate some things and look in to creating aliases for common accessors.
• There is a lingering lifetime hazard between the module and newly added operations. We have the facilities now to solve for this but I will do that in a follow-up.
• We may need to craft a more limited API for safely referencing successors when creating operations. We need more facilities to really prove that out and should defer for now.

[ftynse]

mlir/lib/Bindings/Python/IRModules.cpp

407

https://reviews.llvm.org/D88148

1306

Nit: should this rather be a member of PyMlirContext? Compact .def lists are easier to scan as "documentation".

1357

We can have a RAII cleaner if needed

1367

Nit: reserve space

[stellaraccident]

Address comments.

[stellaraccident]

mlir/lib/Bindings/Python/IRModules.cpp

1306

I switched it and maybe it is marginally better. This is kind of the worst-case scenario: lots of args that need to line up, complicated implementation. It's a burden either way.

1357

I started down that path and then just opted for simplicity since this is the one and only place we are doing this kind of juggling (mental overhead of a cleaner and releasing properly to avoid double free can be bug prone too).

mlir/lib/Bindings/Python/IRModules.h

 #include <pybind11/pybind11.h>
 
 #include "mlir-c/IR.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 namespace python {
 
+class PyBlock;
+class PyLocation;
 class PyMlirContext;
 class PyModule;
 class PyOperation;
+class PyType;
 
 /// Template for a reference to a concrete type which captures a python
 /// reference to its underlying python object.
 template <typename T>
 class PyObjectRef {
 public:
   PyObjectRef(T *referrent, pybind11::object object)
       : referrent(referrent), object(std::move(object)) {
public:
 
   /// Gets the count of live context objects. Used for testing.
   static size_t getLiveCount();
 
   /// Gets the count of live operations associated with this context.
   /// Used for testing.
   size_t getLiveOperationCount();
 
+  /// Creates an operation. See corresponding python docstring.
+  pybind11::object
+  createOperation(std::string name, PyLocation location,
+                  llvm::Optional<std::vector<PyType *>> results,
+                  llvm::Optional<pybind11::dict> attributes,
+                  llvm::Optional<std::vector<PyBlock *>> successors,
+                  int regions);
+
 private:
   PyMlirContext(MlirContext context);
   // Interns the mapping of live MlirContext::ptr to PyMlirContext instances,
   // preserving the relationship that an MlirContext maps to a single
   // PyMlirContext wrapper. This could be replaced in the future with an
   // extension mechanism on the MlirContext for stashing user pointers.
   // Note that this holds a handle, which does not imply ownership.
   // Mappings will be removed when the context is destructed.
public:
   }
 
   PyOperationRef getRef() {
     return PyOperationRef(
         this, pybind11::reinterpret_borrow<pybind11::object>(handle));
   }
 
   bool isAttached() { return attached; }
+  void setAttached() {
+    assert(!attached && "operation already attached");
+    attached = true;
+  }
   void checkValid();
 
 private:
   PyOperation(PyMlirContextRef contextRef, MlirOperation operation);
   static PyOperationRef createInstance(PyMlirContextRef contextRef,
                                        MlirOperation operation,
                                        pybind11::object parentKeepAlive);
 

mlir/lib/Bindings/Python/IRModules.cpp

 using namespace mlir::python;
 
 using llvm::SmallVector;
 
 //------------------------------------------------------------------------------
 // Docstrings (trivial, non-duplicated docstrings are included inline).
 //------------------------------------------------------------------------------
 
+static const char kContextCreateOperationDocstring[] =
+    R"(Creates a new operation.
+
+Args:
+  name: Operation name (e.g. "dialect.operation").
+  location: A Location object.
+  results: Sequence of Type representing op result types.
+  attributes: Dict of str:Attribute.
+  successors: List of Block for the operation's successors.
+  regions: Number of regions to create.
+
+Returns:
+  A new "detached" Operation object. Detached operations can be added
+  to blocks, which causes them to become "attached."
+)";
+
 static const char kContextParseDocstring[] =
     R"(Parses a module's assembly format from a string.
 
 Returns a new MlirModule or raises a ValueError if the parsing fails.
 
 See also: https://mlir.llvm.org/docs/LangRef/
 )";
 

 )";
 
 static const char kTypeStrDunderDocstring[] =
     R"(Prints the assembly form of the type.)";
 
 static const char kDumpDocstring[] =
     R"(Dumps a debug representation of the object to stderr.)";
 
+static const char kAppendBlockDocstring[] =
+    R"(Appends a new block, with argument types as positional args.
+
+Returns:
+  The created block.
+)";
+
 //------------------------------------------------------------------------------
 // Conversion utilities.
 //------------------------------------------------------------------------------
 
 namespace {
 
 /// Accumulates into a python string from a method that accepts an
 /// MlirStringCallback.
while (!mlirBlockIsNull(block)) {
         return PyBlock(operation, block);
       }
       block = mlirBlockGetNextInRegion(block);
       index -= 1;
     }
     throw SetPyError(PyExc_IndexError, "attempt to access out of bounds block");
   }
 
+  PyBlock appendBlock(py::args pyArgTypes) {
+    operation->checkValid();
+    llvm::SmallVector<MlirType, 4> argTypes;
+    argTypes.reserve(pyArgTypes.size());
+    for (auto &pyArg : pyArgTypes) {
+      argTypes.push_back(pyArg.cast<PyType &>().type);
+    }
+
+    MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data());
+    mlirRegionAppendOwnedBlock(region, block);
+    return PyBlock(operation, block);
+  }
+
   static void bind(py::module &m) {
     py::class_<PyBlockList>(m, "BlockList")
         .def("__getitem__", &PyBlockList::dunderGetItem)
         .def("__iter__", &PyBlockList::dunderIter)
-        .def("__len__", &PyBlockList::dunderLen);
+        .def("__len__", &PyBlockList::dunderLen)
+        .def("append", &PyBlockList::appendBlock, kAppendBlockDocstring);
   }
 
 private:
   PyOperationRef operation;
   MlirRegion region;
 };
 
 class PyOperationIterator {
while (!mlirOperationIsNull(childOp)) {
       }
       childOp = mlirOperationGetNextInBlock(childOp);
       index -= 1;
     }
     throw SetPyError(PyExc_IndexError,
                      "attempt to access out of bounds operation");
   }
 
+  void insert(int index, PyOperation &newOperation) {
+    parentOperation->checkValid();
+    newOperation.checkValid();
+    if (index < 0) {
+      throw SetPyError(
+          PyExc_IndexError,
+          "only positive insertion indices are supported for operations");
+    }
+    if (newOperation.isAttached()) {
+      throw SetPyError(
+          PyExc_ValueError,
+          "attempt to insert an operation that has already been inserted");
+    }
+    // TODO: Needing to do this check is unfortunate, especially since it will
+    // be a forward-scan, just like the following call to
+    // mlirBlockInsertOwnedOperation. Switch to insert before/after once
+    // D88148 lands.
+    if (index > dunderLen()) {
+      throw SetPyError(PyExc_IndexError,
+                       "attempt to insert operation past end");
+    }
+    mlirBlockInsertOwnedOperation(block, index, newOperation.get());
+    newOperation.setAttached();
+    // TODO: Rework the parentKeepAlive so as to avoid ownership hazards under
+    // the new ownership.
+  }
+
   static void bind(py::module &m) {
     py::class_<PyOperationList>(m, "OperationList")
         .def("__getitem__", &PyOperationList::dunderGetItem)
         .def("__iter__", &PyOperationList::dunderIter)
-        .def("__len__", &PyOperationList::dunderLen);
+        .def("__len__", &PyOperationList::dunderLen)
+        .def("insert", &PyOperationList::insert, py::arg("index"),
+             py::arg("operation"),
+             "Inserts an operation at an indexed position");
   }
 
 private:
   PyOperationRef parentOperation;
   MlirBlock block;
 };
 
 } // namespace
PyMlirContext::LiveContextMap &PyMlirContext::getLiveContexts() {
   static LiveContextMap liveContexts;
   return liveContexts;
 }
 
 size_t PyMlirContext::getLiveCount() { return getLiveContexts().size(); }
 
 size_t PyMlirContext::getLiveOperationCount() { return liveOperations.size(); }
 
+py::object PyMlirContext::createOperation(
+    std::string name, PyLocation location,
+    llvm::Optional<std::vector<PyType *>> results,
+    llvm::Optional<py::dict> attributes,
+    llvm::Optional<std::vector<PyBlock *>> successors, int regions) {
+  llvm::SmallVector<MlirType, 4> mlirResults;
+  llvm::SmallVector<MlirBlock, 4> mlirSuccessors;
+  llvm::SmallVector<std::pair<std::string, MlirAttribute>, 4> mlirAttributes;
+
+  // General parameter validation.
+  if (regions < 0)
+    throw SetPyError(PyExc_ValueError, "number of regions must be >= 0");
+
+  // Unpack/validate results.
+  if (results) {
+    mlirResults.reserve(results->size());
+    for (PyType *result : *results) {
+      // TODO: Verify result type originate from the same context.
+      if (!result)
+        throw SetPyError(PyExc_ValueError, "result type cannot be None");
+      mlirResults.push_back(result->type);
+    }
+  }
+  // Unpack/validate attributes.
+  if (attributes) {
+    mlirAttributes.reserve(attributes->size());
+    for (auto &it : *attributes) {
+
+      auto name = it.first.cast<std::string>();
+      auto &attribute = it.second.cast<PyAttribute &>();
+      // TODO: Verify attribute originates from the same context.
+      mlirAttributes.emplace_back(std::move(name), attribute.attr);
+    }
+  }
+  // Unpack/validate successors.
+  if (successors) {
+    llvm::SmallVector<MlirBlock, 4> mlirSuccessors;
+    mlirSuccessors.reserve(successors->size());
+    for (auto *successor : *successors) {
+      // TODO: Verify successor originate from the same context.
+      if (!successor)
+        throw SetPyError(PyExc_ValueError, "successor block cannot be None");
+      mlirSuccessors.push_back(successor->get());
+    }
+  }
+
+  // Apply unpacked/validated to the operation state. Beyond this
+  // point, exceptions cannot be thrown or else the state will leak.
+  MlirOperationState state = mlirOperationStateGet(name.c_str(), location.loc);
+  if (!mlirResults.empty())
+    mlirOperationStateAddResults(&state, mlirResults.size(),
+                                 mlirResults.data());
+  if (!mlirAttributes.empty()) {
+    // Note that the attribute names directly reference bytes in
+    // mlirAttributes, so that vector must not be changed from here
+    // on.
+    llvm::SmallVector<MlirNamedAttribute, 4> mlirNamedAttributes;
+    mlirNamedAttributes.reserve(mlirAttributes.size());
+    for (auto &it : mlirAttributes)
+      mlirNamedAttributes.push_back(
+          mlirNamedAttributeGet(it.first.c_str(), it.second));
+    mlirOperationStateAddAttributes(&state, mlirNamedAttributes.size(),
+                                    mlirNamedAttributes.data());
+  }
+  if (!mlirSuccessors.empty())
+    mlirOperationStateAddSuccessors(&state, mlirSuccessors.size(),
+                                    mlirSuccessors.data());
+  if (regions) {
+    llvm::SmallVector<MlirRegion, 4> mlirRegions;
+    mlirRegions.resize(regions);
+    for (int i = 0; i < regions; ++i)
+      mlirRegions[i] = mlirRegionCreate();
+    mlirOperationStateAddOwnedRegions(&state, mlirRegions.size(),
+                                      mlirRegions.data());
+  }
+
+  // Construct the operation.
+  MlirOperation operation = mlirOperationCreate(&state);
+  return PyOperation::createDetached(getRef(), operation).releaseObject();
+}
+
 //------------------------------------------------------------------------------
 // PyModule
 //------------------------------------------------------------------------------
 
 PyModuleRef PyModule::create(PyMlirContextRef contextRef, MlirModule module) {
   PyModule *unownedModule = new PyModule(std::move(contextRef), module);
   // Note that the default return value policy on cast is automatic_reference,
   // which does not take ownership (delete will not be called).
py::class_<PyMlirContext>(m, "Context")
       .def_property(
           "allow_unregistered_dialects",
           [](PyMlirContext &self) -> bool {
             return mlirContextGetAllowUnregisteredDialects(self.get());
           },
           [](PyMlirContext &self, bool value) {
             mlirContextSetAllowUnregisteredDialects(self.get(), value);
           })
+      .def("create_operation", &PyMlirContext::createOperation, py::arg("name"),
+           py::arg("location"), py::arg("results") = py::none(),
+           py::arg("attributes") = py::none(),
+           py::arg("successors") = py::none(), py::arg("regions") = 0,
+           kContextCreateOperationDocstring)
       .def(
           "parse_module",
           [](PyMlirContext &self, const std::string moduleAsm) {
             MlirModule module =
                 mlirModuleCreateParse(self.get(), moduleAsm.c_str());
             // TODO: Rework error reporting once diagnostic engine is exposed
             // in C API.
             if (mlirModuleIsNull(module)) {
py::class_<PyMlirContext>(m, "Context")
             }
             return PyType(self.getRef(), type);
           },
           kContextParseTypeDocstring)
       .def(
           "get_unknown_location",
           [](PyMlirContext &self) {
             return PyLocation(self.getRef(),
                               mlirLocationUnknownGet(self.get()));
           },
           kContextGetUnknownLocationDocstring)
       .def(
           "get_file_location",
           [](PyMlirContext &self, std::string filename, int line, int col) {
             return PyLocation(self.getRef(),
                               mlirLocationFileLineColGet(
                                   self.get(), filename.c_str(), line, col));
           },
           kContextGetFileLocationDocstring, py::arg("filename"),
           py::arg("line"), py::arg("col"));
 
   py::class_<PyLocation>(m, "Location").def("__repr__", [](PyLocation &self) {
     PyPrintAccumulator printAccum;
     mlirLocationPrint(self.loc, printAccum.getCallback(),
                       printAccum.getUserData());
     return printAccum.join();
   });

mlir/lib/Bindings/Python/PybindUtils.h

 //===- PybindUtils.h - Utilities for interop with pybind11 ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_BINDINGS_PYTHON_PYBINDUTILS_H
 #define MLIR_BINDINGS_PYTHON_PYBINDUTILS_H
 
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/Twine.h"
 
+namespace pybind11 {
+namespace detail {
+template <typename T>
+struct type_caster<llvm::Optional<T>> : optional_caster<llvm::Optional<T>> {};
+} // namespace detail
+} // namespace pybind11
+
 namespace mlir {
 namespace python {
 
 // Sets a python error, ready to be thrown to return control back to the
 // python runtime.
 // Correct usage:
 //   throw SetPyError(PyExc_ValueError, "Foobar'd");
 pybind11::error_already_set SetPyError(PyObject *excClass,
                                        const llvm::Twine &message);
 
 } // namespace python
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_PYBINDUTILS_H

mlir/test/Bindings/Python/ir_operation.py

def testTraverseOpRegionBlockIndices():
   # CHECK:       REGION 0:
   # CHECK:         BLOCK 0:
   # CHECK:           OP 0: %0 = "custom.addi"
   # CHECK:           OP 1: return
   # CHECK:    OP 1: "module_terminator"
   walk_operations("", module.operation)
 
 run(testTraverseOpRegionBlockIndices)
+
+
+# CHECK-LABEL: TEST: testDetachedOperation
+def testDetachedOperation():
+  ctx = mlir.ir.Context()
+  ctx.allow_unregistered_dialects = True
+  loc = ctx.get_unknown_location()
+  i32 = mlir.ir.IntegerType.get_signed(ctx, 32)
+  op1 = ctx.create_operation(
+      "custom.op1", loc, results=[i32, i32], regions=1, attributes={
+          "foo": mlir.ir.StringAttr.get(ctx, "foo_value"),
+          "bar": mlir.ir.StringAttr.get(ctx, "bar_value"),
+      })
+  # CHECK: %0:2 = "custom.op1"() ( {
+  # CHECK: }) {bar = "bar_value", foo = "foo_value"} : () -> (si32, si32)
+  print(op1)
+
+  # TODO: Check successors once enough infra exists to do it properly.
+
+run(testDetachedOperation)
+
+
+# CHECK-LABEL: TEST: testOperationInsert
+def testOperationInsert():
+  ctx = mlir.ir.Context()
+  ctx.allow_unregistered_dialects = True
+  module = ctx.parse_module(r"""
+    func @f1(%arg0: i32) -> i32 {
+      %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
+      return %1 : i32
+    }
+  """)
+
+  # Create test op.
+  loc = ctx.get_unknown_location()
+  op1 = ctx.create_operation("custom.op1", loc)
+  op2 = ctx.create_operation("custom.op2", loc)
+
+  func = module.operation.regions[0].blocks[0].operations[0]
+  entry_block = func.regions[0].blocks[0]
+  entry_block.operations.insert(0, op1)
+  entry_block.operations.insert(1, op2)
+  # CHECK: func @f1
+  # CHECK: "custom.op1"()
+  # CHECK: "custom.op2"()
+  # CHECK: %0 = "custom.addi"
+  print(module)
+
+  # Trying to add a previously added op should raise.
+  try:
+    entry_block.operations.insert(0, op1)
+  except ValueError:
+    pass
+  else:
+    assert False, "expected insert of attached op to raise"
+
+run(testOperationInsert)
+
+
+# CHECK-LABEL: TEST: testOperationWithRegion
+def testOperationWithRegion():
+  ctx = mlir.ir.Context()
+  ctx.allow_unregistered_dialects = True
+  loc = ctx.get_unknown_location()
+  i32 = mlir.ir.IntegerType.get_signed(ctx, 32)
+  op1 = ctx.create_operation("custom.op1", loc, regions=1)
+  block = op1.regions[0].blocks.append(i32, i32)
+  # CHECK: "custom.op1"() ( {
+  # CHECK: ^bb0(%arg0: si32, %arg1: si32):  // no predecessors
+  # CHECK:   "custom.terminator"() : () -> ()
+  # CHECK: }) : () -> ()
+  terminator = ctx.create_operation("custom.terminator", loc)
+  block.operations.insert(0, terminator)
+  print(op1)
+
+  # Now add the whole operation to another op.
+  # TODO: Verify lifetime hazard by nulling out the new owning module and
+  # accessing op1.
+  # TODO: Also verify accessing the terminator once both parents are nulled
+  # out.
+  module = ctx.parse_module(r"""
+    func @f1(%arg0: i32) -> i32 {
+      %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
+      return %1 : i32
+    }
+  """)
+  func = module.operation.regions[0].blocks[0].operations[0]
+  entry_block = func.regions[0].blocks[0]
+  entry_block.operations.insert(0, op1)
+  # CHECK: func @f1
+  # CHECK: "custom.op1"()
+  # CHECK:   "custom.terminator"
+  # CHECK: %0 = "custom.addi"
+  print(module)
+
+run(testOperationWithRegion)