[mlir][python] Extend C/Python API to be usable for CFG construction.

[stellaraccident]

Summary

• It is pretty clear that no one has tried this yet since it was both incomplete and broken.
• Fixes a symbol hiding issues keeping even the generic builder from constructing an operation with successors.
• Adds ODS support for successors.
• Adds CAPI mlirBlockGetParentRegion, mlirRegionEqual + tests (and missing test for mlirBlockGetParentOperation).
• Adds Python property: Block.region.
• Adds Python methods: Block.create_before and Block.create_after.
• Adds new blocks.py test to verify a plausible CFG construction case.

[stellaraccident]

Docs and formatting.

mlir/include/mlir-c/IR.h

 MLIR_CAPI_EXPORTED MlirRegion mlirRegionCreate();
 
 /// Takes a region owned by the caller and destroys it.
 MLIR_CAPI_EXPORTED void mlirRegionDestroy(MlirRegion region);
 
 /// Checks whether a region is null.
 static inline bool mlirRegionIsNull(MlirRegion region) { return !region.ptr; }
 
+/// Checks whether two region handles point to the same region. This does not
+/// perform deep comparison.
+MLIR_CAPI_EXPORTED bool mlirRegionEqual(MlirRegion region, MlirRegion other);
+
 /// Gets the first block in the region.
 MLIR_CAPI_EXPORTED MlirBlock mlirRegionGetFirstBlock(MlirRegion region);
 
 /// Takes a block owned by the caller and appends it to the given region.
 MLIR_CAPI_EXPORTED void mlirRegionAppendOwnedBlock(MlirRegion region,
                                                    MlirBlock block);
 
 /// Takes a block owned by the caller and inserts it at `pos` to the given

 
 /// Checks whether two blocks handles point to the same block. This does not
 /// perform deep comparison.
 MLIR_CAPI_EXPORTED bool mlirBlockEqual(MlirBlock block, MlirBlock other);
 
 /// Returns the closest surrounding operation that contains this block.
 MLIR_CAPI_EXPORTED MlirOperation mlirBlockGetParentOperation(MlirBlock);
 
+/// Returns the region that contains this block.
+MLIR_CAPI_EXPORTED MlirRegion mlirBlockGetParentRegion(MlirBlock block);
+
 /// Returns the block immediately following the given block in its parent
 /// region.
 MLIR_CAPI_EXPORTED MlirBlock mlirBlockGetNextInRegion(MlirBlock block);
 
 /// Returns the first operation in the block.
 MLIR_CAPI_EXPORTED MlirOperation mlirBlockGetFirstOperation(MlirBlock block);
 
 /// Returns the terminator operation in the block or null if no terminator.

mlir/lib/Bindings/Python/IRCore.cpp

for (auto &it : *attributes) {
                           "\" when attempting to create the operation \"" +
                           name + "\" (" + err.what() + ")";
         throw py::cast_error(msg);
       }
     }
   }
   // 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());
     }
   }
