[mlir][DialectConversion] Add support for properly tracking replaceUsesOfBlockArgument

[rriddle]

Summary

The current implementation of this method performs the replacement directly, and thus doesn't support proper back tracking.

[GMNGeoffrey]

Thanks!

mlir/lib/Transforms/DialectConversion.cpp

200

Does this mean the mapping parameter is no longer necessary? Seems like there's still something to fix if it shouldn't be necessary

mlir/test/lib/Dialect/Test/TestPatterns.cpp

251

Not this revision, but WTF is with the capitalization on ILLegal?

[rriddle]

mlir/lib/Transforms/DialectConversion.cpp

200

It was only necessary because of the replaceUsesOfBlockArgument bug when I originally wrote this, but now it is necessary for other reasons.

mlir/test/lib/Dialect/Test/TestPatterns.cpp

251

I think I wrote this at like 3 am, i.e., when I was ILL myself.

[GMNGeoffrey]

mlir/test/lib/Dialect/Test/TestPatterns.cpp

251

😁

mlir/include/mlir/IR/Value.h

public:
 
   /// Replace all uses of 'this' value with 'newValue', updating anything in the
   /// IR that uses 'this' to use the other value instead except if the user is
   /// listed in 'exceptions' .
   void
   replaceAllUsesExcept(Value newValue,
                        const SmallPtrSetImpl<Operation *> &exceptions) const;
 
+  /// Replace all uses of 'this' value with 'newValue' if the given callback
+  /// returns true.
+  void replaceUsesWithIf(Value newValue,
+                         function_ref<bool(OpOperand &)> shouldReplace);
+
   //===--------------------------------------------------------------------===//
   // Uses
 
   /// This class implements an iterator over the uses of a value.
   using use_iterator = ValueUseIterator<OpOperand>;
   using use_range = iterator_range<use_iterator>;
 
   use_iterator use_begin() const;

mlir/lib/IR/Value.cpp

 void Value::replaceAllUsesExcept(
     Value newValue, const SmallPtrSetImpl<Operation *> &exceptions) const {
   for (auto &use : llvm::make_early_inc_range(getUses())) {
     if (exceptions.count(use.getOwner()) == 0)
       use.set(newValue);
   }
 }
 
+/// Replace all uses of 'this' value with 'newValue' if the given callback
+/// returns true.
+void Value::replaceUsesWithIf(Value newValue,
+                              function_ref<bool(OpOperand &)> shouldReplace) {
+  for (OpOperand &use : llvm::make_early_inc_range(getUses()))
+    if (shouldReplace(use))
+      use.set(newValue);
+}
+
 //===--------------------------------------------------------------------===//
 // Uses
 
 auto Value::use_begin() const -> use_iterator {
   return getUseList()->use_begin();
 }
 
 /// Returns true if this value has exactly one use.

mlir/lib/Transforms/DialectConversion.cpp

struct ArgConverter {
 
   /// Cleanup and undo any generated conversions for the arguments of block.
   /// This method replaces the new block with the original, reverting the IR to
   /// its original state.
   void discardRewrites(Block *block);
 
   /// Fully replace uses of the old arguments with the new, materializing cast
   /// operations as necessary.
-  // FIXME(riverriddle) The 'mapping' parameter is only necessary because the
-  // implementation of replaceUsesOfBlockArgument is buggy.
   void applyRewrites(ConversionValueMapping &mapping);
 
   //===--------------------------------------------------------------------===//
   // Conversion
   //===--------------------------------------------------------------------===//
 
   /// Attempt to convert the signature of the given block, if successful a new
   /// block is returned containing the new arguments. On failure, nullptr is

 //===----------------------------------------------------------------------===//
 // ConversionPatternRewriterImpl
 //===----------------------------------------------------------------------===//
 namespace {
 /// This class contains a snapshot of the current conversion rewriter state.
 /// This is useful when saving and undoing a set of rewrites.
 struct RewriterState {
   RewriterState(unsigned numCreatedOps, unsigned numReplacements,
-                unsigned numBlockActions, unsigned numIgnoredOperations,
-                unsigned numRootUpdates)
+                unsigned numArgReplacements, unsigned numBlockActions,
+                unsigned numIgnoredOperations, unsigned numRootUpdates)
       : numCreatedOps(numCreatedOps), numReplacements(numReplacements),
+        numArgReplacements(numArgReplacements),
         numBlockActions(numBlockActions),
         numIgnoredOperations(numIgnoredOperations),
         numRootUpdates(numRootUpdates) {}
 
   /// The current number of created operations.
   unsigned numCreatedOps;
 
