[mlir][NFC] accept plain OpBuidler in folded construction helpers

[ftynse]

Summary

A group of functions in the Affine dialect provides a mechanism for
buliding folded-by-construction operations. These functions used to
accept a RewriterBase reference because they may need to erase the
operations that were folded and notify the rewriter when called from
rewrite patterns. Adopt a different approach: postpone the builder
notification of the op creation until we are certain that the op will
not be folded away. This removes the need to notify the rewriter about
op deletion following op construction in case of successful folding, and
removes a bunch of one-off IRRewriter instances in transform code that
may mess up insertion points.

[mravishankar]

I hit the exact same issue, and was embarking on a similar change, but I think there is an issue with erasing operations here.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

787

Actually this is pretty dangerous AFAIK. If you send a pattern rewriter objet and it doesnt know that the operation is being erased that can cause some very subtle errors. My rule of thumb is no method that uses OpBuilder should ever erase an operation. That makes it unusable if a PatternRewriter is passed in.

@rriddle can thorw more light on this.

[ftynse]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

787

This is a good knee-jerk reaction, but I understand the builder code well enough to claim there is no danger here. The rewriter will not get notified about the op being created in the first place because the listener is reset 5 lines above (PatternRewriter derives from OpBuilder::Listener and sets itself as listener). So it doesn't need to know about the op being erased either. The existence of this op will have been completely opaque to the rewriter.

[mravishankar]

Thanks!

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

787

Sorry, didnt mean to suggest that you didnt know. Its something I have hit many times and been burnt by it. Should have worded it better. You obviously know this much better than me. Didnt know about the Listener trick above. Handy!

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

 #include "mlir/IR/Builders.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 
 namespace mlir {
 class AffineApplyOp;
 class AffineBound;
 class AffineValueMap;
-class RewriterBase;
 
 /// TODO: These should be renamed if they are on the mlir namespace.
 ///       Ideally, they should go in a mlir::affine:: namespace.
 
 /// A utility function to check if a value is defined at the top level of an
 /// op with trait `AffineScope` or is a region argument for such an op. A value
 /// of index type defined at the top level is always a valid symbol for all its
 /// uses.
AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map,
                                       ValueRange operands);
 /// Variant of `makeComposedAffineApply` which infers the AffineMap from `e`.
 AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineExpr e,
                                       ValueRange values);
 
 /// Constructs an AffineApplyOp that applies `map` to `operands` after composing
 /// the map with the maps of any other AffineApplyOp supplying the operands,
 /// then immediately attempts to fold it. If folding results in a constant
-/// value, erases all created ops. The `map` must be a single-result affine map.
-OpFoldResult makeComposedFoldedAffineApply(RewriterBase &b, Location loc,
+/// value, no ops are actually created. The `map` must be a single-result affine
+/// map.
+OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
                                            AffineMap map,
                                            ArrayRef<OpFoldResult> operands);
 /// Variant of `makeComposedFoldedAffineApply` that applies to an expression.
-OpFoldResult makeComposedFoldedAffineApply(RewriterBase &b, Location loc,
+OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
                                            AffineExpr expr,
                                            ArrayRef<OpFoldResult> operands);
 /// Variant of `makeComposedFoldedAffineApply` suitable for multi-result maps.
 /// Note that this may create as many affine.apply operations as the map has
 /// results given that affine.apply must be single-result.
-SmallVector<OpFoldResult>
-makeComposedFoldedMultiResultAffineApply(RewriterBase &b, Location loc,
-                                         AffineMap map,
-                                         ArrayRef<OpFoldResult> operands);
+SmallVector<OpFoldResult> makeComposedFoldedMultiResultAffineApply(
+    OpBuilder &b, Location loc, AffineMap map, ArrayRef<OpFoldResult> operands);
 
 /// Returns an AffineMinOp obtained by composing `map` and `operands` with
 /// AffineApplyOps supplying those operands.
 Value makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map,
                             ValueRange operands);
 
 /// Constructs an AffineMinOp that computes a minimum across the results of
 /// applying `map` to `operands`, then immediately attempts to fold it. If
