[Mlir-commits] [mlir] [mlir, python] Expose replaceAllUsesExcept to Python bindings (PR #115850)

Perry Gibson llvmlistbot at llvm.org
Tue Nov 12 06:07:50 PST 2024


================
@@ -3718,6 +3724,32 @@ void mlir::python::populateIRCore(py::module &m) {
             mlirValueReplaceAllUsesOfWith(self.get(), with.get());
           },
           kValueReplaceAllUsesWithDocstring)
+      .def(
+          "replace_all_uses_except",
+          [](PyValue &self, PyValue &with, py::object exceptions) {
+            MlirValue selfValue = self.get();
+            MlirValue withValue = with.get();
+
+            // Check if 'exceptions' is a list
+            if (py::isinstance<py::list>(exceptions)) {
+              // Convert Python list to a vector of MlirOperations
+              std::vector<MlirOperation> exceptionOps;
+              for (py::handle exception : exceptions) {
+                exceptionOps.push_back(exception.cast<PyOperation &>().get());
+              }
+              mlirValueReplaceAllUsesExceptWithSet(
+                  selfValue, withValue, exceptionOps.data(),
+                  static_cast<intptr_t>(exceptionOps.size()));
+            } else {
+              // Assume 'exceptions' is a single Operation
+              MlirOperation exceptedUser =
+                  exceptions.cast<PyOperation &>().get();
+              mlirValueReplaceAllUsesExceptWithSingle(selfValue, withValue,
+                                                      exceptedUser);
+            }
----------------
Wheest wrote:

Good suggestion, that would be clearer and avoids the branch.

I've updated the file with that change.  I've also replaced the `std::vector` with a `llvm::SmallVector` with default size 4.

https://github.com/llvm/llvm-project/pull/115850


More information about the Mlir-commits mailing list