void mlir::python::populateIRCore(py::module &m) {
   py::class_<PyBlock>(m, "Block")
       .def_property_readonly(
           "owner",
           [](PyBlock &self) {
             return self.getParentOperation()->createOpView();
           },
           "Returns the owning operation of this block.")
       .def_property_readonly(
+          "region",
+          [](PyBlock &self) {
+            MlirRegion region = mlirBlockGetParentRegion(self.get());
+            return PyRegion(self.getParentOperation(), region);
+          },
+          "Returns the owning region of this block.")
+      .def_property_readonly(
           "arguments",
           [](PyBlock &self) {
             return PyBlockArgumentList(self.getParentOperation(), self.get());
           },
           "Returns a list of block arguments.")
       .def_property_readonly(
           "operations",
           [](PyBlock &self) {
             return PyOperationList(self.getParentOperation(), self.get());
           },
           "Returns a forward-optimized sequence of operations.")
       .def(
+          "create_before",
+          [](PyBlock &self, py::args pyArgTypes) {
+            self.checkValid();
+            llvm::SmallVector<MlirType, 4> argTypes;
+            argTypes.reserve(pyArgTypes.size());
+            for (auto &pyArg : pyArgTypes) {
+              argTypes.push_back(pyArg.cast<PyType &>());
+            }
+
+            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data());
+            MlirRegion region = mlirBlockGetParentRegion(self.get());
+            mlirRegionInsertOwnedBlockBefore(region, self.get(), block);
+            return PyBlock(self.getParentOperation(), block);
+          },
+          "Creates and returns a new Block before this block "
+          "(with given argument types).")
+      .def(
+          "create_after",
+          [](PyBlock &self, py::args pyArgTypes) {
+            self.checkValid();
+            llvm::SmallVector<MlirType, 4> argTypes;
+            argTypes.reserve(pyArgTypes.size());
+            for (auto &pyArg : pyArgTypes) {
+              argTypes.push_back(pyArg.cast<PyType &>());
+            }
+
+            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data());
+            MlirRegion region = mlirBlockGetParentRegion(self.get());
+            mlirRegionInsertOwnedBlockAfter(region, self.get(), block);
+            return PyBlock(self.getParentOperation(), block);
+          },
+          "Creates and returns a new Block after this block "
+          "(with given argument types).")
+      .def(
           "__iter__",
           [](PyBlock &self) {
             self.checkValid();
             MlirOperation firstOperation =
                 mlirBlockGetFirstOperation(self.get());
             return PyOperationIterator(self.getParentOperation(),
                                        firstOperation);
           },
py::class_<PyInsertionPoint>(m, "InsertionPoint")
           "ValueError if none has been set")
       .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.");
