[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