[mlir][Affine] Delete duplicate code: `applyMapToValues`

[springerm]

Summary

The same functionality is provided by makeComposedFoldedAffineApply.

Depends On: D154196

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

 
 /// 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, no ops are actually created.
 OpFoldResult makeComposedFoldedAffineMax(OpBuilder &b, Location loc,
                                          AffineMap map,
                                          ArrayRef<OpFoldResult> operands);
 
-/// Returns the values obtained by applying `map` to the list of values.
-SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc,
-                                       AffineMap map, ValueRange values);
-
 /// Given an affine map `map` and its input `operands`, this method composes
 /// into `map`, maps of AffineApplyOps whose results are the values in
 /// `operands`, iteratively until no more of `operands` are the result of an
 /// AffineApplyOp. When this function returns, `map` becomes the composed affine
 /// map, and each Value in `operands` is guaranteed to be either a loop IV or a
 /// terminal symbol, i.e., a symbol defined at the top level or a block/function
 /// argument.
 void fullyComposeAffineMapAndOperands(AffineMap *map,

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

 
 OpFoldResult
 mlir::affine::makeComposedFoldedAffineMax(OpBuilder &b, Location loc,
                                           AffineMap map,
                                           ArrayRef<OpFoldResult> operands) {
   return makeComposedFoldedMinMax<AffineMaxOp>(b, loc, map, operands);
 }
 
-/// Fully compose map with operands and canonicalize the result.
-/// Return the `createOrFold`'ed AffineApply op.
-static Value createFoldedComposedAffineApply(OpBuilder &b, Location loc,
-                                             AffineMap map,
-                                             ValueRange operandsRef) {
-  SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
-  fullyComposeAffineMapAndOperands(&map, &operands);
-  canonicalizeMapAndOperands(&map, &operands);
-  return b.createOrFold<AffineApplyOp>(loc, map, operands);
-}
-
-SmallVector<Value, 4> mlir::affine::applyMapToValues(OpBuilder &b, Location loc,
-                                                     AffineMap map,
-                                                     ValueRange values) {
-  SmallVector<Value, 4> res;
-  res.reserve(map.getNumResults());
-  unsigned numDims = map.getNumDims(), numSym = map.getNumSymbols();
-  // For each `expr` in `map`, applies the `expr` to the values extracted from
-  // ranges. If the resulting application can be folded into a Value, the
-  // folding occurs eagerly.
-  for (auto expr : map.getResults()) {
-    AffineMap map = AffineMap::get(numDims, numSym, expr);
-    res.push_back(createFoldedComposedAffineApply(b, loc, map, values));
-  }
-  return res;
-}
-
 // A symbol may appear as a dim in affine.apply operations. This function
 // canonicalizes dims that are valid symbols into actual symbols.
 template <class MapOrSet>
 static void canonicalizePromotedSymbols(MapOrSet *mapOrSet,
                                         SmallVectorImpl<Value> *operands) {
   if (!mapOrSet || operands->empty())
     return;
 

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

tileLinalgOpImpl(RewriterBase &b, LinalgOp op, ArrayRef<OpFoldResult> tileSizes,
 
   // Handle distribution. Create a vector of the same size of loops that are to
   // be tiled.
   SmallVector<linalg::ProcInfo> procInfo;
   if (options.distribution) {
     procInfo.resize(
         iteratorTypes.size(),
         linalg::ProcInfo{nullptr, nullptr, linalg::DistributionMethod::None});
-    // Collect loop ranges of tiled loopss, loops that are parallel.
+    // Collect loop ranges of tiled loops, loops that are parallel.
     SmallVector<Range> parallelLoopRanges;
     for (const auto &iteratorType : llvm::enumerate(iteratorTypes)) {
       if (!isParallelIterator(iteratorType.value()))
         break;
       parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
     }
     auto returnedProcInfo =
         options.distribution->procInfo(b, op.getLoc(), parallelLoopRanges);
auto tiledLoopBodyBuilder =
       [&](OpBuilder &builder, Location loc, ValueRange localIvs,
           ValueRange operandValuesToUse) -> scf::ValueVector {
     ivs.assign(localIvs.begin(), localIvs.end());
 
     // When an `interchangeVector` is present, it has been applied to the
     // loop ranges and the iterator types. Apply its inverse to the
     // resulting loop `ivs` to match the op definition.
     SmallVector<Value, 4> interchangedIvs;
-    if (!options.interchangeVector.empty())
-      interchangedIvs = applyMapToValues(b, loc, invPermutationMap, ivs);
-    else
+    if (!options.interchangeVector.empty()) {
+      for (AffineExpr result : invPermutationMap.getResults())
+        interchangedIvs.push_back(
+            ivs[result.cast<AffineDimExpr>().getPosition()]);
+    } 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->getNumOperands()) &&
            "expect the number of operands and inputs and outputs to match");
     SmallVector<Value> valuesToTile = operandValuesToUse;
     SmallVector<OpFoldResult> sizeBounds =

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

static OpFoldResult getCollapsedOutputDimFromInputShape(
   if (!ShapedType::isDynamic(dstStaticShape[dimIndex])) {
     return builder.getIndexAttr(dstStaticShape[dimIndex]);
   }
   AffineMap map = reassociationMap[dimIndex];
   unsigned startPos =
       map.getResults().front().cast<AffineDimExpr>().getPosition();
   unsigned endPos = map.getResults().back().cast<AffineDimExpr>().getPosition();
   AffineExpr expr;
-  SmallVector<Value, 2> dynamicDims;
+  SmallVector<OpFoldResult> dynamicDims;
   for (auto dim : llvm::seq_inclusive(startPos, endPos)) {
     dynamicDims.push_back(builder.createOrFold<tensor::DimOp>(loc, src, dim));
     AffineExpr currExpr = builder.getAffineSymbolExpr(dim - startPos);
     expr = (expr ? expr * currExpr : currExpr);
   }
-  return affine::applyMapToValues(
+  return affine::makeComposedFoldedAffineApply(
       builder, loc, AffineMap::get(0, endPos - startPos + 1, expr),
-      dynamicDims)[0];
+      dynamicDims);
 }
 
 /// Given the `src` of a collapsing reshape op and its reassociation maps,
 /// compute the shape of the result of the reshape.
 static SmallVector<OpFoldResult, 4> getCollapsedOutputShapeFromInputShape(
     OpBuilder &builder, Location loc, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
   return llvm::to_vector<4>(llvm::map_range(
for (auto d :
        llvm::enumerate(dstStaticShape.slice(startPos, endPos - startPos + 1))) {
     if (d.index() + startPos == static_cast<unsigned>(dimIndex))
       continue;
     assert(!ShapedType::isDynamic(d.value()) &&
            "single dimension cannot be expanded into multiple dynamic "
            "dimensions");
     linearizedStaticDim *= d.value();
   }
-  Value sourceDim = builder.create<tensor::DimOp>(loc, src, sourceDimPos);
-  return affine::applyMapToValues(
+  OpFoldResult sourceDim =
+      builder.create<tensor::DimOp>(loc, src, sourceDimPos).getResult();
+  return affine::makeComposedFoldedAffineApply(
       builder, loc,
       AffineMap::get(
           0, 1, builder.getAffineSymbolExpr(0).floorDiv(linearizedStaticDim)),
-      sourceDim)[0];
+      sourceDim);
 }
 
 /// Given the `src` of an expanding reshape op, the reassociation maps and the
 /// result type, compute the shape of the result of the reshape.
 static SmallVector<OpFoldResult, 4> getExpandedOutputShapeFromInputShape(
     OpBuilder &builder, Location loc, Value src,
     ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
   llvm::DenseMap<int64_t, int64_t> expandedDimToCollapsedDim =
reifyResultShapes(Operation *op, OpBuilder &b,
     SmallVector<OpFoldResult> shapes;
     for (auto dim : llvm::seq<int64_t>(0, padOp.getSourceType().getRank())) {
       if (!padOp.getResultType().isDynamicDim(dim)) {
         shapes.push_back(b.getIndexAttr(padOp.getResultType().getDimSize(dim)));
         continue;
       }
 
       // Shape along each dimension is source dim + low pad + high pad.
-      SmallVector<Value> mapOperands;
+      SmallVector<OpFoldResult> mapOperands;
       mapOperands.push_back(
           b.createOrFold<tensor::DimOp>(loc, padOp.getSource(), dim));
-      AffineExpr expr = b.getAffineDimExpr(0);
-      unsigned numSymbols = 0;
-      auto addOpFoldResult = [&](OpFoldResult valueOrAttr) {
-        if (Value v = llvm::dyn_cast_if_present<Value>(valueOrAttr)) {
-          expr = expr + b.getAffineSymbolExpr(numSymbols++);
-          mapOperands.push_back(v);
-          return;
-        }
-        int64_t staticValue =
-            llvm::cast<IntegerAttr>(valueOrAttr.get<Attribute>()).getInt();
-        expr = expr + staticValue;
-      };
-      addOpFoldResult(lowPad[dim]);
-      addOpFoldResult(highPad[dim]);
-      shapes.push_back(affine::applyMapToValues(
-          b, loc, AffineMap::get(1, numSymbols, expr), mapOperands)[0]);
+      mapOperands.push_back(lowPad[dim]);
+      mapOperands.push_back(highPad[dim]);
+      AffineExpr expr = b.getAffineDimExpr(0) + b.getAffineSymbolExpr(0) +
+                        b.getAffineSymbolExpr(1);
+      shapes.push_back(getValueOrCreateConstantIndexOp(
+          b, loc,
+          affine::makeComposedFoldedAffineApply(
+              b, loc, AffineMap::get(1, 2, expr), mapOperands)));
     }
     reifiedReturnShapes.emplace_back(std::move(shapes));
     return success();
   }
 };
 
 } // namespace
 

mlir/test/Dialect/Linalg/pad_fusion.mlir

%0 = linalg.generic {
     } -> tensor<?x?xf32>
   %1 = tensor.pad %0 low [%arg1, %arg2] high [%arg3, %arg4] {
     ^bb0(%arg6: index, %arg7 : index):
       tensor.yield %arg5 : f32
     } : tensor<?x?xf32> to tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
 
-//  CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s2 + s0 + s1)>
+//  CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 + s1 + s2)>
 //      CHECK: func @dynamic_pad_fusion
 // CHECK-SAME:     %[[ARG0:.+]]: tensor<?x?xf32>
 // CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index
 // CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index
 // CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: index
 // CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index
 // CHECK-SAME:     %[[ARG5:[a-zA-Z0-9]+]]: f32
 //  CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
^bb0(%arg4 : f32, %arg5 : f32):
       linalg.yield %1 : f32
     } -> tensor<42x?xf32>
   %1 = tensor.pad %0 low [3, %arg1] high [4, %arg2] {
     ^bb0(%arg4: index, %arg5 : index):
       tensor.yield %arg3 : f32
     } : tensor<42x?xf32> to tensor<49x?xf32>
   return %1 : tensor<49x?xf32>
 }
-//  CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s2 + s0 + s1)>
+//  CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 + s1 + s2)>
 //      CHECK: func @mixed_pad_fusion
 // CHECK-SAME:     %[[ARG0:.+]]: tensor<?x42xf32>
 // CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index
 // CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index
 // CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: f32
 //  CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
 //  CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
 //  CHECK-DAG:   %[[SOURCE:.+]] = linalg.generic

mlir/test/Dialect/Linalg/resolve-shaped-type-result-dims.mlir

%0 = tensor.pad %arg0 low[%c3, %arg1, %c4] high[7, %c5, %arg2] {
      ^bb0(%arg4: index, %arg5: index, %arg6: index):
        tensor.yield %arg3 : f32
    } : tensor<2x?x?xf32> to tensor<?x?x?xf32>
    %1 = tensor.dim %0, %c0 : tensor<?x?x?xf32>
    %2 = tensor.dim %0, %c1 : tensor<?x?x?xf32>
    %3 = tensor.dim %0, %c2 : tensor<?x?x?xf32>
    return %1, %2, %3 : index, index, index
 }
-//  CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1] -> (s1 + s0 + 5)>
-//  CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s1 + s0 + 4)>
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 5)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 4)>
 //      CHECK: func @dim_of_pad_op
 // CHECK-SAME:   %[[ARG0:[A-Za-z0-9_]+]]: tensor<2x?x?xf32>
 // CHECK-SAME:   %[[ARG1:[A-Za-z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG2:[A-Za-z0-9_]+]]: index
 //  CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
 //  CHECK-DAG:   %[[C2:.+]] = arith.constant 2 : index
 //  CHECK-DAG:   %[[C12:.+]] = arith.constant 12 : index
 //      CHECK:   %[[IN_DIM1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
 //      CHECK:   %[[OUT_DIM1:.+]] = affine.apply #[[MAP0]]()[%[[ARG1]], %[[IN_DIM1]]]
 //      CHECK:   %[[IN_DIM2:.+]] = tensor.dim %[[ARG0]], %[[C2]]
 //      CHECK:   %[[OUT_DIM2:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[IN_DIM2]]]
 //      CHECK:   return %[[C12]], %[[OUT_DIM1]], %[[OUT_DIM2]]