+           "Inserts an operation.")
+      .def_property_readonly(
+          "block", [](PyInsertionPoint &self) { return self.getBlock(); },
+          "Returns the block that this InsertionPoint points to.");
 
   //----------------------------------------------------------------------------
   // Mapping of PyAttribute.
   //----------------------------------------------------------------------------
   py::class_<PyAttribute>(m, "Attribute")
       // Delegate to the PyAttribute copy constructor, which will also lifetime
       // extend the backing context which owns the MlirAttribute.
       .def(py::init<PyAttribute &>(), py::arg("cast_from_type"),

mlir/lib/CAPI/IR/IR.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // Region API.
 //===----------------------------------------------------------------------===//
 
 MlirRegion mlirRegionCreate() { return wrap(new Region); }
 
+bool mlirRegionEqual(MlirRegion region, MlirRegion other) {
+  return unwrap(region) == unwrap(other);
+}
+
 MlirBlock mlirRegionGetFirstBlock(MlirRegion region) {
   Region *cppRegion = unwrap(region);
   if (cppRegion->empty())
     return wrap(static_cast<Block *>(nullptr));
   return wrap(&cppRegion->front());
 }
 
 void mlirRegionAppendOwnedBlock(MlirRegion region, MlirBlock block) {

 bool mlirBlockEqual(MlirBlock block, MlirBlock other) {
   return unwrap(block) == unwrap(other);
 }
 
 MlirOperation mlirBlockGetParentOperation(MlirBlock block) {
   return wrap(unwrap(block)->getParentOp());
 }
 
+MlirRegion mlirBlockGetParentRegion(MlirBlock block) {
+  return wrap(unwrap(block)->getParent());
+}
+
 MlirBlock mlirBlockGetNextInRegion(MlirBlock block) {
   return wrap(unwrap(block)->getNextNode());
 }
 
 MlirOperation mlirBlockGetFirstOperation(MlirBlock block) {
   Block *cppBlock = unwrap(block);
   if (cppBlock->empty())
     return wrap(static_cast<Operation *>(nullptr));

mlir/test/CAPI/ir.c

static void printFirstOfEach(MlirContext ctx, MlirOperation operation) {
   operation = mlirBlockGetFirstOperation(block);
   region = mlirOperationGetRegion(operation, 0);
   MlirOperation parentOperation = operation;
   block = mlirRegionGetFirstBlock(region);
   operation = mlirBlockGetFirstOperation(block);
   assert(mlirModuleIsNull(mlirModuleFromOperation(operation)));
 
   // Verify that parent operation and block report correctly.
+  // CHECK: Parent operation eq: 1
   fprintf(stderr, "Parent operation eq: %d\n",
           mlirOperationEqual(mlirOperationGetParentOperation(operation),
                              parentOperation));
+  // CHECK: Block eq: 1
   fprintf(stderr, "Block eq: %d\n",
           mlirBlockEqual(mlirOperationGetBlock(operation), block));
-  // CHECK: Parent operation eq: 1
-  // CHECK: Block eq: 1
+  // CHECK: Block parent operation eq: 1
+  fprintf(
+      stderr, "Block parent operation eq: %d\n",
+      mlirOperationEqual(mlirBlockGetParentOperation(block), parentOperation));
+  // CHECK: Block parent region eq: 1
+  fprintf(stderr, "Block parent region eq: %d\n",
+          mlirRegionEqual(mlirBlockGetParentRegion(block), region));
 
   // In the module we created, the first operation of the first function is
   // an "memref.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);
static void printFirstOfEach(MlirContext ctx, MlirOperation operation) {
 
   // Add a large attribute to verify printing flags.
   int64_t eltsShape[] = {4};
   int32_t eltsData[] = {1, 2, 3, 4};
   mlirOperationSetAttributeByName(
       operation, mlirStringRefCreateFromCString("elts"),
       mlirDenseElementsAttrInt32Get(
           mlirRankedTensorTypeGet(1, eltsShape, mlirIntegerTypeGet(ctx, 32),
-                                  mlirAttributeGetNull()), 4, eltsData));
+                                  mlirAttributeGetNull()),
+          4, eltsData));
   MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
   mlirOpPrintingFlagsElideLargeElementsAttrs(flags, 2);
   mlirOpPrintingFlagsPrintGenericOpForm(flags);
   mlirOpPrintingFlagsEnableDebugInfo(flags, /*prettyForm=*/0);
   mlirOpPrintingFlagsUseLocalScope(flags);
   fprintf(stderr, "Op print with all flags: ");
   mlirOperationPrintWithFlags(operation, flags, printToStderr, NULL);
   fprintf(stderr, "\n");
int printBuiltinAttributes(MlirContext ctx) {
   MlirAttribute uint8Elements = mlirDenseElementsAttrUInt8Get(
       mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeUnsignedGet(ctx, 8),
                               encoding),
       2, uints8);
   MlirAttribute int8Elements = mlirDenseElementsAttrInt8Get(
       mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 8), encoding),
       2, ints8);
   MlirAttribute uint32Elements = mlirDenseElementsAttrUInt32Get(
-      mlirRankedTensorTypeGet(2, shape,
-                              mlirIntegerTypeUnsignedGet(ctx, 32), encoding),
+      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeUnsignedGet(ctx, 32),
+                              encoding),
       2, uints32);
   MlirAttribute int32Elements = mlirDenseElementsAttrInt32Get(
       mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 32), encoding),
       2, ints32);
   MlirAttribute uint64Elements = mlirDenseElementsAttrUInt64Get(
-      mlirRankedTensorTypeGet(2, shape,
-                              mlirIntegerTypeUnsignedGet(ctx, 64), encoding),
+      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeUnsignedGet(ctx, 64),
+                              encoding),
       2, uints64);
   MlirAttribute int64Elements = mlirDenseElementsAttrInt64Get(
       mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 64), encoding),
       2, ints64);
   MlirAttribute floatElements = mlirDenseElementsAttrFloatGet(
-      mlirRankedTensorTypeGet(2, shape, mlirF32TypeGet(ctx), encoding),
-      2, floats);
+      mlirRankedTensorTypeGet(2, shape, mlirF32TypeGet(ctx), encoding), 2,
+      floats);
   MlirAttribute doubleElements = mlirDenseElementsAttrDoubleGet(
-      mlirRankedTensorTypeGet(2, shape, mlirF64TypeGet(ctx), encoding),
-      2, doubles);
+      mlirRankedTensorTypeGet(2, shape, mlirF64TypeGet(ctx), encoding), 2,
+      doubles);
 
   if (!mlirAttributeIsADenseElements(boolElements) ||
       !mlirAttributeIsADenseElements(uint8Elements) ||
       !mlirAttributeIsADenseElements(int8Elements) ||
       !mlirAttributeIsADenseElements(uint32Elements) ||
       !mlirAttributeIsADenseElements(int32Elements) ||
       !mlirAttributeIsADenseElements(uint64Elements) ||
       !mlirAttributeIsADenseElements(int64Elements) ||
int printBuiltinAttributes(MlirContext ctx) {
   // CHECK: 1.000000e+00 : f64
 
   int64_t indices[] = {4, 7};
   int64_t two = 2;
   MlirAttribute indicesAttr = mlirDenseElementsAttrInt64Get(
       mlirRankedTensorTypeGet(1, &two, mlirIntegerTypeGet(ctx, 64), encoding),
       2, indices);
   MlirAttribute valuesAttr = mlirDenseElementsAttrFloatGet(
-      mlirRankedTensorTypeGet(1, &two, mlirF32TypeGet(ctx), encoding),
-      2, floats);
+      mlirRankedTensorTypeGet(1, &two, mlirF32TypeGet(ctx), encoding), 2,
+      floats);
   MlirAttribute sparseAttr = mlirSparseElementsAttribute(
       mlirRankedTensorTypeGet(2, shape, mlirF32TypeGet(ctx), encoding),
       indicesAttr, valuesAttr);
   mlirAttributeDump(sparseAttr);
   // CHECK: sparse<[4, 7], [0.000000e+00, 1.000000e+00]> : tensor<1x2xf32>
 
   return 0;
 }
