[Mlir-commits] [mlir] [mlir][Transforms] Dialect conversion: Add support for `replaceUsesWithIf` (PR #169606)

[Diego Caballero]

dcaballe wrote:

Thanks, Matthias!

> Additional changes are needed to support this API in rollback mode. In particular, no entries should be added to the ConversionValueMapping for conditional replacements. It's unclear at this point if this API can be supported in rollback mode, so this is deferred to later.

My understanding is that the rollback implementation should perform a conditional in-place replacement for the existing users at the time of the invocation + a conditional differed replacement of future users. Is the problem the evaluation of the condition for the differed users? A couple of points:

* `replaceAllUsesExcept` is *less generic* than `replaceUsesWithIf` so the former might be easier to implement than the latter. For example, for `replaceAllUsesExcept`, the filtering condition is limited to ops (users) that must exist at the time of the invocation (I'm not sure how we could refer in the condition to future ops that haven't been created yet!). This means that the condition only applies to the "in-place" part of the replacement and will always be true for the differed part, right? This turns the replacement into "a conditional in-place replacement for the existing users at the time of the invocation + an **unconditional** differed replacement of future users". This should be easier to implement and would cover the hack in `replaceAllUsesWith`.

* For the generic `replaceUsesWithIf`, could we store the conditions (lambda functions) in the driver and use them to decide if an entry should be added to `ConversionValueMapping`?

> If you are seeing failed assertions due to this change, you are using unsupported API in your dialect conversion. You have 3 options: (1) Migrate to the no-rollback driver. (2) Rewrite your patterns without the unsupported API. (3) Last resort: bypass the rewriter and call replaceUsesWithIf etc. directly on the Value object.

Would it be possible to preserve the existing functionality for the rollback driver under a different name (e.g., `replaceUsesWithIfXYZ`) until we have the proper implementation, then remove it? The problem I see is that this PR could be a dead end for users of the current `replaceUsesWithIf` in rollback mode if the migration to the no-rollback driver is not possible (i.e., in my case, the no-rollback mode would invalidate analysis that it's needed to do the conversion).


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