Copying dependency chains

MLIR
1

Hello All,

I have some code that traces and collects the def-use chain of an operation up to the block arguments of its scope. In other words, given an operation P with Operands o1…oN; I add the operations that define o1…oN to my list and then recurse on the operands of o1…oN, etc. Stopping after we reach a block arg or constant value.

I want to insert a clone of this chain into a modified copy of the block (a different module).

How can I go about doing this?

I currently have something that inserts the cloned operation chain - except that when I clone each operation I think I lose the dependencies between them (results in <<UNKNOWN SSA VALUE>>, as the new op has no reference to it’s cloned operands)

I notice I can maybe remap the operands of an operation using an IRMapping as an argument to the clone method, or use something like cloneInto…

I would appreciate any advice, thanks in advance!

2

Assuming that the region has SSA dominance, I would clone the ops on the chain top-to-bottom, one-by-one and after cloning each op, map the old results to the new results in the IRMapping object (same for block if necessary). I am not aware of a helper function to automate this.

1 Like
3

I will try this, thank you! :slight_smile:

4

I had a go with the top down approach, however I am still getting unknown ssa values when I call clone with a mapping:

I have an IRMapping mapping:

%0 = "hw.constant"() {value = -2 : i2} : () -> i2
| to
%c-2_i2 = hw.constant -2 : i2
,
%1 = "hw.constant"() {value = 0 : i2} : () -> i2
| to
%c0_i2 = hw.constant 0 : i2

And an operation op to clone:

%5 = "comb.mux"(%arg1, %1, %0) : (i1, i2, i2) -> i2

When I call op->clone(mapping), the operands are correctly remapped when I print them, however the result of the cloned operation is still:

%0 = comb.mux <<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>> : i2

Am I forgetting something when I call clone ? or forgetting to add something to the map ?

Here are the printed operands:

<block argument> of type 'i1' at index: 1 
%c0_i2 = hw.constant 0 : i2
%c-2_i2 = hw.constant -2 : i2
5

How are you inserting the cloned Op? It needs to be inserted in a place where the remapped operands are actually defined. I usually use an IRRewriter (if it’s not within a Pattern), do rewriter.setInsertionPointToStart(...) to the block I want the cloned Ops in and then go through the loop described my Matthias, i.e. rewriter.clone() and mapping.map()

6

YES! That seems to work thank you :D, although I am not sure what to do about recursive operands…
(e.g. %word = seq.compreg %word, %clock : i2)

I get

%5 = "seq.compreg"(<<UNKNOWN SSA VALUE>>, %arg0) {name = "word"} : (i2, i1) -> i2

Since I can’t refer to the clone at the time of cloning, do I need to clone and then remap after?

7

Ah, yes in that case it looks like you will need a remapping pass after the cloning. I would just iterate over the cloned ops, iterate over their operands and check if they have been remapped in the IRMapping, if so, you can change the operand in-place.

8

Ahh yes that is a good solution! I am testing a method where I push the indices of recursive args of an operation to a map : DenseMap<Operation *, SmallVector<unsigned>> , while climbing the tree and then lookup the args to replace with itself after the call to clone… Possibly a bit over complicated.

Also, do you mean haven’t been remapped ? Not sure if I have fully understood

9

I think the naive solution would look something like this:

    SmallVector<Operation *> clonedOps;
    for (Operation *opToClone : opsToClone) {
      Operation *cloned = rewriter.clone(*opToClone);
      clonedOps.push_back(cloned);

      for (auto [oldResult, newResult] : llvm::zip(opToClone->getResults(), cloned->getResults()))
        mapping.map(oldResult, newResult);
    }

    for (Operation *cloned : clonedOps) {
      for (OpOperand *opOper : cloned->getOpOperands())
        if (Value m = mapping.lookupOrNull(opOper->get()))
          rewriter.updateRootInPlace(cloned, [&]() {opOper->set(m);});

But I’m sure there is a lot to optimize about that.

1 Like
10

That’s super helpful thank you! I’m going to revise my implementation :nerd_face: