[MLIR] Fix operation clone

[wsmoses]

Summary

Operation clone is currently faulty.

Suppose you have a block like as follows:

(%x0 : i32) {
   %x1 = f(%x0)
   return %x1
}

The test case we have is that we want to "unroll" this, in which we want to change this to compute f(f(x0)) instead of just f(x0). We do so by making a copy of the body at the end of the block and set the uses of the argument in the copy operations with the value returned from the original block.
This is implemented as follows:

1. map to the block arguments to the returned value (map[x0] = x1).
2. clone the body

Now for this small example, this works as intended and we get the following.

(%x0 : i32) {
   %x1 = f(%x0)
   %x2 = f(%x1)
   return %x2
}

This is because the current logic to clone x1 = f(x0) first looks up the arguments in the map (which finds x0 maps to x1 from the initialization), and then sets the map of the result to the cloned result (map[x1] = x2).

However, this fails if x0 is not an argument to the op, but instead used inside the region, like below.

(%x0 : i32) {
   %x1 = f() {
      yield %x0
   }
   return %x1
}

This is because cloning an op currently first looks up the args (none), sets the map of the result (map[%x1] = %x2), and then clones the regions. This results in the following, which is clearly illegal:

(%x0 : i32) {
   %x1 = f() {
      yield %x0
   }
   %x2 = f() {
      yield %x2
   }
   return %x2
}

Diving deeper, this is partially due to the ordering (how this PR fixes it), as well as how region cloning works. Namely it will first clone with the mapping, and then it will remap all operands. Since the ordering above now has a map of x0 -> x1 and x1 -> x2, we end up with the incorrect behavior here.

[ftynse]

This needs a standalone test that must not depend on any affine stuff. It can be implemented as a simple pass defined in test/lib/IR, there's already a bunch of such passes.

However, I am not convinced the standalone test can be constructed. Looking at the semantics of the change, it moves the addition of the operation results to mapper to be performed after cloning the regions of this operation rather than before. In valid IR, SSACFG regions of an operation cannot use its results, so the presence of these results in the mapper should not have any effect on the contents of the cloned regions. Maybe the IR that Operation::clone is being called on is already invalid. Maybe a nastier thing happens: Value is essentially a pointer, so if the operation defining it is erased without also updating the mapper and another operation is created, the pointer may get reused to point to an irrelevant value. In any case, I suspect it's the unroll pass that needs a fix.

mlir/lib/IR/Operation.cpp

528

Nit: update this comment too.

[wsmoses]

In D122531#3410912, @ftynse wrote:

This needs a standalone test that must not depend on any affine stuff. It can be implemented as a simple pass defined in test/lib/IR, there's already a bunch of such passes.

However, I am not convinced the standalone test can be constructed. Looking at the semantics of the change, it moves the addition of the operation results to mapper to be performed after cloning the regions of this operation rather than before. In valid IR, SSACFG regions of an operation cannot use its results, so the presence of these results in the mapper should not have any effect on the contents of the cloned regions. Maybe the IR that Operation::clone is being called on is already invalid. Maybe a nastier thing happens: Value is essentially a pointer, so if the operation defining it is erased without also updating the mapper and another operation is created, the pointer may get reused to point to an irrelevant value. In any case, I suspect it's the unroll pass that needs a fix.

Suppose we have an affine.for which yields some values, and contains some instructions

In particular, what is occurring within unroll is as follows:

1. The map is initialized to have the iter args map to the yielded arguments.
2. Each instruction in the for loop is cloned, using the map.
3. Step 2 is repeated until there are the desired number of copies corresponding to the coarseness.

The intention here is that the cloning of the loop body instructions creates the unrolled copies of the loop, whereas the existing instruction are the 0-th copy.

Thus when the body instruction is being cloned, it isn't in a scenario where it can use its results, but rather it is (correctly) using an operand which is previously defined, and that operand being replaced by a value which the original op yields. So long as the cloned op comes after the original operation, this is semantically valid.

