[mlir] fold more eagerly in structured op splitting

[ftynse]

Summary

Existing implementation of structured op splitting creates several
affine.apply and affine.min operations in its subshape computation.
As these shapes are further used in data slice extraction, this may lead
to slice shapes being dynamic even when the original shapes and the
splitting point are static. This is particularly visible when splitting
is combined with further subsetting transformations such as tiling. Use
composition and folding more aggressively in splitting to avoid this.

In particular, introduce a createComposedAffineMin function that the
affine map used in "min" with the maps used by any affine.apply that
may be feeding the operands to the "min". This enables production of
more static shapes. Also introduce a createComposedFoldedAffineApply
function that combines the existing createComposedAffineApply with
in-place folding to propagate constants produced by zero-input affine
maps. Using these when splitting allows the subsequent canonicalizer
pass to recover static shapes for structured ops.

[bondhugula]

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

738

This looks suspicious. Can you ever pass in a null map here? Even if you don't, if you pass in a map with > 1 results and which is already composed, the assertion message would incorrectly say "could not compose maps". It looks like this should be asserting in a different or perhaps earlier if at all.

819–826

Doc comment.

[ftynse]

Address review.

[ftynse]

Fold even more eagerly and recover static shapes in splitting.

[nicolasvasilache]

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

398

Is there no need for a folded version like in the above ?

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

739

Great, thanks, this has been missing since forever.
I wish we also renamed OpFoldResult to ValueOrAttr or something more descriptive but that was somehow contentious at the time ..

747

This function reinvents non trivial amounts of logic around normalization, dims/sym compression etc. and additionally gets clever with inplace logic.

How about just linearizing and shifting instead:

SmallVector<Value> dims, symbols;
SmallVector<AffineExpr> exprs;
for (unsigned i : llvm::seq<unsigned>(0, map.getNumResults())) {
  SmallVector<Value> allOperands(operands.begin(), operands.end());
  AffineMap map = map.getSubMap({i});
  fullyComposeAffineMapAndOperands(&map, &allOperands);
  unsigned numNewDims = map.getNumDims();
  map = map.shiftDims(dims.size()).shiftSymbols(symbols.size());
  llvm::append_range(dims, ArrayRef<Value>(allOperands).take_front(numNewDims));
  llvm::append_range(symbols, ArrayRef<Value>(allOperands).drop_front(numNewDims));
  exprs.push_back(map.getResult(0));
}

// At this point exprs, dims and symbols are aligned potentially with duplicates.
// Run one last normalization step to get rid of duplicates and then call createOrFold.

?

[ftynse]

Address review.

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

739

I would generally like to have better support for folding IR on construction, but we need only a small subset of what would be required for a full-fledged AutoFoldIRBuilder.

747

Great idea, thanks! I wish I thought about that before spelling out all the logic here.

[nicolasvasilache]

Nice!

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

739

+1

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

 #include "mlir/IR/Builders.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 
 namespace mlir {
 class AffineApplyOp;
 class AffineBound;
 class AffineValueMap;
-class IRRewriter;
+class RewriterBase;
 
 /// TODO: These should be renamed if they are on the mlir namespace.
 ///       Ideally, they should go in a mlir::affine:: namespace.
 
 /// 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.

 /// 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`.
 AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineExpr e,
                                       ValueRange values);
 
