[mlir] more aggressive folding in tiling/fusion transformations

[ftynse]

Summary

Combine the recently added utilities for folded-by-construction affine
operations with the attribute-based Range to enable more folding. This
decreases the amount of emitted code but has little effect on test
precisely because the tests are not checking for the spurious constants.
The difference in the shape of affine maps comes from the internals of
affine folding.

Depends on D129633

[nicolasvasilache]

Mo' folding, fewer problems

[ftynse]

Rebase.

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

 
 /// 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,
                                          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,
+                                         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);
 
 /// Returns the values obtained by applying `map` to the list of values, which
 /// may be known constants.
 SmallVector<OpFoldResult> applyMapToValues(RewriterBase &b, Location loc,
                                            AffineMap map,

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
                             [](AffineMap map) { return map.isProjectedPermutation(); });
       }]
     >
   ];
 
   let extraClassDeclaration = [{
     /// Return the flat list of all operand dimension sizes in the order they
     /// appear in the operands.
-    SmallVector<Value, 4> createFlatListOfOperandDims(OpBuilder &, Location);
+    SmallVector<OpFoldResult> createFlatListOfOperandDims(OpBuilder &, Location);
 
     /// Return the flat list of all operands' static dimension sizes in the
     /// order they appear in the operands. All operand dimension sizes have to
     /// be statically known.
     SmallVector<int64_t, 4> createFlatListOfOperandStaticDims();
 
     /// Create the loop ranges to materialize the computation over the current
     /// operands. This is done by applying `getShapesToLoopsMap` to

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// avoiding affine map manipulations.
 /// The returned ranges correspond to the loop ranges, in the proper order, that
 /// are tiled and for which new loops will be created. Also the function returns
 /// a map from loop indices of the LinalgOp to the corresponding non-empty range
 /// indices of newly created loops.
 using LoopIndexToRangeIndexMap = DenseMap<int, int>;
 std::tuple<SmallVector<Range, 4>, LoopIndexToRangeIndexMap>
 makeTiledLoopRanges(RewriterBase &b, Location loc, AffineMap map,
-                    ValueRange allShapeSizes, ValueRange allTileSizes);
+                    ArrayRef<OpFoldResult> allShapeSizes,
+                    ArrayRef<OpFoldResult> allTileSizes);
 
 /// A description of a multi-size tiling comprising tile sizes and numbers of
 /// tiles, expressed as Values which may or may not be constant. Multi-size
 /// currently means two-size.
 struct MultiSizeSpecification {
   /// Tile sizes.
   Value lowTileSize, highTileSize;
   /// Number of tiles associated with each size.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 
 /// Check if `permutation` is a permutation of the range
 /// `[0, permutation.size())`.
 bool isPermutation(ArrayRef<int64_t> permutation);
 
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
+OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
+                               int64_t dim);
 
 /// Given an operation, retrieves the value of each dynamic dimension through
 /// constructing the necessary DimOp operators.
 SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
 
 /// Computes an upper bound for the result `value` of an index computation.
 /// Translates AffineMinOps and AffineApplyOps along the use-def chains of the
 /// index computation to affine constraints and projects out intermediate

 /// Checks whether fusing the specific `producer` of the `consumedView` is
 /// feasible. This checks `producer` is the last write of `consumedView` and
 /// that no interleaved dependence would be violated (RAW, WAR or WAW).
 bool isFusableInto(const LinalgDependenceGraph &graph, LinalgOp consumer,
                    Value consumedView, LinalgOp producer);
 
 /// Computes tile offsets, given a list of loop `ivs` and `tileSizes`. In case a
 /// tile size is zero (i.e., no tiling), the corresponding offset is also zero.
-SmallVector<Value> computeTileOffsets(OpBuilder &b, Location loc,
-                                      ValueRange ivs, ValueRange tileSizes);
+SmallVector<OpFoldResult> computeTileOffsets(OpBuilder &b, Location loc,
+                                             ArrayRef<OpFoldResult> ivs,
+                                             ArrayRef<OpFoldResult> tileSizes);
 
 /// Computes tile sizes, given a list of `tileSizes` and dimension
 /// sizes (`sizeBounds`). In case a tile size is zero (i.e., no tiling), the
 /// corresponding result size is the corresponding value from `sizeBounds`.
 /// Note: The returned tile sizes are closed intervals.
-SmallVector<Value> computeTileSizes(OpBuilder &b, Location loc,
-                                    ValueRange tileSizes,
-                                    ArrayRef<Value> sizeBounds);
+SmallVector<OpFoldResult> computeTileSizes(OpBuilder &b, Location loc,
+                                           ArrayRef<OpFoldResult> tileSizes,
+                                           ArrayRef<OpFoldResult> sizeBounds);
 
 /// Returns the list of tensor output types produced when the given structured
 /// operation `op` is applied to the given `operands`. Note that `operands` are
 /// not necessarily the actual operands of `op`.
 SmallVector<Type> getTensorOutputTypes(LinalgOp op, ValueRange operands);
 
 /// Creates `insert_slice` ops that insert `results` back into larger tensors
 /// they were originally extracted from with `extract_slice` before being passed
Value materializeOpFoldResult(OpBuilder &b, Location loc,
                               OpFoldResult opFoldResult);
 
 /// Creates an extract_slice/subview op for a single `valueToTile` with
 /// `builder`. This new operation extracts a tile of `valueToTile`, starting
 /// at offsets `lbs` and with sizes `subShapeSizes`. `omitPartialTileCheck`
 /// controls whether to omit the partial/boundary tile condition check in cases
 /// where we statically know that it is unnecessary.
 Value makeTiledShape(OpBuilder &builder, Location loc, Value valueToTile,
-                     ValueRange tileSizes, AffineMap map, ValueRange lbs,
-                     ValueRange ubs, ValueRange subShapeSizes,
+                     ArrayRef<OpFoldResult> tileSizes, AffineMap map,
+                     ArrayRef<OpFoldResult> lbs, ArrayRef<OpFoldResult> ubs,
+                     ArrayRef<OpFoldResult> subShapeSizes,
                      bool omitPartialTileCheck);
 
 /// Creates extract_slice/subview ops for all `valuesToTile` of the given
 /// `linalgOp` with `builder`, assuming `linalgOp` is being fused into a loop
 /// nest for tiling with the given induction variables `ivs` and tile sizes
 /// `tileSizes`. `sizeBounds` are the iteration space bounds for *all* the
 /// implicit loops in `linalgOp`. `omitPartialTileCheck` controls whether to
 /// omit the partial/boundary tile condition check in cases where we statically
 /// know that it is unnecessary.
 ///
 /// Note that a constant zero in `tileSizes` means no tiling at that implicit
 /// loop. The number of non-zero values in `tileSizes` should be equal to the
 /// number of values in `ivs`.
-SmallVector<Value, 4> makeTiledShapes(OpBuilder &builder, Location loc,
-                                      LinalgOp linalgOp,
-                                      ArrayRef<Value> valuesToTile,
-                                      ValueRange ivs, ValueRange tileSizes,
-                                      ArrayRef<Value> sizeBounds,
-                                      bool omitPartialTileCheck);
+SmallVector<Value> makeTiledShapes(OpBuilder &builder, Location loc,
+                                   LinalgOp linalgOp, ValueRange valuesToTile,
+                                   ArrayRef<OpFoldResult> ivs,
+                                   ArrayRef<OpFoldResult> tileSizes,
+                                   ArrayRef<OpFoldResult> sizeBounds,
 
 /// Add the specified offsets to any `linalg.index` ops contained in the given
 /// `linalgOp`. The offsets are provided in the same order as iteration space
 /// dimensions. Null offests are assumed to be zero.
-void offsetIndices(OpBuilder &b, LinalgOp linalgOp, ArrayRef<Value> offests);
-void offsetIndices(RewriterBase &b, LinalgOp linalgOp, ArrayRef<Value> offests);
+void offsetIndices(OpBuilder &b, LinalgOp linalgOp,
+                   ArrayRef<OpFoldResult> offests);
+void offsetIndices(RewriterBase &b, LinalgOp linalgOp,
+                   ArrayRef<OpFoldResult> offests);
 
 using FusableOpDependencesTy = llvm::MapVector<
     Operation *,
     SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 1>>;
 FusableOpDependencesTy
 findAllFusableDependences(ArrayRef<LinalgOp> ops,
                           const LinalgDependenceGraph &dependenceGraph);
 

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

mlir::makeComposedFoldedAffineApply(RewriterBase &b, Location loc,
                                     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);
-  if (result.is<Attribute>()) {
-    for (Operation *op : constants)
-      b.eraseOp(op);
-  }
+
+  // Constants are always folded into affine min/max because they can be
+  // represented as constant expressions, so delete them.
   return result;
 }
 
 OpFoldResult
 mlir::makeComposedFoldedAffineApply(RewriterBase &b, Location loc,
                                     AffineExpr expr,
                                     ArrayRef<OpFoldResult> operands) {
   return makeComposedFoldedAffineApply(

 
 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);
 }
 
-OpFoldResult
-mlir::makeComposedFoldedAffineMin(RewriterBase &b, Location loc, AffineMap map,
+template <typename OpTy>
+static OpFoldResult makeComposedFoldedMinMax(RewriterBase &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<AffineMinOp>(b, loc, actualValues, b.getIndexType(), map);
-  if (result.is<Attribute>()) {
-    for (Operation *op : constants)
-      b.eraseOp(op);
-  }
+      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.
   return result;
 }
 
+OpFoldResult
+mlir::makeComposedFoldedAffineMin(RewriterBase &b, Location loc, AffineMap map,
+                                  ArrayRef<OpFoldResult> operands) {
+  return makeComposedFoldedMinMax<AffineMinOp>(b, loc, map, operands);
+}
+
+OpFoldResult
+mlir::makeComposedFoldedAffineMax(RewriterBase &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,
                                              ValueRange operandsRef) {
   SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
   fullyComposeAffineMapAndOperands(&map, &operands);
   canonicalizeMapAndOperands(&map, &operands);

mlir/lib/Dialect/Linalg/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgDialect
   MLIRLinalgOpsAttributesIncGen
   MLIRLinalgOpsEnumsIncGen
   MLIRLinalgOpsIncGen
   MLIRLinalgStructuredOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRArithmeticDialect
+  MLIRArithmeticUtils
   MLIRBufferizationDialect
   MLIRDialectUtils
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRParser
   MLIRSideEffectInterfaces
   MLIRSparseTensorDialect
   MLIRSCFDialect
   MLIRMathDialect
   MLIRMemRefDialect
   MLIRTensorDialect
   MLIRTilingInterface
   MLIRViewLikeInterface
   )

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

 //===- LinalgInterfaces.cpp - Linalg interfaces implementation ------------===//
 //
 // 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 "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/SmallBitVector.h"
 

 static Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
                                int64_t dim) {
   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
     return b.createOrFold<memref::DimOp>(loc, source, dim);
   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
     return b.createOrFold<tensor::DimOp>(loc, source, dim);
   llvm_unreachable("Expected MemRefType or TensorType");
 }
+static OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
+                                      int64_t dim) {
+  auto shapedType = source.getType().cast<ShapedType>();
+  if (!shapedType.hasRank() || shapedType.isDynamicDim(dim))
+    return createOrFoldDimOp(b, loc, source, dim);
+  return b.getIndexAttr(shapedType.getDimSize(dim));
+}
 