[mehdi_amini]

I didn't quite follow the description of the unroll logic, but can you add a test? It may be easier to figure it out then.

[wsmoses]

Add simple clone test

[wsmoses]

Fix format

[wsmoses]

Fix linking

[mehdi_amini]

I didn't quite follow the description of the unroll logic, but can you add a test? It may be easier to figure it out then.

mlir/test/lib/IR/TestClone.cpp

61

Isn't this simpler to implement this whole thing with a single as a clone (with region) having pre-populated the mapper appropriately?

[wsmoses]

In D122531#3453157, @mehdi_amini wrote:

I didn't quite follow the description of the unroll logic, but can you add a test? It may be easier to figure it out then.

Did you happen to see the mlir/test/IR/test-clone.mlir file I added, or was that not illustrative?

Basically the gist is as follows:

Suppose you have a block like as follows:

(%x0 : i32) {
   %x1 = f(%x0)
   return %x1
}

The test case we have is that we want to "unroll" this, in which we want to change this to compute f(f(x0)) instead of just f(x0). We do so by making a copy of the body at the end of the block and set the uses of the argument in the copy operations with the value returned from the original block.
This is implemented as follows:

1. map to the block arguments to the returned value (map[x0] = x1).
2. clone the body

Now for this small example, this works as intended and we get the following.

(%x0 : i32) {
   %x1 = f(%x0)
   %x2 = f(%x1)
   return %x2
}

This is because the current logic to clone x1 = f(x0) first looks up the arguments in the map (which finds x0 maps to x1 from the initialization), and then sets the map of the result to the cloned result (map[x1] = x2).

However, this fails if x0 is not an argument to the op, but instead used inside the region, like below.

(%x0 : i32) {
   %x1 = f() {
      yield %x0
   }
   return %x1
}

This is because cloning an op currently first looks up the args (none), sets the map of the result (map[%x1] = %x2), and then clones the regions. This results in the following, which is clearly illegal:

(%x0 : i32) {
   %x1 = f() {
      yield %x0
   }
   %x2 = f() {
      yield %x2
   }
   return %x2
}

Diving deeper, this is partially due to the ordering (how this PR fixes it), as well as how region cloning works (reproduced below). Namely it will first clone with the mapping, and then it will remap all operands. Since the ordering above now has a map of x0 -> x1 and x1 -> x2, we end up with the incorrect behavior here.

void Region::cloneInto(Region *dest, Region::iterator destPos,
                       BlockAndValueMapping &mapper) {
  assert(dest && "expected valid region to clone into");
  assert(this != dest && "cannot clone region into itself");

  // If the list is empty there is nothing to clone.
  if (empty())
    return;

  for (Block &block : *this) {
    Block *newBlock = new Block();
    mapper.map(&block, newBlock);

    // Clone the block arguments. The user might be deleting arguments to the
    // block by specifying them in the mapper. If so, we don't add the
    // argument to the cloned block.
    for (auto arg : block.getArguments())
      if (!mapper.contains(arg))
        mapper.map(arg, newBlock->addArgument(arg.getType(), arg.getLoc()));

    // Clone and remap the operations within this block.
    for (auto &op : block)
      newBlock->push_back(op.clone(mapper));

    dest->getBlocks().insert(destPos, newBlock);
  }

  // Now that each of the blocks have been cloned, go through and remap the
  // operands of each of the operations.
  auto remapOperands = [&](Operation *op) {
    for (auto &operand : op->getOpOperands())
      if (auto mappedOp = mapper.lookupOrNull(operand.get()))
        operand.set(mappedOp);
    for (auto &succOp : op->getBlockOperands())
      if (auto *mappedOp = mapper.lookupOrNull(succOp.get()))
        succOp.set(mappedOp);
  };

  for (iterator it(mapper.lookup(&front())); it != destPos; ++it)
    it->walk(remapOperands);
}

[ftynse]