int testClone() {
   fprintf(stderr, "@testClone\n");
   // CHECK-LABEL: @testClone
 
   MlirContext ctx = mlirContextCreate();
   mlirRegisterAllDialects(ctx);
   mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("std"));
   MlirLocation loc = mlirLocationUnknownGet(ctx);
   MlirType indexType = mlirIndexTypeGet(ctx);
-  MlirStringRef valueStringRef =  mlirStringRefCreateFromCString("value");
+  MlirStringRef valueStringRef = mlirStringRefCreateFromCString("value");
 
   MlirAttribute indexZeroLiteral =
       mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("0 : index"));
-  MlirNamedAttribute indexZeroValueAttr = mlirNamedAttributeGet(mlirIdentifierGet(ctx, valueStringRef), indexZeroLiteral);
+  MlirNamedAttribute indexZeroValueAttr = mlirNamedAttributeGet(
+      mlirIdentifierGet(ctx, valueStringRef), indexZeroLiteral);
   MlirOperationState constZeroState = mlirOperationStateGet(
       mlirStringRefCreateFromCString("std.constant"), loc);
   mlirOperationStateAddResults(&constZeroState, 1, &indexType);
   mlirOperationStateAddAttributes(&constZeroState, 1, &indexZeroValueAttr);
   MlirOperation constZero = mlirOperationCreate(&constZeroState);
 
   MlirAttribute indexOneLiteral =
       mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("1 : index"));

mlir/test/mlir-tblgen/op-python-bindings.td

