Skip to content
This repository was archived by the owner on Oct 11, 2025. It is now read-only.

Commit e2a714c

Browse files
authored
[mlir,python] Expose replaceAllUsesExcept to Python bindings (#115850)
Problem originally described in [the forums here](https://discourse.llvm.org/t/mlir-python-expose-replaceallusesexcept/83068/1). Using the MLIR Python bindings, the method [`replaceAllUsesWith`](https://mlir.llvm.org/doxygen/classmlir_1_1Value.html#ac56b0fdb6246bcf7fa1805ba0eb71aa2) for `Value` is exposed, e.g., ```python orig_value.replace_all_uses_with( new_value ) ``` However, in my use-case I am separating a block into multiple blocks, so thus want to exclude certain Operations from having their Values replaced (since I want them to diverge). Within Value, we have [`replaceAllUsesExcept`](https://mlir.llvm.org/doxygen/classmlir_1_1Value.html#a9ec8d5c61f8a6aada4062f609372cce4), where we can pass the Operations which should be skipped. This is not currently exposed in the Python bindings: this PR fixes this. Adds `replace_all_uses_except`, which works with individual Operations, and lists of Operations.
1 parent a61c465 commit e2a714c

File tree

1 file changed

+29
-0
lines changed

1 file changed

+29
-0
lines changed

IRCore.cpp

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,12 @@ static const char kValueReplaceAllUsesWithDocstring[] =
178178
the IR that uses 'self' to use the other value instead.
179179
)";
180180

181+
static const char kValueReplaceAllUsesExceptDocstring[] =
182+
R"("Replace all uses of this value with the 'with' value, except for those
183+
in 'exceptions'. 'exceptions' can be either a single operation or a list of
184+
operations.
185+
)";
186+
181187
//------------------------------------------------------------------------------
182188
// Utilities.
183189
//------------------------------------------------------------------------------
@@ -3718,6 +3724,29 @@ void mlir::python::populateIRCore(py::module &m) {
37183724
mlirValueReplaceAllUsesOfWith(self.get(), with.get());
37193725
},
37203726
kValueReplaceAllUsesWithDocstring)
3727+
.def(
3728+
"replace_all_uses_except",
3729+
[](MlirValue self, MlirValue with, PyOperation &exception) {
3730+
MlirOperation exceptedUser = exception.get();
3731+
mlirValueReplaceAllUsesExcept(self, with, 1, &exceptedUser);
3732+
},
3733+
py::arg("with"), py::arg("exceptions"),
3734+
kValueReplaceAllUsesExceptDocstring)
3735+
.def(
3736+
"replace_all_uses_except",
3737+
[](MlirValue self, MlirValue with, py::list exceptions) {
3738+
// Convert Python list to a SmallVector of MlirOperations
3739+
llvm::SmallVector<MlirOperation> exceptionOps;
3740+
for (py::handle exception : exceptions) {
3741+
exceptionOps.push_back(exception.cast<PyOperation &>().get());
3742+
}
3743+
3744+
mlirValueReplaceAllUsesExcept(
3745+
self, with, static_cast<intptr_t>(exceptionOps.size()),
3746+
exceptionOps.data());
3747+
},
3748+
py::arg("with"), py::arg("exceptions"),
3749+
kValueReplaceAllUsesExceptDocstring)
37213750
.def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR,
37223751
[](PyValue &self) { return self.maybeDownCast(); });
37233752
PyBlockArgument::bind(m);

0 commit comments

Comments
 (0)