MLIR Python: expose `replaceAllUsesExcept`

MLIR
,
1

Using the MLIR Python bindings, the method replaceAllUsesWith for Value is exposed, e.g.,

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, where we can pass the Operations which should be skipped.

However, as far as I can tell this is not exposed in the Python bindings. Under IRCore.cpp, you can see where replace_all_uses_with is exposed here.

.def(
    "replace_all_uses_with",
    [](PyValue &self, PyValue &with) {
      mlirValueReplaceAllUsesOfWith(self.get(), with.get());
    },
    kValueReplaceAllUsesWithDocstring)

However, there is no such equivalent for replaceAllUsesExcept I can see. Iā€™ll update this thread with my attempt, but would adding a similar .def in IRCore.cpp be sufficient, or is there other places I need to add the binding?

.def(
    "replace_all_uses_except",
    [](PyValue &self, PyValue &new_value, py::list exceptions) {
        llvm::SmallPtrSet<Operation *, 4> exceptionSet;
        for (auto &exc : exceptions) {
            exceptionSet.insert(py::cast<PyOperationRef &>(exc).get());
        }
        mlirValueReplaceAllUsesExcept(self.get(), new_value.get(), exceptionSet);
    },
    "Replace all uses of this value with 'new_value' except for the operations in 'exceptions'."
)
1 Like
2

Well mlirValueReplaceAllUsesExcept doesnā€™t exist so youā€™d have to start thereā€¦

But yea thatā€™s generally fine/acceptable - if you send a PR I can review.

3

A pull request for this feature is available here:

2 Likes