def AttrSizedOperandsOp : TestOp<"attr_sized_operands",
                                  [AttrSizedOperandSegments]> {
   // CHECK: def __init__(self, variadic1, non_variadic, variadic2, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   operands.append(variadic1)
   // CHECK:   operands.append(non_variadic)
   // CHECK:   if variadic2 is not None: operands.append(variadic2)
+  // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def variadic1(self):
   // CHECK:   operand_range = _ods_segmented_accessor(
   // CHECK:       self.operation.operands,
   // CHECK:       self.operation.attributes["operand_segment_sizes"], 0)
   // CHECK:   return operand_range
   //
def AttrSizedResultsOp : TestOp<"attr_sized_results",
                                [AttrSizedResultSegments]> {
   // CHECK: def __init__(self, variadic1, non_variadic, variadic2, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   if variadic1 is not None: results.append(variadic1)
   // CHECK:   results.append(non_variadic)
   // CHECK:   if variadic2 is not None: results.append(variadic2)
+  // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def variadic1(self):
   // CHECK:   result_range = _ods_segmented_accessor(
   // CHECK:       self.operation.results,
   // CHECK:       self.operation.attributes["result_segment_sizes"], 0)
   // CHECK:   return result_range[0] if len(result_range) > 0 else None
   //
def AttributedOp : TestOp<"attributed_op"> {
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   attributes["i32attr"] = i32attr
   // CHECK:   if optionalF32Attr is not None: attributes["optionalF32Attr"] = optionalF32Attr
   // CHECK:   if bool(unitAttr): attributes["unitAttr"] = _ods_ir.UnitAttr.get(
   // CHECK:     _ods_get_default_loc_context(loc))
   // CHECK:   attributes["in"] = in_
+  // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def i32attr(self):
   // CHECK:   return _ods_ir.IntegerAttr(self.operation.attributes["i32attr"])
 
   // CHECK: @builtins.property
   // CHECK: def optionalF32Attr(self):
   // CHECK:   if "optionalF32Attr" not in self.operation.attributes:
def AttributedOpWithOperands : TestOp<"attributed_op_with_operands"> {
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   operands.append(_gen_arg_0)
   // CHECK:   operands.append(_gen_arg_2)
   // CHECK:   if bool(in_): attributes["in"] = _ods_ir.UnitAttr.get(
   // CHECK:     _ods_get_default_loc_context(loc))
   // CHECK:   if is_ is not None: attributes["is"] = is_
+  // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def in_(self):
   // CHECK:   return "in" in self.operation.attributes
 
   // CHECK: @builtins.property
   // CHECK: def is_(self):
   // CHECK:   if "is" not in self.operation.attributes:
   // CHECK:     return None
   // CHECK:   return _ods_ir.FloatAttr(self.operation.attributes["is"])
   let arguments = (ins I32, UnitAttr:$in, F32, OptionalAttr<F32Attr>:$is);
 }
 
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class EmptyOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.empty"
 def EmptyOp : TestOp<"empty">;
   // CHECK: def __init__(self, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
+  // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class MissingNamesOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.missing_names"
 def MissingNamesOp : TestOp<"missing_names"> {
   // CHECK: def __init__(self, i32, _gen_res_1, i64, _gen_arg_0, f32, _gen_arg_2, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   results.append(i32)
   // CHECK:   results.append(_gen_res_1)
   // CHECK:   results.append(i64)
   // CHECK:   operands.append(_gen_arg_0)
   // CHECK:   operands.append(f32)
   // CHECK:   operands.append(_gen_arg_2)
+  // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def f32(self):
   // CHECK:   return self.operation.operands[1]
   let arguments = (ins I32, F32:$f32, I64);
 
   // CHECK: @builtins.property
   // CHECK: def i32(self):

 // CHECK-NOT: _ODS_RESULT_SEGMENTS
 def OneVariadicOperandOp : TestOp<"one_variadic_operand"> {
   // CHECK: def __init__(self, non_variadic, variadic, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   operands.append(non_variadic)
   // CHECK:   operands.extend(variadic)
+  // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def non_variadic(self):
   // CHECK:   return self.operation.operands[0]
   //
   // CHECK: @builtins.property
   // CHECK: def variadic(self):
   // CHECK:   _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
   // CHECK:   return self.operation.operands[1:1 + _ods_variadic_group_length]
   let arguments = (ins AnyType:$non_variadic, Variadic<AnyType>:$variadic);
 }
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class OneVariadicResultOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.one_variadic_result"
 // CHECK-NOT: _ODS_OPERAND_SEGMENTS
 // CHECK-NOT: _ODS_RESULT_SEGMENTS
 def OneVariadicResultOp : TestOp<"one_variadic_result"> {
   // CHECK: def __init__(self, variadic, non_variadic, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   results.extend(variadic)
   // CHECK:   results.append(non_variadic)
+  // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def variadic(self):
   // CHECK:   _ods_variadic_group_length = len(self.operation.results) - 2 + 1
   // CHECK:   return self.operation.results[0:0 + _ods_variadic_group_length]
   //
   // CHECK: @builtins.property
   // CHECK: def non_variadic(self):
   // CHECK:   _ods_variadic_group_length = len(self.operation.results) - 2 + 1
   // CHECK:   return self.operation.results[1 + _ods_variadic_group_length - 1]
   let results = (outs Variadic<AnyType>:$variadic, AnyType:$non_variadic);
 }
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class PythonKeywordOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.python_keyword"
 def PythonKeywordOp : TestOp<"python_keyword"> {
   // CHECK: def __init__(self, in_, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   operands.append(in_)
+  // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def in_(self):
   // CHECK:   return self.operation.operands[0]
   let arguments = (ins AnyType:$in);
 }
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
def SimpleOp : TestOp<"simple"> {
   // CHECK: def __init__(self, i64, f64, i32, f32, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   results.append(i64)
   // CHECK:   results.append(f64)
   // CHECK:   operands.append(i32)
   // CHECK:   operands.append(f32)
+  // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
-  // CHECK:     loc=loc, ip=ip))
+  // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def i32(self):
   // CHECK:   return self.operation.operands[0]
   //
   // CHECK: @builtins.property
   // CHECK: def f32(self):
   // CHECK:   return self.operation.operands[1]
   let arguments = (ins I32:$i32, F32:$f32);
 
   // CHECK: @builtins.property
   // CHECK: def i64(self):
   // CHECK:   return self.operation.results[0]
   //
   // CHECK: @builtins.property
   // CHECK: def f64(self):
   // CHECK:   return self.operation.results[1]
   let results = (outs I64:$i64, F64:$f64);
 }
+
+// CHECK: @_ods_cext.register_operation(_Dialect)
+// CHECK: class WithSuccessorsOp(_ods_ir.OpView):
+// CHECK-LABEL: OPERATION_NAME = "test.with_successors"
+def WithSuccessorsOp : TestOp<"with_successors"> {
+  // CHECK-NOT:  _ods_successors = None
+  // CHECK:      _ods_successors = []
+  // CHECK-NEXT: _ods_successors.append(successor)
+  // CHECK-NEXT: _ods_successors.extend(successors)
+  let successors = (successor AnySuccessor:$successor,
+                              VariadicSuccessor<AnySuccessor>:$successors);
+}

mlir/test/python/ir/blocks.py

+# RUN: %PYTHON %s | FileCheck %s
+
+import gc
+import io
+import itertools
+from mlir.ir import *
+from mlir.dialects import builtin
+# Note: std dialect needed for terminators.
+from mlir.dialects import std
+
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+  gc.collect()
+  assert Context._get_live_count() == 0
+  return f
+
+
+# CHECK-LABEL: TEST: testBlockCreation
+# CHECK: func @test(%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i16)
+# CHECK:   br ^bb1(%[[ARG1]] : i16)
+# CHECK: ^bb1(%[[PHI0:.*]]: i16):
+# CHECK:   br ^bb2(%[[ARG0]] : i32)
+# CHECK: ^bb2(%[[PHI1:.*]]: i32):
+# CHECK:   return
+@run
+def testBlockCreation():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      f_type = FunctionType.get(
+          [IntegerType.get_signless(32),
+           IntegerType.get_signless(16)], [])
+      f_op = builtin.FuncOp("test", f_type)
+      entry_block = f_op.add_entry_block()
+      i32_arg, i16_arg = entry_block.arguments
+      successor_block = entry_block.create_after(i32_arg.type)
+      with InsertionPoint(successor_block) as successor_ip:
+        assert successor_ip.block == successor_block
+        std.ReturnOp([])
+      middle_block = successor_block.create_before(i16_arg.type)
+
+      with InsertionPoint(entry_block) as entry_ip:
+        assert entry_ip.block == entry_block
+        std.BranchOp([i16_arg], dest=middle_block)
+
+      with InsertionPoint(middle_block) as middle_ip:
+        assert middle_ip.block == middle_block
+        std.BranchOp([i32_arg], dest=successor_block)
+    print(module.operation)
+    # Ensure region back references are coherent.
+    assert entry_block.region == middle_block.region == successor_block.region

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

if (namedAttr.attr.isOptional()) {
       // Non-optional attributes cannot be deleted.
     }
   }
 }
 
 /// Template for the default auto-generated builder.
 ///   {0} is a comma-separated list of builder arguments, including the trailing
 ///       `loc` and `ip`;
-///   {1} is the code populating `operands`, `results` and `attributes` fields.
+///   {1} is the code populating `operands`, `results` and `attributes`,
+///       `successors` fields.
 constexpr const char *initTemplate = R"Py(
   def __init__(self, {0}):
     operands = []
     results = []
     attributes = {{}
     {1}
     super().__init__(self.build_generic(
       attributes=attributes, results=results, operands=operands,
-      loc=loc, ip=ip))
+      successors=_ods_successors, loc=loc, ip=ip))
 )Py";
 
 /// Template for appending a single element to the operand/result list.
 ///   {0} is either 'operand' or 'result';
 ///   {1} is the field name.
 constexpr const char *singleElementAppendTemplate = "{0}s.append({1})";
 
 /// Template for appending an optional element to the operand/result list.

 ///   {1} is the builder argument name.
 constexpr const char *initOptionalAttributeTemplate =
     R"Py(if {1} is not None: attributes["{0}"] = {1})Py";
 
 constexpr const char *initUnitAttributeTemplate =
     R"Py(if bool({1}): attributes["{0}"] = _ods_ir.UnitAttr.get(
       _ods_get_default_loc_context(loc)))Py";
 