-/// folding results in a constant value, erases all created ops.
-OpFoldResult makeComposedFoldedAffineMin(RewriterBase &b, Location loc,
+/// folding results in a constant value, no ops are actually created.
+OpFoldResult makeComposedFoldedAffineMin(OpBuilder &b, Location loc,
                                          AffineMap map,
                                          ArrayRef<OpFoldResult> operands);
 
 /// Constructs an AffineMinOp that computes a maximum across the results of
 /// applying `map` to `operands`, then immediately attempts to fold it. If
-/// folding results in a constant value, erases all created ops.
-OpFoldResult makeComposedFoldedAffineMax(RewriterBase &b, Location loc,
+/// folding results in a constant value, no ops are actually created.
+OpFoldResult makeComposedFoldedAffineMax(OpBuilder &b, Location loc,
                                          AffineMap map,
                                          ArrayRef<OpFoldResult> operands);
 
 /// Returns the values obtained by applying `map` to the list of values.
 SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc,
                                        AffineMap map, ValueRange values);
 
 /// Given an affine map `map` and its input `operands`, this method composes

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/InliningUtils.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "affine-analysis"
 
void mlir::fullyComposeAffineMapAndOperands(AffineMap *map,
   })) {
     composeAffineMapAndOperands(map, operands);
   }
 }
 
 /// Given a list of `OpFoldResult`, build the necessary operations to populate
 /// `actualValues` with values produced by operations. In particular, for any
 /// attribute-typed element in `values`, call the constant materializer
