[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

py::class_<PyOperationList>(m, "OperationList", py::module_local())
         .def("__len__", &PyOperationList::dunderLen);
   }
 
 private:
   PyOperationRef parentOperation;
   MlirBlock block;
 };
 
+class PyOpOperand {
+public:
+  PyOpOperand(MlirOpOperand opOperand) : opOperand(opOperand) {}
+
+  py::object getOwner() {
+    MlirOperation owner = mlirOpOperandGetOwner(opOperand);
+    PyMlirContextRef context =
+        PyMlirContext::forContext(mlirOperationGetContext(owner));
+    return PyOperation::forOperation(context, owner)->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:
+  MlirOpOperand opOperand;
+};
+
+class PyOpOperandIterator {
+public:
+  PyOpOperandIterator(MlirOpOperand opOperand) : opOperand(opOperand) {}
+
+  PyOpOperandIterator &dunderIter() { return *this; }
+
+  PyOpOperand dunderNext() {
+    if (mlirOpOperandIsNull(opOperand))
+      throw py::stop_iteration();
+
+    PyOpOperand returnOpOperand(opOperand);
+    opOperand = mlirOpOperandGetNextUse(opOperand);
+    return returnOpOperand;
+  }
+
+  static void bind(py::module &m) {
+    py::class_<PyOpOperandIterator>(m, "OpOperandIterator", py::module_local())
+        .def("__iter__", &PyOpOperandIterator::dunderIter)
+        .def("__next__", &PyOpOperandIterator::dunderNext);
+  }
+
+private:
+  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
+@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])
+
+  # CHECK: Use owner: "custom.op2"
+  # CHECK: Use operand_number: 0
+  # CHECK: Use owner: "custom.op2"
+  # CHECK: Use operand_number: 0
+  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}")