+/// Template to initialize the successors list in the builder if there are any
+/// successors.
+///   {0} is the value to initialize the successors list to.
+constexpr const char *initSuccessorsTemplate = R"Py(_ods_successors = {0})Py";
+
+/// Template to append or extend the list of successors in the builder.
+///   {0} is the list method ('append' or 'extend');
+///   {1} is the value to add.
+constexpr const char *addSuccessorTemplate = R"Py(_ods_successors.{0}({1}))Py";
+
 /// Populates `builderArgs` with the Python-compatible names of builder function
 /// arguments, first the results, then the intermixed attributes and operands in
 /// the same order as they appear in the `arguments` field of the op definition.
 /// Additionally, `operandNames` is populated with names of operands in their
 /// order of appearance.
 static void
 populateBuilderArgs(const Operator &op,
                     llvm::SmallVectorImpl<std::string> &builderArgs,
-                    llvm::SmallVectorImpl<std::string> &operandNames) {
+                    llvm::SmallVectorImpl<std::string> &operandNames,
+                    llvm::SmallVectorImpl<std::string> &successorArgNames) {
   for (int i = 0, e = op.getNumResults(); i < e; ++i) {
     std::string name = op.getResultName(i).str();
     if (name.empty()) {
       if (op.getNumResults() == 1) {
         // Special case for one result, make the default name be 'result'
         // to properly match the built-in result accessor.
         name = "result";
       } else {
         name = llvm::formatv("_gen_res_{0}", i);
       }
     }
     name = sanitizeName(name);
     builderArgs.push_back(name);
   }
   for (int i = 0, e = op.getNumArgs(); i < e; ++i) {
     std::string name = op.getArgName(i).str();
     if (name.empty())
       name = llvm::formatv("_gen_arg_{0}", i);
     name = sanitizeName(name);
     builderArgs.push_back(name);
     if (!op.getArg(i).is<NamedAttribute *>())
       operandNames.push_back(name);
   }
+
+  for (int i = 0, e = op.getNumSuccessors(); i < e; ++i) {
+    NamedSuccessor successor = op.getSuccessor(i);
+    std::string name = std::string(successor.name);
+    if (name.empty())
+      name = llvm::formatv("_gen_successor_{0}", i);
+    name = sanitizeName(name);
+    builderArgs.push_back(name);
+    successorArgNames.push_back(name);
+  }
 }
 
 /// Populates `builderLines` with additional lines that are required in the
 /// builder to set up operation attributes. `argNames` is expected to contain
 /// the names of builder arguments that correspond to op arguments, i.e. to the
 /// operands and attributes in the same order as they appear in the `arguments`
 /// field.
 static void