-/// associated with the Affine dialect to produce an operation.
+/// associated with the Affine dialect to produce an operation. Do NOT notify
+/// the builder listener about the constant ops being created as they are
+/// intended to be removed after being folded into affine constructs; this is
+/// not suitable for use beyond the Affine dialect.
 static void materializeConstants(OpBuilder &b, Location loc,
                                  ArrayRef<OpFoldResult> values,
                                  SmallVectorImpl<Operation *> &constants,
                                  SmallVectorImpl<Value> &actualValues) {
+  OpBuilder::Listener *listener = b.getListener();
+  b.setListener(nullptr);
+  auto listenerResetter =
+      llvm::make_scope_exit([listener, &b] { b.setListener(listener); });
+
   actualValues.reserve(values.size());
   auto *dialect = b.getContext()->getLoadedDialect<AffineDialect>();
   for (OpFoldResult ofr : values) {
     if (auto value = ofr.dyn_cast<Value>()) {
       actualValues.push_back(value);
       continue;
     }
     // Since we are directly specifying `index` as the result type, we need to

 /// arbitrary `leadingArguments`, followed by a list of Value-typed `operands`.
 /// The operation is also expected to always produce a single result. Return an
 /// `OpFoldResult` containing the Attribute representing the folded constant if
 /// complete folding was possible and a Value produced by the created operation
 /// otherwise.
 template <typename OpTy, typename... Args>
 static std::enable_if_t<OpTy::template hasTrait<OpTrait::OneResult>(),
                         OpFoldResult>
-createOrFold(RewriterBase &b, Location loc, ValueRange operands,
+createOrFold(OpBuilder &b, Location loc, ValueRange operands,
              Args &&...leadingArguments) {
   // Identify the constant operands and extract their values as attributes.
   // Note that we cannot use the original values directly because the list of
   // operands may have changed due to canonicalization and composition.
   SmallVector<Attribute> constantOperands;
   constantOperands.reserve(operands.size());
   for (Value operand : operands) {
     IntegerAttr attr;
     if (matchPattern(operand, m_Constant(&attr)))
       constantOperands.push_back(attr);
     else
       constantOperands.push_back(nullptr);
   }
 
   // Create the operation and immediately attempt to fold it. On success,
   // delete the operation and prepare the (unmaterialized) value for being
-  // returned. On failure, return the operation result value.
+  // returned. On failure, return the operation result value. Temporarily remove
+  // the listener to avoid notifying it when the op is created as it may be
+  // removed immediately and there is no way of notifying the caller about that
+  // without resorting to RewriterBase.
+  //
   // TODO: arguably, the main folder (createOrFold) API should support this use
   // case instead of indiscriminately materializing constants.
+  OpBuilder::Listener *listener = b.getListener();
+  b.setListener(nullptr);
+  auto listenerResetter =
+      llvm::make_scope_exit([listener, &b] { b.setListener(listener); });
   OpTy op =
       b.create<OpTy>(loc, std::forward<Args>(leadingArguments)..., operands);
   SmallVector<OpFoldResult, 1> foldResults;
   if (succeeded(op->fold(constantOperands, foldResults)) &&
       !foldResults.empty()) {
-    b.eraseOp(op);
+    op->erase();
     return foldResults.front();
   }
+
+  // Notify the listener now that we definitely know that the operation will
+  // persist. Use the original listener stored in the variable.
+  if (listener)
+    listener->notifyOperationInserted(op);
   return op->getResult(0);
 }
 
 AffineApplyOp mlir::makeComposedAffineApply(OpBuilder &b, Location loc,
                                             AffineMap map,
                                             ValueRange operands) {
   AffineMap normalizedMap = map;
   SmallVector<Value, 8> normalizedOperands(operands.begin(), operands.end());
static void composeMultiResultAffineMap(AffineMap &map,
   // Canonicalize the map created from composed expressions to deduplicate the
   // dimension and symbol operands.
   operands = llvm::to_vector(llvm::concat<Value>(dims, symbols));
   map = AffineMap::get(dims.size(), symbols.size(), exprs, map.getContext());
   canonicalizeMapAndOperands(&map, &operands);
 }
 
 OpFoldResult
-mlir::makeComposedFoldedAffineApply(RewriterBase &b, Location loc,
-                                    AffineMap map,
+mlir::makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map,
                                     ArrayRef<OpFoldResult> operands) {
   assert(map.getNumResults() == 1 && "building affine.apply with !=1 result");
 
   SmallVector<Operation *> constants;
   SmallVector<Value> actualValues;
   materializeConstants(b, loc, operands, constants, actualValues);
   composeAffineMapAndOperands(&map, &actualValues);
   OpFoldResult result = createOrFold<AffineApplyOp>(b, loc, actualValues, map);
 
   // Constants are always folded into affine min/max because they can be
   // represented as constant expressions, so delete them.
   for (Operation *op : constants)
-    b.eraseOp(op);
+    op->erase();
   return result;
 }
 
 OpFoldResult
-mlir::makeComposedFoldedAffineApply(RewriterBase &b, Location loc,
-                                    AffineExpr expr,
+mlir::makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineExpr expr,
                                     ArrayRef<OpFoldResult> operands) {
   return makeComposedFoldedAffineApply(
       b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{expr}).front(),
       operands);
 }
 
 SmallVector<OpFoldResult> mlir::makeComposedFoldedMultiResultAffineApply(
-    RewriterBase &b, Location loc, AffineMap map,
+    OpBuilder &b, Location loc, AffineMap map,
     ArrayRef<OpFoldResult> operands) {
   return llvm::to_vector(llvm::map_range(
       llvm::seq<unsigned>(0, map.getNumResults()), [&](unsigned i) {
         return makeComposedFoldedAffineApply(b, loc, map.getSubMap({i}),
                                              operands);
       }));
 }
 
 Value mlir::makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map,
                                   ValueRange operands) {
   SmallVector<Value> allOperands = llvm::to_vector(operands);
   composeMultiResultAffineMap(map, allOperands);
   return b.createOrFold<AffineMinOp>(loc, b.getIndexType(), map, allOperands);
 }
 
 template <typename OpTy>
-static OpFoldResult makeComposedFoldedMinMax(RewriterBase &b, Location loc,
+static OpFoldResult makeComposedFoldedMinMax(OpBuilder &b, Location loc,
                                              AffineMap map,
                                              ArrayRef<OpFoldResult> operands) {
   SmallVector<Operation *> constants;
   SmallVector<Value> actualValues;
   materializeConstants(b, loc, operands, constants, actualValues);
   composeMultiResultAffineMap(map, actualValues);
   OpFoldResult result =
       createOrFold<OpTy>(b, loc, actualValues, b.getIndexType(), map);
 
   // Constants are always folded into affine min/max because they can be
   // represented as constant expressions, so delete them.
   for (Operation *op : constants)
-    b.eraseOp(op);
+    op->erase();
   return result;
 }
 
 OpFoldResult
-mlir::makeComposedFoldedAffineMin(RewriterBase &b, Location loc, AffineMap map,
+mlir::makeComposedFoldedAffineMin(OpBuilder &b, Location loc, AffineMap map,
                                   ArrayRef<OpFoldResult> operands) {
   return makeComposedFoldedMinMax<AffineMinOp>(b, loc, map, operands);
 }
 
 OpFoldResult
-mlir::makeComposedFoldedAffineMax(RewriterBase &b, Location loc, AffineMap map,
+mlir::makeComposedFoldedAffineMax(OpBuilder &b, Location loc, AffineMap map,
                                   ArrayRef<OpFoldResult> operands) {
   return makeComposedFoldedMinMax<AffineMaxOp>(b, loc, map, operands);
 }
 
 /// Fully compose map with operands and canonicalize the result.
 /// Return the `createOrFold`'ed AffineApply op.
 static Value createFoldedComposedAffineApply(OpBuilder &b, Location loc,
                                              AffineMap map,

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

mlir::linalg::computeMultiTileSizes(OpBuilder &builder, LinalgOp op,
   Value targetSizeValue = materializeOpFoldResult(b, targetSize);
   Value divisorValue = materializeOpFoldResult(b, divisor);
 
   // Find the trip count of the iteration space dimension for which the tile
   // sizes are computed.
   SmallVector<OpFoldResult> allShapes =
       op.createFlatListOfOperandDims(b, b.getLoc());
   AffineMap shapesToLoops = op.getShapesToLoopsMap();
-  IRRewriter rewriter(b);
   SmallVector<OpFoldResult> loopRanges =
-      makeComposedFoldedMultiResultAffineApply(rewriter, op.getLoc(),
-                                               shapesToLoops, allShapes);
+      makeComposedFoldedMultiResultAffineApply(b, op.getLoc(), shapesToLoops,
+                                               allShapes);
   Value tripCount =
-      materializeOpFoldResult(rewriter, op.getLoc(), loopRanges[dimension]);
+      materializeOpFoldResult(b, op.getLoc(), loopRanges[dimension]);
 
   // Compute the tile sizes and the respective numbers of tiles.
   AffineExpr s0 = b.getAffineSymbolExpr(0);
   AffineExpr s1 = b.getAffineSymbolExpr(1);
   AffineExpr s2 = b.getAffineSymbolExpr(2);
   auto apply = [&](AffineExpr expr, ValueRange values) -> Value {
     return makeComposedAffineApply(b, b.getLoc(), expr, values);
   };
static bool canOmitTileOffsetInBoundsCheck(OpFoldResult tileSize,
   if (!tileSizeConst || !numThreadsConst || !iterSizeConst)
     return false;
   return *tileSizeConst * (*numThreadsConst - 1) < *iterSizeConst;
 }
 
 /// Build an `affine_max` of all the `vals`.
 static OpFoldResult buildMax(OpBuilder &b, Location loc,
                              ArrayRef<OpFoldResult> vals) {
-  IRRewriter rewriter(b);
   return makeComposedFoldedAffineMax(
-      rewriter, loc,
-      AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()), vals);
+      b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
+      vals);
 }
 
 /// Build an `affine_min` of all the `vals`.
 static OpFoldResult buildMin(OpBuilder &b, Location loc,
                              ArrayRef<OpFoldResult> vals) {
-  IRRewriter rewriter(b);
   return makeComposedFoldedAffineMin(
-      rewriter, loc,
-      AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()), vals);
+      b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
+      vals);
 }
 
 /// Rewrite a TilingInterface `op` to a tiled `scf.foreach_thread`. The
 /// tiling is specified by the number of tiles/threads `numThreads` and the
 /// optional nominal tile size `nominalTileSizes`. If `nominalTilSizes` is
 /// not specified, then  it is derived from `numThreads` as `ceilDiv(dimSize[i],
 /// numThreads[i])`. If non-empty, the `threadDimMapping` is added as an
 /// attribute to the resulting `scf.foreach_thread`. A zero tile sizes indicate
return tileToForeachThreadOpImpl(b, op, numThreads,
                                    /*nominalTileSizes=*/tileSizes,
                                    threadDimMapping,
                                    /*omitTileOffsetBoundsCheck=*/true);
 }
 
 // Insert a tile `source` into the destination tensor `dest`. The position at
 // which the tile is inserted (as well as size of tile) is taken from a given
 // ExtractSliceOp `sliceOp`.
-static Value insertSliceIntoTensor(RewriterBase &b, Location loc,
+static Value insertSliceIntoTensor(OpBuilder &b, Location loc,
                                    tensor::ExtractSliceOp sliceOp, Value source,
                                    Value dest) {
   return b.create<tensor::InsertSliceOp>(
       loc, sliceOp.getSource().getType(), source, dest, sliceOp.getOffsets(),
       sliceOp.getSizes(), sliceOp.getStrides(), sliceOp.getStaticOffsets(),
       sliceOp.getStaticSizes(), sliceOp.getStaticStrides());
 }
 
else
       interchangedIvs.assign(ivs.begin(), ivs.end());
 
     // Tile the `operandValuesToUse` that either match the `op` operands
     // themselves or the tile loop arguments forwarding them.
     assert(operandValuesToUse.size() ==
                static_cast<size_t>(op.getNumInputsAndOutputs()) &&
            "expect the number of operands and inputs and outputs to match");
     SmallVector<Value> valuesToTile = operandValuesToUse;
-    IRRewriter rewriter(b);
     SmallVector<OpFoldResult> sizeBounds =
-        makeComposedFoldedMultiResultAffineApply(
-            rewriter, loc, shapeSizesToLoopsMap, allShapeSizes);
+        makeComposedFoldedMultiResultAffineApply(b, loc, shapeSizesToLoopsMap,
+                                                 allShapeSizes);
     SmallVector<Value> tiledOperands = makeTiledShapes(
         b, loc, op, valuesToTile, getAsOpFoldResult(interchangedIvs), tileSizes,
         sizeBounds,
         /*omitPartialTileCheck=*/false);
 
     SmallVector<Type> resultTensorTypes =
         getTensorOutputTypes(op, tiledOperands);
     res = op.clone(b, loc, resultTensorTypes, tiledOperands);
loopNest = mlir::scf::buildLoopNest(
         // later moved inside by ExtractSliceOfPadTensorSwapPattern.
         auto map = AffineMap::getMultiDimIdentityMap(rank, b.getContext());
         Value tiledOutput = makeTiledShape(
             b, loc, newPadOp->getResult(0), tileSizes, map, offsets, allDims,
             sizes, /*omitPartialTileCheck=*/false);
         auto sliceOp = tiledOutput.getDefiningOp<tensor::ExtractSliceOp>();
         assert(sliceOp && "expected ExtractSliceOp");
         // Insert the tile into the output tensor.
-        // TODO: Propagate RewriterBase everywhere.
-        IRRewriter rewriter(b);
         Value yieldValue =
-            insertSliceIntoTensor(rewriter, loc, sliceOp, sliceOp, iterArgs[0]);
+            insertSliceIntoTensor(b, loc, sliceOp, sliceOp, iterArgs[0]);
         return scf::ValueVector({yieldValue});
       });
   return success();
 }
 
 namespace {
 struct PadOpTilingPattern : public OpRewritePattern<tensor::PadOp> {
   PadOpTilingPattern(MLIRContext *ctx, LinalgTilingOptions opt)

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
     Location loc = op->getLoc();
     LinalgOp linalgOp = cast<LinalgOp>(op);
     SmallVector<OpFoldResult> allShapesSizes =
         linalgOp.createFlatListOfOperandDims(b, loc);
     AffineMap map = linalgOp.getShapesToLoopsMap();
 
-    IRRewriter rewriter(b);
     return llvm::to_vector(
         llvm::map_range(map.getResults(), [&](AffineExpr loopExpr) {
-          OpFoldResult ofr = makeComposedFoldedAffineApply(
-              rewriter, loc, loopExpr, allShapesSizes);
+          OpFoldResult ofr =
+              makeComposedFoldedAffineApply(b, loc, loopExpr, allShapesSizes);
           return Range{b.getIndexAttr(0), ofr, b.getIndexAttr(1)};
         }));
   }
 
   // Instantiate the tiled implementation of the operation.
   SmallVector<Operation *>
   getTiledImplementation(Operation *op, OpBuilder &b, ValueRange dest,
                          ArrayRef<OpFoldResult> offsets,
getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
                         ArrayRef<OpFoldResult> sizes,
                         SmallVector<OpFoldResult> &resultOffsets,
                         SmallVector<OpFoldResult> &resultSizes) const {
     Location loc = op->getLoc();
     LinalgOp linalgOp = cast<LinalgOp>(op);
 
     AffineExpr d0;
     bindDims(b.getContext(), d0);
-    IRRewriter rewriter(b);
     SmallVector<OpFoldResult> subShapeSizes =
         llvm::to_vector(llvm::map_range(sizes, [&](OpFoldResult ofr) {
-          return makeComposedFoldedAffineApply(rewriter, loc, d0 - 1, ofr);
+          return makeComposedFoldedAffineApply(b, loc, d0 - 1, ofr);
         }));
 
     OpOperand *outOperand = linalgOp.getOutputOperand(resultNumber);
     Value sliceOpResult =
         makeTiledShape(b, loc, outOperand->get(), sizes,
                        linalgOp.getTiedIndexingMap(outOperand), offsets,
                        /*ubs*/ {}, subShapeSizes, true);
     auto sliceOp = sliceOpResult.getDefiningOp<tensor::ExtractSliceOp>();