Add support for `replace_uses_with_if`

[erick-xanadu]

Please assign this issue to me, I'll try to work on this later this week.

See here

[superlopuh]

Added in #3190 :)

[erick-xanadu]

@superlopuh should this not be accessible from the rewriter class?

[superlopuh]

I'm not sure, on one hand it would make sense to have all mutations go through a rewriter, on the other it would be nice to have a concrete use-case for where that would be useful, like for adding the modified operations to a worklist. Is that the case for you?

[erick-xanadu]

@superlopuh , I'm writing a small tutorial to teach xDSL internally to a couple of colleagues. The tutorial is a simple implementation of applying finite-difference to a python function:

def two_times(x):
  return 2 * x

The goal of the tutorial is to introduce them to the clone API as I am cloning the call operation (and just getting people more used to the API) I.e.,

# pseudo-code
arg = matchedCallOp.results[0]
arg_plus_h = addOp.results[0]
valueMapper = {arg: arg_plus_h}
cloneCallOp = matchedCallOp.clone(valueMapper)

At some point I need to also do this:

#pseudo-code
# In the IR
new_value_to_return = (matchedCallOp.results[0] + cloneCallOp.results[0]) / 2 

I would like now to replace all uses of matchedCallOp.results[0] with new_value_to_return except for the use in computing new_value_to_return.

I don't have any strong arguments for implementing this in the rewriter except that it more closely resembles what is implemented in MLIR and it saves one extra call.

Without this interface in the rewriter I have:

        # Notice the use of replace_by_if
        callOp.results[0].replace_by_if(division.results[0], lambda use: use.operation != addition)
        rewriter.notify_op_modified(callOp)

Which could be modified tom something like:

  rewriter.replace_uses_with_if(callOp.results[0], division.results[0], lambda use: use.operation != addition)

I'm happy leaving it as is, but I just wanted to double check about whether it should be implemented through the rewriter or not. :)

EDIT: Thinking about this, I believe that there is no reason to write rewriter.notify_op_modified(callOp) but technically all other operations which had their use replace should be added to the operation list.

[superlopuh]

This is reasonable, does MLIR do this also? CC @math-fehr

[math-fehr]

AFAIK, MLIR has a function for it!
Also, if we can pass it through the rewriter that would be best, I'm trying to fit everything in the rewriter so all modifications go through it

[superlopuh]

Sounds like we need two functions then, replace_by, and replace_by_if. I would prefer for the API to be consistent across SSAValue and Rewriter, so if we want to have replace_uses_with_if then we would first need to rename the SSAValue API IMO.