Add InsertionPoint and context managers to the Python API.

[stellaraccident]

Summary

• Removes index based insertion. All insertion now happens through the insertion point.
• Introduces thread local context managers for implicit creation relative to an insertion point.
• Introduces (but does not yet use) binding the Context to the thread local context stack. Intent is to refactor all methods to take context optionally and have them use the default if available.
• Adds C APIs for mlirOperationGetParentOperation(), mlirOperationGetBlock() and mlirBlockGetTerminator().
• Removes an assert in PyOperation creation that was incorrectly constraining. There is already a TODO to rework the keepAlive field that it was guarding and without the assert, it is no worse than the current state.

[mehdi_amini]

Very nice!

mlir/lib/Bindings/Python/IRModules.cpp

647

Nit: remove trivial brace.

685

This API does not have a user in this patch right now.

905

(trivial braces twice here)

967

(trivial brace)

996

(trivial brace)

1023

(trivial brace, and above)

mlir/lib/Bindings/Python/IRModules.h

100

Why the difference between the two APIs? (one is returning nullptr the other throwing an exception)?

110

Why a std::list instead of a vector?

[stellaraccident]

Address comments.

[stellaraccident]

Thanks. PTAL.

mlir/lib/Bindings/Python/IRModules.cpp

685

I'm going to leave it (it is leaking from the next patch and has an implementation in line with the one for default insertion points).

mlir/lib/Bindings/Python/IRModules.h

100

Just added them as needed (here and leaking from the followup). I've normalized them with a required flag. I would like the accessors for the default context and insertion point centralized because the next patch wires them in various ways for converting args when not explicitly provided, and the exception raising should be in one place.

110

I suppose because "that's what I always do for these context manager holders" isn't a great answer, considering that we all know that list is inferior if not used for specific things :)

I've switch to vector, and in my defense, I think I was reaching for it because of previous designs that needed to store non-movable types.

mlir/include/mlir-c/IR.h

 
 /** Checks whether the underlying operation is null. */
 static inline int mlirOperationIsNull(MlirOperation op) { return !op.ptr; }
 
 /** Checks whether two operation handles point to the same operation. This does
  * not perform deep comparison. */
 int mlirOperationEqual(MlirOperation op, MlirOperation other);
 
+/** Gets the block that owns this operation, returning null if the operation is
+ * not owned. */
+MlirBlock mlirOperationGetBlock(MlirOperation op);
+
+/** Gets the operation that owns this operation, returning null if the operation
+ * is not owned. */
+MlirOperation mlirOperationGetParentOperation(MlirOperation op);
+
 /** Returns the number of regions attached to the given operation. */
 intptr_t mlirOperationGetNumRegions(MlirOperation op);
 
 /** Returns `pos`-th region attached to the operation. */
 MlirRegion mlirOperationGetRegion(MlirOperation op, intptr_t pos);
 
 /** Returns an operation immediately following the given operation it its
  * enclosing block. */

 
 /** Returns the block immediately following the given block in its parent
  * region. */
 MlirBlock mlirBlockGetNextInRegion(MlirBlock block);
 
 /** Returns the first operation in the block. */
 MlirOperation mlirBlockGetFirstOperation(MlirBlock block);
 
+/** Returns the terminator operation in the block or null if no terminator. */
+MlirOperation mlirBlockGetTerminator(MlirBlock block);
+
 /** Takes an operation owned by the caller and appends it to the block. */
 void mlirBlockAppendOwnedOperation(MlirBlock block, MlirOperation operation);
 
 /** Takes an operation owned by the caller and inserts it as `pos` to the block.
    This is an expensive operation that scans the block linearly, prefer
    insertBefore/After instead. */
 void mlirBlockInsertOwnedOperation(MlirBlock block, intptr_t pos,
                                    MlirOperation operation);

mlir/lib/Bindings/Python/IRModules.h

 //===- IRModules.h - IR Submodules of pybind module -----------------------===//
 //
 // 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_IRMODULES_H
 #define MLIR_BINDINGS_PYTHON_IRMODULES_H
 
+#include <vector>
+
 #include <pybind11/pybind11.h>
 
 #include "mlir-c/IR.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 namespace python {
 
 class PyBlock;
