[mlir][python] Provide some methods and properties for API completeness

[ftynse]

Summary

When writing the user-facing documentation, I noticed several inconsistencies
and asymmetries in the Python API we provide. Fix them by adding:

• the owner property to regions, similarly to blocks;
• the isinstance method to any class derived from PyConcreteAttr, PyConcreteValue and PyConreteAffineExpr, similar to PyConcreteType to enable isa-like calls without having to handle exceptions;
• a mechanism to create the first block in the region as we could only create blocks relative to other blocks, with is impossible in an empty region.

[gysit]

LGTM!

mlir/lib/Bindings/Python/IRAffine.cpp

102–104

is this working similar to attributes, i.e. one can distinguish const / dim expressions etc? then a small test could make sense as well?

mlir/docs/Bindings/Python.md

 ```
 
 Regions are created for an operation when constructing it on the C++ side. They
 are not constructible in Python and are not expected to exist outside of
 operations (unlike in C++ that supports detached regions).
 
 Blocks can be created within a given region and inserted before or after another
 block of the same region using `create_before()`, `create_after()` methods of
-the `Block` class. They are not expected to exist outside of regions (unlike in
-C++ that supports detached blocks).
+the `Block` class, or the `create_at_start()` static method of the same class.
+They are not expected to exist outside of regions (unlike in C++ that supports
+detached blocks).
+
+```python
+from mlir.ir import Block, Context, Operation
+
+with Context():
+  op = Operation.create("generic.op", regions=1)
+
+  # Create the first block in the region.
+  entry_block = Block.create_at_start(op.regions[0])
+
+  # Create further blocks.
+  other_block = entry_block.create_after()
+```
 
 Blocks can be used to create `InsertionPoint`s, which can point to the beginning
 or the end of the block, or just before its terminator. It is common for
 `OpView` subclasses to provide a `.body` property that can be used to construct
 an `InsertionPoint`. For example, builtin `Module` and `FuncOp` provide a
 `.body` and `.add_entry_blocK()`, respectively.
 
 #### Attributes and Types

mlir/lib/Bindings/Python/IRAffine.cpp

if (!DerivedTy::isaFunction(orig)) {
                            DerivedTy::pyClassName + " (from " + origRepr + ")");
     }
     return orig;
   }
 
   static void bind(py::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, py::module_local());
     cls.def(py::init<PyAffineExpr &>());
+    cls.def_static("isinstance", [](PyAffineExpr &otherAffineExpr) -> bool {
+      return DerivedTy::isaFunction(otherAffineExpr);
+    });
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
 class PyAffineConstantExpr : public PyConcreteAffineExpr<PyAffineConstantExpr> {

mlir/lib/Bindings/Python/IRCore.cpp

static MlirValue castFrom(PyValue &orig) {
     }
     return orig.get();
   }
 
   /// Binds the Python module objects to functions of this class.
   static void bind(py::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, py::module_local());
     cls.def(py::init<PyValue &>(), py::keep_alive<0, 1>());
+    cls.def_static("isinstance", [](PyValue &otherValue) -> bool {
+      return DerivedTy::isaFunction(otherValue);
+    });
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
 /// Python wrapper for MlirBlockArgument.
void mlir::python::populateIRCore(py::module &m) {
   //----------------------------------------------------------------------------
   py::class_<PyRegion>(m, "Region", py::module_local())
       .def_property_readonly(
           "blocks",
           [](PyRegion &self) {
             return PyBlockList(self.getParentOperation(), self.get());
           },
           "Returns a forward-optimized sequence of blocks.")
+      .def_property_readonly(
+          "owner",
+          [](PyRegion &self) {
+            return self.getParentOperation()->createOpView();
+          },
+          "Returns the operation owning this region.")
       .def(
           "__iter__",
           [](PyRegion &self) {
             self.checkValid();
             MlirBlock firstBlock = mlirRegionGetFirstBlock(self.get());
             return PyBlockIterator(self.getParentOperation(), firstBlock);
           },
           "Iterates over blocks in the region.")
py::class_<PyBlock>(m, "Block", py::module_local())
           },
           "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_static(
+          "create_at_start",
+          [](PyRegion &parent, py::list pyArgTypes) {
+            parent.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());
+            mlirRegionInsertOwnedBlock(parent, 0, block);
+            return PyBlock(parent.getParentOperation(), block);
+          },
+          py::arg("parent"), py::arg("pyArgTypes") = py::list(),
+          "Creates and returns a new Block at the beginning of the given "
+          "region (with given argument types).")
       .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 &>());

mlir/lib/Bindings/Python/IRModule.h

 /// Regions are managed completely by their containing operation. Unlike the
 /// C++ API, the python API does not support detached regions.
 class PyRegion {
 public:
   PyRegion(PyOperationRef parentOperation, MlirRegion region)
       : parentOperation(std::move(parentOperation)), region(region) {
     assert(!mlirRegionIsNull(region) && "python region cannot be null");
   }
+  operator MlirRegion() const { return region; }
 
   MlirRegion get() { return region; }
   PyOperationRef &getParentOperation() { return parentOperation; }
 
   void checkValid() { return parentOperation->checkValid(); }
 
 private:
   PyOperationRef parentOperation;
static MlirAttribute castFrom(PyAttribute &orig) {
     }
     return orig;
   }
 
   static void bind(pybind11::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, pybind11::buffer_protocol(),
                        pybind11::module_local());
     cls.def(pybind11::init<PyAttribute &>(), pybind11::keep_alive<0, 1>());
+    cls.def_static("isinstance", [](PyAttribute &otherAttr) -> bool {
+      return DerivedTy::isaFunction(otherAttr);
+    });
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
 /// Wrapper around the generic MlirType.

 /// definition. For op result value, this is the operation that defines the
 /// value. For block argument values, this is the operation that contains the
 /// block to which the value is an argument (blocks cannot be detached in Python
 /// bindings so such operation always exists).
 class PyValue {
 public:
   PyValue(PyOperationRef parentOperation, MlirValue value)
       : parentOperation(parentOperation), value(value) {}
+  operator MlirValue() const { return value; }
 
   MlirValue get() { return value; }
   PyOperationRef &getParentOperation() { return parentOperation; }
 
   void checkValid() { return parentOperation->checkValid(); }
 
   /// Gets a capsule wrapping the void* within the MlirValue.
   pybind11::object getCapsule();

mlir/test/python/ir/affine_expr.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
+  return f
 
 
 # CHECK-LABEL: TEST: testAffineExprCapsule
+@run
 def testAffineExprCapsule():
   with Context() as ctx:
     affine_expr = AffineExpr.get_constant(42)
 
   affine_expr_capsule = affine_expr._CAPIPtr
   # CHECK: capsule object
   # CHECK: mlir.ir.AffineExpr._CAPIPtr
   print(affine_expr_capsule)
 
   affine_expr_2 = AffineExpr._CAPICreate(affine_expr_capsule)
   assert affine_expr == affine_expr_2
   assert affine_expr_2.context == ctx
 
-run(testAffineExprCapsule)
-
 
 # CHECK-LABEL: TEST: testAffineExprEq
+@run
 def testAffineExprEq():
   with Context():
     a1 = AffineExpr.get_constant(42)
     a2 = AffineExpr.get_constant(42)
     a3 = AffineExpr.get_constant(43)
     # CHECK: True
     print(a1 == a1)
     # CHECK: True
     print(a1 == a2)
     # CHECK: False
     print(a1 == a3)
     # CHECK: False
     print(a1 == None)
     # CHECK: False
     print(a1 == "foo")
 
-run(testAffineExprEq)
-
 
 # CHECK-LABEL: TEST: testAffineExprContext
+@run
 def testAffineExprContext():
   with Context():
     a1 = AffineExpr.get_constant(42)
   with Context():
     a2 = AffineExpr.get_constant(42)
 
   # CHECK: False
   print(a1 == a2)
 
 run(testAffineExprContext)
 
 
 # CHECK-LABEL: TEST: testAffineExprConstant
+@run
 def testAffineExprConstant():
   with Context():
     a1 = AffineExpr.get_constant(42)
     # CHECK: 42
     print(a1.value)
     # CHECK: 42
     print(a1)
 
     a2 = AffineConstantExpr.get(42)
     # CHECK: 42
     print(a2.value)
     # CHECK: 42
     print(a2)
 
     assert a1 == a2
 
-run(testAffineExprConstant)
-
 
 # CHECK-LABEL: TEST: testAffineExprDim
+@run
 def testAffineExprDim():
   with Context():
     d1 = AffineExpr.get_dim(1)
     d11 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
 
     # CHECK: 1
     print(d1.position)
     # CHECK: d1
     print(d1)
 
     # CHECK: 2
     print(d2.position)
     # CHECK: d2
     print(d2)
 
     assert d1 == d11
     assert d1 != d2
 
-run(testAffineExprDim)
-
 
 # CHECK-LABEL: TEST: testAffineExprSymbol
+@run
 def testAffineExprSymbol():
   with Context():
     s1 = AffineExpr.get_symbol(1)
     s11 = AffineSymbolExpr.get(1)
     s2 = AffineSymbolExpr.get(2)
 
     # CHECK: 1
     print(s1.position)
     # CHECK: s1
     print(s1)
 
     # CHECK: 2
     print(s2.position)
     # CHEKC: s2
     print(s2)
 
     assert s1 == s11
     assert s1 != s2
 
-run(testAffineExprSymbol)
-
 
 # CHECK-LABEL: TEST: testAffineAddExpr
+@run
 def testAffineAddExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
     d12 = AffineExpr.get_add(d1, d2)
     # CHECK: d1 + d2
     print(d12)
 
     d12op = d1 + d2
     # CHECK: d1 + d2
     print(d12op)
 
     assert d12 == d12op
     assert d12.lhs == d1
     assert d12.rhs == d2
 
-run(testAffineAddExpr)
-
 
 # CHECK-LABEL: TEST: testAffineMulExpr
+@run
 def testAffineMulExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_mul(d1, c2)
     # CHECK: d1 * 2
     print(expr)
 
     # CHECK: d1 * 2
     op = d1 * c2
     print(op)
 
     assert expr == op
     assert expr.lhs == d1
     assert expr.rhs == c2
 
-run(testAffineMulExpr)
-
 
 # CHECK-LABEL: TEST: testAffineModExpr
+@run
 def testAffineModExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_mod(d1, c2)
     # CHECK: d1 mod 2
     print(expr)
 
     # CHECK: d1 mod 2
     op = d1 % c2
     print(op)
 
     assert expr == op
     assert expr.lhs == d1
     assert expr.rhs == c2
 
-run(testAffineModExpr)
-
 
 # CHECK-LABEL: TEST: testAffineFloorDivExpr
+@run
 def testAffineFloorDivExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_floor_div(d1, c2)
     # CHECK: d1 floordiv 2
     print(expr)
 
     assert expr.lhs == d1
     assert expr.rhs == c2
 
-run(testAffineFloorDivExpr)
-
 
 # CHECK-LABEL: TEST: testAffineCeilDivExpr
+@run
 def testAffineCeilDivExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_ceil_div(d1, c2)
     # CHECK: d1 ceildiv 2
     print(expr)
 
     assert expr.lhs == d1
     assert expr.rhs == c2
 
-run(testAffineCeilDivExpr)
-
 
 # CHECK-LABEL: TEST: testAffineExprSub
+@run
 def testAffineExprSub():
   with Context():
     d1 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
     expr = d1 - d2
     # CHECK: d1 - d2
     print(expr)
 
     assert expr.lhs == d1
     rhs = AffineMulExpr(expr.rhs)
     # CHECK: d2
     print(rhs.lhs)
     # CHECK: -1
     print(rhs.rhs)
 
-run(testAffineExprSub)
-
-
+# CHECK-LABEL: TEST: testClassHierarchy
+@run
 def testClassHierarchy():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     add = AffineAddExpr.get(d1, c2)
     mul = AffineMulExpr.get(d1, c2)
     mod = AffineModExpr.get(d1, c2)
     floor_div = AffineFloorDivExpr.get(d1, c2)
except ValueError as e:
       print(e)
 
     try:
       AffineBinaryExpr(c2)
     except ValueError as e:
       # CHECK: Cannot cast affine expression to AffineBinaryExpr
       print(e)
 
-run(testClassHierarchy)
+# CHECK-LABEL: TEST: testIsInstance
+@run
+def testIsInstance():
+  with Context():
+    d1 = AffineDimExpr.get(1)
+    c2 = AffineConstantExpr.get(2)
+    add = AffineAddExpr.get(d1, c2)
+    mul = AffineMulExpr.get(d1, c2)
+
+    # CHECK: True
+    print(AffineDimExpr.isinstance(d1))
+    # CHECK: False
+    print(AffineConstantExpr.isinstance(d1))
+    # CHECK: True
+    print(AffineConstantExpr.isinstance(c2))
+    # CHECK: False
+    print(AffineMulExpr.isinstance(c2))
+    # CHECK: True
+    print(AffineAddExpr.isinstance(add))
+    # CHECK: False
+    print(AffineMulExpr.isinstance(add))
+    # CHECK: True
+    print(AffineMulExpr.isinstance(mul))
+    # CHECK: False
+    print(AffineAddExpr.isinstance(mul))

mlir/test/python/ir/attributes.py

 def testAttrCast():
   with Context():
     a1 = Attribute.parse('"attr1"')
     a2 = Attribute(a1)
     # CHECK: a1 == a2: True
     print("a1 == a2:", a1 == a2)
 
 
+# CHECK-LABEL: TEST: testAttrIsInstance
+@run
+def testAttrIsInstance():
+  with Context():
+    a1 = Attribute.parse("42")
+    a2 = Attribute.parse("[42]")
+    assert IntegerAttr.isinstance(a1)
+    assert not IntegerAttr.isinstance(a2)
+    assert not ArrayAttr.isinstance(a1)
+    assert ArrayAttr.isinstance(a2)
+
+
 # CHECK-LABEL: TEST: testAttrEqDoesNotRaise
 @run
 def testAttrEqDoesNotRaise():
   with Context():
     a1 = Attribute.parse('"attr1"')
     not_an_attr = "foo"
     # CHECK: False
     print(a1 == not_an_attr)

mlir/test/python/ir/blocks.py

with InsertionPoint(module.body):
         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
+
+
+# CHECK-LABEL: TEST: testFirstBlockCreation
+# CHECK: func @test(%{{.*}}: f32)
+# CHECK:   return
+@run
+def testFirstBlockCreation():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    f32 = F32Type.get()
+    with InsertionPoint(module.body):
+      func = builtin.FuncOp("test", ([f32], []))
+      entry_block = Block.create_at_start(func.operation.regions[0], [f32])
+      with InsertionPoint(entry_block):
+        std.ReturnOp([])
+
+    print(module)
+    assert module.operation.verify()
+    assert func.body.blocks[0] == entry_block

mlir/test/python/ir/operation.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
+  return f
 
 
 # Verify iterator based traversal of the op/region/block hierarchy.
 # CHECK-LABEL: TEST: testTraverseOpRegionBlockIterators
+@run
 def testTraverseOpRegionBlockIterators():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(
       r"""
     func @f1(%arg0: i32) -> i32 {
       %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
       return %1 : i32
def testTraverseOpRegionBlockIterators():
   # CHECK:     OP 0: func
   # CHECK:       REGION 0:
   # CHECK:         BLOCK 0:
   # CHECK:           OP 0: %0 = "custom.addi"
   # CHECK:           OP 1: return
   walk_operations("", op)
 
 
-run(testTraverseOpRegionBlockIterators)
-
-
 # Verify index based traversal of the op/region/block hierarchy.
 # CHECK-LABEL: TEST: testTraverseOpRegionBlockIndices
+@run
 def testTraverseOpRegionBlockIndices():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(
       r"""
     func @f1(%arg0: i32) -> i32 {
       %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
       return %1 : i32
def testTraverseOpRegionBlockIndices():
   # CHECK:         BLOCK 0:
   # CHECK:           OP 0: %0 = "custom.addi"
   # CHECK:           OP 0: parent builtin.func
   # CHECK:           OP 1: return
   # CHECK:           OP 1: parent builtin.func
   walk_operations("", module.operation)
 
 
-run(testTraverseOpRegionBlockIndices)
+# CHECK-LABEL: TEST: testBlockAndRegionOwners
+@run
+def testBlockAndRegionOwners():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  module = Module.parse(
+      r"""
+    builtin.module {
+      builtin.func @f() {
+        std.return
+      }
+    }
+  """, ctx)
+
+  assert module.operation.regions[0].owner == module.operation
+  assert module.operation.regions[0].blocks[0].owner == module.operation
+
+  func = module.body.operations[0]
+  assert func.operation.regions[0].owner == func
+  assert func.operation.regions[0].blocks[0].owner == func
 
 
 # CHECK-LABEL: TEST: testBlockArgumentList
+@run
 def testBlockArgumentList():
   with Context() as ctx:
     module = Module.parse(
         r"""
       func @f1(%arg0: i32, %arg1: f64, %arg2: index) {
         return
       }
     """, ctx)
with Context() as ctx:
 
     # CHECK: Type: i8
     # CHECK: Type: i16
     # CHECK: Type: i24
     for t in entry_block.arguments.types:
       print("Type: ", t)
 
 
-run(testBlockArgumentList)
-
-
 # CHECK-LABEL: TEST: testOperationOperands
+@run
 def testOperationOperands():
   with Context() as ctx:
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
       func @f1(%arg0: i32) {
         %0 = "test.producer"() : () -> i64
         "test.consumer"(%arg0, %0) : (i32, i64) -> ()
         return
       }""")
     func = module.body.operations[0]
     entry_block = func.regions[0].blocks[0]
     consumer = entry_block.operations[1]
     assert len(consumer.operands) == 2
     # CHECK: Operand 0, type i32
     # CHECK: Operand 1, type i64
     for i, operand in enumerate(consumer.operands):
       print(f"Operand {i}, type {operand.type}")
 
 
-run(testOperationOperands)
 
 
 # CHECK-LABEL: TEST: testOperationOperandsSlice
+@run
 def testOperationOperandsSlice():
   with Context() as ctx:
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
       func @f1() {
         %0 = "test.producer0"() : () -> i64
         %1 = "test.producer1"() : () -> i64
         %2 = "test.producer2"() : () -> i64
for operand in even:
       print(operand)
 
     # CHECK: test.producer2
     fourth = consumer.operands[::2][1::2]
     for operand in fourth:
       print(operand)
 
 
-run(testOperationOperandsSlice)
 
 
 # CHECK-LABEL: TEST: testOperationOperandsSet
+@run
 def testOperationOperandsSet():
   with Context() as ctx, Location.unknown(ctx):
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
       func @f1() {
         %0 = "test.producer0"() : () -> i64
         %1 = "test.producer1"() : () -> i64
         %2 = "test.producer2"() : () -> i64
with Context() as ctx, Location.unknown(ctx):
     consumer.operands[0] = producer1.result
     print(consumer.operands[0])
 
     # CHECK: test.producer2
     consumer.operands[-1] = producer2.result
     print(consumer.operands[0])
 
 
-run(testOperationOperandsSet)
 
 
 # CHECK-LABEL: TEST: testDetachedOperation
+@run
 def testDetachedOperation():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     i32 = IntegerType.get_signed(32)
     op1 = Operation.create(
         "custom.op1",
         results=[i32, i32],
         regions=1,
         attributes={
             "foo": StringAttr.get("foo_value"),
             "bar": StringAttr.get("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: testOperationInsertionPoint
+@run
 def testOperationInsertionPoint():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(
       r"""
     func @f1(%arg0: i32) -> i32 {
       %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
       return %1 : i32
def testOperationInsertionPoint():
   try:
     ip.insert(op1)
   except ValueError:
     pass
   else:
     assert False, "expected insert of attached op to raise"
 
 
-run(testOperationInsertionPoint)
-
-
 # CHECK-LABEL: TEST: testOperationWithRegion
+@run
 def testOperationWithRegion():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     i32 = IntegerType.get_signed(32)
     op1 = Operation.create("custom.op1", regions=1)
     block = op1.regions[0].blocks.append(i32, i32)
     # CHECK: "custom.op1"() ( {
with Location.unknown(ctx):
     ip.insert(op1)
     # CHECK: func @f1
     # CHECK: "custom.op1"()
     # CHECK:   "custom.terminator"
     # CHECK: %0 = "custom.addi"
     print(module)
 
 
-run(testOperationWithRegion)
-
-
 # CHECK-LABEL: TEST: testOperationResultList
+@run
 def testOperationResultList():
   ctx = Context()
   module = Module.parse(
       r"""
     func @f1() {
       %0:3 = call @f2() : () -> (i32, f64, index)
       return
     }
def testOperationResultList():
 
   # CHECK: Result type i32
   # CHECK: Result type f64
   # CHECK: Result type index
   for t in call.results.types:
     print(f"Result type {t}")
 
 
-run(testOperationResultList)
 
 
 # CHECK-LABEL: TEST: testOperationResultListSlice
+@run
 def testOperationResultListSlice():
   with Context() as ctx:
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
       func @f1() {
         "some.op"() : () -> (i1, i2, i3, i4, i5)
         return
       }
with Context() as ctx:
 
     # CHECK: Result 3, type i4
     # CHECK: Result 1, type i2
     inverted_middle = producer.results[-2:0:-2]
     for res in inverted_middle:
       print(f"Result {res.result_number}, type {res.type}")
 
 
-run(testOperationResultListSlice)
 
 
 # CHECK-LABEL: TEST: testOperationAttributes
+@run
 def testOperationAttributes():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(
       r"""
     "some.op"() { some.attribute = 1 : i8,
                   other.attribute = 3.0,
                   dependent = "text" } : () -> ()
def testOperationAttributes():
   try:
     op.attributes[42]
   except IndexError:
     pass
   else:
     assert False, "expected IndexError on accessing an out-of-bounds attribute"
 
 
-run(testOperationAttributes)
 
 
 # CHECK-LABEL: TEST: testOperationPrint
+@run
 def testOperationPrint():
   ctx = Context()
   module = Module.parse(
       r"""
     func @f1(%arg0: i32) -> i32 {
       %0 = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
       return %arg0 : i32
     }
def testOperationPrint():
   module.operation.print(
       large_elements_limit=2,
       enable_debug_info=True,
       pretty_debug_info=True,
       print_generic_op_form=True,
       use_local_scope=True)
 
 
-run(testOperationPrint)
 
 
 # CHECK-LABEL: TEST: testKnownOpView
+@run
 def testKnownOpView():
   with Context(), Location.unknown():
     Context.current.allow_unregistered_dialects = True
     module = Module.parse(r"""
       %1 = "custom.f32"() : () -> f32
       %2 = "custom.f32"() : () -> f32
       %3 = arith.addf %1, %2 : f32
     """)
with Context(), Location.unknown():
     print(repr(custom))
 
     # Check again to make sure negative caching works.
     custom = module.body.operations[0]
     # CHECK: OpView object
     print(repr(custom))
 
 
-run(testKnownOpView)
-
-
 # CHECK-LABEL: TEST: testSingleResultProperty
+@run
 def testSingleResultProperty():
   with Context(), Location.unknown():
     Context.current.allow_unregistered_dialects = True
     module = Module.parse(r"""
       "custom.no_result"() : () -> ()
       %0:2 = "custom.two_result"() : () -> (f32, f32)
       %1 = "custom.one_result"() : () -> f32
     """)
except ValueError as e:
     print(e)
   else:
     assert False, "Expected exception"
 
   # CHECK: %1 = "custom.one_result"() : () -> f32
   print(module.body.operations[2])
 
 
-run(testSingleResultProperty)
-
-
 # CHECK-LABEL: TEST: testPrintInvalidOperation
+@run
 def testPrintInvalidOperation():
   ctx = Context()
   with Location.unknown(ctx):
     module = Operation.create("builtin.module", regions=2)
     # This module has two region and is invalid verify that we fallback
     # to the generic printer for safety.
     block = module.regions[0].blocks.append()
     # CHECK: // Verification failed, printing generic form
     # CHECK: "builtin.module"() ( {
     # CHECK: }) : () -> ()
     print(module)
     # CHECK: .verify = False
     print(f".verify = {module.operation.verify()}")
 
 
-run(testPrintInvalidOperation)
-
-
 # CHECK-LABEL: TEST: testCreateWithInvalidAttributes
+@run
 def testCreateWithInvalidAttributes():
   ctx = Context()
   with Location.unknown(ctx):
     try:
       Operation.create(
           "builtin.module", attributes={None: StringAttr.get("name")})
     except Exception as e:
       # CHECK: Invalid attribute key (not a string) when attempting to create the operation "builtin.module"
except Exception as e:
       print(e)
     try:
       Operation.create("builtin.module", attributes={"some_key": None})
     except Exception as e:
       # CHECK: Found an invalid (`None`?) attribute value for the key "some_key" when attempting to create the operation "builtin.module"
       print(e)
 
 
-run(testCreateWithInvalidAttributes)
-
-
 # CHECK-LABEL: TEST: testOperationName
+@run
 def testOperationName():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(
       r"""
     %0 = "custom.op1"() : () -> f32
     %1 = "custom.op2"() : () -> i32
     %2 = "custom.op1"() : () -> f32
   """, ctx)
 
   # CHECK: custom.op1
   # CHECK: custom.op2
   # CHECK: custom.op1
   for op in module.body.operations:
     print(op.operation.name)
 
 
-run(testOperationName)
-
-
 # CHECK-LABEL: TEST: testCapsuleConversions
+@run
 def testCapsuleConversions():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     m = Operation.create("custom.op1").operation
     m_capsule = m._CAPIPtr
     assert '"mlir.ir.Operation._CAPIPtr"' in repr(m_capsule)
     m2 = Operation._CAPICreate(m_capsule)
     assert m2 is m
 
 
-run(testCapsuleConversions)
-
-
 # CHECK-LABEL: TEST: testOperationErase
+@run
 def testOperationErase():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     m = Module.create()
     with InsertionPoint(m.body):
       op = Operation.create("custom.op1")
 
       # CHECK: "custom.op1"
       print(m)
 
       op.operation.erase()
 
       # CHECK-NOT: "custom.op1"
       print(m)
 
       # Ensure we can create another operation
       Operation.create("custom.op2")
-
-
-run(testOperationErase)

mlir/test/python/ir/value.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
+  return f
 
 
 # CHECK-LABEL: TEST: testCapsuleConversions
+@run
 def testCapsuleConversions():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     i32 = IntegerType.get_signless(32)
     value = Operation.create("custom.op1", results=[i32]).result
     value_capsule = value._CAPIPtr
     assert '"mlir.ir.Value._CAPIPtr"' in repr(value_capsule)
     value2 = Value._CAPICreate(value_capsule)
     assert value2 == value
 
 
-run(testCapsuleConversions)
-
-
 # CHECK-LABEL: TEST: testOpResultOwner
+@run
 def testOpResultOwner():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     i32 = IntegerType.get_signless(32)
     op = Operation.create("custom.op1", results=[i32])
     assert op.result.owner == op
 
 
-run(testOpResultOwner)
+# CHECK-LABEL: TEST: testValueIsInstance
+@run
+def testValueIsInstance():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  module = Module.parse(
+      r"""
+    func @foo(%arg0: f32) {
+      %0 = "some_dialect.some_op"() : () -> f64
+      return
+    }""", ctx)
+  func = module.body.operations[0]
+  assert BlockArgument.isinstance(func.regions[0].blocks[0].arguments[0])
+  assert not OpResult.isinstance(func.regions[0].blocks[0].arguments[0])
+
+  op = func.regions[0].blocks[0].operations[0]
+  assert not BlockArgument.isinstance(op.results[0])
+  assert OpResult.isinstance(op.results[0])