You have a point, this doesn't look like a misuse of the API and one should expect it to produce valid IR. Thanks for iterating! Please update the commit description to contain the more abstract version that does not refer to affine or unrolling based on your test.

mlir/test/lib/IR/TestClone.cpp

54–59

Nit: drop trivial braces.

[mehdi_amini]

In D122531#3453220, @wsmoses wrote:

In D122531#3453157, @mehdi_amini wrote:

I didn't quite follow the description of the unroll logic, but can you add a test? It may be easier to figure it out then.

Did you happen to see the mlir/test/IR/test-clone.mlir file I added, or was that not illustrative?

Sorry I didn't intend to send this now, I think I typed this last week and didn't send it apparently, I just meant to send my inline comment today...

Diving deeper, this is partially due to the ordering (how this PR fixes it), as well as how region cloning works (reproduced below). Namely it will first clone with the mapping, and then it will remap all operands. Since the ordering above now has a map of x0 -> x1 and x1 -> x2, we end up with the incorrect behavior here.

Uh, this transitive mapping seems wrong to me, isn't it? If it is: should we fix this instead?

[mehdi_amini]

mlir/lib/IR/Operation.cpp

528

You missed this comment by the way

[wsmoses]

Uh, this transitive mapping seems wrong to me, isn't it? If it is: should we fix this instead?

I agree the transitive mapping also seems wrong and is seemingly due to how region cloning works. That said the previous cloning mechanism also had the region cloning occur in a different order from the arguments, so I think this is a good change regardless.

I'm less familiar with the region cloning infrastructure and how it aims to handle block arguments/etc and it seems that the transitivity is needed to handle these. Regardless, making this not happen for anything except block arguments in the region seems like a good thing to do.

mlir/lib/IR/Operation.cpp

528

Whoops, thanks for the heads up.

[NathanielMcVicar]

I posted this on the commit, sorry! Meant to do it here.
This appears to have broken a number of tests on the Windows buildbot: https://lab.llvm.org/buildbot/#/builders/13/builds/19731

LLVM ERROR: TypeID::get<`anonymous-namespace'::ClonePass>(): Using TypeID on a class with an anonymous namespace requires an explicit TypeID definition. The implicit fallback uses string name, which does not guarantee uniqueness in anonymous contexts. Define an explicit TypeID instantiation for this type using `MLIR_DECLARE_EXPLICIT_TYPE_ID`/`MLIR_DEFINE_EXPLICIT_TYPE_ID` or `DEFINE_EXPLICIT_PRIVATE_INLINE_TYPE_ID`.

[wsmoses]

In D122531#3454481, @NathanielMcVicar wrote:

I posted this on the commit, sorry! Meant to do it here.
This appears to have broken a number of tests on the Windows buildbot: https://lab.llvm.org/buildbot/#/builders/13/builds/19731

LLVM ERROR: TypeID::get<`anonymous-namespace'::ClonePass>(): Using TypeID on a class with an anonymous namespace requires an explicit TypeID definition. The implicit fallback uses string name, which does not guarantee uniqueness in anonymous contexts. Define an explicit TypeID instantiation for this type using `MLIR_DECLARE_EXPLICIT_TYPE_ID`/`MLIR_DEFINE_EXPLICIT_TYPE_ID` or `DEFINE_EXPLICIT_PRIVATE_INLINE_TYPE_ID`.

Well that's exciting, let me take a look and see if I can make a quick fix, and if not I'll revert and reapply when done. Also this is odd, since I thought CI passed prior to landing and I'm surprised this wasn't picked up.

[mehdi_amini]

You likely didn’t rebase for a while and the CI may not do it itself

[wsmoses]

In D122531#3454495, @mehdi_amini wrote:

You likely didn’t rebase for a while and the CI may not do it itself

Yup that was it (and the change occurred between then).

It appears what was missing was (testing locally on Ubuntu and trying to find a windows machine before pushing):

 struct ClonePass
     : public PassWrapper<ClonePass, InterfacePass<FunctionOpInterface>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ClonePass)

