[mlir][affine] More efficient `makeComposedFolded...` helpers

[springerm]

Summary

The old code used to materialize constants as ops, immediately folded them into the resulting affine map and then deleted the constant ops again. Instead, directly fold the attributes into the affine map. Furthermore, all helpers accept OpFoldResult instead of Value now. This makes the code at call sites more efficient, because it is no longer necessary to materialize a Value, just to be able to use these helper functions.

Note: The API has changed (accepts OpFoldResult instead of Value), otherwise this change is NFC.

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

 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 
 namespace mlir {
 namespace affine {
 
 class AffineApplyOp;
 class AffineBound;
+class AffineMaxOp;
+class AffineMinOp;
 class AffineValueMap;
 
 /// A utility function to check if a value is defined at the top level of an
 /// op with trait `AffineScope` or is a region argument for such an op. A value
 /// of index type defined at the top level is always a valid symbol for all its
 /// uses.
 bool isTopLevelValue(Value value);
 

 /// for affine maps.
 void canonicalizeSetAndOperands(IntegerSet *set,
                                 SmallVectorImpl<Value> *operands);
 
 /// Returns a composed AffineApplyOp by composing `map` and `operands` with
 /// other AffineApplyOps supplying those operands. The operands of the resulting
 /// AffineApplyOp do not change the length of  AffineApplyOp chains.
 AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map,
-                                      ValueRange operands);
-/// Variant of `makeComposedAffineApply` which infers the AffineMap from `e`.
+                                      ArrayRef<OpFoldResult> operands);
 AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineExpr e,
-                                      ValueRange values);
+                                      ArrayRef<OpFoldResult> operands);
 
 /// Constructs an AffineApplyOp that applies `map` to `operands` after composing
 /// the map with the maps of any other AffineApplyOp supplying the operands,
 /// then immediately attempts to fold it. If folding results in a constant
 /// value, no ops are actually created. The `map` must be a single-result affine
 /// map.
 OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
                                            AffineMap map,
                                            ArrayRef<OpFoldResult> operands);
 /// Variant of `makeComposedFoldedAffineApply` that applies to an expression.
 OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
                                            AffineExpr expr,
                                            ArrayRef<OpFoldResult> operands);
 /// Variant of `makeComposedFoldedAffineApply` suitable for multi-result maps.
 /// Note that this may create as many affine.apply operations as the map has
 /// results given that affine.apply must be single-result.
 SmallVector<OpFoldResult> makeComposedFoldedMultiResultAffineApply(
     OpBuilder &b, Location loc, AffineMap map, ArrayRef<OpFoldResult> operands);
 
 /// Returns an AffineMinOp obtained by composing `map` and `operands` with
 /// AffineApplyOps supplying those operands.
-Value makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map,
-                            ValueRange operands);
+AffineMinOp makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map,
+                                  ArrayRef<OpFoldResult> operands);
 
 /// Constructs an AffineMinOp that computes a minimum across the results of
 /// applying `map` to `operands`, then immediately attempts to fold it. If
 /// folding results in a constant value, no ops are actually created.
 OpFoldResult makeComposedFoldedAffineMin(OpBuilder &b, Location loc,
                                          AffineMap map,
                                          ArrayRef<OpFoldResult> operands);
 

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

static void getXferIndices(RewriterBase &rewriter, TransferOpType xferOp,
                            AffineMap offsetMap, ArrayRef<Value> dimValues,
                            SmallVector<Value, 4> &indices) {
   indices.append(xferOp.getIndices().begin(), xferOp.getIndices().end());
   Location loc = xferOp.getLoc();
   unsigned offsetsIdx = 0;
   for (auto expr : xferOp.getPermutationMap().getResults()) {
     if (auto dim = expr.template dyn_cast<AffineDimExpr>()) {
       Value prevIdx = indices[dim.getPosition()];
-      SmallVector<Value, 3> dims(dimValues.begin(), dimValues.end());
+      SmallVector<OpFoldResult, 3> dims(dimValues.begin(), dimValues.end());
       dims.push_back(prevIdx);
       AffineExpr d0 = rewriter.getAffineDimExpr(offsetMap.getNumDims());
       indices[dim.getPosition()] = affine::makeComposedAffineApply(
           rewriter, loc, d0 + offsetMap.getResult(offsetsIdx++), dims);
       continue;
     }
   }
 }

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

