[mlir] Expose value replacement methods through CAPI & Python bindings

[jdd]

As we discussed offline, I don't think this is the appropriate level to expose through the C API. Some sort of getUses function would be more appropriate IMO.

[mortbopet]

Looking into this, it seems more difficult than initially thought to expose low-level operations to perform value replacement through the CAPI. We'd expect getUses to return OpOperands (by inheritance, IrOperand) which in turn is what we would call replace on.
Main pain points which i see here are:

1. OpOperand is heavily encapsulated to protect how and who can construct an operand. https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/Value.h#L265
2. By the above, there's no immediate way that we can roundtrip the value through CAPI, which would require that an OpOperand can be constructed statically through an opaque pointer.

There may be a way around this if the constructors are made public. However, that design choice was made for a reason so that would not be a sound decision. Suggestions are welcome, but I don't immediately see any non-intrusive options apart from the one presented in this PR.

[jdd]

I'm pretty sure that OpOperand just needs to implement getAsOpaquePointer/getFromOpaquePointer and it is safe for those implementations to just cast from/to this.

[jdd]

Or just represent MlirOpOperand as (Operation*, operand index).

mlir/include/mlir-c/IR.h

 //===-- mlir-c/IR.h - C API to Core MLIR IR classes ---------------*- C -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//

 MLIR_CAPI_EXPORTED void mlirValueDump(MlirValue value);
 
 /// Prints a value by sending chunks of the string representation and
 /// forwarding `userData to `callback`. Note that the callback may be called
 /// several times with consecutive chunks of the string.
 MLIR_CAPI_EXPORTED void
 mlirValuePrint(MlirValue value, MlirStringCallback callback, void *userData);
 
+/// Replace all uses of 'this' value with the new value, updating anything in
+/// the IR that uses 'this' to use the other value instead.  When this returns
+/// there are zero uses of 'this'.
+MLIR_CAPI_EXPORTED void mlirValueReplaceAllUsesWith(MlirValue value,
+                                                    MlirValue newValue);
+
+/// Replace all uses of 'this' value with 'newValue', updating anything in the
+/// IR that uses 'this' to use the other value instead except if the user is
+/// listed in 'exceptions' .
+MLIR_CAPI_EXPORTED void mlirValueReplaceAllUsesExcept(MlirValue value,
+                                                      MlirValue newValue,
+                                                      MlirOperation except);
+
 //===----------------------------------------------------------------------===//
 // Type API.
 //===----------------------------------------------------------------------===//
 
 /// Parses a type. The type is owned by the context.
 MLIR_CAPI_EXPORTED MlirType mlirTypeParseGet(MlirContext context,
                                              MlirStringRef type);
 

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_<PyValue>(m, "Value", py::module_local())
             PyPrintAccumulator printAccum;
             printAccum.parts.append("Value(");
             mlirValuePrint(self.get(), printAccum.getCallback(),
                            printAccum.getUserData());
             printAccum.parts.append(")");
             return printAccum.join();
           },
           kValueDunderStrDocstring)
-      .def_property_readonly("type", [](PyValue &self) {
-        return PyType(self.getParentOperation()->getContext(),
+      .def_property_readonly("type",
+                             [](PyValue &self) {
+                               return PyType(
+                                   self.getParentOperation()->getContext(),
                                    mlirValueGetType(self.get()));
+                             })
+      .def("replace_all_uses_with",
+           [](PyValue &self, PyValue &newValue) {
+             mlirValueReplaceAllUsesWith(self.get(), newValue.get());
+           })
+      .def("replace_all_uses_except",
+           [](PyValue &self, PyValue &newValue, PyOperation &operation) {
+             mlirValueReplaceAllUsesExcept(self.get(), newValue.get(),
+                                           operation.get());
            });
+
   PyBlockArgument::bind(m);
   PyOpResult::bind(m);
 
   //----------------------------------------------------------------------------
   // Mapping of SymbolTable.
   //----------------------------------------------------------------------------
   py::class_<PySymbolTable>(m, "SymbolTable", py::module_local())
       .def(py::init<PyOperationBase &>())

mlir/lib/CAPI/IR/IR.cpp

 //===- IR.cpp - C Interface for Core MLIR APIs ----------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 

 void mlirValueDump(MlirValue value) { unwrap(value).dump(); }
 
 void mlirValuePrint(MlirValue value, MlirStringCallback callback,
                     void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(value).print(stream);
 }
 
+void mlirValueReplaceAllUsesWith(MlirValue value, MlirValue newValue) {
+  unwrap(value).replaceAllUsesWith(unwrap(newValue));
+}
+
+void mlirValueReplaceAllUsesExcept(MlirValue value, MlirValue newValue,
+                                   MlirOperation except) {
+  unwrap(value).replaceAllUsesExcept(unwrap(newValue), unwrap(except));
+}
+
 //===----------------------------------------------------------------------===//
 // Type API.
 //===----------------------------------------------------------------------===//
 
 MlirType mlirTypeParseGet(MlirContext context, MlirStringRef type) {
   return wrap(mlir::parseType(unwrap(type), unwrap(context)));
 }