[python] Expose replaceUsesOfWith C API (#171892)

This PR exposes the `replaceUsesOfWith` C API to Python

Author: chhzh123

mlir/include/mlir-c/IR.h

@@ -858,6 +858,10 @@ MLIR_CAPI_EXPORTED
 void mlirOperationWalk(MlirOperation op, MlirOperationWalkCallback callback,
                        void *userData, MlirWalkOrder walkOrder);
 
+/// Replace uses of 'of' value with the 'with' value inside the 'op' operation.
+MLIR_CAPI_EXPORTED void
+mlirOperationReplaceUsesOfWith(MlirOperation op, MlirValue of, MlirValue with);
+
 //===----------------------------------------------------------------------===//
 // Region API.
 //===----------------------------------------------------------------------===//

mlir/lib/Bindings/Python/IRCore.cpp

@@ -3939,6 +3939,14 @@ void mlir::python::populateIRCore(nb::module_ &m) {
             return PyOpSuccessors(self.getOperation().getRef());
           },
           "Returns the list of Operation successors.")
+      .def(
+          "replace_uses_of_with",
+          [](PyOperation &self, PyValue &of, PyValue &with) {
+            mlirOperationReplaceUsesOfWith(self.get(), of.get(), with.get());
+          },
+          "of"_a, "with_"_a,
+          "Replaces uses of the 'of' value with the 'with' value inside the "
+          "operation.")
       .def("_set_invalid", &PyOperation::setInvalid,
            "Invalidate the operation.");
 

mlir/lib/CAPI/IR/IR.cpp

@@ -902,6 +902,11 @@ void mlirOperationWalk(MlirOperation op, MlirOperationWalkCallback callback,
   }
 }
 
+void mlirOperationReplaceUsesOfWith(MlirOperation op, MlirValue oldValue,
+                                    MlirValue newValue) {
+  unwrap(op)->replaceUsesOfWith(unwrap(oldValue), unwrap(newValue));
+}
+
 //===----------------------------------------------------------------------===//
 // Region API.
 //===----------------------------------------------------------------------===//

mlir/test/python/ir/operation.py

@@ -1208,6 +1208,40 @@ def callback(op):
         print("Exception raised")
 
 
+# CHECK-LABEL: TEST: testOpReplaceUsesWith
+@run
+def testOpReplaceUsesWith():
+    ctx = Context()
+    ctx.allow_unregistered_dialects = True
+    with Location.unknown(ctx):
+        m = Module.create()
+        i32 = IntegerType.get_signless(32)
+        with InsertionPoint(m.body):
+            value = Operation.create("custom.op1", results=[i32]).results[0]
+            value2 = Operation.create("custom.op2", results=[i32]).results[0]
+            op = Operation.create("custom.op3", operands=[value])
+            op2 = Operation.create("custom.op4", operands=[value])
+            op.replace_uses_of_with(value, value2)
+
+    assert len(list(value.uses)) == 1
+
+    # CHECK: Use owner: "custom.op4"
+    # CHECK: Use operand_number: 0
+    for use in value.uses:
+        assert use.owner in [op2]
+        print(f"Use owner: {use.owner}")
+        print(f"Use operand_number: {use.operand_number}")
+
+    assert len(list(value2.uses)) == 1
+
+    # CHECK: Use owner: "custom.op3"
+    # CHECK: Use operand_number: 0
+    for use in value2.uses:
+        assert use.owner in [op]
+        print(f"Use owner: {use.owner}")
+        print(f"Use operand_number: {use.operand_number}")
+
+
 # CHECK-LABEL: TEST: testGetOwnerConcreteOpview
 @run
 def testGetOwnerConcreteOpview():