This missing macro was put in place in between (git blame on SymbolUses):

b8cd0c1486449 (Tres Popp     2019-12-05 03:56:18 -0800  17) /// provided by the symbol table along with erasing from the symbol table.
80aca1eaf778a (River Riddle  2020-04-07 13:56:16 -0700  18) struct SymbolUsesPass
80aca1eaf778a (River Riddle  2020-04-07 13:56:16 -0700  19)     : public PassWrapper<SymbolUsesPass, OperationPass<ModuleOp>> {
5e50dd048e3a2 (River Riddle  2022-03-30 17:00:37 -0700  20)   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SymbolUsesPass)
5e50dd048e3a2 (River Riddle  2022-03-30 17:00:37 -0700  21)

mlir/include/mlir/IR/Operation.h

 //===- Operation.h - MLIR Operation Class -----------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
public:
   Operation *clone(BlockAndValueMapping &mapper);
   Operation *clone();
 
   /// Create a partial copy of this operation without traversing into attached
   /// regions. The new operation will have the same number of regions as the
   /// original one, but they will be left empty.
   /// Operands are remapped using `mapper` (if present), and `mapper` is updated
   /// to contain the results.
-  Operation *cloneWithoutRegions(BlockAndValueMapping &mapper);
+  /// The `mapResults` argument specifies whether the results of the operation
+  /// should also be mapped.
+  Operation *cloneWithoutRegions(BlockAndValueMapping &mapper,
+                                 bool mapResults = true);
 
   /// Create a partial copy of this operation without traversing into attached
   /// regions. The new operation will have the same number of regions as the
   /// original one, but they will be left empty.
   Operation *cloneWithoutRegions();
 
   /// Returns the operation block that contains this operation.
   Block *getBlock() { return block; }

mlir/lib/IR/Operation.cpp

 //===- Operation.cpp - Operation support code -----------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 

 }
 
 //===----------------------------------------------------------------------===//
 // Operation Cloning
 //===----------------------------------------------------------------------===//
 
 /// Create a deep copy of this operation but keep the operation regions empty.
 /// Operands are remapped using `mapper` (if present), and `mapper` is updated
 /// to contain the results.