   /// The current number of replacements queued.
   unsigned numReplacements;
 
+  /// The current number of argument replacements queued.
+  unsigned numArgReplacements;
+
   /// The current number of block actions performed.
   unsigned numBlockActions;
 
   /// The current number of ignored operations.
   unsigned numIgnoredOperations;
 
   /// The current number of operations that were updated in place.
   unsigned numRootUpdates;
struct ConversionPatternRewriterImpl {
   ArgConverter argConverter;
 
   /// Ordered vector of all of the newly created operations during conversion.
   std::vector<Operation *> createdOps;
 
   /// Ordered vector of any requested operation replacements.
   SmallVector<OpReplacement, 4> replacements;
 
+  /// Ordered vector of any requested block argument replacements.
+  SmallVector<BlockArgument, 4> argReplacements;
+
   /// Ordered list of block operations (creations, splits, motions).
   SmallVector<BlockAction, 4> blockActions;
 
   /// A set of operations that have been erased/replaced/etc that should no
   /// longer be considered for legalization. This is not meant to be an
   /// exhaustive list of all operations, but the minimal set that can be used to
   /// detect if a given operation should be `ignored`. For example, we may add
   /// the operations that define non-empty regions to the set, but not any of
#ifndef NDEBUG
   llvm::ScopedPrinter logger{llvm::dbgs()};
 #endif
 };
 } // end namespace detail
 } // end namespace mlir
 
 RewriterState ConversionPatternRewriterImpl::getCurrentState() {
   return RewriterState(createdOps.size(), replacements.size(),
-                       blockActions.size(), ignoredOps.size(),
-                       rootUpdates.size());
+                       argReplacements.size(), blockActions.size(),
+                       ignoredOps.size(), rootUpdates.size());
 }
 
 void ConversionPatternRewriterImpl::resetState(RewriterState state) {
   // Reset any operations that were updated in place.
   for (unsigned i = state.numRootUpdates, e = rootUpdates.size(); i != e; ++i)
     rootUpdates[i].resetOperation();
   rootUpdates.resize(state.numRootUpdates);
 
+  // Reset any replaced arguments.
+  for (BlockArgument replacedArg :
+       llvm::drop_begin(argReplacements, state.numArgReplacements))
+    mapping.erase(replacedArg);
+  argReplacements.resize(state.numArgReplacements);
+
   // Undo any block actions.
   undoBlockActions(state.numBlockActions);
 
   // Reset any replaced operations and undo any saved mappings.
   for (auto &repl : llvm::drop_begin(replacements, state.numReplacements))
     for (auto result : repl.op->getResults())
       mapping.erase(result);
   replacements.resize(state.numReplacements);
for (auto &repl : replacements) {
     }
 
     // If this operation defines any regions, drop any pending argument
     // rewrites.
     if (argConverter.typeConverter && repl.op->getNumRegions())
       argConverter.notifyOpRemoved(repl.op);
   }
 
+  // Apply all of the requested argument replacements.
+  for (BlockArgument arg : argReplacements) {
+    Value repl = mapping.lookupOrDefault(arg);
+    if (repl.isa<BlockArgument>()) {
+      arg.replaceAllUsesWith(repl);
+      continue;
+    }
+
+    // If the replacement value is an operation, we check to make sure that we
+    // don't replace uses that are within the parent operation of the
+    // replacement value.
+    Operation *replOp = repl.cast<OpResult>().getOwner();
+    Block *replBlock = replOp->getBlock();
+    arg.replaceUsesWithIf(repl, [&](OpOperand &operand) {
+      Operation *user = operand.getOwner();
+      return user->getBlock() != replBlock || replOp->isBeforeInBlock(user);
+    });
+  }
+
   // In a second pass, erase all of the replaced operations in reverse. This
   // allows processing nested operations before their parent region is
   // destroyed.
   for (auto &repl : llvm::reverse(replacements))
     repl.op->erase();
 
   argConverter.applyRewrites(mapping);
 }

 /// Apply a signature conversion to the entry block of the given region.
 Block *ConversionPatternRewriter::applySignatureConversion(
     Region *region, TypeConverter::SignatureConversion &conversion) {
   return impl->applySignatureConversion(region, conversion);
 }
 
 void ConversionPatternRewriter::replaceUsesOfBlockArgument(BlockArgument from,
                                                            Value to) {
-  for (auto &u : from.getUses()) {
-    if (u.getOwner() == to.getDefiningOp())
-      continue;
-    u.getOwner()->replaceUsesOfWith(from, to);
-  }
+  LLVM_DEBUG({
+    Operation *parentOp = from.getOwner()->getParentOp();
+    impl->logger.startLine() << "** Replace Argument : '" << from
+                             << "'(in region of '" << parentOp->getName()
+                             << "'(" << from.getOwner()->getParentOp() << ")\n";
+  });
+  impl->argReplacements.push_back(from);
   impl->mapping.map(impl->mapping.lookupOrDefault(from), to);
 }
 
 /// Return the converted value that replaces 'key'. Return 'key' if there is
 /// no such a converted value.
 Value ConversionPatternRewriter::getRemappedValue(Value key) {
   return impl->mapping.lookupOrDefault(key);
 }