void mlir::affine::fullyComposeAffineMapAndOperands(
     AffineMap *map, SmallVectorImpl<Value> *operands) {
   while (llvm::any_of(*operands, [](Value v) {
     return isa_and_nonnull<AffineApplyOp>(v.getDefiningOp());
   })) {
     composeAffineMapAndOperands(map, operands);
   }
 }
 
-/// Given a list of `OpFoldResult`, build the necessary operations to populate
-/// `actualValues` with values produced by operations. In particular, for any
-/// attribute-typed element in `values`, call the constant materializer
-/// associated with the Affine dialect to produce an operation. Do NOT notify
-/// the builder listener about the constant ops being created as they are
-/// intended to be removed after being folded into affine constructs; this is
-/// not suitable for use beyond the Affine dialect.
-static void materializeConstants(OpBuilder &b, Location loc,
-                                 ArrayRef<OpFoldResult> values,
-                                 SmallVectorImpl<Operation *> &constants,
-                                 SmallVectorImpl<Value> &actualValues) {
-  OpBuilder::Listener *listener = b.getListener();
-  b.setListener(nullptr);
-  auto listenerResetter =
-      llvm::make_scope_exit([listener, &b] { b.setListener(listener); });
-
-  actualValues.reserve(values.size());
-  auto *dialect = b.getContext()->getLoadedDialect<AffineDialect>();
-  for (OpFoldResult ofr : values) {
-    if (auto value = llvm::dyn_cast_if_present<Value>(ofr)) {
-      actualValues.push_back(value);
-      continue;
+/// Fold all attributes among the given operands into the affine map and return
+/// all remaining values. The affine map is modified in-place.
+static SmallVector<Value>
+foldAttributesIntoMap(Builder &b, AffineMap *map,
+                      ArrayRef<OpFoldResult> operands) {
+  SmallVector<AffineExpr> dimReplacements, symReplacements;
+  SmallVector<Value, 8> valueOperands;
+  int64_t numDims = 0;
+  for (int64_t i = 0; i < map->getNumDims(); ++i) {
+    if (auto attr = operands[i].dyn_cast<Attribute>()) {
+      dimReplacements.push_back(
+          b.getAffineConstantExpr(attr.cast<IntegerAttr>().getInt()));
+    } else {
+      dimReplacements.push_back(b.getAffineDimExpr(numDims++));
+      valueOperands.push_back(operands[i].get<Value>());
     }
-    // Since we are directly specifying `index` as the result type, we need to
-    // ensure the provided attribute is also an index type. Otherwise, the
-    // AffineDialect materializer will create invalid `arith.constant`
-    // operations if the provided Attribute is any other kind of integer.
-    constants.push_back(dialect->materializeConstant(
-        b,
-        b.getIndexAttr(llvm::cast<IntegerAttr>(ofr.get<Attribute>()).getInt()),
-        b.getIndexType(), loc));
-    actualValues.push_back(constants.back()->getResult(0));
   }
+  int64_t numSymbols = 0;
+  for (int64_t i = 0; i < map->getNumSymbols(); ++i) {
+    if (auto attr = operands[i + map->getNumDims()].dyn_cast<Attribute>()) {
+      symReplacements.push_back(
+          b.getAffineConstantExpr(attr.cast<IntegerAttr>().getInt()));
+    } else {
+      symReplacements.push_back(b.getAffineSymbolExpr(numSymbols++));
+      valueOperands.push_back(operands[i + map->getNumDims()].get<Value>());
     }
-
-/// Create an operation of the type provided as template argument and attempt to
-/// fold it immediately. The operation is expected to have a builder taking
-/// arbitrary `leadingArguments`, followed by a list of Value-typed `operands`.
-/// The operation is also expected to always produce a single result. Return an
-/// `OpFoldResult` containing the Attribute representing the folded constant if
-/// complete folding was possible and a Value produced by the created operation
-/// otherwise.
-template <typename OpTy, typename... Args>
-static std::enable_if_t<OpTy::template hasTrait<OpTrait::OneResult>(),
-                        OpFoldResult>
-createOrFold(OpBuilder &b, Location loc, ValueRange operands,
-             Args &&...leadingArguments) {
-  // Identify the constant operands and extract their values as attributes.
-  // Note that we cannot use the original values directly because the list of
-  // operands may have changed due to canonicalization and composition.
-  SmallVector<Attribute> constantOperands;
-  constantOperands.reserve(operands.size());
-  for (Value operand : operands) {
-    IntegerAttr attr;
-    if (matchPattern(operand, m_Constant(&attr)))
-      constantOperands.push_back(attr);
-    else
-      constantOperands.push_back(nullptr);
   }
-
-  // Create the operation and immediately attempt to fold it. On success,
-  // delete the operation and prepare the (unmaterialized) value for being
-  // returned. On failure, return the operation result value. Temporarily remove
-  // the listener to avoid notifying it when the op is created as it may be
-  // removed immediately and there is no way of notifying the caller about that
-  // without resorting to RewriterBase.
-  //
-  // TODO: arguably, the main folder (createOrFold) API should support this use
-  // case instead of indiscriminately materializing constants.
-  OpBuilder::Listener *listener = b.getListener();
-  b.setListener(nullptr);
-  auto listenerResetter =
-      llvm::make_scope_exit([listener, &b] { b.setListener(listener); });
-  OpTy op =
-      b.create<OpTy>(loc, std::forward<Args>(leadingArguments)..., operands);
-  SmallVector<OpFoldResult, 1> foldResults;
-  if (succeeded(op->fold(constantOperands, foldResults)) &&
-      !foldResults.empty()) {
-    op->erase();
-    return foldResults.front();
-  }
-
-  // Notify the listener now that we definitely know that the operation will
-  // persist. Use the original listener stored in the variable.
-  if (listener)
-    listener->notifyOperationInserted(op);
-  return op->getResult(0);
+  *map = map->replaceDimsAndSymbols(dimReplacements, symReplacements, numDims,
+                                    numSymbols);
+  return valueOperands;
 }
 
-AffineApplyOp mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc,
-                                                    AffineMap map,
-                                                    ValueRange operands) {
-  AffineMap normalizedMap = map;
-  SmallVector<Value, 8> normalizedOperands(operands.begin(), operands.end());
-  composeAffineMapAndOperands(&normalizedMap, &normalizedOperands);
-  assert(normalizedMap);
-  return b.create<AffineApplyOp>(loc, normalizedMap, normalizedOperands);
+AffineApplyOp
+mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map,
+                                      ArrayRef<OpFoldResult> operands) {
+  SmallVector<Value> valueOperands = foldAttributesIntoMap(b, &map, operands);
+  composeAffineMapAndOperands(&map, &valueOperands);
+  assert(map);
+  return b.create<AffineApplyOp>(loc, map, valueOperands);
 }
 
-AffineApplyOp mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc,
-                                                    AffineExpr e,
-                                                    ValueRange values) {
+AffineApplyOp
+mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc, AffineExpr e,
+                                      ArrayRef<OpFoldResult> operands) {
   return makeComposedAffineApply(
       b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{e}).front(),
-      values);
+      operands);
 }
 
 /// Composes the given affine map with the given list of operands, pulling in
 /// the maps from any affine.apply operations that supply the operands.
 static void composeMultiResultAffineMap(AffineMap &map,
                                         SmallVectorImpl<Value> &operands) {
   // Compose and canonicalize each expression in the map individually because
   // composition only applies to single-result maps, collecting potentially

 }
 
 OpFoldResult
 mlir::affine::makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
                                             AffineMap map,
                                             ArrayRef<OpFoldResult> operands) {
   assert(map.getNumResults() == 1 && "building affine.apply with !=1 result");
 
-  SmallVector<Operation *> constants;
-  SmallVector<Value> actualValues;
-  materializeConstants(b, loc, operands, constants, actualValues);
-  composeAffineMapAndOperands(&map, &actualValues);
-  OpFoldResult result = createOrFold<AffineApplyOp>(b, loc, actualValues, map);
-
-  // Constants are always folded into affine min/max because they can be
-  // represented as constant expressions, so delete them.
-  for (Operation *op : constants)
-    op->erase();
-  return result;
+  // Temporarily disconnect the listener, so that no notification is triggered
+  // if the op is folded.
+  // TODO: OpBuilder::createOrFold should return OpFoldResults, then this
+  // workaround is no longer needed.
+  OpBuilder::Listener *listener = b.getListener();
+  b.setListener(nullptr);
+  auto listenerResetter =
+      llvm::make_scope_exit([listener, &b] { b.setListener(listener); });
+
+  // Create op.
+  AffineApplyOp applyOp = makeComposedAffineApply(b, loc, map, operands);
+
+  // Get constant operands.
+  SmallVector<Attribute> constOperands(applyOp->getNumOperands());
+  for (unsigned i = 0, e = constOperands.size(); i != e; ++i)
+    matchPattern(applyOp->getOperand(i), m_Constant(&constOperands[i]));
+
+  // Try to fold the operation.
+  SmallVector<OpFoldResult> foldResults;
+  if (failed(applyOp->fold(constOperands, foldResults)) ||
+      foldResults.empty()) {
+    if (listener)
+      listener->notifyOperationInserted(applyOp);
+    return applyOp.getResult();
+  }
+
+  applyOp->erase();
+  assert(foldResults.size() == 1 && "expected 1 folded result");
+  return foldResults.front();
 }
 
 OpFoldResult
 mlir::affine::makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
                                             AffineExpr expr,
                                             ArrayRef<OpFoldResult> operands) {
   return makeComposedFoldedAffineApply(
       b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{expr}).front(),
       operands);
 }
 
 SmallVector<OpFoldResult>
 mlir::affine::makeComposedFoldedMultiResultAffineApply(
     OpBuilder &b, Location loc, AffineMap map,
     ArrayRef<OpFoldResult> operands) {
   return llvm::map_to_vector(llvm::seq<unsigned>(0, map.getNumResults()),
                              [&](unsigned i) {
                                return makeComposedFoldedAffineApply(
                                    b, loc, map.getSubMap({i}), operands);
                              });
 }
 
