[mlir] Make ViewLikeInterface Range work with attributes

[ftynse]

Summary

While most of methods in ViewLikeInterface accept an OpFoldResult for
the offset/size/stride that may be static, represented as Attribute,
or dynamic, represented as Value, the Range abstraction only
accepted Values. This can often lead to known-constant
offset/size/strides being materialized into constant operations and
hinder further constant propagation without explicitly running the
constant folding pass. This often leads to a more complicated than
necessary addressing code being emitted. Switch Range to use
OpFoldResult. Code that uses Range currently keeps materializing the
constants to minimize the effect of this change on the IR. Further
commits will make use of this.

[ftynse]

Rebase.

[mravishankar]

Nice! Always wanted this, but it also means changes to downstream usages, but worth it for me!

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

1413

Could we check for this condition and fail early. So move it all the way to the top. Sorry, I know this is not part of the change really, but then you wont need the #ifndef NDEBUG eyesore...

[ftynse]

Address review.

mlir/include/mlir/Dialect/Linalg/IR/Linalg.h

 #include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 namespace mlir {
 namespace linalg {
 
 class LinalgOp;
 
-// TOFO: allow an extra ValueRange to specify an indexing and allow
-// non-hyperrectangular shapes.
-using LoopRangeBuilder =
-    std::function<SmallVector<Range, 4>(ImplicitLocOpBuilder)>;
-
-/// Provide a very simple inference procedure to build the loop ranges from the
-/// op and its operands. This only works with permutation affine maps and
-/// patterns of the form `(m, n)[s] -> (m + n - s floordiv 2)`.
-/// A more advanced Tensor-Comprehension like inference is possible but has
-/// proven to be ambiguous in unfavorable case.
-/// As a consequence, we relax the default behavior very conservatively and
-/// provide an op-specified hook so that Linalg ops may override the behavior.
-LoopRangeBuilder defaultLoopRangesBuilder(LinalgOp op);
-
 /// Returns the name mangled library call name to disambiguate between different
 /// overloads at the C level. The name mangling scheme is basic and uses MLIR
 /// type names:
 ///   1. form a string which is the concatenation of the linalg op name with all
 ///      the operand type names, separate by underscores;
 ///   2. drop the `linalg.` prefix, and the `<`, `>`, `?` symbols from the type.
 /// Assumes `op` is a LinalgOp.
 ///

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

 SmallVector<Value> insertSlicesBack(OpBuilder &builder, Location loc,
                                     LinalgOp op, ValueRange operands,
                                     ValueRange results);
 
 /// Turns an OpFoldResult into a value, creating an index-typed constant if
 /// necessary.
 Value materializeOpFoldResult(ImplicitLocOpBuilder &builder,
                               OpFoldResult opFoldResult);
+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,

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpImplementation.h"
 
 namespace mlir {
 /// Auxiliary range data structure to unpack the offset, size and stride
 /// operands into a list of triples. Such a list can be more convenient to
 /// manipulate.
 struct Range {
-  Value offset;
-  Value size;
-  Value stride;
+  OpFoldResult offset;
+  OpFoldResult size;
+  OpFoldResult stride;
 };
 
 class OffsetSizeAndStrideOpInterface;
 
 /// Return a vector of all the static or dynamic offsets of the op from provided
 /// external static and dynamic offsets.
 SmallVector<OpFoldResult, 4> getMixedOffsets(OffsetSizeAndStrideOpInterface op,
                                              ArrayAttr staticOffsets,

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

static FailureOr<SmallVector<Value>> collapseGenericOpIterationDims(
 
   CollapsingInfo collapsingInfo;
   if (failed(collapsingInfo.initialize(genericOp.getNumLoops(),
                                        foldedIterationDims))) {
     return rewriter.notifyMatchFailure(
         genericOp, "illegal to collapse specified dimensions");
   }
 
+  // Bail on non-canonical ranges.
+  SmallVector<Range> loopRanges =
+      cast<LinalgOp>(genericOp.getOperation())
+          .createLoopRanges(rewriter, genericOp.getLoc());
+  auto opFoldIsConstantValue = [](OpFoldResult ofr, int64_t value) {
+    if (auto attr = ofr.dyn_cast<Attribute>())
+      return attr.cast<IntegerAttr>().getInt() == value;
+    llvm::APInt actual;
+    return matchPattern(ofr.get<Value>(), m_ConstantInt(&actual)) &&
+           actual.getSExtValue() == value;
+  };
+  if (!llvm::all_of(loopRanges, [&](Range range) {
+        return opFoldIsConstantValue(range.offset, 0) &&
+               opFoldIsConstantValue(range.stride, 1);
+      })) {
+    return rewriter.notifyMatchFailure(
+        genericOp,
+        "expected all loop ranges to have zero start and unit stride");
+  }
+
   // Get the iterator types for the operand.
   SmallVector<StringRef> iteratorTypes = getCollapsedOpIteratorTypes(
       genericOp.iterator_types().getValue(), collapsingInfo);
 
   // Get the indexing maps.
   auto indexingMaps = llvm::to_vector(
       llvm::map_range(genericOp.getIndexingMapsArray(), [&](AffineMap map) {
         return getCollapsedOpIndexingMap(map, collapsingInfo);
static FailureOr<SmallVector<Value>> collapseGenericOpIterationDims(
   Block *collapsedOpBlock = &collapsedGenericOp->getRegion(0).front();
   rewriter.mergeBlocks(origOpBlock, collapsedOpBlock,
                        collapsedOpBlock->getArguments());
 
   if (collapsedGenericOp.hasIndexSemantics()) {
     // Collect the loop range of the generic op.
     OpBuilder::InsertionGuard g(rewriter);
     rewriter.setInsertionPoint(collapsedGenericOp);
-    SmallVector<Range> loopRanges =
-        cast<LinalgOp>(genericOp.getOperation())
-            .createLoopRanges(rewriter, genericOp.getLoc());
-    assert(llvm::all_of(loopRanges,
-                        [](Range range) {
-                          return matchPattern(range.offset, m_Zero()) &&
-                                 matchPattern(range.stride, m_One());
-                        }) &&
-           "expected all loop ranges to have zero start and unit stride");
-    SmallVector<Value> loopBound = llvm::to_vector(
-        llvm::map_range(loopRanges, [](Range range) { return range.size; }));
+    SmallVector<Value> loopBound =
+        llvm::to_vector(llvm::map_range(loopRanges, [&](Range range) {
+          return materializeOpFoldResult(rewriter, loc, range.size);
+        }));
     generateCollapsedIndexingRegion(loc,
                                     &collapsedGenericOp->getRegion(0).front(),
                                     collapsingInfo, loopBound, rewriter);
   }
 
   // Insert expanding reshape for the result to get back the original result
   // type.
   SmallVector<Value> results;

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

 /// 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;
   Location loc = producer.getLoc();
   auto zero = b.create<arith::ConstantIndexOp>(loc, 0);
-  auto one = b.create<arith::ConstantIndexOp>(loc, 1);
 
   for (unsigned i = 0, e = producer.getNumLoops(); i < e; ++i) {
     auto shapeDim = getShapeDefiningLoopRange(producer, i);
     Value dim = createOrFoldDimOp(b, loc, shapeDim.shape, shapeDim.dimension);
     sizeBounds.push_back(dim);
     auto it = fusedLoopsAndRanges.find(i);
     if (it != fusedLoopsAndRanges.end()) {
-      ivs.push_back(it->second.offset);
-      tileSizes.push_back(it->second.size);
+      ivs.push_back(materializeOpFoldResult(b, loc, it->second.offset));
+      tileSizes.push_back(materializeOpFoldResult(b, loc, 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);
-      loopRanges.push_back(Range{zero, dim, one});
+      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());
 
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 range.offset; });
+  llvm::transform(loopRanges, std::back_inserter(allIvs), [&](Range range) {
+    return materializeOpFoldResult(b, loc, 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,
 
   // Get the producer.
   LinalgOp producerOp = producerResult.getOwner();
   Location loc = producerOp.getLoc();
 
   // Obtain the `producerOp` loop bounds and the `sliceOp` ranges.
   SmallVector<Value> producerLoopBounds;
   llvm::transform(producerOp.createLoopRanges(b, loc),
-                  std::back_inserter(producerLoopBounds),
-                  [](Range range) { return range.size; });
+                  std::back_inserter(producerLoopBounds), [&](Range range) {
+                    return materializeOpFoldResult(b, loc, 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);
   for (auto it : zip(tiledSliceDimIndices, tiledProducerLoopIndices)) {
     int64_t tiledSliceDim = std::get<0>(it);
     int64_t tiledProducerLoop = std::get<1>(it);
-    tileIvs[tiledProducerLoop] = sliceOpRanges[tiledSliceDim].offset;
-    tileSizes[tiledProducerLoop] = sliceOpRanges[tiledSliceDim].size;
+    tileIvs[tiledProducerLoop] =
+        materializeOpFoldResult(b, loc, sliceOpRanges[tiledSliceDim].offset);
+    tileSizes[tiledProducerLoop] =
+        materializeOpFoldResult(b, loc, sliceOpRanges[tiledSliceDim].size);
     allIvs[tiledProducerLoop] = tileIvs[tiledProducerLoop];
   }
   erase_value(tileIvs, nullptr);
   SmallVector<Value> tiledOperands = producerOp.getInputAndOutputOperands();
   tiledOperands = makeTiledShapes(b, loc, producerOp, tiledOperands, tileIvs,
                                   tileSizes, producerLoopBounds,
                                   /**omitPartialTileCheck=*/false);
 

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

FailureOr<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer(
   fullSizes.reserve(rank);
   partialSizes.reserve(rank);
   llvm::SmallBitVector droppedDims = subView.getDroppedDims();
   int64_t resultDimIdx = 0;
   for (const auto &en : llvm::enumerate(subView.getOrCreateRanges(b, loc))) {
     if (droppedDims[en.index()])
       continue;
     auto rangeValue = en.value();
-    // Try to extract a tight constant.
+    // Try to extract a tight constant. If the size is known statically, no need
+    // to look for the bound.
     LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n");
+    Value size;
+    if (auto attr = rangeValue.size.dyn_cast<Attribute>()) {
+      size = materializeOpFoldResult(b, loc, rangeValue.size);
+    } else {
+      Value materializedSize = materializeOpFoldResult(b, loc, rangeValue.size);
       FailureOr<int64_t> upperBound =
-        getConstantUpperBoundForIndex(rangeValue.size);
-    Value size =
-        failed(upperBound)
-            ? rangeValue.size
-            : b.create<arith::ConstantIndexOp>(loc, upperBound.value());
+          getConstantUpperBoundForIndex(materializedSize);
+      size = failed(upperBound)
+                 ? materializedSize
+                 : b.create<arith::ConstantIndexOp>(loc, upperBound.getValue());
+    }
     LLVM_DEBUG(llvm::dbgs() << "Extracted tightest: " << size << "\n");
     fullSizes.push_back(size);
     partialSizes.push_back(
         b.createOrFold<memref::DimOp>(loc, subView, resultDimIdx++));
   }
   SmallVector<int64_t, 4> dynSizes(fullSizes.size(), -1);
   // If a callback is not specified, then use the default implementation for
   // allocating the promoted buffer.

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

linalg::splitOp(RewriterBase &rewriter, TilingInterface op, unsigned dimension,
                 OpFoldResult splitPoint) {
   // Compute the iteration space.
   SmallVector<Range> iterationSpace = op.getIterationDomain(rewriter);
 
   // Bail out on dimension overflow.
   if (dimension >= iterationSpace.size())
     return std::make_pair(op, TilingInterface());
 
-  SmallVector<OpFoldResult> offsets =
-      getAsOpFoldResult(llvm::to_vector(llvm::map_range(
-          iterationSpace, [](const Range &range) { return range.offset; })));
-  SmallVector<OpFoldResult> sizes =
-      getAsOpFoldResult(llvm::to_vector(llvm::map_range(
-          iterationSpace, [](const Range &range) { return range.size; })));
+  SmallVector<OpFoldResult> offsets = llvm::to_vector(llvm::map_range(
+      iterationSpace, [](const Range &range) { return range.offset; }));
+  SmallVector<OpFoldResult> sizes = llvm::to_vector(llvm::map_range(
+      iterationSpace, [](const Range &range) { return range.size; }));
 
   // Adjust the split point so that it doesn't overflow the size.
   AffineExpr d0, d1, d2;
   bindDims(rewriter.getContext(), d0, d1, d2);
   OpFoldResult minSplitPoint = makeComposedFoldedAffineMin(
       rewriter, op.getLoc(),
       AffineMap::inferFromExprList(ArrayRef<AffineExpr>{d0, d1 + d2}).front(),
       {splitPoint, offsets[dimension], sizes[dimension]});
if (attr.cast<IntegerAttr>().getValue().isZero())
       return {op, TilingInterface()};
   }
 
   // Create the first part.
   SmallVector<Value> firstResults;
   TilingInterface firstPart = createSplitPart(
       rewriter, op.getLoc(), op, offsets, sizes,
       op.getDestinationOperands(rewriter), dimension, minSplitPoint,
-      getAsOpFoldResult(iterationSpace[dimension].offset), firstResults);
+      iterationSpace[dimension].offset, firstResults);
 
   // Need to pretend that the original op now takes as operands firstResults,
   // otherwise tiling interface implementation will take the wrong value to
   // produce data tiles.
   rewriter.updateRootInPlace(op, [&]() {
     unsigned numTotalOperands = op->getNumOperands();
     unsigned numOutputOperands = firstResults.size();
     op->setOperands(numTotalOperands - numOutputOperands, numOutputOperands,

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

if (isZero(tileSizes[idx - zerosCount])) {
       continue;
     }
     loopIndexToRangeIndex[idx] = idx - zerosCount;
   }
 
   // Create a new range with the applied tile sizes.
   SmallVector<Range, 4> res;
   for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx)
-    res.push_back(Range{b.create<arith::ConstantIndexOp>(loc, 0),
-                        shapeSizes[idx], tileSizes[idx]});
+    res.push_back(Range{b.getIndexAttr(0), shapeSizes[idx], tileSizes[idx]});
   return std::make_tuple(res, loopIndexToRangeIndex);
 }
 
 void mlir::linalg::transformIndexOps(
     RewriterBase &b, LinalgOp op, SmallVectorImpl<Value> &ivs,
     const LoopIndexToRangeIndexMap &loopIndexToRangeIndex) {
   SmallVector<Value> allIvs(op.getNumLoops(), nullptr);
   for (auto &en : enumerate(allIvs)) {
static LogicalResult tilePadOp(RewriterBase &builder, tensor::PadOp op,
   Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
   tileSizes.append(rank - tileSizes.size(), zero);
   // 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;
   for (int64_t i = 0; i < rank; ++i) {
-    allDims.push_back(ranges[i].size);
+    Value materializedSize =
+        materializeOpFoldResult(builder, loc, ranges[i].size);
+    allDims.push_back(materializedSize);
     if (!isZero(tileSizes[i])) {
-      lbs.push_back(ranges[i].offset);
-      dims.push_back(ranges[i].size);
+      lbs.push_back(materializeOpFoldResult(builder, loc, ranges[i].offset));
+      dims.push_back(materializedSize);
       steps.push_back(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),

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

 
 /// Explicit instantiation of loop nest generator for different loop types.
 template struct mlir::linalg::GenerateLoopNest<scf::ForOp>;
 template struct mlir::linalg::GenerateLoopNest<scf::ParallelOp>;
 template struct mlir::linalg::GenerateLoopNest<AffineForOp>;
 
 /// Given a list of subview ranges, extract individual values for lower, upper
 /// bounds and steps and put them into the corresponding vectors.
-static void unpackRanges(ArrayRef<Range> ranges, SmallVectorImpl<Value> &lbs,
+static void unpackRanges(OpBuilder &builder, Location loc,
+                         ArrayRef<Range> ranges, SmallVectorImpl<Value> &lbs,
                          SmallVectorImpl<Value> &ubs,
                          SmallVectorImpl<Value> &steps) {
   for (Range range : ranges) {
-    lbs.emplace_back(range.offset);
-    ubs.emplace_back(range.size);
-    steps.emplace_back(range.stride);
+    lbs.emplace_back(materializeOpFoldResult(builder, loc, range.offset));
+    ubs.emplace_back(materializeOpFoldResult(builder, loc, range.size));
+    steps.emplace_back(materializeOpFoldResult(builder, loc, range.stride));
   }
 }
 
 namespace mlir {
 namespace linalg {
 
 bool allIndexingsAreProjectedPermutation(LinalgOp op) {
   return llvm::all_of(op.getIndexingMapsArray(), [](AffineMap m) {
if (distributionOptions) {
     // Get their distribution schemes.
     distributionMethod = distributionOptions->distributionMethod;
     if (distributionMethod.size() < parallelLoopRanges.size())
       parallelLoopRanges.resize(distributionMethod.size());
     procInfo = distributionOptions->procInfo(b, loc, parallelLoopRanges);
   }
 
   SmallVector<Value, 4> lbs, ubs, steps;
-  unpackRanges(loopRanges, lbs, ubs, steps);
+  unpackRanges(b, loc, loopRanges, lbs, ubs, steps);
   LoopNest loopNest = mlir::scf::buildLoopNest(
       b, loc, lbs, ubs, steps, iterArgInitValues,
       [&](OpBuilder &b, Location loc, ValueRange ivs, ValueRange iterArgs) {
         assert(iterArgs.size() == linalgOp.getOutputTensorOperands().size() &&
                "expect the number of output tensors and iter args to match");
         SmallVector<Value> operandValuesToUse =
             linalgOp.getInputAndOutputOperands();
         if (!iterArgs.empty()) {
void GenerateLoopNest<AffineForOp>::doit(
     ArrayRef<Attribute> iteratorTypes,
     function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
                                   ValueRange)>
         bodyBuilderFn,
     Optional<LinalgLoopDistributionOptions>, ArrayRef<StringRef>) {
   SmallVector<Value> iterArgInitValues = linalgOp.getOutputTensorOperands();
   assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
   SmallVector<Value, 4> lbs, ubs, steps;
-  unpackRanges(loopRanges, lbs, ubs, steps);
+  unpackRanges(b, loc, loopRanges, lbs, ubs, steps);
 
   // Affine loops require constant steps.
   SmallVector<int64_t, 4> constantSteps;
   constantSteps.reserve(steps.size());
   for (Value v : steps) {
     auto op = v.getDefiningOp<arith::ConstantIndexOp>();
     assert(op && "Affine loops require constant steps");
     constantSteps.push_back(op.value());
void GenerateLoopNest<scf::ParallelOp>::doit(
   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
   unsigned numLoops = iteratorTypes.size();
   ivs.reserve(numLoops);
   lbsStorage.reserve(numLoops);
   ubsStorage.reserve(numLoops);
   stepsStorage.reserve(numLoops);
 
   // Get the loop lb, ub, and step.
-  unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
+  unpackRanges(b, loc, loopRanges, lbsStorage, ubsStorage, stepsStorage);
 
   // Modify the lb, ub, and step based on the distribution options.
   SmallVector<DistributionMethod, 0> distributionMethod;
   if (distributionOptions) {
     auto &options = *distributionOptions;
     distributionMethod.assign(distributionOptions->distributionMethod.begin(),
                               distributionOptions->distributionMethod.end());
     SmallVector<Range, 2> parallelLoopRanges;

 Value materializeOpFoldResult(ImplicitLocOpBuilder &builder,
                               OpFoldResult opFoldResult) {
   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) {
   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
                            llvm::make_range(tileSizes.begin(), tileSizes.end()),

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

 //===- Tiling.cpp - Implementation of tiling using TilingInterface -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the tiling using TilingInterface.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.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/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/TilingInterface.h"
 #include "llvm/Support/Debug.h"
 
generateTileLoopNest(OpBuilder &builder, Location loc,
   // `tileSize` and `ub - iv`, where `iv` is the induction variable
   // of the tiled loop.
   AffineExpr s0, s1, d0;
   bindDims(builder.getContext(), d0);
   bindSymbols(builder.getContext(), s0, s1);
   AffineMap minMap = AffineMap::get(1, 2, {s0, s1 - d0}, builder.getContext());
 
   for (auto loopRange : llvm::enumerate(loopRanges)) {
+    Value offset =
+        getValueOrCreateConstantIndexOp(builder, loc, loopRange.value().offset);
+    Value size =
+        getValueOrCreateConstantIndexOp(builder, loc, loopRange.value().size);
     // No loops if tile size is zero. Set offset and size to the loop
     // offset and size.
     if (matchPattern(tileSizeVals[loopRange.index()], m_Zero())) {
-      offsets[loopRange.index()] = loopRange.value().offset;
-      sizes[loopRange.index()] = loopRange.value().size;
+      offsets[loopRange.index()] = offset;
+      sizes[loopRange.index()] = size;
       continue;
     }
 
     auto loop = builder.create<scf::ForOp>(
-        loc, loopRange.value().offset, loopRange.value().size,
-        tileSizeVals[loopRange.index()], ValueRange{},
+        loc, offset, size, tileSizeVals[loopRange.index()], ValueRange{},
         [&](OpBuilder &bodyBuilder, Location bodyLoc, Value iv,
             ValueRange /*iterArgs*/) {
           Value boundedTileSize = builder.create<AffineMinOp>(
               bodyLoc, minMap,
-              ValueRange{iv, tileSizeVals[loopRange.index()],
-                         loopRange.value().size});
+              ValueRange{iv, tileSizeVals[loopRange.index()], size});
           sizes[loopRange.index()] = boundedTileSize;
           builder.create<scf::YieldOp>(loc);
         });
     offsets[loopRange.index()] = loop.getInductionVar();
     loops.push_back(loop);
     builder.setInsertionPoint(loop.getBody()->getTerminator());
   }
   return loops;