[mlir][Python] Add a simple PyOpOperand iterator for PyValue uses.

[mikeurbach]

Summary

This adds a simple PyOpOperand based on MlirOpOperand, which can has
properties for the owner op and operation number.

This also adds a PyOpOperandIterator that defines methods for iter
and next so PyOpOperands can be iterated over using the the
MlirOpOperand C API.

Finally, a uses psuedo-container is added to PyValue so the uses can
generically be iterated.

Depends on D139596

[mikeurbach]

Fixup unrelated whitespace change.

[mikeurbach]

Switch to the new C API for OpOperand.

[mikeurbach]

Rebase main after C API landed.

[mikeurbach]

I'm not sure who to tag on Python API additions these days, maybe @mehdi_amini, @jpienaar, or @jdd could take a look?

[jdd]

Looks correct to me!

mlir/lib/Bindings/Python/IRCore.cpp

452

It should be possible to get the context from the operand via the forContext method on PyMlirContextRef.

486–487

Why reassign? It'd be clearer to collapse 486 through 488 into one line IMO.

[stellaraccident]

Lgtm modulo trivial comments

mlir/test/python/ir/value.py

94

Can you move these checks down to the for loop for locality?

[mikeurbach]

mlir/lib/Bindings/Python/IRCore.cpp

486–487

I'm not sure I follow how to avoid the reassignment. We need to update some state so the next call to dunderNext either stops the iteration or returns the next op operand. This is the same pattern used for blocks, for example: https://github.com/llvm/llvm-project/blob/3c9f479a5e47d5f3de07ccac365589dd928bfd20/mlir/lib/Bindings/Python/IRCore.cpp#L268-L277

[jdd]

mlir/lib/Bindings/Python/IRCore.cpp

486–487

Sorry... I misread. You're right. I don't read so good.

[mikeurbach]

Some clean ups suggested in review.

[mikeurbach]

Thanks for the reviews!

mlir/lib/Bindings/Python/IRCore.cpp

 //===- IRModules.cpp - 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
 //
 //===----------------------------------------------------------------------===//
 
py::class_<PyOperationList>(m, "OperationList", py::module_local())
         .def("__len__", &PyOperationList::dunderLen);
   }
 
 private:
   PyOperationRef parentOperation;
   MlirBlock block;
 };
 
+class PyOpOperand {
+public:
+  PyOpOperand(PyMlirContextRef context, MlirOpOperand opOperand)
+      : context(std::move(context)), opOperand(opOperand) {}
+
+  py::object getOwner() {
+    return PyOperation::forOperation(context, mlirOpOperandGetOwner(opOperand))
+        ->createOpView();
+  }
+
+  size_t getOperandNumber() { return mlirOpOperandGetOperandNumber(opOperand); }
+
+  static void bind(py::module &m) {
+    py::class_<PyOpOperand>(m, "OpOperand", py::module_local())
+        .def_property_readonly("owner", &PyOpOperand::getOwner)
+        .def_property_readonly("operand_number",
+                               &PyOpOperand::getOperandNumber);
+  }
+
+private:
+  PyMlirContextRef context;
+  MlirOpOperand opOperand;
+};
+
+class PyOpOperandIterator {
+public:
+  PyOpOperandIterator(MlirOpOperand opOperand)
+      : context(PyMlirContext::forContext(
+            mlirOperationGetContext(mlirOpOperandGetOwner(opOperand)))),
+        opOperand(opOperand) {}
+
+  PyOpOperandIterator &dunderIter() { return *this; }
+
+  PyOpOperand dunderNext() {
+    if (mlirOpOperandIsNull(opOperand))
+      throw py::stop_iteration();
+    MlirOpOperand returnOpOperand = opOperand;
+    opOperand = mlirOpOperandGetNextUse(opOperand);
+    return PyOpOperand(context, returnOpOperand);
+  }
+
+  static void bind(py::module &m) {
+    py::class_<PyOpOperandIterator>(m, "OpOperandIterator", py::module_local())
+        .def("__iter__", &PyOpOperandIterator::dunderIter)
+        .def("__next__", &PyOpOperandIterator::dunderNext);
+  }
+
+private:
+  PyMlirContextRef context;
+  MlirOpOperand opOperand;
+};
+
 } // namespace
 
 //------------------------------------------------------------------------------
 // PyMlirContext
 //------------------------------------------------------------------------------
 
 PyMlirContext::PyMlirContext(MlirContext context) : context(context) {
   py::gil_scoped_acquire acquire;
py::class_<PyValue>(m, "Value", py::module_local())
             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");
             return py::none();
           })
+      .def_property_readonly("uses",
+                             [](PyValue &self) {
+                               return PyOpOperandIterator(
+                                   mlirValueGetFirstUse(self.get()));
+                             })
       .def("__eq__",
            [](PyValue &self, PyValue &other) {
              return self.get().ptr == other.get().ptr;
            })
       .def("__eq__", [](PyValue &self, py::object other) { return false; })
       .def("__hash__",
            [](PyValue &self) {
              return static_cast<size_t>(llvm::hash_value(self.get().ptr));
py::class_<PyValue>(m, "Value", py::module_local())
           },
           kValueDunderStrDocstring)
       .def_property_readonly("type", [](PyValue &self) {
         return PyType(self.getParentOperation()->getContext(),
                       mlirValueGetType(self.get()));
       });
   PyBlockArgument::bind(m);
   PyOpResult::bind(m);
+  PyOpOperand::bind(m);
 
   //----------------------------------------------------------------------------
   // Mapping of SymbolTable.
   //----------------------------------------------------------------------------
   py::class_<PySymbolTable>(m, "SymbolTable", py::module_local())
       .def(py::init<PyOperationBase &>())
       .def("__getitem__", &PySymbolTable::dunderGetItem)
       .def("insert", &PySymbolTable::insert, py::arg("operation"))

 
   // Container bindings.
   PyBlockArgumentList::bind(m);
   PyBlockIterator::bind(m);
   PyBlockList::bind(m);
   PyOperationIterator::bind(m);
   PyOperationList::bind(m);
   PyOpAttributeMap::bind(m);
+  PyOpOperandIterator::bind(m);
   PyOpOperandList::bind(m);
   PyOpResultList::bind(m);
   PyRegionIterator::bind(m);
   PyRegionList::bind(m);
 
   // Debug bindings.
   PyGlobalDebugFlag::bind(m);
 }

mlir/test/python/ir/value.py

func.func @foo(%arg0: f32) -> f32 {
       return %0 : f32
     }""", ctx)
 
   [func] = module.body.operations
   block = func.entry_block
   op, ret = block.operations
   assert hash(block.arguments[0]) == hash(op.operands[0])
   assert hash(op.result) == hash(ret.operands[0])
+
+# CHECK-LABEL: TEST: testValueUses
+# CHECK: Use owner: "custom.op2"
+# CHECK: Use operand_number: 0
+# CHECK: Use owner: "custom.op2"
+# CHECK: Use operand_number: 0
+@run
+def testValueUses():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with Location.unknown(ctx):
+    i32 = IntegerType.get_signless(32)
+    module = Module.create()
+    with InsertionPoint(module.body):
+      value = Operation.create("custom.op1", results=[i32]).results[0]
+      op1 = Operation.create("custom.op2", operands=[value])
+      op2 = Operation.create("custom.op2", operands=[value])
+
+  for use in value.uses:
+    assert use.owner in [op1, op2]
+    print(f"Use owner: {use.owner}")
+    print(f"Use operand_number: {use.operand_number}")