[mlir][Tensor] Use folded evaluators in tiling implementation of `tensor.pad`.

[mravishankar]

Fix bug.

[springerm]

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

62–63

Maybe this should be an assertion?

mlir/include/mlir/Dialect/Tensor/Utils/Utils.h

 PadOp createPadHighOp(RankedTensorType type, Value source, Value pad,
                       bool nofold, Location loc, OpBuilder &builder);
 
 // Creates dim ops for each dynamic dimension of the ranked tensor argument and
 // returns these as values.
 SmallVector<Value> createDynamicDimValues(OpBuilder &b, Location loc,
                                           Value rankedTensor);
 
+// Returns the tensor extent along dimension `dim` if `rankedTensor` is of
+// `RankedTensorType`. Returns `failure()` otherwise.
+FailureOr<OpFoldResult> createDimValue(OpBuilder &b, Location loc,
+                                       Value rankedTensor, int64_t dim);
+
 // Creates dim ops or constant ops for each dimension of the ranked tensor
 // argument and returns these as values.
 SmallVector<OpFoldResult> createDimValues(OpBuilder &b, Location loc,
                                           Value rankedTensor);
 
 /// Returns the transposed `rankedTensorType` if `transposeVector` is non-empty.
 /// Fail if `transposeVector` is not a permutation matching the tensor rank.
 FailureOr<RankedTensorType>
 computeTransposedType(RankedTensorType rankedTensorType,
                       ArrayRef<int64_t> transposeVector);
 
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_UTILS_UTILS_H_

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

Operation *tensor::bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp,
 
   // Helper variables and functions for various arithmetic operations. These
   // are used extensively for computing new offset/length and padding values.
   Location loc = padOp->getLoc();
   AffineExpr dim0, dim1;
   bindDims(b.getContext(), dim0, dim1);
   // Add two integers.
   auto addMap = AffineMap::get(2, 0, {dim0 + dim1});