for (int i = 0, e = op.getNumArgs(); i < e; ++i) {
     builderLines.push_back(llvm::formatv(attribute->attr.isOptional()
                                              ? initOptionalAttributeTemplate
                                              : initAttributeTemplate,
                                          attribute->name, argNames[i]));
   }
 }
 
 /// Populates `builderLines` with additional lines that are required in the
+/// builder to set up successors. successorArgNames is expected to correspond
+/// to the Python argument name for each successor on the op.
+static void populateBuilderLinesSuccessors(
+    const Operator &op, llvm::ArrayRef<std::string> successorArgNames,
+    llvm::SmallVectorImpl<std::string> &builderLines) {
+  if (successorArgNames.empty()) {
+    builderLines.push_back(llvm::formatv(initSuccessorsTemplate, "None"));
+    return;
+  }
+
+  builderLines.push_back(llvm::formatv(initSuccessorsTemplate, "[]"));
+  for (int i = 0, e = successorArgNames.size(); i < e; ++i) {
+    auto &argName = successorArgNames[i];
+    const NamedSuccessor &successor = op.getSuccessor(i);
+    builderLines.push_back(
+        llvm::formatv(addSuccessorTemplate,
+                      successor.isVariadic() ? "extend" : "append", argName));
+  }
+}
+
+/// Populates `builderLines` with additional lines that are required in the
 /// builder. `kind` must be either "operand" or "result". `names` contains the
 /// names of init arguments that correspond to the elements.
 static void populateBuilderLines(
     const Operator &op, const char *kind, llvm::ArrayRef<std::string> names,
     llvm::SmallVectorImpl<std::string> &builderLines,
     llvm::function_ref<int(const Operator &)> getNumElements,
     llvm::function_ref<const NamedTypeConstraint &(const Operator &, int)>
         getElement) {

 
 /// Emits a default builder constructing an operation from the list of its
 /// result types, followed by a list of its operands.
 static void emitDefaultOpBuilder(const Operator &op, raw_ostream &os) {
   // If we are asked to skip default builders, comply.
   if (op.skipDefaultBuilders())
     return;
 
-  llvm::SmallVector<std::string, 8> builderArgs;
-  llvm::SmallVector<std::string, 8> builderLines;
-  llvm::SmallVector<std::string, 4> operandArgNames;
+  llvm::SmallVector<std::string> builderArgs;
+  llvm::SmallVector<std::string> builderLines;
+  llvm::SmallVector<std::string> operandArgNames;
+  llvm::SmallVector<std::string> successorArgNames;
   builderArgs.reserve(op.getNumOperands() + op.getNumResults() +
-                      op.getNumNativeAttributes());
-  populateBuilderArgs(op, builderArgs, operandArgNames);
+                      op.getNumNativeAttributes() + op.getNumSuccessors());
+  populateBuilderArgs(op, builderArgs, operandArgNames, successorArgNames);
+
   populateBuilderLines(
       op, "result",
       llvm::makeArrayRef(builderArgs).take_front(op.getNumResults()),
       builderLines, getNumResults, getResult);
   populateBuilderLines(op, "operand", operandArgNames, builderLines,
                        getNumOperands, getOperand);
   populateBuilderLinesAttr(
       op, llvm::makeArrayRef(builderArgs).drop_front(op.getNumResults()),
       builderLines);
+  populateBuilderLinesSuccessors(op, successorArgNames, builderLines);
 
   builderArgs.push_back("*");
   builderArgs.push_back("loc=None");
   builderArgs.push_back("ip=None");
   os << llvm::formatv(initTemplate, llvm::join(builderArgs, ", "),
                       llvm::join(builderLines, "\n    "));
 }