+class PyInsertionPoint;
 class PyLocation;
 class PyMlirContext;
 class PyModule;
 class PyOperation;
 class PyType;
 class PyValue;
 
 /// Template for a reference to a concrete type which captures a python
public:
   /// the reference. Note that this does not work from initializers.
   pybind11::object releaseObject() {
     assert(referrent && object);
     referrent = nullptr;
     auto stolen = std::move(object);
     return stolen;
   }
 
+  T *get() { return referrent; }
   T *operator->() {
     assert(referrent && object);
     return referrent;
   }
   pybind11::object getObject() {
     assert(referrent && object);
     return object;
   }
   operator bool() const { return referrent && object; }
 
 private:
   T *referrent;
   pybind11::object object;
 };
 
-using PyMlirContextRef = PyObjectRef<PyMlirContext>;
+/// Tracks an entry in the thread context stack. New entries are pushed onto
+/// here for each with block that activates a new InsertionPoint or Context.
+/// Pushing either a context or an insertion point resets the other:
+///   - a new context activates a new entry with a null insertion point.
+///   - a new insertion point activates a new entry with the context that the
+///     insertion point is bound to.
+class PyThreadContextEntry {
+public:
+  PyThreadContextEntry(pybind11::object context,
+                       pybind11::object insertionPoint)
+      : context(std::move(context)), insertionPoint(std::move(insertionPoint)) {
+  }
+
+  /// Gets the top of stack context and return nullptr if not defined.
+  /// If required is true and there is no default, a nice user-facing exception
+  /// is raised.
+  static PyMlirContext *getDefaultContext(bool required);
+
+  /// Gets the top of stack insertion point and return nullptr if not defined.
+  /// If required is true and there is no default, a nice user-facing exception
+  /// is raised.
+  static PyInsertionPoint *getDefaultInsertionPoint(bool required);
+
+  PyMlirContext *getContext();
+  PyInsertionPoint *getInsertionPoint();
+
+  /// Stack management.
+  static PyThreadContextEntry *getTos();
+  static void push(pybind11::object context, pybind11::object insertionPoint);
+
+  /// Gets the thread local stack.
+  static std::vector<PyThreadContextEntry> &getStack();
+
+private:
+  /// An object reference to the PyContext.
+  pybind11::object context;
+  /// An object reference to the current insertion point.
+  pybind11::object insertionPoint;
+};
 
 /// Wrapper around MlirContext.
+using PyMlirContextRef = PyObjectRef<PyMlirContext>;
 class PyMlirContext {
 public:
   PyMlirContext() = delete;
   PyMlirContext(const PyMlirContext &) = delete;
   PyMlirContext(PyMlirContext &&) = delete;
 
   /// 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

 /// 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 {
 public:
   ~PyOperation();
   /// Returns a PyOperation for the given MlirOperation, optionally associating
-  /// it with a parentKeepAlive (which must match on all such calls for the
-  /// same operation).
+  /// 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.
   static PyOperationRef
   createDetached(PyMlirContextRef contextRef, MlirOperation operation,
public:
   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();
+
 private:
   PyOperation(PyMlirContextRef contextRef, MlirOperation operation);
   static PyOperationRef createInstance(PyMlirContextRef contextRef,
                                        MlirOperation operation,
                                        pybind11::object parentKeepAlive);
 
   MlirOperation operation;
   pybind11::handle handle;
public:
 
   void checkValid() { return parentOperation->checkValid(); }
 
 private:
   PyOperationRef parentOperation;
   MlirBlock block;
 };
 
+/// An insertion point maintains a pointer to a Block and a reference operation.
+/// Calls to insert() will insert a new operation before the
+/// 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);
+
+  /// 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);
+
+  /// Enter and exit the context manager.
+  pybind11::object contextEnter();
+  void contextExit(pybind11::object excType, pybind11::object excVal,
+                   pybind11::object excTb);
+
+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)) {}
+
+  PyBlock block;
+  llvm::Optional<PyOperationRef> refOperation;
+};
+
 /// 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) {}
   bool operator==(const PyAttribute &other);
 