-  auto add = [&](Value v1, Value v2) {
-    return b.createOrFold<AffineApplyOp>(loc, addMap, ValueRange{v1, v2});
+  auto add = [&](OpFoldResult v1, OpFoldResult v2) {
+    return makeComposedFoldedAffineApply(b, loc, addMap, {v1, v2});
   };
   // Subtract two integers.
   auto subMap = AffineMap::get(2, 0, {dim0 - dim1});
-  auto sub = [&](Value v1, Value v2) {
-    return b.createOrFold<AffineApplyOp>(loc, subMap, ValueRange{v1, v2});
+  auto sub = [&](OpFoldResult v1, OpFoldResult v2) {
+    return makeComposedFoldedAffineApply(b, loc, subMap, {v1, v2});
   };
   // Take the minimum of two integers.
   auto idMap = AffineMap::getMultiDimIdentityMap(2, b.getContext());
-  auto min = [&](Value v1, Value v2) {
-    return b.createOrFold<AffineMinOp>(loc, idMap, ValueRange{v1, v2});
+  auto min = [&](OpFoldResult v1, OpFoldResult v2) {
+    return makeComposedFoldedAffineMin(b, loc, idMap, {v1, v2});
   };
   // Take the maximum of two integers.
-  auto max = [&](Value v1, Value v2) {
-    return b.createOrFold<AffineMaxOp>(loc, idMap, ValueRange{v1, v2});
+  auto max = [&](OpFoldResult v1, OpFoldResult v2) {
+    return makeComposedFoldedAffineMax(b, loc, idMap, {v1, v2});
   };
   // Zero index-typed integer.
-  auto zero = b.create<arith::ConstantIndexOp>(loc, 0);
+  OpFoldResult zero = b.getIndexAttr(0);
 
   // Helper function for filling static/dynamic low/high padding indices
   // vectors of PadOp.
   auto appendIndex = [&](Value val, SmallVector<Value> &dynIndices,
                          SmallVector<int64_t> &staticIndices) {
     if (auto constInt = getConstantIntValue(val)) {
       staticIndices.push_back(*constInt);
     } else {
       staticIndices.push_back(ShapedType::kDynamic);
       dynIndices.push_back(val);
     }
   };
 
   // Compute new offsets, lengths, low padding, high padding.
   SmallVector<OpFoldResult> newOffsets, newLengths, newStrides;
-  SmallVector<Value> newLows, newHighs;
-  SmallVector<int64_t> staticNewLows, staticNewHighs;
+  SmallVector<OpFoldResult> newLows, newHighs;
   // Set to true if the original data source is not read at all.
   bool hasZeroLen = false;
   // Same as hasZeroLen, but for dynamic dimension sizes. This condition
   // is true if the original data source turns out to be unused at runtime.
   Value dynHasZeroLenCond;
 
   int64_t rank = padOp.getSourceType().getRank();
   for (unsigned dim = 0; dim < rank; ++dim) {
-    auto low =
-        getValueOrCreateConstantIndexOp(b, loc, padOp.getMixedLowPad()[dim]);
-    bool hasLowPad = getConstantIntValue(low) != static_cast<int64_t>(0);
-    auto high =
-        getValueOrCreateConstantIndexOp(b, loc, padOp.getMixedHighPad()[dim]);
-    bool hasHighPad = getConstantIntValue(high) != static_cast<int64_t>(0);
-    auto offset = getValueOrCreateConstantIndexOp(b, loc, offsets[dim]);
-    auto length = getValueOrCreateConstantIndexOp(b, loc, sizes[dim]);
-    auto srcSize = b.createOrFold<tensor::DimOp>(loc, padOp.getSource(), dim);
+    auto low = padOp.getMixedLowPad()[dim];
+    bool hasLowPad = !isConstantIntValue(low, 0);
+    auto high = padOp.getMixedHighPad()[dim];
+    bool hasHighPad = !isConstantIntValue(high, 0);
+    auto offset = offsets[dim];
+    auto length = sizes[dim];
+    auto srcSize =
+        tensor::createDimValue(b, loc, padOp.getSource(), dim).value();
 
     // The new amount of low padding is `low - offset`. Except for the case
     // where none of the low padding is read. In that case, the new amount of
     // low padding is zero.
     //
     // Optimization: If low = 0, then newLow = 0.
-    Value newLow = hasLowPad ? max(zero, sub(low, offset)) : zero;
-    appendIndex(newLow, newLows, staticNewLows);
+    OpFoldResult newLow = hasLowPad ? max(zero, sub(low, offset)) : zero;
+    newLows.push_back(newLow);
 
     // Start reading the data from position `offset - low`. Since the original
     // read may have started in the low padding zone, this value could be
     // negative. Therefore, start reading from:
     //
     // max(offset - low, 0)
     //
     // The original read could also have started in the high padding zone.
     // In that case, set the offset to the end of source tensor. The new
     // ExtractSliceOp length will be zero in that case. (Effectively reading
     // no data from the source.)
     //
     // Optimization: If low = 0, then the formula can be simplified.
-    Value newOffset = hasLowPad ? min(max(sub(offset, low), zero), srcSize)
+    OpFoldResult newOffset = hasLowPad
+                                 ? min(max(sub(offset, low), zero), srcSize)
                                  : min(offset, srcSize);
-    newOffsets.push_back(getAsOpFoldResult(newOffset));
+    newOffsets.push_back(newOffset);
 
     // The original ExtractSliceOp was reading until position `offset +
     // length`. Therefore, the corresponding position within the source tensor
     // is:
     //
     // offset + length - low
     //
     // In case the original ExtractSliceOp stopped reading within the low
     // padding zone, this value can be negative. In that case, the end
     // position of the read should be zero. (Similar to newOffset.)
     //
     // The original read could also have stopped in the high padding zone.
     // In that case, set the end positition of the read should be the end of
     // the source tensor. (Similar to newOffset.)
     //
     // endLoc = min(max(offset - low + length, 0), srcSize)
     //
     // The new ExtractSliceOp length is `endLoc - newOffset`.
     //
     // Optimization: If low = 0, then the formula can be simplified.
-    Value endLoc = hasLowPad
-                       ? min(max(add(sub(offset, low), length), zero), srcSize)
-                       : min(add(offset, length), srcSize);
-    Value newLength = sub(endLoc, newOffset);
-    newLengths.push_back(getAsOpFoldResult(newLength));
+    OpFoldResult endLoc =
+        hasLowPad ? min(max(add(sub(offset, low), length), zero), srcSize)
+    OpFoldResult newLength = sub(endLoc, newOffset);
+    newLengths.push_back(newLength);
 
     // Check if newLength is zero. In that case, no SubTensorOp should be
     // executed.
-    if (auto newLengthInt = getConstantIntValue(newLength)) {
-      hasZeroLen |= *newLengthInt == 0;
-    } else {
-      Value check = b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
-                                            newLength, zero);
+    if (isConstantIntValue(newLength, 0)) {
+      hasZeroLen = true;
+    } else if (!hasZeroLen) {
+      Value check = b.create<arith::CmpIOp>(
+          loc, arith::CmpIPredicate::eq,
+          getValueOrCreateConstantIndexOp(b, loc, newLength),
+          getValueOrCreateConstantIndexOp(b, loc, zero));
       dynHasZeroLenCond =
           dynHasZeroLenCond
               ? b.create<arith::OrIOp>(loc, check, dynHasZeroLenCond)
               : check;
     }
 
     // The amount of high padding is simply the number of elements remaining,
     // so that the result has the same length as the original ExtractSliceOp.
     // As an optimization, if the original high padding is zero, then the new
     // high padding must also be zero.
-    Value newHigh = hasHighPad ? sub(sub(length, newLength), newLow) : zero;
-    appendIndex(newHigh, newHighs, staticNewHighs);
+    OpFoldResult newHigh =
+        hasHighPad ? sub(sub(length, newLength), newLow) : zero;
+    newHighs.push_back(newHigh);
 
     // Only unit stride supported.
     newStrides.push_back(b.getIndexAttr(1));
   }
 
   // The shape of the result can be obtained from the sizes passed in.
   SmallVector<Value> dynDims;
   SmallVector<int64_t> shape;
   dispatchIndexOpFoldResults(sizes, dynDims, shape);
   RankedTensorType resultType =
       RankedTensorType::get(shape, padOp.getResultType().getElementType());
 
   // Insert cast to ensure that types match. (May be folded away.)
-  auto castResult = [&](Value val) -> Operation * {
+  auto castResult = [&](Operation *op) -> Operation * {
+    Value val = op->getResult(0);
+    if (resultType == val.getType())
+      return op;
     return b.create<tensor::CastOp>(loc, resultType, val);
   };
 
   // In cases where the original data source is unused: Emit a GenerateOp and
   // do not generate a SliceOp. (The result shape of the SliceOp would
   // have a dimension of size 0, the semantics of which is unclear.)
   auto createGenerateOp = [&]() {
     // Create GenerateOp.
     auto generateOp = b.create<tensor::GenerateOp>(
         loc, resultType, dynDims,
         [&](OpBuilder &builder, Location gLoc, ValueRange indices) {
           builder.create<tensor::YieldOp>(gLoc, padValue);
         });
     return castResult(generateOp);
   };
 
   // Emit a SliceOp and a PadOp. Should not be used in cases where
   // the result shape of the new SliceOp has a zero dimension.
   auto createPadOfExtractSlice = [&]() {
     // Create pad(extract_slice(x)).
-    auto newSliceOp = b.create<tensor::ExtractSliceOp>(
+    Value newSliceOp = b.create<tensor::ExtractSliceOp>(
         loc, padOp.getSource(), newOffsets, newLengths, newStrides);
-    auto newPadOp = b.create<PadOp>(loc, newSliceOp, staticNewLows,
-                                    staticNewHighs, newLows, newHighs);
+    auto newPadOp = b.create<PadOp>(loc, Type(), newSliceOp, newLows, newHighs);
 
     // Copy region to new PadOp.
     IRMapping bvm;
     padOp.getRegion().cloneInto(&newPadOp.getRegion(), bvm);
 
     // Cast result and return.
     return castResult(newPadOp);
   };

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
   for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
     if (en.value() == ShapedType::kDynamic)
       dynamicDims.push_back(
           b.create<tensor::DimOp>(loc, rankedTensor, en.index()));
   }
   return dynamicDims;
 }
 
+FailureOr<OpFoldResult> mlir::tensor::createDimValue(OpBuilder &b, Location loc,
+                                                     Value rankedTensor,
+                                                     int64_t dim) {
+  auto tensorTy = rankedTensor.getType().dyn_cast<RankedTensorType>();
+  if (!tensorTy)
+    return failure();
+  auto shape = tensorTy.getShape();
+  if (dim >= shape.size())
+    return failure();
+  if (ShapedType::isDynamic(shape[dim]))
+    return OpFoldResult(b.createOrFold<tensor::DimOp>(loc, rankedTensor, dim));
+  return OpFoldResult(b.getIndexAttr(shape[dim]));
+}
+
 SmallVector<OpFoldResult>
 mlir::tensor::createDimValues(OpBuilder &b, Location loc, Value rankedTensor) {
   auto tensorTy = rankedTensor.getType().cast<RankedTensorType>();
   SmallVector<OpFoldResult> dims;
   for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
     if (ShapedType::isDynamic(en.value())) {
       dims.push_back(
           b.createOrFold<tensor::DimOp>(loc, rankedTensor, en.index()));