-Operation *Operation::cloneWithoutRegions(BlockAndValueMapping &mapper) {
+Operation *Operation::cloneWithoutRegions(BlockAndValueMapping &mapper,
+                                          bool mapResults) {
   SmallVector<Value, 8> operands;
   SmallVector<Block *, 2> successors;
 
   // Remap the operands.
   operands.reserve(getNumOperands());
   for (auto opValue : getOperands())
     operands.push_back(mapper.lookupOrDefault(opValue));
 
   // Remap the successors.
   successors.reserve(getNumSuccessors());
   for (Block *successor : getSuccessors())
     successors.push_back(mapper.lookupOrDefault(successor));
 
   // Create the new operation.
   auto *newOp = create(getLoc(), getName(), getResultTypes(), operands, attrs,
                        successors, getNumRegions());
 
   // Remember the mapping of any results.
+  if (mapResults) {
     for (unsigned i = 0, e = getNumResults(); i != e; ++i)
       mapper.map(getResult(i), newOp->getResult(i));
+  }
 
   return newOp;
 }
 
 Operation *Operation::cloneWithoutRegions() {
   BlockAndValueMapping mapper;
   return cloneWithoutRegions(mapper);
 }
 
 /// Create a deep copy of this operation, remapping any operands that use
 /// values outside of the operation using the map that is provided (leaving
 /// them alone if no entry is present).  Replaces references to cloned
 /// sub-operations to the corresponding operation that is copied, and adds
 /// those mappings to the map.
 Operation *Operation::clone(BlockAndValueMapping &mapper) {
-  auto *newOp = cloneWithoutRegions(mapper);
+  auto *newOp = cloneWithoutRegions(mapper, /*mapResults=*/false);
 
   // Clone the regions.
   for (unsigned i = 0; i != numRegions; ++i)
     getRegion(i).cloneInto(&newOp->getRegion(i), mapper);
 
+  for (unsigned i = 0, e = getNumResults(); i != e; ++i)
+    mapper.map(getResult(i), newOp->getResult(i));
+
   return newOp;
 }
 
 Operation *Operation::clone() {
   BlockAndValueMapping mapper;
   return clone(mapper);
 }
 

mlir/test/IR/test-clone.mlir

+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="func.func(test-clone)" -split-input-file 
+
+module {
+  func @fixpoint(%arg1 : i32) -> i32 {
+    %r = "test.use"(%arg1) ({
+       "test.yield"(%arg1) : (i32) -> ()
+    }) : (i32) -> i32
+    return %r : i32
+  }
+}
+
+// CHECK:   func @fixpoint(%[[arg0:.+]]: i32) -> i32 {
+// CHECK-NEXT:     %[[i0:.+]] = "test.use"(%[[arg0]]) ({
+// CHECK-NEXT:       "test.yield"(%arg0) : (i32) -> ()
+// CHECK-NEXT:     }) : (i32) -> i32
+// CHECK-NEXT:     %[[i1:.+]] = "test.use"(%[[i0]]) ({
+// CHECK-NEXT:       "test.yield"(%[[i0]]) : (i32) -> ()
+// CHECK-NEXT:     }) : (i32) -> i32
+// CHECK-NEXT:     return %[[i1]] : i32
+// CHECK-NEXT:   }

mlir/test/lib/IR/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
   TestBuiltinAttributeInterfaces.cpp
+  TestClone.cpp
   TestDiagnostics.cpp
   TestDominance.cpp
   TestFunc.cpp
   TestInterfaces.cpp
   TestMatchers.cpp
   TestOpaqueLoc.cpp
   TestOperationEquals.cpp
   TestPrintDefUse.cpp

mlir/test/lib/IR/TestClone.cpp

+//===- TestSymbolUses.cpp - Pass to test symbol uselists ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestDialect.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+
+/// This is a test pass which clones the body of a function. Specifically
+/// this pass replaces f(x) to instead return f(f(x)) in which the cloned body
+/// takes the result of the first operation return as an input.
+struct ClonePass
+    : public PassWrapper<ClonePass, InterfacePass<FunctionOpInterface>> {
+  StringRef getArgument() const final { return "test-clone"; }
+  StringRef getDescription() const final {
+    return "Test replacement of clone";
+  }
+  void runOnOperation() override {
+    FunctionOpInterface op = getOperation();
+
+    // Limit testing to ops with only one region.
+    if (op->getNumRegions() != 1)
+      return;
+
+    Region &region = op->getRegion(0);
+    if (!region.hasOneBlock())
+      return;
+
+    Block &regionEntry = region.front();
+    auto terminator = regionEntry.getTerminator();
+
+    // Only handle functions whose returns match the inputs.
+  if (terminator->getNumOperands() != regionEntry.getNumArguments())
+      return;
+
+    BlockAndValueMapping map;
+    for (auto tup : llvm::zip(terminator->getOperands(), regionEntry.getArguments())) {
+      if (std::get<0>(tup).getType() != std::get<1>(tup).getType()) return;
+      map.map(std::get<1>(tup), std::get<0>(tup));
+    }
+
+    OpBuilder B(op->getContext());
+    B.setInsertionPointToEnd(&regionEntry);
+    SmallVector<Operation*> toClone;
+    for (Operation& inst : regionEntry) {
+        toClone.push_back(&inst);    
+    }
+    for (Operation* inst : toClone) {
+       B.clone(*inst, map);
+    }
+    terminator->erase();
+  }
+};
+} // namespace
+
+namespace mlir {
+void registerSymbolTestPasses() {
+  PassRegistration<ClonePass>();
+}
+} // namespace mlir