mlir/lib/Bindings/Python/IRModules.cpp

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("insert", &PyOperationList::insert, py::arg("index"),
-             py::arg("operation"),
-             "Inserts an operation at an indexed position");
+        .def("__len__", &PyOperationList::dunderLen);
   }
 
 private:
   PyOperationRef parentOperation;
   MlirBlock block;
 };
 
 } // namespace
if (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();
+  PyOperationRef created = PyOperation::createDetached(getRef(), operation);
+
+  // InsertPoint active?
+  PyInsertionPoint *ip =
+      PyThreadContextEntry::getDefaultInsertionPoint(/*required=*/false);
+  if (ip)
+    ip->insert(*created.get());
+
+  return created.releaseObject();
+}
+
+//------------------------------------------------------------------------------
+// PyThreadContextEntry management
+//------------------------------------------------------------------------------
+
+std::vector<PyThreadContextEntry> &PyThreadContextEntry::getStack() {
+  static thread_local std::vector<PyThreadContextEntry> stack;
+  return stack;
+}
+
+PyThreadContextEntry *PyThreadContextEntry::getTos() {
+  auto &stack = getStack();
+  if (stack.empty())
+    return nullptr;
+  return &stack.back();
+}
+
+void PyThreadContextEntry::push(pybind11::object context,
+                                pybind11::object insertionPoint) {
+  auto &stack = getStack();
+  stack.emplace_back(std::move(context), std::move(insertionPoint));
+}
+
+PyMlirContext *PyThreadContextEntry::getContext() {
+  if (!context)
+    return nullptr;
+  return py::cast<PyMlirContext *>(context);
+}
+
+PyInsertionPoint *PyThreadContextEntry::getInsertionPoint() {
+  if (!insertionPoint)
+    return nullptr;
+  return py::cast<PyInsertionPoint *>(insertionPoint);
+}
+
+PyMlirContext *PyThreadContextEntry::getDefaultContext(bool required) {
+  auto *tos = getTos();
+  PyMlirContext *context = tos ? tos->getContext() : nullptr;
+  if (required && !context) {
+    throw SetPyError(
+        PyExc_RuntimeError,
+        "A default context is required for this call but is not provided. "
+        "Establish a default by surrounding the code with "
+        "'with context:'");
+  }
+  return context;
+}
+
+PyInsertionPoint *
+PyThreadContextEntry::getDefaultInsertionPoint(bool required) {
+  auto *tos = getTos();
+  PyInsertionPoint *ip = tos ? tos->getInsertionPoint() : nullptr;
+  if (required && !ip)
+    throw SetPyError(PyExc_RuntimeError,
+                     "A default insertion point is required for this call but "
+                     "is not provided. "
+                     "Establish a default by surrounding the code with "
+                     "'with InsertionPoint(...):'");
+  return ip;
 }
 
 //------------------------------------------------------------------------------
 // PyDialect, PyDialectDescriptor, PyDialects
 //------------------------------------------------------------------------------
 
 MlirDialect PyDialects::getDialectForKey(const std::string &key,
                                          bool attrError) {
PyOperationRef PyOperation::forOperation(PyMlirContextRef contextRef,
   auto it = liveOperations.find(operation.ptr);
   if (it == liveOperations.end()) {
     // Create.
     return createInstance(std::move(contextRef), operation,
                           std::move(parentKeepAlive));
   }
   // Use existing.
   PyOperation *existing = it->second.second;
-  assert(existing->parentKeepAlive.is(parentKeepAlive));
   py::object pyRef = py::reinterpret_borrow<py::object>(it->second.first);
   return PyOperationRef(existing, std::move(pyRef));
 }
 
 PyOperationRef PyOperation::createDetached(PyMlirContextRef contextRef,
                                            MlirOperation operation,
                                            py::object parentKeepAlive) {
   auto &liveOperations = contextRef->liveOperations;
print(fileObject, /*binary=*/binary,
         /*enableDebugInfo=*/enableDebugInfo,
         /*prettyDebugInfo=*/prettyDebugInfo,
         /*printGenericOpForm=*/printGenericOpForm,
         /*useLocalScope=*/useLocalScope);
 
   return fileObject.attr("getvalue")();
 }
 
+PyOperationRef PyOperation::getParentOperation() {
+  if (!isAttached())
+    throw SetPyError(PyExc_ValueError, "Detached operations have no parent");
+  MlirOperation operation = mlirOperationGetParentOperation(get());
+  if (mlirOperationIsNull(operation))
+    throw SetPyError(PyExc_ValueError, "Operation has no parent.");
+  return PyOperation::forOperation(getContext(), operation);
+}
+
+PyBlock PyOperation::getBlock() {
+  PyOperationRef parentOperation = getParentOperation();
+  MlirBlock block = mlirOperationGetBlock(get());
+  assert(!mlirBlockIsNull(block) && "Attached operation has null parent");
+  return PyBlock{std::move(parentOperation), block};
+}
+
 PyOpView::PyOpView(py::object operation)
     : operationObject(std::move(operation)),
       operation(py::cast<PyOperation *>(this->operationObject)) {}
 
 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):
