[MLIR] [Python] Fix Value.owner to handle BlockArgs

Previously, calling `Value.owner()` would C++ assert in debug builds if
`Value` was a block argument. Additionally, the behavior was just wrong
in release builds. This patch adds support for BlockArg Values.

mlir/lib/Bindings/Python/IRCore.cpp

@@ -3118,11 +3118,22 @@ void mlir::python::populateIRCore(py::module &m) {
       .def_property_readonly(
           "owner",
           [](PyValue &self) {
-            assert(mlirOperationEqual(self.getParentOperation()->get(),
-                                      mlirOpResultGetOwner(self.get())) &&
-                   "expected the owner of the value in Python to match that in "
-                   "the IR");
-            return self.getParentOperation().getObject();
+            MlirValue v = self.get();
+            if (mlirValueIsAOpResult(v)) {
+              assert(
+                  mlirOperationEqual(self.getParentOperation()->get(),
+                                     mlirOpResultGetOwner(self.get())) &&
+                  "expected the owner of the value in Python to match that in "
+                  "the IR");
+              return self.getParentOperation().getObject();
+            }
+
+            if (mlirValueIsABlockArgument(v)) {
+              MlirBlock block = mlirBlockArgumentGetOwner(self.get());
+              return py::cast(PyBlock(self.getParentOperation(), block));
+            }
+
+            assert(false && "Value must be a block argument or an op result");
           })
       .def("__eq__",
            [](PyValue &self, PyValue &other) {

mlir/test/python/ir/value.py

@@ -37,6 +37,21 @@ def testOpResultOwner():
     assert op.result.owner == op
 
 
+# CHECK-LABEL: TEST: testBlockArgOwner
+@run
+def testBlockArgOwner():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  module = Module.parse(
+      r"""
+    func.func @foo(%arg0: f32) {
+      return
+    }""", ctx)
+  func = module.body.operations[0]
+  block = func.regions[0].blocks[0]
+  assert block.arguments[0].owner == block
+
+
 # CHECK-LABEL: TEST: testValueIsInstance
 @run
 def testValueIsInstance():