mlir/test/Dialect/Linalg/vectorization-masked.mlir

 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
   transform.structured.masked_vectorize %0 vector_sizes [2, 4] : !transform.any_op
 }
 
 // -----
 
-//       CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s1 + s0)>
+//       CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>
 //       CHECK: func @test_masked_vectorize_dynamic_pad
 func.func @test_masked_vectorize_dynamic_pad(
   %0 : tensor<?x?xf32>, %h0 : index, %h1 : index)
     -> tensor<?x?xf32>
 {
   //  CHECK-DAG: %[[c42:.*]] = arith.constant 4.243000e+01 : f32
   //  CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
   //  CHECK-DAG: %[[res_d0:.+]] = affine.apply #[[MAP]]()

mlir/test/Dialect/Tensor/bufferize.mlir

func.func @tensor.reshape(%t1: tensor<?x10xf32>) -> tensor<2x2x5xf32> {
 
   // CHECK: %[[r:.*]] = bufferization.to_tensor %[[reshaped]]
   // CHECK: return %[[r]]
   return %reshaped : tensor<2x2x5xf32>
 }
 
 // -----
 
-// CHECK:       #[[$sum_map_1:.+]] = affine_map<()[s0, s1] -> (s1 + s0 + 5)>
+// CHECK:       #[[$sum_map_1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 5)>
 // CHECK:       #[[$sum_map_2:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 10)>
 // CHECK-LABEL: func @tensor.pad(
 //  CHECK-SAME:   %[[t1:.*]]: tensor<?x10xindex>, %[[l2:.*]]: index, %[[h1:.*]]: index, %[[h2:.*]]: index
 func.func @tensor.pad(%t1: tensor<?x10xindex>, %l2: index, %h1: index,
                       %h2: index) -> tensor<?x?xindex> {
   // CHECK-DAG: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?x10xindex>
   // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
   // CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index