py::object PyOpView::createRawSubclass(py::object userClass) {
   auto opViewType = py::detail::get_type_handle(typeid(PyOpView), true);
   attributes["__init__"] = opViewType.attr("__init__");
   py::str origName = userClass.attr("__name__");
   py::str newName = py::str("_") + origName;
   return parentMetaclass(newName, py::make_tuple(userClass), attributes);
 }
 
 //------------------------------------------------------------------------------
+// PyInsertionPoint.
+//------------------------------------------------------------------------------
+
+PyInsertionPoint::PyInsertionPoint(PyBlock &block) : block(block) {}
+
+PyInsertionPoint::PyInsertionPoint(PyOperation &beforeOperation)
+    : block(beforeOperation.getBlock()),
+      refOperation(beforeOperation.getRef()) {}
+
+void PyInsertionPoint::insert(PyOperation &operation) {
+  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();
+    beforeOp = (*refOperation)->get();
+  }
+  mlirBlockInsertOwnedOperationBefore(block.get(), beforeOp, operation.get());
+  operation.setAttached();
+}
+
+PyInsertionPoint PyInsertionPoint::atBlockBegin(PyBlock &block) {
+  MlirOperation firstOp = mlirBlockGetFirstOperation(block.get());
+  if (mlirOperationIsNull(firstOp)) {
+    // Just insert at end.
+    return PyInsertionPoint(block);
+  }
+
+  // Insert before first op.
+  PyOperationRef firstOpRef = PyOperation::forOperation(
+      block.getParentOperation()->getContext(), firstOp);
+  return PyInsertionPoint{block, std::move(firstOpRef)};
+}
+
+PyInsertionPoint PyInsertionPoint::atBlockTerminator(PyBlock &block) {
+  MlirOperation terminator = mlirBlockGetTerminator(block.get());
+  if (mlirOperationIsNull(terminator))
+    throw SetPyError(PyExc_ValueError, "Block has no terminator");
+  PyOperationRef terminatorOpRef = PyOperation::forOperation(
+      block.getParentOperation()->getContext(), terminator);
+  return PyInsertionPoint{block, std::move(terminatorOpRef)};
+}
+
+py::object PyInsertionPoint::contextEnter() {
+  auto context = block.getParentOperation()->getContext().getObject();
+  py::object self = py::cast(this);
+  PyThreadContextEntry::push(/*context=*/std::move(context),
+                             /*insertionPoint=*/self);
+  return self;
+}
+
+void PyInsertionPoint::contextExit(pybind11::object excType,
+                                   pybind11::object excVal,
+                                   pybind11::object excTb) {
+  auto &stack = PyThreadContextEntry::getStack();
+  if (stack.empty())
+    throw SetPyError(PyExc_RuntimeError,
+                     "Unbalanced insertion point enter/exit");
+  auto &tos = stack.back();
+  PyInsertionPoint *current = tos.getInsertionPoint();
+  if (current != this)
+    throw SetPyError(PyExc_RuntimeError,
+                     "Unbalanced insertion point enter/exit");
+  stack.pop_back();
+}
+
+//------------------------------------------------------------------------------
 // PyAttribute.
 //------------------------------------------------------------------------------
 
 bool PyAttribute::operator==(const PyAttribute &other) {
   return mlirAttributeEqual(attr, other.attr);
 }
 
 //------------------------------------------------------------------------------