-SmallVector<Value, 4> LinalgOp::createFlatListOfOperandDims(OpBuilder &b,
+SmallVector<OpFoldResult> LinalgOp::createFlatListOfOperandDims(OpBuilder &b,
                                                                 Location loc) {
-  SmallVector<Value, 4> res;
+  SmallVector<OpFoldResult> res;
   for (OpOperand *opOperand : getInputAndOutputOperands()) {
     for (int64_t i = 0, e = getRank(opOperand); i < e; ++i)
-      res.push_back(createOrFoldDimOp(b, loc, opOperand->get(), i));
+      res.push_back(createFoldedDimOp(b, loc, opOperand->get(), i));
   }
   return res;
 }
 
 SmallVector<int64_t, 4> LinalgOp::createFlatListOfOperandStaticDims() {
   SmallVector<int64_t, 4> res;
   assert(!hasDynamicShape() && "expected operands to have static shapes");
   for (OpOperand *opOperand : getInputAndOutputOperands())
     llvm::append_range(res, getShape(opOperand));
   return res;
 }
 
 SmallVector<Range, 4> LinalgOp::createLoopRanges(OpBuilder &b, Location loc) {
   AffineMap map = getLoopsToShapesMap();
   unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
   auto viewSizes = createFlatListOfOperandDims(b, loc);
   SmallVector<Range, 4> res(numDims);
-  Value zeroVal = b.create<arith::ConstantIndexOp>(loc, 0);
-  Value oneVal = b.create<arith::ConstantIndexOp>(loc, 1);
   for (unsigned idx = 0; idx < numRes; ++idx) {
     auto result = map.getResult(idx);
     if (auto d = result.dyn_cast<AffineDimExpr>()) {
       if (res[d.getPosition()].offset)
         continue;
-      res[d.getPosition()] = Range{zeroVal, viewSizes[idx], oneVal};
+      res[d.getPosition()] =
+          Range{b.getIndexAttr(0), viewSizes[idx], b.getIndexAttr(1)};
     }
   }
   return res;
 }
 
 SmallVector<int64_t, 4> LinalgOp::computeStaticLoopSizes() {
   AffineMap map = getLoopsToShapesMap();
   unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
AffineMap resultShapesFromInputShapesMap =
       loopToResultsShapeMap.compose(getShapesToLoopsMap());
 
   // Check that the result dim map does not contain the positions corresponding
   // to the outputs.
   llvm::SmallBitVector outputDims(resultShapesFromInputShapesMap.getNumDims());
   outputDims.set(resultShapesSubMapPos.first, resultShapesSubMapPos.second);
   HasAffineDimExprVisitor checkDimExpr(std::move(outputDims));
   Location loc = getOperation()->getLoc();
-  auto allResultDimValues =
-      applyMapToValues(b, loc, resultShapesFromInputShapesMap,
+  IRRewriter rewriter(b);
+  SmallVector<OpFoldResult> allResultDimValues =
+      applyMapToValues(rewriter, loc, resultShapesFromInputShapesMap,
                        createFlatListOfOperandDims(b, loc));
   int64_t pos = 0;
   ArrayRef<AffineExpr> shapeExprs = resultShapesFromInputShapesMap.getResults();
   for (OpOperand *opOperand : getOutputOperands()) {
     SmallVector<Value> shapes;
     for (int64_t dim : llvm::seq<int64_t>(0, getRank(opOperand))) {
       if (checkDimExpr.visit(shapeExprs[pos]))
         shapes.push_back(createOrFoldDimOp(b, loc, opOperand->get(), dim));
       else
-        shapes.push_back(allResultDimValues[pos]);
+        shapes.push_back(
+            getValueOrCreateConstantIndexOp(b, loc, allResultDimValues[pos]));
       pos++;
     }
     reifiedReturnShapes.emplace_back(std::move(shapes));
   }
   return success();
 }
 
 LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {

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

 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 namespace {
 /// Bubble up extract_slice above Linalg operation.
 ///
LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
     OpOperand *outOperand = linalgOp.getOutputOperand(0);
     AffineMap indexingMap = linalgOp.getTiedIndexingMap(outOperand);
     if (!indexingMap.isProjectedPermutation()) {
       return rewriter.notifyMatchFailure(
           sliceOp, "expected a projected permutation for output");
     }
 
     auto linalgLoc = linalgOp.getLoc();
-    auto allShapeSizes =
+    SmallVector<OpFoldResult> allShapeSizes =
         linalgOp.createFlatListOfOperandDims(rewriter, linalgLoc);
     AffineMap shapeSizesToLoopsMap = linalgOp.getShapesToLoopsMap();
     if (!shapeSizesToLoopsMap) {
       return rewriter.notifyMatchFailure(
           linalgOp, "failed to get loops map from shape sizes");
     }
-    auto sizeBounds = applyMapToValues(rewriter, linalgLoc,
-                                       shapeSizesToLoopsMap, allShapeSizes);
-
-    auto sliceLoc = sliceOp.getLoc();
-    auto offsetVals = getValueOrCreateConstantIndexOp(
-        rewriter, sliceLoc, sliceOp.getMixedOffsets());
-    auto sizeVals = getValueOrCreateConstantIndexOp(rewriter, sliceLoc,
-                                                    sliceOp.getMixedSizes());
+    SmallVector<OpFoldResult> sizeBounds = applyMapToValues(
+        rewriter, linalgLoc, shapeSizesToLoopsMap, allShapeSizes);
 
     // The offsets and sizes from the slice operation only give you the tile
     // size of the output. Use that compute the tile sizes and offsets of the
     // loops. For loops not used to access the output, set the tile sizes to
     // loop bounds and set the offset to 0.
-    Value zero = rewriter.create<arith::ConstantIndexOp>(linalgLoc, 0);
-    SmallVector<Value, 4> tileOffsets(sizeBounds.size(), zero);
-    SmallVector<Value, 4> tileSizes = sizeBounds;
+    SmallVector<OpFoldResult> tileOffsets(sizeBounds.size(),
+                                          rewriter.getIndexAttr(0));
+    SmallVector<OpFoldResult> tileSizes = sizeBounds;
     for (auto const &result : enumerate(indexingMap.getResults())) {
       unsigned position = result.value().cast<AffineDimExpr>().getPosition();
-      tileOffsets[position] = offsetVals[result.index()];
-      tileSizes[position] = sizeVals[result.index()];
+      tileOffsets[position] = sliceOp.getMixedOffsets()[result.index()];
+      tileSizes[position] = sliceOp.getMixedSizes()[result.index()];
     }
 
     SmallVector<Value> valuesToTile = linalgOp.getInputAndOutputOperands();
-
-    SmallVector<Value, 4> tiledOperands = makeTiledShapes(
-        rewriter, linalgLoc, linalgOp, valuesToTile, tileOffsets, tileSizes,
-        sizeBounds, /*omitPartialTileCheck=*/true);
+    SmallVector<Value> tiledOperands =
+        makeTiledShapes(rewriter, linalgLoc, linalgOp, valuesToTile,
+                        tileOffsets, tileSizes, sizeBounds,
+                        /*omitPartialTileCheck=*/true);
 
     SmallVector<Type, 4> resultTensorTypes;
     for (OpOperand *opOperand : linalgOp.getOutputTensorOperands())
       resultTensorTypes.push_back(
           tiledOperands[opOperand->getOperandNumber()].getType());
 
     Operation *newOp =
         linalgOp.clone(rewriter, linalgLoc, resultTensorTypes, tiledOperands);

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

 static SmallVector<OpFoldResult> getGenericOpLoopRange(OpBuilder &b,
                                                        GenericOp op) {
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(op);
   Location loc = op.getLoc();
   auto allShapesSizes =
       cast<LinalgOp>(op.getOperation()).createFlatListOfOperandDims(b, loc);
   AffineMap map = op.getShapesToLoopsMap();
-  return getAsOpFoldResult(applyMapToValues(b, loc, map, allShapesSizes));
+  IRRewriter rewriter(b);
+  return applyMapToValues(rewriter, loc, map, allShapesSizes);
 }
 
 /// Helper method to permute the list of `values` based on the `map`.
 SmallVector<OpFoldResult> permuteValues(ArrayRef<OpFoldResult> values,
                                         AffineMap map) {
   assert(map.isPermutation());
   SmallVector<OpFoldResult> permutedValues(values.size());
   for (auto position :

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

static SmallVector<Value> getTiledOperands(LinalgOp producer) {
   return producer.getInputAndOutputOperands();
 }
 
 /// Fuses the producer by cloning the `producer`. The `fusedLoopsAndRanges`
 /// provides the loop range information for the fused loops. The rest are
 /// obtained from the producer itself, since they are not tiled + fused.
 static LinalgOp fuse(OpBuilder &b, LinalgOp producer,
                      const DenseMap<unsigned, Range> &fusedLoopsAndRanges) {
-  SmallVector<Value, 8> ivs, tileSizes, sizeBounds;
-  SmallVector<Range, 8> loopRanges;
+  SmallVector<OpFoldResult> ivs, tileSizes, sizeBounds;
+  SmallVector<Range> loopRanges;
   Location loc = producer.getLoc();
-  auto zero = b.create<arith::ConstantIndexOp>(loc, 0);
 
   for (unsigned i = 0, e = producer.getNumLoops(); i < e; ++i) {
     auto shapeDim = getShapeDefiningLoopRange(producer, i);
-    Value dim = createOrFoldDimOp(b, loc, shapeDim.shape, shapeDim.dimension);
+    OpFoldResult dim =
+        createFoldedDimOp(b, loc, shapeDim.shape, shapeDim.dimension);
     sizeBounds.push_back(dim);
     auto it = fusedLoopsAndRanges.find(i);
     if (it != fusedLoopsAndRanges.end()) {
-      ivs.push_back(materializeOpFoldResult(b, loc, it->second.offset));
-      tileSizes.push_back(materializeOpFoldResult(b, loc, it->second.size));
+      ivs.push_back(it->second.offset);
+      tileSizes.push_back(it->second.size);
       loopRanges.push_back(it->second);
       LLVM_DEBUG(llvm::dbgs() << "tiled loop#" << i << " with LoopRange "
                               << loopRanges.back() << "\n");
     } else {
-      tileSizes.push_back(zero);
+      tileSizes.push_back(b.getIndexAttr(0));
       loopRanges.push_back(Range{b.getIndexAttr(0), dim, b.getIndexAttr(1)});
       LLVM_DEBUG(llvm::dbgs() << "full loop#" << i << " with LoopRange "
                               << loopRanges.back() << "\n");
     }
   }
 
   SmallVector<Value, 8> clonedShapes;
   clonedShapes.reserve(producer.getNumInputsAndOutputs());
for (RankedTensorType t : producer.getOutputTensorTypes()) {
     resultTypes.push_back(tensor::ExtractSliceOp::inferResultType(
         t.cast<RankedTensorType>(), staticOffsetsVector, staticSizesVector,
         staticStridesVector));
   }
 
   Operation *clonedOp = producer.clone(b, loc, resultTypes, clonedShapes);
 
   // Shift all IndexOp results by the tile offset.
-  SmallVector<Value> allIvs;
-  llvm::transform(loopRanges, std::back_inserter(allIvs), [&](Range range) {
-    return materializeOpFoldResult(b, loc, range.offset);
-  });
+  SmallVector<OpFoldResult> allIvs = llvm::to_vector(
+      llvm::map_range(loopRanges, [&](Range range) { return range.offset; }));
   offsetIndices(b, clonedOp, allIvs);
 
   return clonedOp;
 }
 
 /// Get the loop range for a dimension `dim` based on the `shapedOperand`. It is
 /// expected to be defined by a subview op or an extract_slice op.
 static Range getRangeFromOperandShape(OpBuilder &b, Location loc,

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

static LinalgOp getTiledProducer(OpBuilder &b, OpResult producerResult,
   OpBuilder::InsertionGuard guard(b);
   b.setInsertionPointAfter(sliceOp);
 
   // Get the producer.
   LinalgOp producerOp = producerResult.getOwner();
   Location loc = producerOp.getLoc();
 
   // Obtain the `producerOp` loop bounds and the `sliceOp` ranges.
-  SmallVector<Value> producerLoopBounds;
+  SmallVector<OpFoldResult> producerLoopBounds;
   llvm::transform(producerOp.createLoopRanges(b, loc),
-                  std::back_inserter(producerLoopBounds), [&](Range range) {
-                    return materializeOpFoldResult(b, loc, range.size);
-                  });
+                  std::back_inserter(producerLoopBounds),
+                  [&](Range range) { return range.size; });
   SmallVector<Range> sliceOpRanges = sliceOp.getOrCreateRanges(b, loc);
 
   // Tile the producer operands given the `sliceOp` ranges. Iterate the
   // `tiledSliceDimIndices` and store the tile offset and size for the tiled
   // slice dimension.
-  auto zero = b.create<arith::ConstantIndexOp>(loc, 0);
-  SmallVector<Value> tileIvs(producerOp.getNumLoops(), nullptr);
-  SmallVector<Value> tileSizes(producerOp.getNumLoops(), zero);
-  SmallVector<Value> allIvs(producerOp.getNumLoops(), nullptr);
+  SmallVector<OpFoldResult> tileIvs(producerOp.getNumLoops(), nullptr);
+  SmallVector<OpFoldResult> tileSizes(producerOp.getNumLoops(),
+                                      b.getIndexAttr(0));
+  SmallVector<OpFoldResult> allIvs(producerOp.getNumLoops(), nullptr);
   for (auto it : zip(tiledSliceDimIndices, tiledProducerLoopIndices)) {
     int64_t tiledSliceDim = std::get<0>(it);
     int64_t tiledProducerLoop = std::get<1>(it);
-    tileIvs[tiledProducerLoop] =
-        materializeOpFoldResult(b, loc, sliceOpRanges[tiledSliceDim].offset);
-    tileSizes[tiledProducerLoop] =
-        materializeOpFoldResult(b, loc, sliceOpRanges[tiledSliceDim].size);
+    tileIvs[tiledProducerLoop] = sliceOpRanges[tiledSliceDim].offset;
+    tileSizes[tiledProducerLoop] = sliceOpRanges[tiledSliceDim].size;
     allIvs[tiledProducerLoop] = tileIvs[tiledProducerLoop];
   }
-  erase_value(tileIvs, nullptr);
+  erase_value(tileIvs, OpFoldResult());
   SmallVector<Value> tiledOperands = producerOp.getInputAndOutputOperands();
   tiledOperands = makeTiledShapes(b, loc, producerOp, tiledOperands, tileIvs,
                                   tileSizes, producerLoopBounds,
                                   /**omitPartialTileCheck=*/false);
 
   // Output fusion has to update the iteration arguments of the tile loop nest.
   // In particular, the iteration argument of the outermost tile loop needs to
   // be set to the producer output instead of the producer result and `clonedOp`

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

 #include "llvm/Support/CommandLine.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::scf;
 
 #define DEBUG_TYPE "linalg-tiling"
 
-static bool isZero(Value v) {
-  if (auto cst = v.getDefiningOp<arith::ConstantIndexOp>())
+static bool isZero(OpFoldResult v) {
+  if (!v)
+    return false;
+  if (auto attr = v.dyn_cast<Attribute>()) {
+    IntegerAttr intAttr = attr.dyn_cast<IntegerAttr>();
+    return intAttr && intAttr.getValue().isZero();
+  }
+  if (auto cst = v.get<Value>().getDefiningOp<arith::ConstantIndexOp>())
     return cst.value() == 0;
   return false;
 }
 
 std::tuple<SmallVector<Range, 4>, LoopIndexToRangeIndexMap>
 mlir::linalg::makeTiledLoopRanges(RewriterBase &b, Location loc, AffineMap map,
-                                  ValueRange allShapeSizes,
-                                  ValueRange allTileSizes) {
+                                  ArrayRef<OpFoldResult> allShapeSizes,
+                                  ArrayRef<OpFoldResult> allTileSizes) {
   assert(allTileSizes.size() == map.getNumResults());
   // Apply `map` to get shape sizes in loop order.
-  auto shapeSizes = applyMapToValues(b, loc, map, allShapeSizes);
-  SmallVector<Value, 4> tileSizes(allTileSizes.begin(), allTileSizes.end());
+  SmallVector<OpFoldResult> shapeSizes =
+      applyMapToValues(b, loc, map, allShapeSizes);
+  SmallVector<OpFoldResult> tileSizes(allTileSizes.begin(), allTileSizes.end());
 
   // Traverse the tile sizes, which are in loop order, erase zeros everywhere.
   LoopIndexToRangeIndexMap loopIndexToRangeIndex;
   for (int idx = 0, e = tileSizes.size(), zerosCount = 0; idx < e; ++idx) {
     if (isZero(tileSizes[idx - zerosCount])) {
       shapeSizes.erase(shapeSizes.begin() + idx - zerosCount);
       tileSizes.erase(tileSizes.begin() + idx - zerosCount);
       ++zerosCount;
void mlir::linalg::transformIndexOps(
     const LoopIndexToRangeIndexMap &loopIndexToRangeIndex) {
   SmallVector<Value> allIvs(op.getNumLoops(), nullptr);
   for (auto &en : enumerate(allIvs)) {
     auto rangeIndex = loopIndexToRangeIndex.find(en.index());
     if (rangeIndex == loopIndexToRangeIndex.end())
       continue;
     en.value() = ivs[rangeIndex->second];
   }
-  offsetIndices(b, op, allIvs);
+  offsetIndices(b, op, getAsOpFoldResult(allIvs));
 }
 
 /// Asserts that the given index-typed value is strictly positive. If the value
 /// is an attribute, asserts at compile time, otherwise emits an assertion
 /// checked at runtime.
 static void emitIsPositiveIndexAssertion(ImplicitLocOpBuilder &b,
                                          OpFoldResult value) {
   if (auto attr = value.dyn_cast<Attribute>()) {
if (emitAssertions) {
     emitIsPositiveIndexAssertion(b, targetSize);
     emitIsPositiveIndexAssertion(b, divisor);
   }
   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.
-  // TODO: update createFlatListOfOperandDims to return OpFoldResults and avoid
-  // littering by useless constant materialization.
-  SmallVector<Value, 4> allShapes =
+  SmallVector<OpFoldResult> allShapes =
       op.createFlatListOfOperandDims(b, b.getLoc());
   AffineMap shapesToLoops = op.getShapesToLoopsMap();
-  SmallVector<Value, 4> loopRanges =
-      applyMapToValues(b, op.getLoc(), shapesToLoops, allShapes);
-  Value tripCount = loopRanges[dimension];
+  IRRewriter rewriter(b);
+  SmallVector<OpFoldResult> loopRanges =
+      applyMapToValues(rewriter, op.getLoc(), shapesToLoops, allShapes);
+  Value tripCount =
+      materializeOpFoldResult(rewriter, 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);
   };
createMatchingParallelSubsetInsertOp(OpBuilder &b, Location loc,
                                      tensor::ExtractSliceOp subsetExtractOp,
                                      Value source, Value dest) {
   b.create<tensor::ParallelInsertSliceOp>(
       loc, source, dest, subsetExtractOp.getMixedOffsets(),
       subsetExtractOp.getMixedSizes(), subsetExtractOp.getMixedStrides());
 }
 
 /// Build an `affine_max` of all the `vals`.
-static Value buildMax(OpBuilder &b, Location loc, ValueRange vals) {
-  return b.createOrFold<AffineMaxOp>(
-      loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
-      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);
 }
 
 /// Build an `affine_min` of all the `vals`.
-static Value buildMin(OpBuilder &b, Location loc, ValueRange vals) {
-  return b.createOrFold<AffineMinOp>(
-      loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
-      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);
 }
 
 FailureOr<ForeachThreadTilingResult>
 linalg::tileToForeachThreadOp(OpBuilder &b, TilingInterface op,
                               ArrayRef<OpFoldResult> numThreads,
                               ArrayRef<int64_t> threadDimMapping) {
   Location loc = op->getLoc();
   OpBuilder::InsertionGuard g(b);
SmallVector<OpFoldResult> nonZeroNumThreads =
         return !isConstantIntValue(ofr, 0);
       }));
   SmallVector<Value> materializedNonZeroNumThreads =
       llvm::to_vector(llvm::map_range(nonZeroNumThreads, [&](OpFoldResult ofr) {
         ImplicitLocOpBuilder ilocb(loc, b);
         return materializeOpFoldResult(ilocb, ofr);
       }));
 
-  Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
   Operation *tiledOp = nullptr;
   scf::ForeachThreadOp foreachThreadOp = b.create<scf::ForeachThreadOp>(
       loc, materializedNonZeroNumThreads, threadDimMapping,
       [&](OpBuilder &b, Location loc, ValueRange threadIds) {
         int64_t nLoops = loopRanges.size();
         SmallVector<OpFoldResult> tiledOffsets, tiledSizes;
         tiledOffsets.reserve(nLoops);
         tiledSizes.reserve(nLoops);
         for (unsigned loopIdx = 0, threadIdIdx = 0; loopIdx < nLoops;
              ++loopIdx) {
           bool overflow = loopIdx >= numThreads.size();
           bool isZero = !overflow && isConstantIntValue(numThreads[loopIdx], 0);
           // Degenerate case: take the whole domain.
           if (overflow || isZero) {
             tiledOffsets.push_back(loopRanges[loopIdx].offset);
             tiledSizes.push_back(loopRanges[loopIdx].size);
             continue;
           }
 
           // Tiled case: compute the offset and size.
           AffineExpr i, j, M, N, O;
           bindDims(b.getContext(), i, j);
           bindSymbols(b.getContext(), M, N, O);
-          Value size = loopRanges[loopIdx].size;
-          Value offset = loopRanges[loopIdx].offset;
-          Value threadId = threadIds[threadIdIdx];
-          // TODO: more aggressive foldings.
+          OpFoldResult size = loopRanges[loopIdx].size;
+          OpFoldResult offset = loopRanges[loopIdx].offset;
+          OpFoldResult threadId = threadIds[threadIdIdx];
           // Symbolic fixed max size per thread.
           // TODO: floor + 0/1 depending on case for better load-balancing.
-          Value maxSizePerThread = b.createOrFold<AffineApplyOp>(
-              loc, M.ceilDiv(N),
-              ValueRange{size, materializedNonZeroNumThreads[threadIdIdx]});
+          IRRewriter rewriter(b);
+          OpFoldResult maxSizePerThread = makeComposedFoldedAffineApply(
+              rewriter, loc, M.ceilDiv(N),
+              {size, nonZeroNumThreads[threadIdIdx]});
           // Dynamic offset shifted by threadId * maxSizePerThread.
-          Value offsetPerThread = b.createOrFold<AffineApplyOp>(
-              loc, i + j * M, ValueRange{offset, threadId, maxSizePerThread});
+          OpFoldResult offsetPerThread = makeComposedFoldedAffineApply(
+              rewriter, loc, i + j * M, {offset, threadId, maxSizePerThread});
           // Dynamic upper-bound depending on the threadId.
-          Value sizeMinusOffsetPerThread = b.createOrFold<AffineApplyOp>(
-              loc, -i + M, ValueRange{offsetPerThread, size});
-          Value tileSizePerThread = buildMin(
-              b, loc, ValueRange{sizeMinusOffsetPerThread, maxSizePerThread});
+          OpFoldResult sizeMinusOffsetPerThread = makeComposedFoldedAffineApply(
+              rewriter, loc, -i + M, {offsetPerThread, size});
+          OpFoldResult tileSizePerThread =
+              buildMin(b, loc, {sizeMinusOffsetPerThread, maxSizePerThread});
           tiledOffsets.push_back(offsetPerThread);
           // TODO: if tileSizePerThread <= 0 early exit.
           tiledSizes.push_back(
-              buildMax(b, loc, ValueRange{zero, tileSizePerThread}));
+              buildMax(b, loc, {b.getIndexAttr(0), tileSizePerThread}));
           ++threadIdIdx;
         }
 
         SmallVector<Operation *> tiledOps =
             op.getTiledImplementation(b, destOperands, tiledOffsets, tiledSizes,
                                       /*tileDestOperands=*/true);
         assert(tiledOps.size() == 1 && "expected a single produced tiled op");
         tiledOp = tiledOps.front();
static Value insertSliceIntoTensor(RewriterBase &b, Location loc,
   return b.create<tensor::InsertSliceOp>(
       loc, sliceOp.getSource().getType(), source, dest, sliceOp.getOffsets(),
       sliceOp.getSizes(), sliceOp.getStrides(), sliceOp.getStaticOffsets(),
       sliceOp.getStaticSizes(), sliceOp.getStaticStrides());
 }
 
 template <typename LoopTy>
 static FailureOr<TiledLinalgOp>
-tileLinalgOpImpl(RewriterBase &b, LinalgOp op, ValueRange tileSizes,
+tileLinalgOpImpl(RewriterBase &b, LinalgOp op, ArrayRef<OpFoldResult> tileSizes,
                  const LinalgTilingOptions &options) {
   auto nLoops = op.getNumLoops();
   // Initial tile sizes may be too big, only take the first nLoops.
   tileSizes = tileSizes.take_front(nLoops);
 
   if (llvm::all_of(tileSizes, isZero)) {
     TiledLinalgOp tiledOp;
     tiledOp.op = cast<LinalgOp>(b.clone(*op.getOperation()));
     tiledOp.tensorResults.assign(tiledOp.op->result_begin(),
                                  tiledOp.op->result_end());
     return tiledOp;
   }
 
   // 1. Build the tiled loop ranges.
-  auto allShapeSizes = op.createFlatListOfOperandDims(b, op.getLoc());
+  SmallVector<OpFoldResult> allShapeSizes =
+      op.createFlatListOfOperandDims(b, op.getLoc());
   AffineMap shapeSizesToLoopsMap = op.getShapesToLoopsMap();
   if (!shapeSizesToLoopsMap)
     return failure();
 
   SmallVector<Range, 4> loopRanges;
   LoopIndexToRangeIndexMap loopIndexToRangeIndex;
   std::tie(loopRanges, loopIndexToRangeIndex) = makeTiledLoopRanges(
       b, op.getLoc(), shapeSizesToLoopsMap, allShapeSizes, tileSizes);
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;
-    auto sizeBounds =
-        applyMapToValues(b, loc, shapeSizesToLoopsMap, allShapeSizes);
-    SmallVector<Value, 4> tiledOperands =
-        makeTiledShapes(b, loc, op, valuesToTile, interchangedIvs, tileSizes,
-                        sizeBounds, /*omitPartialTileCheck=*/false);
+    IRRewriter rewriter(b);
+    SmallVector<OpFoldResult> sizeBounds =
+        applyMapToValues(rewriter, 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);
     tensorResults =
         insertSlicesBack(builder, loc, op, tiledOperands, res->getResults());
     return scf::ValueVector(tensorResults.begin(), tensorResults.end());
   };
FailureOr<TiledLinalgOp> static tileLinalgOpImpl(
 
   if (!options.tileSizeComputationFunction)
     return failure();
 
   // Enforce the convention that "tiling by zero" skips tiling a particular
   // dimension. This convention is significantly simpler to handle instead of
   // adjusting affine maps to account for missing dimensions.
   auto nLoops = op.getNumLoops();
-  SmallVector<Value, 4> tileSizeVector =
-      options.tileSizeComputationFunction(b, op);
+  SmallVector<OpFoldResult> tileSizeVector =
+      getAsOpFoldResult(options.tileSizeComputationFunction(b, op));
   if (tileSizeVector.size() < nLoops) {
-    auto zero = b.create<arith::ConstantIndexOp>(op.getLoc(), 0);
-    tileSizeVector.append(nLoops - tileSizeVector.size(), zero);
+    tileSizeVector.append(nLoops - tileSizeVector.size(), b.getIndexAttr(0));
   }
 
   return tileLinalgOpImpl<LoopTy>(b, op, tileSizeVector, options);
 }
 
 FailureOr<TiledLinalgOp>
 mlir::linalg::tileLinalgOp(RewriterBase &b, LinalgOp op,
                            const LinalgTilingOptions &options) {
static LogicalResult tilePadOp(RewriterBase &builder, tensor::PadOp op,
   Location loc = op.getLoc();
   OpBuilder::InsertionGuard g(builder);
   builder.setInsertionPoint(op);
 
   // Clone tensor::PadOp so that the existing op can be replaced more easily.
   newPadOp = cast<tensor::PadOp>(builder.clone(*op.getOperation()));
   // Get rank and tile sizes.
   int64_t rank = op.getResultType().getRank();
-  SmallVector<Value> tileSizes =
-      options.tileSizeComputationFunction(builder, op);
+  SmallVector<OpFoldResult> tileSizes =
+      getAsOpFoldResult(options.tileSizeComputationFunction(builder, op));
   // Normalize untiled padding dimensions to 0.
-  Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
-  tileSizes.append(rank - tileSizes.size(), zero);
+  tileSizes.append(rank - tileSizes.size(), builder.getIndexAttr(0));
   // Compute lower and upper bounds of the loop nest.
   TilingInterface tilingInterface =
       dyn_cast<TilingInterface>(op.getOperation());
   SmallVector<Range> ranges = tilingInterface.getIterationDomain(builder);
-  SmallVector<Value> lbs, dims, allDims, steps;
+  SmallVector<Value> lbs, dims, steps;
+  SmallVector<OpFoldResult> allDims;
   for (int64_t i = 0; i < rank; ++i) {
-    Value materializedSize =
-        materializeOpFoldResult(builder, loc, ranges[i].size);
-    allDims.push_back(materializedSize);
+    allDims.push_back(ranges[i].size);
     if (!isZero(tileSizes[i])) {
       lbs.push_back(materializeOpFoldResult(builder, loc, ranges[i].offset));
-      dims.push_back(materializedSize);
-      steps.push_back(tileSizes[i]);
+      dims.push_back(materializeOpFoldResult(builder, loc, ranges[i].size));
+      steps.push_back(materializeOpFoldResult(builder, loc, tileSizes[i]));
     }
   }
   // Generate loop nest: One loop per dimension.
   SmallVector<Value> destOperand =
       tilingInterface.getDestinationOperands(builder);
   loopNest = mlir::scf::buildLoopNest(
       builder, loc, lbs, /*ubs=*/dims, steps, ValueRange(destOperand),
       [&](OpBuilder &b, Location loc, ValueRange localIvs,
           ValueRange iterArgs) -> scf::ValueVector {
         // Compute offsets and sizes of ExtractSliceOp.
-        SmallVector<Value> offsets =
-            computeTileOffsets(b, loc, localIvs, tileSizes);
-        SmallVector<Value> sizes = computeTileSizes(b, loc, tileSizes, allDims);
+        SmallVector<Value> localIVVector = llvm::to_vector(localIvs);
+        SmallVector<OpFoldResult> offsets = computeTileOffsets(
+            b, loc, getAsOpFoldResult(localIVVector), tileSizes);
+        SmallVector<OpFoldResult> sizes =
+            computeTileSizes(b, loc, tileSizes, allDims);
         // Create ExtractSliceOp: Extract a tile from the tensor::PadOp.
         // Note: The tensor::PadOp is located outside of the loop nest. It is
         // 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>();

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

 /// maybe possible to move this into the op-definition (though there are
 /// advantages to leaving it as an external model)
 template <typename LinalgOpTy>
 struct LinalgOpTilingInterface
     : public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
                                             LinalgOpTy> {
   /// Return the destination operands.
   SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
-    return llvm::cast<LinalgOp>(op).getOutputOperands();
+    return cast<LinalgOp>(op).getOutputOperands();
   }
 
   /// Return the loop iterator type.
   SmallVector<StringRef> getLoopIteratorTypes(Operation *op) const {
     LinalgOpTy concreteOp = cast<LinalgOpTy>(op);
     return llvm::to_vector(
         llvm::map_range(concreteOp.iterator_types(), [](Attribute strAttr) {
           return strAttr.cast<StringAttr>().getValue();
         }));
   }
 
   /// Return the iteration domain range.
   SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
     Location loc = op->getLoc();
     LinalgOp linalgOp = cast<LinalgOp>(op);
-    auto allShapesSizes = linalgOp.createFlatListOfOperandDims(b, loc);
+    SmallVector<OpFoldResult> allShapesSizes =
+        linalgOp.createFlatListOfOperandDims(b, loc);
     AffineMap map = linalgOp.getShapesToLoopsMap();
-    Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
-    Value one = b.create<arith::ConstantIndexOp>(loc, 1);
+
+    IRRewriter rewriter(b);
     return llvm::to_vector(llvm::map_range(
-        applyMapToValues(b, loc, map, allShapesSizes), [&](Value v) {
-          return Range{zero, v, one};
+        applyMapToValues(rewriter, loc, map, allShapesSizes),
+        [&](OpFoldResult ofr) {
+          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,
                          ArrayRef<OpFoldResult> sizes,
                          bool tileDestOperands) const {
     // Leave the `sizeBounds` value empty. That is only needed when the `sizes`
     // specified could lead to out of bounds accesses.
     Location loc = op->getLoc();
     LinalgOp linalgOp = cast<LinalgOp>(op);
     SmallVector<Value> valuesToTile = linalgOp.getInputAndOutputOperands();
-    SmallVector<Value> offsetValues =
-        getValueOrCreateConstantIndexOp(b, loc, offsets);
     SmallVector<Value, 4> tiledOperands = makeTiledShapes(
-        b, loc, linalgOp, valuesToTile, offsetValues,
-        getValueOrCreateConstantIndexOp(b, loc, sizes), {}, true);
+        b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
 
     SmallVector<Type> resultTensorTypes = llvm::to_vector(llvm::map_range(
         linalgOp.getOutputTensorOperands(), [&](OpOperand *opOperand) {
           return tiledOperands[opOperand->getOperandNumber()].getType();
         }));
 
     Operation *tiledOp =
         linalgOp.clone(b, loc, resultTensorTypes, tiledOperands);
-    offsetIndices(b, cast<LinalgOp>(tiledOp), offsetValues);
+    offsetIndices(b, cast<LinalgOp>(tiledOp), offsets);
 
     return {tiledOp};
   }
 
   // Return the details of the output tile generated by the tiled
   // implementation.
   LogicalResult
   getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
                         ArrayRef<OpFoldResult> offsets,
                         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);
-
-    auto fullyComposeAffineMapAndOperands = [](OpBuilder &builder, Location loc,
-                                               AffineExpr expr,
-                                               ValueRange operands) -> Value {
-      AffineMap map = AffineMap::inferFromExprList({expr}).front();
-      SmallVector<Value> normalizedOperands(operands.begin(), operands.end());
-      mlir::fullyComposeAffineMapAndOperands(&map, &normalizedOperands);
-      canonicalizeMapAndOperands(&map, &normalizedOperands);
-      return builder.createOrFold<AffineApplyOp>(loc, map, normalizedOperands);
-    };
-
-    SmallVector<Value> sizeVals =
-        getValueOrCreateConstantIndexOp(b, loc, sizes);
-    SmallVector<Value> subShapeSizes =
-        llvm::to_vector(llvm::map_range(sizeVals, [&](Value v) {
-          return fullyComposeAffineMapAndOperands(b, loc, d0 - 1, v);
+    IRRewriter rewriter(b);
+    SmallVector<OpFoldResult> subShapeSizes =
+        llvm::to_vector(llvm::map_range(sizes, [&](OpFoldResult ofr) {
+          return makeComposedFoldedAffineApply(rewriter, loc, d0 - 1, ofr);
         }));
+
     OpOperand *outOperand = linalgOp.getOutputOperand(resultNumber);
     Value sliceOpResult =
-        makeTiledShape(b, loc, outOperand->get(), sizeVals,
-                       linalgOp.getTiedIndexingMap(outOperand),
-                       getValueOrCreateConstantIndexOp(b, loc, offsets),
+        makeTiledShape(b, loc, outOperand->get(), sizes,
+                       linalgOp.getTiedIndexingMap(outOperand), offsets,
                        /*ubs*/ {}, subShapeSizes, true);
     auto sliceOp = sliceOpResult.getDefiningOp<tensor::ExtractSliceOp>();
     if (!sliceOp)
       return failure();
     resultOffsets = sliceOp.getMixedOffsets();
     resultSizes = sliceOp.getMixedSizes();
     return success();
   }

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

 LinalgTilingOptions &mlir::linalg::LinalgTilingOptions::scalarizeDynamicDims() {
   assert(!tileSizeComputationFunction && "tile sizes already set");
   tileSizeComputationFunction = [](OpBuilder &b, Operation *op) {
     SmallVector<Value, 4> tileSizes;
     auto linalgOp = dyn_cast<LinalgOp>(op);
     if (!linalgOp)
       return tileSizes;
     Location loc = linalgOp.getLoc();
-    auto allShapeSizes = linalgOp.createFlatListOfOperandDims(b, loc);
+    SmallVector<OpFoldResult> allShapeSizes =
+        linalgOp.createFlatListOfOperandDims(b, loc);
     AffineMap map = linalgOp.getShapesToLoopsMap();
     if (!map)
       return tileSizes;
-    auto shapeSizes = applyMapToValues(b, loc, map, allShapeSizes);
+    IRRewriter rewriter(b);
+    SmallVector<OpFoldResult> shapeSizes =
+        applyMapToValues(rewriter, loc, map, allShapeSizes);
     // If the shape size is dynamic, tile by 1. Otherwise, do not tile (tile
     // size 0).
-    for (Value shapeSize : shapeSizes)
+    for (OpFoldResult shapeSize : shapeSizes)
       tileSizes.push_back(getConstantIntValue(shapeSize)
                               ? b.create<arith::ConstantIndexOp>(loc, 0)
                               : b.create<arith::ConstantIndexOp>(loc, 1));
     return tileSizes;
   };
   return *this;
 }
 

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

 
 #define DEBUG_TYPE "linalg-utils"
 
 using namespace mlir;
 using namespace presburger;
 using namespace mlir::linalg;
 using namespace mlir::scf;
 
-static bool isZero(Value v) {
-  if (auto cst = v.getDefiningOp<arith::ConstantIndexOp>())
+static bool isZero(OpFoldResult v) {
+  if (!v)
+    return false;
+  if (auto attr = v.dyn_cast<Attribute>()) {
+    IntegerAttr intAttr = attr.dyn_cast<IntegerAttr>();
+    return intAttr && intAttr.getValue().isZero();
+  }
+  if (auto cst = v.get<Value>().getDefiningOp<arith::ConstantIndexOp>())
     return cst.value() == 0;
   return false;
 }
 
 namespace {
 
 // Helper visitor to determine whether an AffineExpr is tiled.
 // This is achieved by traversing every AffineDimExpr with position `pos` and
 // checking whether the corresponding `tileSizes[pos]` is non-zero.
 // This also enforces only positive coefficients occur in multiplications.
 //
 // Example:
 //   `d0 + 2 * d1 + d3` is tiled by [0, 0, 0, 2] but not by [0, 0, 2, 0]
 //
 struct TileCheck : public AffineExprVisitor<TileCheck> {
-  TileCheck(ValueRange tileSizes) : tileSizes(tileSizes) {}
+  TileCheck(ArrayRef<OpFoldResult> tileSizes) : tileSizes(tileSizes) {}
 
   void visitDimExpr(AffineDimExpr expr) {
     isTiled |= !isZero(tileSizes[expr.getPosition()]);
   }
   void visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) {
     visit(expr.getLHS());
     visit(expr.getRHS());
     if (expr.getKind() == mlir::AffineExprKind::Mul)
       assert(expr.getRHS().cast<AffineConstantExpr>().getValue() > 0 &&
              "nonpositive multiplying coefficient");
   }
   bool isTiled = false;
-  ValueRange tileSizes;
+  ArrayRef<OpFoldResult> tileSizes;
 };
 
 } // namespace
 
-static bool isTiled(AffineExpr expr, ValueRange tileSizes) {
+static bool isTiled(AffineExpr expr, ArrayRef<OpFoldResult> tileSizes) {
   if (!expr)
     return false;
   TileCheck t(tileSizes);
   t.visit(expr);
   return t.isTiled;
 }
 
 // Checks whether the `map  varies with respect to a non-zero `tileSize`.
-static bool isTiled(AffineMap map, ValueRange tileSizes) {
+static bool isTiled(AffineMap map, ArrayRef<OpFoldResult> tileSizes) {
   if (!map)
     return false;
   for (unsigned r = 0; r < map.getNumResults(); ++r)
     if (isTiled(map.getResult(r), tileSizes))
       return true;
   return false;
 }
 

 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim) {
   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
     return b.createOrFold<memref::DimOp>(loc, source, dim);
   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
     return b.createOrFold<tensor::DimOp>(loc, source, dim);
   llvm_unreachable("Expected MemRefType or TensorType");
 }
 
+OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
+                               int64_t dim) {
+  auto shapedType = source.getType().cast<ShapedType>();
+  if (!shapedType.hasRank() || shapedType.isDynamicDim(dim))
+    return createOrFoldDimOp(b, loc, source, dim);
+  return b.getIndexAttr(shapedType.getDimSize(dim));
+}
+
 /// Given an operation, retrieves the value of each dynamic dimension through
 /// constructing the necessary DimOp operators.
 SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b) {
   SmallVector<Value, 4> dynOperands;
   auto shapedType = val.getType().cast<ShapedType>();
   for (const auto &dim : llvm::enumerate(shapedType.getShape())) {
     if (dim.value() == ShapedType::kDynamicSize)
       dynOperands.push_back(createOrFoldDimOp(b, loc, val, dim.index()));
generateParallelLoopNest(
             linalgOp.getInputAndOutputOperands();
         bodyBuilderFn(b, loc, ivs, operandValuesToUse);
       },
       ivs, distributionMethod);
 
   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
 }
 
-static Value fullyComposeAndAffineApply(OpBuilder &b, Location loc,
-                                        AffineExpr expr, ValueRange operands) {
-  AffineMap map = AffineMap::inferFromExprList({expr}).front();
-  SmallVector<Value> normalizedOperands(operands.begin(), operands.end());
-  mlir::fullyComposeAffineMapAndOperands(&map, &normalizedOperands);
-  canonicalizeMapAndOperands(&map, &normalizedOperands);
-  return b.createOrFold<AffineApplyOp>(loc, map, normalizedOperands);
-}
-
 Value makeTiledShape(OpBuilder &builder, Location loc, Value valueToTile,
-                     ValueRange tileSizes, AffineMap map, ValueRange lbs,
-                     ValueRange ubs, ValueRange subShapeSizes,
+                     ArrayRef<OpFoldResult> tileSizes, AffineMap map,
+                     ArrayRef<OpFoldResult> lbs, ArrayRef<OpFoldResult> ubs,
+                     ArrayRef<OpFoldResult> subShapeSizes,
                      bool omitPartialTileCheck) {
   auto shapedType = valueToTile.getType().dyn_cast<ShapedType>();
   assert(shapedType && "only shaped types can be tiled");
   ArrayRef<int64_t> shape = shapedType.getShape();
   int64_t rank = shapedType.getRank();
 
   // Construct a new subview / extract_slice for the tile.
   SmallVector<OpFoldResult, 4> offsets, sizes, strides;
if (!isTiled(map.getSubMap({r}), tileSizes)) {
       continue;
     }
     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n");
 
     // Tiling creates a new slice at the proper index, the slice step is 1
     // (i.e. the op does not subsample, stepping occurs in the loop).
     auto m = map.getSubMap({r});
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: submap: " << m << "\n");
-    auto offset = applyMapToValues(builder, loc, m, lbs).front();
-    offsets.push_back(getAsOpFoldResult(offset));
-    auto closedIntSize =
-        applyMapToValues(builder, loc, m, subShapeSizes).front();
+    IRRewriter rewriter(builder);
+    OpFoldResult offset = applyMapToValues(rewriter, loc, m, lbs).front();
+    offsets.push_back(offset);
+    OpFoldResult closedIntSize =
+        applyMapToValues(rewriter, loc, m, subShapeSizes).front();
     // Resulting size needs to be made half open interval again.
     AffineExpr s0 = getAffineSymbolExpr(0, builder.getContext());
-    Value size =
-        fullyComposeAndAffineApply(builder, loc, s0 + 1, closedIntSize);
+    OpFoldResult size =
+        makeComposedFoldedAffineApply(rewriter, loc, s0 + 1, closedIntSize);
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: raw size: " << size << "\n");
     LLVM_DEBUG(llvm::dbgs()
                << "makeTiledShape: new offset: " << offset << "\n");
     strides.push_back(builder.getIndexAttr(1));
 
     if (omitPartialTileCheck) {
       // We statically know that the partial/boundary tile condition is
       // unnecessary.
       LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: new size: " << size << "\n");
-      sizes.push_back(getAsOpFoldResult(size));
+      sizes.push_back(size);
       continue;
     }
 
     // The size of the subview / extract_slice should be trimmed to avoid
     // out-of-bounds accesses, unless:
     // a. We statically know the subshape size divides the shape size evenly.
     // b. The subshape size is 1. According to the way the loops are set up,
     //    tensors with "0" dimensions would never be constructed.
     int64_t shapeSize = shape[r];
-    auto sizeCst = size.getDefiningOp<arith::ConstantIndexOp>();
-    auto hasTileSizeOne = sizeCst && sizeCst.value() == 1;
+    Optional<int64_t> sizeCst = getConstantIntValue(size);
+    auto hasTileSizeOne = sizeCst && *sizeCst == 1;
     auto dividesEvenly = sizeCst && !ShapedType::isDynamic(shapeSize) &&
-                         ((shapeSize % sizeCst.value()) == 0);
+                         ((shapeSize % *sizeCst) == 0);
     if (!hasTileSizeOne && !dividesEvenly) {
       LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: shapeSize=" << shapeSize
                               << ", size: " << size
                               << ": make sure in bound with affine.min\n");
 
       AffineExpr dim0, dim1, dim2;
       bindDims(builder.getContext(), dim0, dim1, dim2);
 
       // Get the dimension size for this dimension. We need to first calculate
       // the max index and then plus one. This is important because for
       // convolution ops, we have its input window dimension's affine map of the
       // form `(d0 * s0 + d1)`, where `d0`/`d1 is an output/filter window
       // dimension and `s0` is stride. Directly use the dimension size of
       // output/filer window dimensions will cause incorrect calculation.
       AffineMap minusOneMap =
           AffineMap::inferFromExprList({ArrayRef<AffineExpr>{dim0 - 1}})
               .front();
       AffineMap plusOneMap =
           AffineMap::inferFromExprList({ArrayRef<AffineExpr>{dim0 + 1}})
               .front();
-      auto maxIndices = llvm::to_vector<8>(llvm::map_range(ubs, [&](Value ub) {
-        return makeComposedAffineApply(builder, loc, minusOneMap, {ub})
-            .getResult();
+      SmallVector<OpFoldResult> maxIndices =
+          llvm::to_vector(llvm::map_range(ubs, [&](OpFoldResult ub) {
+            return makeComposedFoldedAffineApply(rewriter, loc, minusOneMap,
+                                                 {ub});
           }));
-      Value maxIndex = applyMapToValues(builder, loc, m, maxIndices).front();
-      Value d = makeComposedAffineApply(builder, loc, plusOneMap, {maxIndex});
+      OpFoldResult maxIndex =
+          applyMapToValues(rewriter, loc, m, maxIndices).front();
+      OpFoldResult d =
+          makeComposedFoldedAffineApply(rewriter, loc, plusOneMap, {maxIndex});
 
       // Compute min(dim - offset, size) to avoid out-of-bounds accesses.
       AffineMap minMap = AffineMap::inferFromExprList(
                              {ArrayRef<AffineExpr>{dim1 - dim2, dim0}})
                              .front();
-      SmallVector<Value, 4> operands{size, d, offset};
-      fullyComposeAffineMapAndOperands(&minMap, &operands);
-      canonicalizeMapAndOperands(&minMap, &operands);
-      size = builder.create<AffineMinOp>(loc, builder.getIndexType(), minMap,
-                                         operands);
+      size =
+          makeComposedFoldedAffineMin(rewriter, loc, minMap, {size, d, offset});
     }
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: new size: " << size << "\n");
-    sizes.push_back(getAsOpFoldResult(size));
+    sizes.push_back(size);
   }
 
   auto *sliceOp = TypeSwitch<ShapedType, Operation *>(shapedType)
                       .Case([&](MemRefType) {
                         return builder.create<memref::SubViewOp>(
                             loc, valueToTile, offsets, sizes, strides);
                       })
                       .Case([&](RankedTensorType) {
                         return makeComposedExtractSliceOp(
                             builder, loc, valueToTile, offsets, sizes, strides);
                       })
                       .Default([](ShapedType) -> Operation * {
                         llvm_unreachable("Unexpected shaped type");
                       });
   return sliceOp->getResult(0);
 }
 
-SmallVector<Value> computeTileOffsets(OpBuilder &b, Location loc,
-                                      ValueRange ivs, ValueRange tileSizes) {
-  SmallVector<Value> offsets;
+SmallVector<OpFoldResult> computeTileOffsets(OpBuilder &b, Location loc,
+                                             ArrayRef<OpFoldResult> ivs,
+                                             ArrayRef<OpFoldResult> tileSizes) {
+  SmallVector<OpFoldResult> offsets;
   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n");
     bool isTiled = !isZero(tileSizes[idx]);
-    offsets.push_back(
-        isTiled ? ivs[idxIvs++]
-                : b.create<arith::ConstantIndexOp>(loc, 0).getResult());
+    offsets.push_back(isTiled ? ivs[idxIvs++] : b.getIndexAttr(0));
     LLVM_DEBUG(llvm::dbgs()
                << "computeTileOffsets: " << offsets.back() << "\n");
   }
   return offsets;
 }
 
-SmallVector<Value> computeTileSizes(OpBuilder &b, Location loc,
-                                    ValueRange tileSizes,
-                                    ArrayRef<Value> sizeBounds) {
-  SmallVector<Value> sizes;
+SmallVector<OpFoldResult> computeTileSizes(OpBuilder &b, Location loc,
+                                           ArrayRef<OpFoldResult> tileSizes,
+                                           ArrayRef<OpFoldResult> sizeBounds) {
+  SmallVector<OpFoldResult> sizes;
   for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx) {
     bool isTiled = !isZero(tileSizes[idx]);
     // Before composing, we need to make range a closed interval.
-    Value size = isTiled ? tileSizes[idx] : sizeBounds[idx];
+    OpFoldResult size = isTiled ? tileSizes[idx] : sizeBounds[idx];
     AffineExpr d0 = getAffineDimExpr(0, b.getContext());
-    sizes.push_back(fullyComposeAndAffineApply(b, loc, d0 - 1, size));
+    IRRewriter rewriter(b);
+    sizes.push_back(makeComposedFoldedAffineApply(rewriter, loc, d0 - 1, size));
     LLVM_DEBUG(llvm::dbgs() << "computeTileSizes: " << sizes.back() << "\n");
   }
   return sizes;
 }
 
 SmallVector<Type> getTensorOutputTypes(LinalgOp op, ValueRange operands) {
   // TODO: use an interface/adaptor to avoid leaking position in
   // `tiledOperands`.
for (OpOperand *opOperand : op.getOutputTensorOperands()) {
     }
     ++resultIdx;
   }
   return tensorResults;
 }
 
 Value materializeOpFoldResult(ImplicitLocOpBuilder &builder,
                               OpFoldResult opFoldResult) {
+  if (!opFoldResult)
+    return nullptr;
+
   if (auto value = opFoldResult.dyn_cast<Value>())
     return value;
   auto attr = opFoldResult.get<Attribute>().cast<IntegerAttr>();
   return builder.create<arith::ConstantIndexOp>(attr.getValue().getSExtValue());
 }
 
 Value materializeOpFoldResult(OpBuilder &builder, Location loc,
                               OpFoldResult opFoldResult) {
   ImplicitLocOpBuilder b(loc, builder);
   return materializeOpFoldResult(b, opFoldResult);
 }
 
-SmallVector<Value, 4> makeTiledShapes(OpBuilder &b, Location loc,
-                                      LinalgOp linalgOp,
-                                      ArrayRef<Value> valuesToTile,
-                                      ValueRange ivs, ValueRange tileSizes,
-                                      ArrayRef<Value> sizeBounds,
-                                      bool omitPartialTileCheck) {
+SmallVector<Value> makeTiledShapes(OpBuilder &b, Location loc,
+                                   LinalgOp linalgOp, ValueRange valuesToTile,
+                                   ArrayRef<OpFoldResult> ivs,
+                                   ArrayRef<OpFoldResult> tileSizes,
+                                   ArrayRef<OpFoldResult> sizeBounds,
   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
                            llvm::make_range(tileSizes.begin(), tileSizes.end()),
-                           [](Value v) { return !isZero(v); })) &&
+                           [](OpFoldResult v) { return !isZero(v); })) &&
          "expected as many ivs as non-zero sizes");
 
   // Construct (potentially temporary) mins and maxes on which to apply maps
   // that define tile subshapes.
-  SmallVector<Value> lbs = computeTileOffsets(b, loc, ivs, tileSizes);
-  SmallVector<Value> subShapeSizes =
+  SmallVector<OpFoldResult> lbs = computeTileOffsets(b, loc, ivs, tileSizes);
+  SmallVector<OpFoldResult> subShapeSizes =
       computeTileSizes(b, loc, tileSizes, sizeBounds);
 
   assert(static_cast<int64_t>(valuesToTile.size()) ==
              linalgOp.getNumInputsAndOutputs() &&
          "expected one value to tile for every operand");
-  SmallVector<Value, 4> tiledShapes;
+  SmallVector<Value> tiledShapes;
   tiledShapes.reserve(valuesToTile.size());
   for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
     Value shapedOp = valuesToTile[opOperand->getOperandNumber()];
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);
     AffineMap map = linalgOp.getTiedIndexingMap(opOperand);
     // Use `opOperand` as is if it is not tiled and not an output tensor. Having
     // an extract/insert slice pair for all output tensors simplifies follow up
     // transformations such as padding and bufferization since the
for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
     tiledShapes.push_back(makeTiledShape(b, loc, shapedOp, tileSizes, map, lbs,
                                          sizeBounds, subShapeSizes,
                                          omitPartialTileCheck));
   }
 
   return tiledShapes;
 }
 
-void offsetIndices(OpBuilder &b, LinalgOp linalgOp, ArrayRef<Value> offsets) {
+void offsetIndices(OpBuilder &b, LinalgOp linalgOp,
+                   ArrayRef<OpFoldResult> offsets) {
   IRRewriter rewriter(b);
   offsetIndices(rewriter, linalgOp, offsets);
 }
 
 void offsetIndices(RewriterBase &b, LinalgOp linalgOp,
-                   ArrayRef<Value> offsets) {
+                   ArrayRef<OpFoldResult> offsets) {
   if (!linalgOp.hasIndexSemantics())
     return;
 
   for (IndexOp indexOp : linalgOp.getBlock()->getOps<IndexOp>()) {
-    if (indexOp.dim() >= offsets.size() || offsets[indexOp.dim()] == nullptr)
+    if (indexOp.dim() >= offsets.size() || !offsets[indexOp.dim()])
       continue;
     OpBuilder::InsertionGuard guard(b);
     b.setInsertionPointAfter(indexOp);
     AffineExpr index, offset;
     bindDims(b.getContext(), index, offset);
-    AffineApplyOp applyOp = makeComposedAffineApply(
+    OpFoldResult applied = makeComposedFoldedAffineApply(
         b, indexOp.getLoc(), index + offset,
-        ValueRange{indexOp.getResult(), offsets[indexOp.dim()]});
-    b.replaceOpWithIf(indexOp, applyOp.getResult(), [&](OpOperand &use) {
-      return use.getOwner() != applyOp;
+        {getAsOpFoldResult(indexOp.getResult()), offsets[indexOp.dim()]});
+    Value materialized = materializeOpFoldResult(b, indexOp.getLoc(), applied);
+    b.replaceOpWithIf(indexOp, materialized, [&](OpOperand &use) {
+      return use.getOwner() != materialized.getDefiningOp();
     });
   }
 }
 
 /// Get the reassociation maps to fold the result of a extract_slice (or source
 /// of a insert_slice) operation with given offsets, and sizes to its
 /// rank-reduced version. This is only done for the cases where the size is 1
 /// and offset is 0. Strictly speaking the offset 0 is not required in general,

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

return llvm::to_vector<4>(
       llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
         return a.cast<IntegerAttr>().getInt();
       }));
 }
 
 /// Given a value, try to extract a constant Attribute. If this fails, return
 /// the original value.
 OpFoldResult getAsOpFoldResult(Value val) {
+  if (!val)
+    return OpFoldResult();
   Attribute attr;
   if (matchPattern(val, m_Constant(&attr)))
     return attr;
   return val;
 }
 
 /// Given an array of values, try to extract a constant Attribute from each
 /// value. If this fails, return the original value.

mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir

func.func @conv_tensors_dynamic(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?x?x?xf32>, %elementwise: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
   }
   return %for0 : tensor<?x?x?x?xf32>
 }
 
 // CHECK: #[[BOUND8_MAP:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 8)>
 // CHECK: #[[BOUND8_MAP_2:.+]] = affine_map<(d0)[s0, s1] -> (-d0 + s1, -d0 + s0, 8)>
 // CHECK: #[[BOUND16_MAP:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 16)>
 // CHECK: #[[X2_MAP:.+]] = affine_map<(d0) -> (d0 * 2)>
-// CHECK: #[[INPUT_BOUND:.+]] = affine_map<(d0, d1)[s0, s1] -> (d1 * -2 + s0 + s1 * 2 - 2, d0 * 2 + s0 - 2)>
+// CHECK: #[[INPUT_BOUND:.+]] = affine_map<(d0, d1)[s0, s1] -> (d0 * -2 + s0 * 2 + s1 - 2, d1 * 2 + s1 - 2)>
 // CHECK: #[[BOUND16_MAP_2:.+]] = affine_map<(d0)[s0, s1] -> (-d0 + s1, -d0 + s0, 16)>
 // CHECK: #[[BOUND4_MAP:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 4)>
 // CHECK: #[[BOUND2_MAP:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)>
 // CHECK: #[[BOUND4_MAP_2:.+]] = affine_map<(d0)[s0, s1] -> (-d0 + s1, -d0 + s0, 4)>
 // CHECK: #[[BOUND2_MAP_2:.+]] = affine_map<(d0, d1)[s0, s1] -> (-d0 + s0, -d1 + s1, 2)>
 
 //      CHECK: func @conv_tensors_dynamic
 // CHECK-SAME: (%[[INPUT]]: tensor<?x?x?x?xf32>, %[[FILTER]]: tensor<?x?x?x?xf32>, %[[ELEM]]: tensor<?x?x?x?xf32>)

 //  CHECK-DAG:   %[[FILL_W:.+]] = tensor.dim %[[FILL]], %[[C2]] : tensor<?x?x?x?xf32>
 
 //      CHECK:   scf.for %[[IV0:.+]] = %{{.+}} to %[[ELEM_N]] step %{{.+}} iter_args(%{{.+}} = %[[FILL]])
 // CHECK-NEXT:     %[[SIZE_ELEM_N:.+]] = affine.min #[[BOUND8_MAP]](%[[IV0]])[%[[ELEM_N]]]
 // CHECK-NEXT:     %[[SIZE_INPUT_N:.+]] = affine.min #[[BOUND8_MAP_2]](%[[IV0]])[%[[INPUT_N]], %[[ELEM_N]]]
 // CHECK-NEXT:     scf.for %[[IV1:.+]] = %{{.+}} to %[[ELEM_OH]]
 // CHECK-NEXT:       %[[SIZE_ELEM_OH:.+]] = affine.min #[[BOUND16_MAP]](%[[IV1]])[%[[ELEM_OH]]]
 // CHECK-NEXT:       %[[OFFSET_OH:.+]] = affine.apply #[[X2_MAP]](%[[IV1]])
-// CHECK-NEXT:       %[[SIZE_INPUT_H:.+]] = affine.min #[[INPUT_BOUND]](%[[SIZE_ELEM_OH]], %[[IV1]])[%[[FILTER_H]], %[[FILL_H]]]
+// CHECK-NEXT:       %[[SIZE_INPUT_H:.+]] = affine.min #[[INPUT_BOUND]](%[[IV1]], %[[SIZE_ELEM_OH]])[%[[FILL_H]], %[[FILTER_H]]]
 // CHECK-NEXT:       %[[SIZE_ELEM_OH_2:.+]] = affine.min #[[BOUND16_MAP_2]](%[[IV1]])[%[[FILL_H]], %[[ELEM_OH]]]
 // CHECK-NEXT:       scf.for %[[IV2:.+]] = %{{.+}} to %[[ELEM_OW]]
 // CHECK-NEXT:         %[[SIZE_ELEM_OW:.+]] = affine.min #[[BOUND4_MAP]](%[[IV2]])[%[[ELEM_OW]]]
 // CHECK-NEXT:         %[[SIZE_ELEM_OC:.+]] = affine.min #[[BOUND2_MAP]](%[[IV2]])[%[[ELEM_OC]]]
 // CHECK-NEXT:         %[[OFFSET_OW:.+]] = affine.apply #[[X2_MAP]](%[[IV2]])
-// CHECK-NEXT:         %[[SIZE_INPUT_W:.+]] = affine.min #[[INPUT_BOUND]](%[[SIZE_ELEM_OW]], %[[IV2]])[%[[FILTER_W]], %[[FILL_W]]]
+// CHECK-NEXT:         %[[SIZE_INPUT_W:.+]] = affine.min #[[INPUT_BOUND]](%[[IV2]], %[[SIZE_ELEM_OW]])[%[[FILL_W]], %[[FILTER_W]]]
 // CHECK-NEXT:         %[[ST_INPUT:.+]] = tensor.extract_slice %[[INPUT]][%[[IV0]], %[[OFFSET_OH]], %[[OFFSET_OW]], 0]
 // CHECK-SAME:               [%[[SIZE_INPUT_N]], %[[SIZE_INPUT_H]], %[[SIZE_INPUT_W]], %[[INPUT_C]]]
 // CHECK-NEXT:         %[[SIZE_ELEM_OW_2:.+]] = affine.min #[[BOUND4_MAP_2]](%[[IV2]])[%[[FILL_W]], %[[ELEM_OW]]]
 // CHECK-NEXT:         scf.for %[[IV3:.+]] = %{{.+}} to %[[ELEM_OC]] step %{{.+}} iter_args(%[[ARG:[a-z0-9]+]]
 // CHECK-NEXT:           %[[ST_ELEM:.+]] = tensor.extract_slice %[[ELEM]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]]
 // CHECK-SAME:                 [%[[SIZE_ELEM_N]], %[[SIZE_ELEM_OH]], %[[SIZE_ELEM_OW]], %[[SIZE_ELEM_OC]]]
 // CHECK-NEXT:           %[[ST_ARG:.+]] = tensor.extract_slice %[[ARG]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]]
 // CHECK-SAME:                 [%[[SIZE_ELEM_N]], %[[SIZE_ELEM_OH]], %[[SIZE_ELEM_OW]], %[[SIZE_ELEM_OC]]]

mlir/test/Dialect/Linalg/tile-conv.mlir

 // RUN: mlir-opt %s -linalg-tile="tile-sizes=2,3" | FileCheck %s
 
-//  CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 - 1, s0 + 1)>
-//  CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 - 1, s0 + 2)>
+//  CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 - 1, s1 + 1)>
+//  CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 - 1, s1 + 2)>
 //  CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)>
 //  CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 3)>
 
 func.func @conv(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>, %arg2 : memref<?x?xf32>) {
   linalg.conv_2d ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>) outs(%arg2 : memref<?x?xf32>)
   return
 }
 
 //       CHECK: func @conv
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?xf32>
 //  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?xf32>
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
 //   CHECK-DAG:   %[[T0:.*]] = memref.dim %[[ARG1]], %[[C0]]
 //   CHECK-DAG:   %[[T1:.*]] = memref.dim %[[ARG1]], %[[C1]]
 //   CHECK-DAG:   %[[T2:.*]] = memref.dim %[[ARG2]], %[[C0]]
 //   CHECK-DAG:   %[[T3:.*]] = memref.dim %[[ARG2]], %[[C1]]
 //       CHECK:   scf.for %[[ARG3:.*]] = %[[C0]] to %[[T2]] step %[[C2]]
 //       CHECK:     scf.for %[[ARG4:.*]] = %[[C0]] to %[[T3]] step %[[C3]]
-//       CHECK:       %[[T4:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T0]], %[[T2]]]
-//       CHECK:       %[[T5:.*]] = affine.min #[[MAP1]](%[[ARG4]])[%[[T1]], %[[T3]]]
+//       CHECK:       %[[T4:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T2]], %[[T0]]]
+//       CHECK:       %[[T5:.*]] = affine.min #[[MAP1]](%[[ARG4]])[%[[T3]], %[[T1]]]
 //       CHECK:       %[[SV1:.*]] = memref.subview %[[ARG0]][%[[ARG3]], %[[ARG4]]] [%[[T4]], %[[T5]]]
 //       CHECK:       %[[T6:.*]] = affine.min #[[MAP2]](%[[ARG3]])[%[[T2]]
 //       CHECK:       %[[T7:.*]] = affine.min #[[MAP3]](%[[ARG4]])[%[[T3]]]
 //       CHECK:       %[[SV2:.*]] = memref.subview %[[ARG2]][%[[ARG3]], %[[ARG4]]] [%[[T6]], %[[T7]]]
 //       CHECK:       linalg.conv_2d
 //  CHECK-SAME:         ins(%[[SV1]], %[[ARG1]]
 //  CHECK-SAME:         outs(%[[SV2]]

mlir/test/Dialect/Linalg/transform-op-fuse.mlir

 // -----
 
 // CHECK-LABEL: func.func @interchange_reduction
 //  CHECK-SAME: (%[[INPUT:.+]]: tensor<12x7x25xf32>)
 func.func @interchange_reduction(%input: tensor<12x7x25xf32>) -> tensor<12x25xf32> {
   %five = arith.constant 5.0 : f32
   %init = linalg.init_tensor [12, 25] : tensor<12x25xf32>
 
-//   CHECK-DAG:  %[[C4:.+]] = arith.constant 4 : index
-//   CHECK-DAG:  %[[C5:.+]] = arith.constant 5 : index
-//   CHECK-DAG:  %[[C7:.+]] = arith.constant 7 : index
-
-//       CHECK: %[[INIT:.+]] = linalg.init_tensor [12, 25]
+//       CHECK: %[[INIT:.+]] = linalg.init_tensor [12, 25]
+//   CHECK-DAG: %[[C5:.+]] = arith.constant 5 : index
+//   CHECK-DAG: %[[C7:.+]] = arith.constant 7 : index
 //       CHECK: scf.for %[[IV0:.+]] = %{{.+}} to %{{.+}} step %[[C5]] iter_args(%[[FOR_ARG0:.+]] = %[[INIT]])
 //       CHECK:   scf.for %[[IV1:.+]] = %{{.+}} to %{{.+}} step %[[C7]] iter_args(%[[FOR_ARG1:.+]] = %[[FOR_ARG0]])
 //       CHECK:     %[[OUT_SLICE0:.+]] = tensor.extract_slice %[[FOR_ARG1]][%[[IV0]], %[[IV1]]]
 //       CHECK:     %[[FILL:.+]] = linalg.fill {{.+}} outs(%[[OUT_SLICE0]] : tensor<?x?xf32>)
+//       CHECK:     %[[C4:.+]] = arith.constant 4 : index
 //       CHECK:     scf.for %[[IV2:.+]] = %{{.+}} to %{{.+}} step %[[C4]] iter_args(%[[FOR_ARG2:.+]] = %[[FILL]])
 //       CHECK:       %[[IN_SLICE:.+]] = tensor.extract_slice %[[INPUT]]
 //       CHECK:       %[[OUT_SLICE2:.+]] = tensor.extract_slice %[[FOR_ARG2]][0, 0]
 //       CHECK:       linalg.generic {{.+}} ins(%[[IN_SLICE]] : tensor<?x?x?xf32>) outs(%[[OUT_SLICE2]] : tensor<?x?xf32>)
 
   %fill = linalg.fill ins(%five : f32) outs(%init : tensor<12x25xf32>) -> tensor<12x25xf32>
   %0 = linalg.generic {
     indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d2)>],