mlir/test/Transforms/test-legalizer.mlir

"test.container"() ({
     // Check that we can undo block creation, i.e. that the block was removed.
     // CHECK: test.create_illegal_block
     // CHECK-NOT: ^{{.*}}(%{{.*}}: i32, %{{.*}}: i32):
     "test.create_illegal_block"() : () -> ()
     "test.finish"() : () -> ()
   }) : () -> ()
   return
 }
+
+// -----
+
+// CHECK-LABEL: @undo_block_arg_replace
+func @undo_block_arg_replace() {
+  "test.undo_block_arg_replace"() ({
+  ^bb0(%arg0: i32):
+    // CHECK: ^bb0(%[[ARG:.*]]: i32):
+    // CHECK-NEXT: "test.return"(%[[ARG]]) : (i32)
+
+    "test.return"(%arg0) : (i32) -> ()
+  }) : () -> ()
+  return
+}

mlir/test/lib/Dialect/Test/TestPatterns.cpp

LogicalResult matchAndRewrite(Operation *op,
     // Create an illegal op to ensure the conversion fails.
     rewriter.create<ILLegalOpF>(op->getLoc(), i32Type);
     rewriter.create<TerminatorOp>(op->getLoc());
     rewriter.replaceOp(op, {});
     return success();
   }
 };
 
+/// A simple pattern that tests the undo mechanism when replacing the uses of a
+/// block argument.
+struct TestUndoBlockArgReplace : public ConversionPattern {
+  TestUndoBlockArgReplace(MLIRContext *ctx)
+      : ConversionPattern("test.undo_block_arg_replace", /*benefit=*/1, ctx) {}
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    auto illegalOp =
+        rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getF32Type());
+    rewriter.replaceUsesOfBlockArgument(op->getRegion(0).front().getArgument(0),
+                                        illegalOp);
+    rewriter.updateRootInPlace(op, [] {});
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // Type-Conversion Rewrite Testing
 
 /// This patterns erases a region operation that has had a type conversion.
 struct TestDropOpSignatureConversion : public ConversionPattern {
   TestDropOpSignatureConversion(MLIRContext *ctx, TypeConverter &converter)
       : ConversionPattern("test.drop_region_op", 1, ctx), converter(converter) {
   }
struct TestLegalizePatternDriver
   enum class ConversionMode { Analysis, Full, Partial };
 
   TestLegalizePatternDriver(ConversionMode mode) : mode(mode) {}
 
   void runOnOperation() override {
     TestTypeConverter converter;
     mlir::OwningRewritePatternList patterns;
     populateWithGenerated(&getContext(), &patterns);
-    patterns.insert<
-        TestRegionRewriteBlockMovement, TestRegionRewriteUndo, TestCreateBlock,
-        TestCreateIllegalBlock, TestPassthroughInvalidOp, TestSplitReturnType,
-        TestChangeProducerTypeI32ToF32, TestChangeProducerTypeF32ToF64,
+    patterns.insert<TestRegionRewriteBlockMovement, TestRegionRewriteUndo,
+                    TestCreateBlock, TestCreateIllegalBlock,
+                    TestUndoBlockArgReplace, TestPassthroughInvalidOp,
+                    TestSplitReturnType, TestChangeProducerTypeI32ToF32,
+                    TestChangeProducerTypeF32ToF64,
                     TestChangeProducerTypeF32ToInvalid, TestUpdateConsumerType,
-        TestNonRootReplacement, TestBoundedRecursiveRewrite>(&getContext());
+                    TestNonRootReplacement, TestBoundedRecursiveRewrite>(
+        &getContext());
     patterns.insert<TestDropOpSignatureConversion>(&getContext(), converter);
     mlir::populateFuncOpTypeConversionPattern(patterns, &getContext(),
                                               converter);
     mlir::populateCallOpTypeConversionPattern(patterns, &getContext(),
                                               converter);
 
     // Define the conversion target used for the test.
     ConversionTarget target(getContext());