+/// Returns an AffineMinOp obtained by composing `map` and `operands` with
+/// AffineApplyOps supplying those operands.
+Value makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map,
+                            ValueRange 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(IRRewriter &b, Location loc,
+SmallVector<OpFoldResult> applyMapToValues(RewriterBase &b, Location loc,
                                            AffineMap map,
                                            ArrayRef<OpFoldResult> 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

 
 /// Replace all occurrences of AffineExpr at position `pos` in `map` by the
 /// defining AffineApplyOp expression and operands.
 /// When `dimOrSymbolPosition < dims.size()`, AffineDimExpr@[pos] is replaced.
 /// When `dimOrSymbolPosition >= dims.size()`,
 /// AffineSymbolExpr@[pos - dims.size()] is replaced.
 /// Mutate `map`,`dims` and `syms` in place as follows:
 ///   1. `dims` and `syms` are only appended to.
-///   2. `map` dim and symbols are gradually shifted to higer positions.
+///   2. `map` dim and symbols are gradually shifted to higher positions.
 ///   3. Old `dim` and `sym` entries are replaced by nullptr
 /// This avoids the need for any bookkeeping.
 static LogicalResult replaceDimOrSym(AffineMap *map,
                                      unsigned dimOrSymbolPosition,
                                      SmallVectorImpl<Value> &dims,
                                      SmallVectorImpl<Value> &syms) {
   bool isDimReplacement = (dimOrSymbolPosition < dims.size());
   unsigned pos = isDimReplacement ? dimOrSymbolPosition

 
 AffineApplyOp mlir::makeComposedAffineApply(OpBuilder &b, Location loc,
                                             AffineExpr e, ValueRange values) {
   return makeComposedAffineApply(
       b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{e}).front(),
       values);
 }
 
+Value mlir::makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map,
+                                  ValueRange operands) {
+  // 1. Compose maps and operands for individual expressions. Store results as
+  // single-expression maps in order to keep track of the numbers of dimensions
+  // and symbols used in each.
+  SmallVector<SmallVector<Value>, 2> exprOperands;
+  SmallVector<AffineMap, 2> exprs;
+  exprOperands.reserve(map.getNumResults());
+  exprs.reserve(map.getNumResults());
+  for (unsigned i : llvm::seq<unsigned>(0, map.getNumResults())) {
+    AffineMap &exprMap = exprs.emplace_back(map.getSubMap({i}));
+    fullyComposeAffineMapAndOperands(
+        &exprMap, &exprOperands.emplace_back(llvm::to_vector(operands)));
+  }
+
+  // 2. Collect unique dimension and symbol operands used across individual
+  // expressions. We don't expect these sets to overlap because composition
+  // implies canonicalization that would have promoted the same operand to be a
+  // symbol based on its affine category (that only depends on where the operand
+  // value is defined) identically across all expressions.
+  SetVector<Value> normalizedDimOperandSet;
+  SetVector<Value> normalizedSymbolOperandSet;
+  for (unsigned i : llvm::seq<unsigned>(0, map.getNumResults())) {
+    auto *it = exprOperands[i].begin() + exprs[i].getNumDims();
+    normalizedDimOperandSet.insert(exprOperands[i].begin(), it);
+    normalizedSymbolOperandSet.insert(it, exprOperands[i].end());
+  }
+
+  // 3. Create a single list of unique operands containing dimensions followed
+  // by symbols.
+  unsigned numDims = normalizedDimOperandSet.size();
+  unsigned numSymbols = normalizedSymbolOperandSet.size();
+  SmallVector<Value> normalizedOperands =
+      llvm::to_vector(normalizedDimOperandSet);
+  llvm::append_range(normalizedOperands, normalizedSymbolOperandSet);
+
+  // 4. For each composed expression, remap its dimension and symbols to other
+  // dimensions and symbols that correspond to the positions of the operands
+  // specific to this expression in the normalized operand list.
+  SmallVector<AffineExpr> normalizedExprs;
+  normalizedExprs.reserve(exprs.size());
+  for (unsigned i : llvm::seq<unsigned>(0, exprs.size())) {
+    SmallVector<AffineExpr> dimReplacements;
+    SmallVector<AffineExpr> symReplacements;
+    for (const auto &en : llvm::enumerate(exprOperands[i])) {
+      Value operand = en.value();
+      unsigned position = std::distance(
+          normalizedOperands.begin(), llvm::find(normalizedOperands, operand));
+      AffineExpr replacement = position < numDims
+                                   ? b.getAffineDimExpr(position)
+                                   : b.getAffineSymbolExpr(position - numDims);
+      if (en.index() < exprs[i].getNumDims())
+        dimReplacements.push_back(replacement);
+      else
+        symReplacements.push_back(replacement);
+    }
+    normalizedExprs.push_back(exprs[i].getResult(0).replaceDimsAndSymbols(
+        dimReplacements, symReplacements));
+  }
+
+  // 5. Construct an affine.min with normalized expressions and operands. Note
+  // that it may even fold to a constant thanks to normalization.
+  auto normalizedMap =
+      AffineMap::get(numDims, numSymbols, normalizedExprs, b.getContext());
+  return b.createOrFold<AffineMinOp>(loc, b.getIndexType(), normalizedMap,
+                                     normalizedOperands);
+}
+
 /// 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);
SmallVector<Value, 4> mlir::applyMapToValues(OpBuilder &b, Location loc,
   // 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;
 }
 
-SmallVector<OpFoldResult>
-mlir::applyMapToValues(IRRewriter &b, Location loc, AffineMap map,
-                       ArrayRef<OpFoldResult> values) {
-  // Materialize constants and keep track of produced operations so we can clean
-  // them up later.
-  SmallVector<Operation *> constants;
-  SmallVector<Value> actualValues;
+/// 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.
+static void materializeConstants(OpBuilder &b, Location loc,
+                                 ArrayRef<OpFoldResult> values,
+                                 SmallVectorImpl<Operation *> &constants,
+                                 SmallVectorImpl<Value> &actualValues) {
   actualValues.reserve(values.size());
   auto *dialect = b.getContext()->getLoadedDialect<AffineDialect>();
   for (OpFoldResult ofr : values) {
     if (auto value = ofr.dyn_cast<Value>()) {
       actualValues.push_back(value);
       continue;
     }
     constants.push_back(dialect->materializeConstant(b, ofr.get<Attribute>(),
                                                      b.getIndexType(), loc));
     actualValues.push_back(constants.back()->getResult(0));
   }
+}
 
-  // Compose, fold and construct maps for each result independently because they
-  // may simplify more effectively.
-  SmallVector<OpFoldResult> results;
-  results.reserve(map.getNumResults());
-  bool foldedAll = true;
-  for (auto i : llvm::seq<unsigned>(0, map.getNumResults())) {
-    AffineMap submap = map.getSubMap({i});
-    SmallVector<Value> operands = actualValues;
-    fullyComposeAffineMapAndOperands(&submap, &operands);
-    canonicalizeMapAndOperands(&submap, &operands);
-
-    // 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 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`.
+template <typename OpTy, typename... Args>
+static std::enable_if_t<OpTy::template hasTrait<OpTrait::OneResult>(),
+                        OpFoldResult>
+createOrFold(RewriterBase &b, Location loc, ValueRange operands,
+             Args &&...leadingArguments) {
 
-    // Create an apply operation and immediately attempt to fold it. On sucess,
-    // delete the operation and prepare the (unmaterialized) value for being
-    // returned. On failure, return the function result.
-    // TODO: arguably, the main folder (createOrFold) API should support this
-    // use case instead of indiscriminately materializing constants.
-    auto apply = b.create<AffineApplyOp>(loc, submap, operands);
-    SmallVector<OpFoldResult, 1> foldResult;
-    if (succeeded(apply->fold(constantOperands, foldResult))) {
-      assert(foldResult.size() == 1 && "expected single-result map");
-      b.eraseOp(apply);
-      results.push_back(foldResult.front());
-    } else {
-      results.push_back(apply.getResult());
-      foldedAll = false;
-    }
+  // Create the operation and immediately attempt to fold it. On success,
+  // returned. On failure, return the operation result value.
+  // TODO: arguably, the main folder (createOrFold) API should support this use
+  // case instead of indiscriminately materializing constants.
+  OpTy op =
+      b.create<OpTy>(loc, std::forward<Args>(leadingArguments)..., operands);
+  SmallVector<OpFoldResult, 1> foldResults;
+  if (succeeded(op->fold(constantOperands, foldResults)) &&
+      !foldResults.empty()) {
+    b.eraseOp(op);
+    return foldResults.front();
+  return op->getResult(0);
+}
+
+SmallVector<OpFoldResult>
+mlir::applyMapToValues(RewriterBase &b, Location loc, AffineMap map,
+                       ArrayRef<OpFoldResult> values) {
+  // Materialize constants and keep track of produced operations so we can clean
+  // them up later.
+  SmallVector<Operation *> constants;
+  SmallVector<Value> actualValues;
+  materializeConstants(b, loc, values, constants, actualValues);
+
+  // Compose, fold and construct maps for each result independently because they
+  // may simplify more effectively.
+  SmallVector<OpFoldResult> results;
+  results.reserve(map.getNumResults());
+  bool foldedAll = true;
+  for (auto i : llvm::seq<unsigned>(0, map.getNumResults())) {
+    AffineMap submap = map.getSubMap({i});
+    SmallVector<Value> operands = actualValues;
+    fullyComposeAffineMapAndOperands(&submap, &operands);
+    canonicalizeMapAndOperands(&submap, &operands);
+    results.push_back(createOrFold<AffineApplyOp>(b, loc, operands, submap));
+    if (!results.back().is<Attribute>())
+      foldedAll = false;
   }
 
   // If the entire map could be folded, remove the constants that were used in
   // the initial ops.
   if (foldedAll) {
     for (Operation *constant : constants)
       b.eraseOp(constant);
   }

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

DiagnosedSilenceableFailure transform::MultiTileSizesOp::applyToOne(
   OpFoldResult targetSize = builder.getIndexAttr(getTargetSize());
   OpFoldResult divisor = builder.getIndexAttr(getDivisor());
   FailureOr<MultiSizeSpecification> spec = computeMultiTileSizes(
       builder, target, getDimension(), targetSize, divisor);
   if (failed(spec)) {
     return emitSilenceableError() << "could not generate tile size computation";
   }
 
+  AffineExpr s0 = builder.getAffineSymbolExpr(0);
+  AffineExpr s1 = builder.getAffineSymbolExpr(1);
   Operation *splitPoint =
-      builder
-          .createOrFold<arith::MulIOp>(target.getLoc(), spec->lowTileSize,
-                                       spec->lowTripCount)
-          .getDefiningOp();
+      makeComposedAffineApply(builder, target.getLoc(), s0 * s1,
+                              {spec->lowTileSize, spec->lowTripCount});
   Operation *lowTileSize = spec->lowTileSize.getDefiningOp();
   Operation *highTileSize = spec->highTileSize.getDefiningOp();
   assert(lowTileSize && highTileSize && splitPoint &&
          "tile sizes are not produced by operations");
   results.reserve(results.size() + 3);
   results.push_back(lowTileSize);
   results.push_back(highTileSize);
   results.push_back(splitPoint);

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

 //===- Split.cpp - Structured op splitting --------------------------------===//
 //
 // 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/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 
 #include "llvm/ADT/STLExtras.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 /// Extract the slices of `operands` supplied to the given operation `op` such
 /// that they are sufficient to execute the op for the subset of its iteration
for (OpOperand *opOperand : op.getInputAndOutputOperands()) {
     // dimension that is not used in indexing this type, just use the entire
     // operand.
     if (!type || dimension >= indexing.getNumDims() ||
         !indexing.isFunctionOfDim(dimension)) {
       slices.push_back(opOperand->get());
       continue;
     }
 
-    SmallVector<Value, 4> sizes =
-        applyMapToValues(builder, op.getLoc(), indexing, splitIterationSpace);
+    SmallVector<Value> sizes;
+    sizes.reserve(indexing.getNumResults());
+    for (AffineExpr dimIndexing : indexing.getResults()) {
+      AffineApplyOp sizeOp = makeComposedAffineApply(
+          builder, op.getLoc(), dimIndexing, splitIterationSpace);
+      assert(sizeOp->getNumResults() == 1 &&
+             "expected single-result affine map");
+      sizes.push_back(sizeOp.getResult());
+    }
     SmallVector<OpFoldResult> offsets(type.getRank(), builder.getIndexAttr(0));
     SmallVector<OpFoldResult> strides(type.getRank(), builder.getIndexAttr(1));
 
     if (offset) {
       offsets[dimension] = offset;
       IRRewriter rewriter(builder);
       offsets = applyMapToValues(rewriter, builder.getLoc(), indexing, offsets);
     }
std::pair<LinalgOp, LinalgOp> linalg::splitOp(RewriterBase &rewriter,
   SmallVector<Value, 4> iterationSpaceShapes =
       applyMapToValues(builder, op.getLoc(), shapesToLoops, allShapes);
 
   // Update the iteration space to have `splitPoint` as the size of `dimension`
   // and use it to slice operands and results for a new, smaller instance of the
   // `op`. Adjust the size if necessary to prevent overflows. Insert the partial
   // results back.
   Value splitPointValue = materializeOpFoldResult(builder, splitPoint);
-  splitPointValue = builder.createOrFold<AffineMinOp>(
-      builder.getIndexType(),
+  splitPointValue = makeComposedAffineMin(
+      builder, builder.getLoc(),
       AffineMap::getMultiDimIdentityMap(/*numDims=*/2, builder.getContext()),
       ValueRange({splitPointValue, iterationSpaceShapes[dimension]}));
   SmallVector<Value> splitIterationSpace =
       llvm::to_vector(iterationSpaceShapes);
   SmallVector<Value> originalResults = llvm::to_vector(
       llvm::map_range(op.getOutputOperands(),
                       [](OpOperand *opOperand) { return opOperand->get(); }));
   SmallVector<Value> firstResults;

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

func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
 
   // CHECK: return %[[RES]]
   return %0 : tensor<100xf32>
 }
 
 // CHECK-LABEL: @one_d_static_overflow
 // CHECK-SAME:  %[[IN:.+]]: tensor<10xf32>, %[[OUT:.+]]: tensor<10xf32>
 func.func @one_d_static_overflow(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> {
-  // CHECK: %[[IN_SLICE_LOW:.+]] = tensor.extract_slice %[[IN]][0] [10] [1] : tensor<10xf32> to tensor<10xf32>
-  // CHECK: %[[OUT_SLICE_LOW:.+]] = tensor.extract_slice %[[OUT]][0] [10] [1] : tensor<10xf32> to tensor<10xf32>
-  // CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic
-  // CHECK:   ins(%[[IN_SLICE_LOW]]
-  // CHECK:   outs(%[[OUT_SLICE_LOW]]
+  // Due to overflow, the first part of the split computes everything and the
+  // insert/extract slices are folded away by the canonicalizer.
+  // CHECK: %[[RES_PARTIAL:.+]] = linalg.generic
+  // CHECK:   ins(%[[IN]]
+  // CHECK:   outs(%[[OUT]]
   // CHECK:   linalg.index 0
   // CHECK:   func.call @elem
-  // CHECK: %[[RES_PARTIAL:.+]] = tensor.insert_slice %[[RES_SLICE_LOW]] into %[[OUT]][0] [10] [1]
   //
+  // The second part operates on zero-sized slices that are not currently
+  // folded away.
   // CHECK: %[[IN_SLICE_HIGH:.+]] = tensor.extract_slice %[[IN]][10] [0] [1] : tensor<10xf32> to tensor<0xf32>
   // CHECK: %[[OUT_SLICE_HIGH:.+]] = tensor.extract_slice %[[RES_PARTIAL]][10] [0] [1] : tensor<10xf32> to tensor<0xf32>
   // CHECK: %[[RES_SLICE_HIGH:.+]] = linalg.generic
   // CHECK:   ins(%[[IN_SLICE_HIGH]]
   // CHECK:   outs(%[[OUT_SLICE_HIGH]]
   // CHECK:   %[[IDX:.+]] = linalg.index 0
   // CHECK:   affine.apply #[[$ADD_10_MAP]](%[[IDX]])
   // CHECK:   func.call @elem
^bb1(%arg1: !pdl.operation):
     %0 = transform.pdl_match @linalg_generic in %arg1
     %1 = transform.pdl_match @func_call in %arg1
     transform.structured.split %0 after %1 { dimension = 0 }
   }
 }
 
 func.func private @get_size() -> index
 
-// CHECK: #[[$MAP_MIN_100:.+]] = affine_map<(d0, d1) -> (d0, 100)>
-// CHECK: #[[$MAP_S_MINUS_100:.+]] = affine_map<()[s0] -> (-s0 + 100)>
+// CHECK-DAG: #[[$MAP_MIN_100:.+]] = affine_map<()[s0] -> (100, s0)>
+// CHECK-DAG: #[[$MAP_S_MINUS_100:.+]] = affine_map<()[s0] -> (-s0 + 100)>
 
 // CHECK-LABEL: @dynamic
 func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
   // CHECK: %[[SPLIT:.+]] = call @get_size
-  // CHECK: %[[SPLIT_LOW:.+]] = affine.min #[[$MAP_MIN_100]](%[[SPLIT]]
+  // CHECK: %[[SPLIT_LOW:.+]] = affine.min #[[$MAP_MIN_100]]()[%[[SPLIT]]]
   // CHECK: %[[IN_SLICE_LOW:.+]] = tensor.extract_slice %[[IN:.+]][0] [%[[SPLIT_LOW]]] [1] : tensor<100xf32> to tensor<?xf32>
   // CHECK: %[[OUT_SLICE_LOW:.+]] = tensor.extract_slice %[[OUT:.+]][0] [%[[SPLIT_LOW]]] [1] : tensor<100xf32> to tensor<?xf32>
   // CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic
   // CHECK:   ins(%[[IN_SLICE_LOW]]
   // CHECK:   outs(%[[OUT_SLICE_LOW]]
   // CHECK: %[[PARTIAL:.+]] = tensor.insert_slice %[[RES_SLICE_LOW]] into %[[OUT]][0] [%[[SPLIT_LOW]]] [1]
   //
   // CHECK: %[[SPLIT_HIGH_1:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]]
+  // CHECK: %[[IN_SLICE_HIGH:.+]] = tensor.extract_slice %[[IN:.+]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_1]]] [1] : tensor<100xf32> to tensor<?xf32>
   // CHECK: %[[SPLIT_HIGH_2:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]]
-  // CHECK: %[[IN_SLICE_HIGH:.+]] = tensor.extract_slice %[[IN:.+]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_2]]] [1] : tensor<100xf32> to tensor<?xf32>
-  // CHECK: %[[SPLIT_HIGH_3:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]]
-  // CHECK: %[[OUT_SLICE_HIGH:.+]] = tensor.extract_slice %[[PARTIAL:.+]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_3]]] [1] : tensor<100xf32> to tensor<?xf32>
+  // CHECK: %[[OUT_SLICE_HIGH:.+]] = tensor.extract_slice %[[PARTIAL:.+]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_2]]] [1] : tensor<100xf32> to tensor<?xf32>
   // CHECK: %[[RES_SLICE_HIGH:.+]] = linalg.generic
   // CHECK:   ins(%[[IN_SLICE_HIGH]]
   // CHECK:   outs(%[[OUT_SLICE_HIGH]]
-  // CHECK: tensor.insert_slice %[[RES_SLICE_HIGH]] into %[[PARTIAL]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_3]]] [1]
+  // CHECK: tensor.insert_slice %[[RES_SLICE_HIGH]] into %[[PARTIAL]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_2]]] [1]
   %0 = func.call @get_size() : () -> index
   %1 = linalg.generic {
     indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
     iterator_types = ["parallel"]
   }
   ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) {
   ^bb0(%3: f32, %4: f32):
-    linalg.yield %3 : f32
+    %5 = arith.addf %3, %4 : f32
+    linalg.yield %5 : f32
   } -> tensor<100xf32>
   return %1 : tensor<100xf32>
 }
 
 // -----
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):