py::class_<PyBlock>(m, "Block")
             PyPrintAccumulator printAccum;
             mlirBlockPrint(self.get(), printAccum.getCallback(),
                            printAccum.getUserData());
             return printAccum.join();
           },
           "Returns the assembly form of the block.");
 
   //----------------------------------------------------------------------------
+  // Mapping of PyInsertionPoint.
+  //----------------------------------------------------------------------------
+
+  py::class_<PyInsertionPoint>(m, "InsertionPoint")
+      .def(py::init<PyBlock &>(), py::arg("block"),
+           "Inserts after the last operation but still inside the block.")
+      .def("__enter__", &PyInsertionPoint::contextEnter)
+      .def("__exit__", &PyInsertionPoint::contextExit)
+      .def(py::init<PyOperation &>(), 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.");
+
+  //----------------------------------------------------------------------------
   // Mapping of PyAttribute.
   //----------------------------------------------------------------------------
   py::class_<PyAttribute>(m, "Attribute")
       .def_property_readonly(
           "context",
           [](PyAttribute &self) { return self.getContext().getObject(); },
           "Context that owns the Attribute")
       .def_property_readonly("type",

mlir/lib/CAPI/IR/IR.cpp

 }
 
 void mlirOperationDestroy(MlirOperation op) { unwrap(op)->erase(); }
 
 int mlirOperationEqual(MlirOperation op, MlirOperation other) {
   return unwrap(op) == unwrap(other);
 }
 
+MlirBlock mlirOperationGetBlock(MlirOperation op) {
+  return wrap(unwrap(op)->getBlock());
+}
+
+MlirOperation mlirOperationGetParentOperation(MlirOperation op) {
+  return wrap(unwrap(op)->getParentOp());
+}
+
 intptr_t mlirOperationGetNumRegions(MlirOperation op) {
   return static_cast<intptr_t>(unwrap(op)->getNumRegions());
 }
 
 MlirRegion mlirOperationGetRegion(MlirOperation op, intptr_t pos) {
   return wrap(&unwrap(op)->getRegion(static_cast<unsigned>(pos)));
 }
 

 
 MlirOperation mlirBlockGetFirstOperation(MlirBlock block) {
   Block *cppBlock = unwrap(block);
   if (cppBlock->empty())
     return wrap(static_cast<Operation *>(nullptr));
   return wrap(&cppBlock->front());
 }
 
+MlirOperation mlirBlockGetTerminator(MlirBlock block) {
+  Block *cppBlock = unwrap(block);
+  if (cppBlock->empty())
+    return wrap(static_cast<Operation *>(nullptr));
+  Operation &back = cppBlock->back();
+  if (!back.isKnownTerminator())
+    return wrap(static_cast<Operation *>(nullptr));
+  return wrap(&back);
+}
+
 void mlirBlockAppendOwnedOperation(MlirBlock block, MlirOperation operation) {
   unwrap(block)->push_back(unwrap(operation));
 }
 
 void mlirBlockInsertOwnedOperation(MlirBlock block, intptr_t pos,
                                    MlirOperation operation) {
   auto &opList = unwrap(block)->getOperations();
   opList.insert(std::next(opList.begin(), pos), unwrap(operation));

mlir/test/Bindings/Python/dialects.py

 # TODO: Op creation and access is still quite verbose: simplify this test as
 # additional capabilities come online.
 def testCustomOpView():
   ctx = mlir.ir.Context()
   ctx.allow_unregistered_dialects = True
   f32 = mlir.ir.F32Type.get(ctx)
   loc = ctx.get_unknown_location()
   m = ctx.create_module(loc)
-  m_block = m.body
-  # TODO: Remove integer insertion in favor of InsertionPoint and/or op-based.
-  ip = [0]
+
   def createInput():
     op = ctx.create_operation("pytest_dummy.intinput", loc, results=[f32])
-    m_block.operations.insert(ip[0], op)
-    ip[0] += 1
     # TODO: Auto result cast from operation
     return op.results[0]
 
+  with mlir.ir.InsertionPoint.at_block_terminator(m.body):
     # Create via dialects context collection.
     input1 = createInput()
     input2 = createInput()
     op1 = ctx.dialects.std.AddFOp(loc, input1, input2)
-  # TODO: Auto operation cast from OpView
-  # TODO: Context manager insertion point
-  m_block.operations.insert(ip[0], op1.operation)
-  ip[0] += 1
 
     # Create via an import
     from mlir.dialects.std import AddFOp
-  op2 = AddFOp(loc, input1, op1.result)
-  m_block.operations.insert(ip[0], op2.operation)
-  ip[0] += 1
+    AddFOp(loc, input1, op1.result)
 
   # CHECK: %[[INPUT0:.*]] = "pytest_dummy.intinput"
   # CHECK: %[[INPUT1:.*]] = "pytest_dummy.intinput"
   # CHECK: %[[R0:.*]] = addf %[[INPUT0]], %[[INPUT1]] : f32
   # CHECK: %[[R1:.*]] = addf %[[INPUT0]], %[[R0]] : f32
   m.operation.print()
 
 run(testCustomOpView)

mlir/test/Bindings/Python/insertion_point.py

+# RUN: %PYTHON %s | FileCheck %s
+
+import gc
+import io
+import itertools
+from mlir.ir import *
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+  gc.collect()
+  assert Context._get_live_count() == 0
+
+
+# CHECK-LABEL: TEST: test_insert_at_block_end
+def test_insert_at_block_end():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  loc = ctx.get_unknown_location()
+  module = ctx.parse_module(r"""
+    func @foo() -> () {
+      "custom.op1"() : () -> ()
+    }
+  """)
+  entry_block = module.body.operations[0].regions[0].blocks[0]
+  ip = InsertionPoint(entry_block)
+  ip.insert(ctx.create_operation("custom.op2", loc))
+  # CHECK: "custom.op1"
+  # CHECK: "custom.op2"
+  module.operation.print()
+
+run(test_insert_at_block_end)
+
+
+# CHECK-LABEL: TEST: test_insert_before_operation
+def test_insert_before_operation():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  loc = ctx.get_unknown_location()
+  module = ctx.parse_module(r"""
+    func @foo() -> () {
+      "custom.op1"() : () -> ()
+      "custom.op2"() : () -> ()
+    }
+  """)
+  entry_block = module.body.operations[0].regions[0].blocks[0]
+  ip = InsertionPoint(entry_block.operations[1])
+  ip.insert(ctx.create_operation("custom.op3", loc))
+  # CHECK: "custom.op1"
+  # CHECK: "custom.op3"
+  # CHECK: "custom.op2"
+  module.operation.print()
+
+run(test_insert_before_operation)
+
+
+# CHECK-LABEL: TEST: test_insert_at_block_begin
+def test_insert_at_block_begin():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  loc = ctx.get_unknown_location()
+  module = ctx.parse_module(r"""
+    func @foo() -> () {
+      "custom.op2"() : () -> ()
+    }
+  """)
+  entry_block = module.body.operations[0].regions[0].blocks[0]
+  ip = InsertionPoint.at_block_begin(entry_block)
+  ip.insert(ctx.create_operation("custom.op1", loc))
+  # CHECK: "custom.op1"
+  # CHECK: "custom.op2"
+  module.operation.print()
+
+run(test_insert_at_block_begin)
+
+
+# CHECK-LABEL: TEST: test_insert_at_block_begin_empty
+def test_insert_at_block_begin_empty():
+  # TODO: Write this test case when we can create such a situation.
+  pass
+
+run(test_insert_at_block_begin_empty)
+
+
+# CHECK-LABEL: TEST: test_insert_at_terminator
+def test_insert_at_terminator():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  loc = ctx.get_unknown_location()
+  module = ctx.parse_module(r"""
+    func @foo() -> () {
+      "custom.op1"() : () -> ()
+      return
+    }
+  """)
+  entry_block = module.body.operations[0].regions[0].blocks[0]
+  ip = InsertionPoint.at_block_terminator(entry_block)
+  ip.insert(ctx.create_operation("custom.op2", loc))
+  # CHECK: "custom.op1"
+  # CHECK: "custom.op2"
+  module.operation.print()
+
+run(test_insert_at_terminator)
+
+
+# CHECK-LABEL: TEST: test_insert_at_block_terminator_missing
+def test_insert_at_block_terminator_missing():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  loc = ctx.get_unknown_location()
+  module = ctx.parse_module(r"""
+    func @foo() -> () {
+      "custom.op1"() : () -> ()
+    }
+  """)
+  entry_block = module.body.operations[0].regions[0].blocks[0]
+  try:
+    ip = InsertionPoint.at_block_terminator(entry_block)
+  except ValueError as e:
+    # CHECK: Block has no terminator
+    print(e)
+  else:
+    assert False, "Expected exception"
+
+run(test_insert_at_block_terminator_missing)
+
+
+# CHECK-LABEL: TEST: test_insertion_point_context
+def test_insertion_point_context():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  loc = ctx.get_unknown_location()
+  module = ctx.parse_module(r"""
+    func @foo() -> () {
+      "custom.op1"() : () -> ()
+    }
+  """)
+  entry_block = module.body.operations[0].regions[0].blocks[0]
+  with InsertionPoint(entry_block):
+    ctx.create_operation("custom.op2", loc)
+    with InsertionPoint.at_block_begin(entry_block):
+      ctx.create_operation("custom.opa", loc)
+      ctx.create_operation("custom.opb", loc)
+    ctx.create_operation("custom.op3", loc)
+  # CHECK: "custom.opa"
+  # CHECK: "custom.opb"
+  # CHECK: "custom.op1"
+  # CHECK: "custom.op2"
+  # CHECK: "custom.op3"
+  module.operation.print()
+
+run(test_insertion_point_context)

mlir/test/Bindings/Python/ir_operation.py

 # CHECK-LABEL: TEST: testBlockArgumentList
 def testBlockArgumentList():
   ctx = mlir.ir.Context()
   module = ctx.parse_module(r"""
     func @f1(%arg0: i32, %arg1: f64, %arg2: index) {
       return
     }
   """)
-  func = module.operation.regions[0].blocks[0].operations[0]
+  func = module.body.operations[0]
   entry_block = func.regions[0].blocks[0]
   assert len(entry_block.arguments) == 3
   # CHECK: Argument 0, type i32
   # CHECK: Argument 1, type f64
   # CHECK: Argument 2, type index
   for arg in entry_block.arguments:
     print(f"Argument {arg.arg_number}, type {arg.type}")
     new_type = mlir.ir.IntegerType.get_signless(ctx, 8 * (arg.arg_number + 1))
def testDetachedOperation():
   # 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():
+# CHECK-LABEL: TEST: testOperationInsertionPoint
+def testOperationInsertionPoint():
   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]
+  func = module.body.operations[0]
   entry_block = func.regions[0].blocks[0]
-  entry_block.operations.insert(0, op1)
-  entry_block.operations.insert(1, op2)
+  ip = mlir.ir.InsertionPoint.at_block_begin(entry_block)
+  ip.insert(op1)
+  ip.insert(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)
+    ip.insert(op1)
   except ValueError:
     pass
   else:
     assert False, "expected insert of attached op to raise"
 
-run(testOperationInsert)
+run(testOperationInsertionPoint)
 
 
 # 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)
+  ip = mlir.ir.InsertionPoint(block)
+  ip.insert(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]
+  func = module.body.operations[0]
   entry_block = func.regions[0].blocks[0]
-  entry_block.operations.insert(0, op1)
+  ip = mlir.ir.InsertionPoint.at_block_begin(entry_block)
+  ip.insert(op1)
   # CHECK: func @f1
   # CHECK: "custom.op1"()
   # CHECK:   "custom.terminator"
   # CHECK: %0 = "custom.addi"
   print(module)
 
 run(testOperationWithRegion)
 
 
 # CHECK-LABEL: TEST: testOperationResultList
 def testOperationResultList():
   ctx = mlir.ir.Context()
   module = ctx.parse_module(r"""
     func @f1() {
       %0:3 = call @f2() : () -> (i32, f64, index)
       return
     }
     func @f2() -> (i32, f64, index)
   """)
-  caller = module.operation.regions[0].blocks[0].operations[0]
+  caller = module.body.operations[0]
   call = caller.regions[0].blocks[0].operations[0]
   assert len(call.results) == 3
   # CHECK: Result 0, type i32
   # CHECK: Result 1, type f64
   # CHECK: Result 2, type index
   for res in call.results:
     print(f"Result {res.result_number}, type {res.type}")
 

mlir/test/CAPI/ir.c

 
 static void printFirstOfEach(MlirContext ctx, MlirOperation operation) {
   // Assuming we are given a module, go to the first operation of the first
   // function.
   MlirRegion region = mlirOperationGetRegion(operation, 0);
   MlirBlock block = mlirRegionGetFirstBlock(region);
   operation = mlirBlockGetFirstOperation(block);
   region = mlirOperationGetRegion(operation, 0);
+  MlirOperation parentOperation = operation;
   block = mlirRegionGetFirstBlock(region);
   operation = mlirBlockGetFirstOperation(block);
 
-  // In the module we created, the first operation of the first function is an
-  // "std.dim", which has an attribute and a single result that we can use to
-  // test the printing mechanism.
+  // Verify that parent operation and block report correctly.
+  fprintf(stderr, "Parent operation eq: %d\n",
+          mlirOperationEqual(mlirOperationGetParentOperation(operation),
+                             parentOperation));
+  fprintf(stderr, "Block eq: %d\n",
+          mlirBlockEqual(mlirOperationGetBlock(operation), block));
+
+  // In the module we created, the first operation of the first function is
+  // an "std.dim", which has an attribute and a single result that we can
+  // use to test the printing mechanism.
   mlirBlockPrint(block, printToStderr, NULL);
   fprintf(stderr, "\n");
   fprintf(stderr, "First operation: ");
   mlirOperationPrint(operation, printToStderr, NULL);
   fprintf(stderr, "\n");
 
+  // Get the block terminator and print it.
+  MlirOperation terminator = mlirBlockGetTerminator(block);
+  fprintf(stderr, "Terminator: ");
+  mlirOperationPrint(terminator, printToStderr, NULL);
+  fprintf(stderr, "\n");
+
   // Get the attribute by index.
   MlirNamedAttribute namedAttr0 = mlirOperationGetAttribute(operation, 0);
   fprintf(stderr, "Get attr 0: ");
   mlirAttributePrint(namedAttr0.attribute, printToStderr, NULL);
   fprintf(stderr, "\n");
 
   // Now re-get the attribute by name.
   MlirAttribute attr0ByName =
int main() {
   // CHECK: Number of values: 9
   // CHECK: Number of block arguments: 3
   // CHECK: Number of op results: 6
   // CHECK: 0
   // clang-format on
 
   printFirstOfEach(ctx, module);
   // clang-format off
+  // CHECK: Parent operation eq: 1
+  // CHECK: Block eq: 1
   // CHECK:   %[[C0:.*]] = constant 0 : index
   // CHECK:   %[[DIM:.*]] = dim %{{.*}}, %[[C0]] : memref<?xf32>
   // CHECK:   %[[C1:.*]] = constant 1 : index
   // CHECK:   scf.for %[[I:.*]] = %[[C0]] to %[[DIM]] step %[[C1]] {
   // CHECK:     %[[LHS:.*]] = load %{{.*}}[%[[I]]] : memref<?xf32>
   // CHECK:     %[[RHS:.*]] = load %{{.*}}[%[[I]]] : memref<?xf32>
   // CHECK:     %[[SUM:.*]] = addf %[[LHS]], %[[RHS]] : f32
   // CHECK:     store %[[SUM]], %{{.*}}[%[[I]]] : memref<?xf32>
   // CHECK:   }
   // CHECK: return
   // CHECK: First operation: {{.*}} = constant 0 : index
+  // CHECK: Terminator: return
   // CHECK: Get attr 0: 0 : index
   // CHECK: Get attr 0 by name: 0 : index
   // CHECK: does_not_exist is null: 1
   // CHECK: Result 0: {{.*}} = constant 0 : index
   // CHECK: Value is null: 0
   // CHECK: Result 0 type: index
   // CHECK: Op with set attr: {{.*}} {custom_attr = true}
   // CHECK: Remove attr: 1