-Value mlir::affine::makeComposedAffineMin(OpBuilder &b, Location loc,
-                                          AffineMap map, ValueRange operands) {
-  SmallVector<Value> allOperands = llvm::to_vector(operands);
-  composeMultiResultAffineMap(map, allOperands);
-  return b.createOrFold<AffineMinOp>(loc, b.getIndexType(), map, allOperands);
+template <typename OpTy>
+static OpTy makeComposedMinMax(OpBuilder &b, Location loc, AffineMap map,
+                               ArrayRef<OpFoldResult> operands) {
+  SmallVector<Value> valueOperands = foldAttributesIntoMap(b, &map, operands);
+  composeMultiResultAffineMap(map, valueOperands);
+  return b.create<OpTy>(loc, b.getIndexType(), map, valueOperands);
+}
+
+AffineMinOp
+mlir::affine::makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map,
+                                    ArrayRef<OpFoldResult> operands) {
+  return makeComposedMinMax<AffineMinOp>(b, loc, map, operands);
 }
 
 template <typename OpTy>
 static OpFoldResult makeComposedFoldedMinMax(OpBuilder &b, Location loc,
                                              AffineMap map,
                                              ArrayRef<OpFoldResult> operands) {
-  SmallVector<Operation *> constants;
-  SmallVector<Value> actualValues;
-  materializeConstants(b, loc, operands, constants, actualValues);
-  composeMultiResultAffineMap(map, actualValues);
-  OpFoldResult result =
-      createOrFold<OpTy>(b, loc, actualValues, b.getIndexType(), map);
-
-  // Constants are always folded into affine min/max because they can be
-  // represented as constant expressions, so delete them.
-  for (Operation *op : constants)
-    op->erase();
-  return result;
+  // Temporarily disconnect the listener, so that no notification is triggered
+  // if the op is folded.
+  // TODO: OpBuilder::createOrFold should return OpFoldResults, then this
+  // workaround is no longer needed.
+  OpBuilder::Listener *listener = b.getListener();
+  b.setListener(nullptr);
+  auto listenerResetter =
+      llvm::make_scope_exit([listener, &b] { b.setListener(listener); });
+
+  // Create op.
+  auto minMaxOp = makeComposedMinMax<OpTy>(b, loc, map, operands);
+
+  // Get constant operands.
+  SmallVector<Attribute> constOperands(minMaxOp->getNumOperands());
+  for (unsigned i = 0, e = constOperands.size(); i != e; ++i)
+    matchPattern(minMaxOp->getOperand(i), m_Constant(&constOperands[i]));
+
+  // Try to fold the operation.
+  SmallVector<OpFoldResult> foldResults;
+  if (failed(minMaxOp->fold(constOperands, foldResults)) ||
+      foldResults.empty()) {
+    if (listener)
+      listener->notifyOperationInserted(minMaxOp);
+    return minMaxOp.getResult();
+  }
+
+  minMaxOp->erase();
+  assert(foldResults.size() == 1 && "expected 1 folded result");
+  return foldResults.front();
 }
 
 OpFoldResult
 mlir::affine::makeComposedFoldedAffineMin(OpBuilder &b, Location loc,
                                           AffineMap map,
                                           ArrayRef<OpFoldResult> operands) {
   return makeComposedFoldedMinMax<AffineMinOp>(b, loc, map, operands);
 }

mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp

idBuilder = [this](RewriterBase &rewriter, Location loc,
       SmallVector<int64_t> reverseBasisSizes(
           llvm::reverse(this->availableMappingSizes));
       SmallVector<int64_t> strides = computeStrides(reverseBasisSizes);
       SmallVector<AffineExpr> delinearizingExprs = delinearize(d0, strides);
       SmallVector<Value> ids;
       // Reverse back to be in [x, y, z] order.
       for (AffineExpr e : llvm::reverse(delinearizingExprs))
         ids.push_back(
-            affine::makeComposedAffineApply(rewriter, loc, e, warpId));
+            affine::makeComposedAffineApply(rewriter, loc, e, {warpId}));
 
       // clang-format off
       LDBG("----linearId: " << linearId);
           LDBG("----warpId: " << warpId);
       LLVM_DEBUG(llvm::interleaveComma(reverseBasisSizes,
                                        DBGS() << "--delinearization basis: ");
                  llvm::dbgs() << "\n";
                  llvm::interleaveComma(strides,
idBuilder = [this](RewriterBase &rewriter, Location loc,
       SmallVector<int64_t> strides = computeStrides(reverseBasisSizes);
       AffineExpr d0;
       bindDims(rewriter.getContext(), d0);
       SmallVector<AffineExpr> delinearizingExprs = delinearize(d0, strides);
       SmallVector<Value> ids;
       // Reverse back to be in [x, y, z] order.
       for (AffineExpr e : llvm::reverse(delinearizingExprs))
         ids.push_back(
-            affine::makeComposedAffineApply(rewriter, loc, e, linearId));
+            affine::makeComposedAffineApply(rewriter, loc, e, {linearId}));
 
       // clang-format off
       LLVM_DEBUG(llvm::interleaveComma(reverseBasisSizes,
                                        DBGS() << "--delinearization basis: ");
                  llvm::dbgs() << "\n";
                  llvm::interleaveComma(strides,
                                        DBGS() << "--delinearization strides: ");
                  llvm::dbgs() << "\n";

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

 // Given indices corresponding to iterators in the output (oIndex) and filter
 // (fIndex) for a convolution, compute the convolved index for the
 // input as `oIndex * stride + fIndex`.
 static Value getConvolvedIndex(OpBuilder &b, Location loc, Value oIndex,
                                Value fIndex, int64_t stride) {
   AffineExpr oExpr, fExpr;
   bindSymbols(b.getContext(), oExpr, fExpr);
   AffineMap convMap = AffineMap::get(0, 2, stride * oExpr + fExpr);
-  return affine::makeComposedAffineApply(b, loc, convMap,
-                                         ValueRange{oIndex, fIndex});
+  return affine::makeComposedAffineApply(b, loc, convMap, {oIndex, fIndex});
 }
 
 FailureOr<std::pair<Operation *, Operation *>>
 rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcHwcfOp convOp) {
   auto inputType = cast<ShapedType>(convOp.getInputs()[0].getType());
   auto filterType = cast<ShapedType>(convOp.getInputs()[1].getType());
   auto outputType = cast<ShapedType>(convOp.getOutputs()[0].getType());
 

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

SmallVector<OpFoldResult> loopRanges =
                                                allShapes);
   Value tripCount =
       getValueOrCreateConstantIndexOp(b, op.getLoc(), loopRanges[dimension]);
 
   // Compute the tile sizes and the respective numbers of tiles.
   AffineExpr s0 = b.getAffineSymbolExpr(0);
   AffineExpr s1 = b.getAffineSymbolExpr(1);
   AffineExpr s2 = b.getAffineSymbolExpr(2);
-  auto apply = [&](AffineExpr expr, ValueRange values) -> Value {
-    return affine::makeComposedAffineApply(b, b.getLoc(), expr, values);
+  auto apply = [&](AffineExpr expr, ArrayRef<OpFoldResult> ofrs) -> Value {
+    return affine::makeComposedAffineApply(b, b.getLoc(), expr, ofrs);
   };
   Value a = apply(s0.floorDiv(s1), {tripCount, divisorValue});
   Value t = apply((s0 + s1 - 1).floorDiv(s1), {targetSizeValue, divisorValue});
   Value d = apply((s0 + s1 - 1).floorDiv(s1), {a, t});
   Value s = apply(s0.floorDiv(s1) * s2, {a, d, divisorValue});
   Value v = apply(s0 % s1, {a, d});
   Value u = apply(s0 - s1, {d, v});
 

mlir/lib/Dialect/Tensor/Transforms/ExtractSliceFromReshapeUtils.cpp

static DimAndIndex invertSliceIndexing(OpBuilder &b, Location loc,
                                        ArrayRef<Range> sliceParams,
                                        const DimAndIndex &dimAndIndex) {
   AffineExpr d0, s0, s1;
   bindDims(b.getContext(), d0);
   bindSymbols(b.getContext(), s0, s1);
   auto [dim, indexValue] = dimAndIndex;
   assert(dim < sliceParams.size() && "slice should be non rank-reducing");
   return std::make_pair(
-      dim,
-      affine::makeComposedAffineApply(
+      dim, affine::makeComposedAffineApply(
                b, loc, s0 + d0 * s1,
-          {indexValue,
-           getValueOrCreateConstantIndexOp(b, loc, sliceParams[dim].offset),
-           getValueOrCreateConstantIndexOp(b, loc, sliceParams[dim].stride)}));
+               {indexValue, sliceParams[dim].offset, sliceParams[dim].stride}));
 }
 
 /// Transform `dimAndIndex` from the result tensor index space of a
 /// CollapseShapeOp to the source tensor index space.
 static ValueRange invertCollapseShapeIndexing(
     OpBuilder &b, Location loc, ArrayRef<ReassociationIndices> reassociation,
     ArrayRef<OpFoldResult> reshapeSourceShape, const DimAndIndex &dimAndIndex) {
   const auto &[dim, indexValue] = dimAndIndex;

mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp

xferOp.zipResultAndIndexing([&](int64_t resultIdx, int64_t indicesIdx) {
     if (xferOp.isDimInBounds(resultIdx))
       return;
     // Fold or create the check that `index + vector_size` <= `memref_size`.
     Location loc = xferOp.getLoc();
     int64_t vectorSize = xferOp.getVectorType().getDimSize(resultIdx);
     auto d0 = getAffineDimExpr(0, xferOp.getContext());
     auto vs = getAffineConstantExpr(vectorSize, xferOp.getContext());
     Value sum = affine::makeComposedAffineApply(b, loc, d0 + vs,
-                                                xferOp.indices()[indicesIdx]);
+                                                {xferOp.indices()[indicesIdx]});
     Value cond = createFoldedSLE(
         b, sum, vector::createOrFoldDimOp(b, loc, xferOp.source(), indicesIdx));
     if (!cond)
       return;
     // Conjunction over all dims for which we are in-bounds.
     if (inBoundsCond)
       inBoundsCond = b.create<arith::AndIOp>(loc, inBoundsCond, cond);
     else