[MLIR] Fix affine.for unroll for multi-result upper bound maps

[bondhugula]

Summary

Fix affine.for unroll for multi-result upper bound maps: these can't be
unrolled/unroll-and-jammed in cases where the trip count isn't known to
be a multiple of the unroll factor.

Fix and clean up repeated/unnecessary checks/comments at helper callees.

Also, fix clang-tidy variable naming warnings and redundant includes.

[dcaballe]

Thanks, Uday. LGTM but I'm not super familiar with the loop unrolling pass. I'll defer the approval to someone with more expertise in the pass.

[vinayaka-polymage]

Looks good, minor request to complicate the test case, thanks!

mlir/lib/Transforms/Utils/LoopUtils.cpp

1177–1178

Comment needs an update ?

mlir/test/Dialect/Affine/unroll.mlir

585

Can you please complicate the test case more ?

func @unroll_multi_upper_bound(%arg0 : index) {
  affine.for %i = 0 to min affine_map<()[s0] -> (8 * s0, 12 * s0)>()[%arg0] {
    "test.foo"(%i) : (index) -> ()
  }
  // UNROLL-BY-4: affine.for %{{.*}} = 0 to min #map{{.*}}() step 4
  // UNROLL-BY-4-NOT: affine.for
  return
}

[ayzhuang]

Looks great, could you add an unrolling test where the lower bound map has multiple results and the trip count is a multiple of the unroll factor? Thanks.

[bondhugula]

In D114662#3159058, @ayzhuang wrote:

Looks great, could you add an unrolling test where the lower bound map has multiple results and the trip count is a multiple of the unroll factor? Thanks.

This is a good point. Although this revision enables go past that unnecessary initial check, the unrolling of loops with multi-result lower bound maps where the trip count is a multiple of unroll factor still won't happen since getTripCountMapAndOperands isn't yet extended for multi-result lower bound maps. There is this in LoopAnalysis.cpp:

53   if (lbMap.getNumResults() != 1) {
54     *tripCountMap = AffineMap();
55     return;
56   }

There isn't a reason for this restriction and this has been a TODO.

But I've added a test case and marked it TODO.

mlir/lib/Transforms/Utils/LoopUtils.cpp

1177–1178

This is already covered in the comment below and so I didn't want to duplicate the specific subset where it isn't possible.

mlir/test/Dialect/Affine/unroll.mlir

585

Good point, done.

[bondhugula]

Address review comments.

[ayzhuang]

LGTM, thanks.

mlir/lib/Transforms/Utils/LoopUtils.cpp

 #include "mlir/Transforms/LoopUtils.h"
 
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
-#include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BlockAndValueMapping.h"
-#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "mlir/Transforms/Utils.h"
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
-#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "LoopUtils"
 
 using namespace mlir;
 using llvm::SmallMapVector;
 
 namespace {
 // This structure is to pass and return sets of loop parameters without
 // confusing the order.
 struct LoopParams {
   Value lowerBound;
   Value upperBound;
   Value step;
 };
 } // namespace
 
 /// Computes the cleanup loop lower bound of the loop being unrolled with
 /// the specified unroll factor; this bound will also be upper bound of the main
 /// part of the unrolled loop. Computes the bound as an AffineMap with its
 /// operands or a null map when the trip count can't be expressed as an affine
 /// expression.
-static void getCleanupLoopLowerBound(AffineForOp forOp, unsigned unrollFactor,
-                                     AffineMap &map,
-                                     SmallVectorImpl<Value> &operands) {
-  auto lbMap = forOp.getLowerBoundMap();
-
-  // Single result lower bound map only.
-  if (lbMap.getNumResults() != 1) {
-    map = AffineMap();
-    return;
-  }
-
+static void
+getCleanupLoopLowerBound(AffineForOp forOp, unsigned unrollFactor,
+                         AffineMap &cleanupLbMap,
+                         SmallVectorImpl<Value> &cleanupLbOperands) {
   AffineMap tripCountMap;
   SmallVector<Value, 4> tripCountOperands;
   getTripCountMapAndOperands(forOp, &tripCountMap, &tripCountOperands);
-
-  // Sometimes the trip count cannot be expressed as an affine expression.
+  // Trip count can't be computed.
   if (!tripCountMap) {
-    map = AffineMap();
+    cleanupLbMap = AffineMap();
     return;
   }
 
   OpBuilder b(forOp);
+  auto lbMap = forOp.getLowerBoundMap();
   auto lb = b.create<AffineApplyOp>(forOp.getLoc(), lbMap,
                                     forOp.getLowerBoundOperands());
 
   // For each upper bound expr, get the range.
   // Eg: affine.for %i = lb to min (ub1, ub2),
   // where tripCountExprs yield (tr1, tr2), we create affine.apply's:
   // lb + tr1 - tr1 % ufactor, lb + tr2 - tr2 % ufactor; the results of all
   // these affine.apply's make up the cleanup loop lower bound.
   SmallVector<AffineExpr, 4> bumpExprs(tripCountMap.getNumResults());
   SmallVector<Value, 4> bumpValues(tripCountMap.getNumResults());
   int64_t step = forOp.getStep();
   for (unsigned i = 0, e = tripCountMap.getNumResults(); i < e; i++) {
     auto tripCountExpr = tripCountMap.getResult(i);
     bumpExprs[i] = (tripCountExpr - tripCountExpr % unrollFactor) * step;
     auto bumpMap = AffineMap::get(tripCountMap.getNumDims(),
                                   tripCountMap.getNumSymbols(), bumpExprs[i]);
     bumpValues[i] =
         b.create<AffineApplyOp>(forOp.getLoc(), bumpMap, tripCountOperands);
   }
 
   SmallVector<AffineExpr, 4> newUbExprs(tripCountMap.getNumResults());
   for (unsigned i = 0, e = bumpExprs.size(); i < e; i++)
     newUbExprs[i] = b.getAffineDimExpr(0) + b.getAffineDimExpr(i + 1);
 
-  operands.clear();
-  operands.push_back(lb);
-  operands.append(bumpValues.begin(), bumpValues.end());
-  map = AffineMap::get(1 + tripCountMap.getNumResults(), 0, newUbExprs,
+  cleanupLbOperands.clear();
+  cleanupLbOperands.push_back(lb);
+  cleanupLbOperands.append(bumpValues.begin(), bumpValues.end());
+  cleanupLbMap = AffineMap::get(1 + tripCountMap.getNumResults(), 0, newUbExprs,
                                 b.getContext());
-  // Simplify the map + operands.
-  fullyComposeAffineMapAndOperands(&map, &operands);
-  map = simplifyAffineMap(map);
-  canonicalizeMapAndOperands(&map, &operands);
+  // Simplify the cleanupLbMap + cleanupLbOperands.
+  fullyComposeAffineMapAndOperands(&cleanupLbMap, &cleanupLbOperands);
+  cleanupLbMap = simplifyAffineMap(cleanupLbMap);
+  canonicalizeMapAndOperands(&cleanupLbMap, &cleanupLbOperands);
   // Remove any affine.apply's that became dead from the simplification above.
   for (auto v : bumpValues)
     if (v.use_empty())
       v.getDefiningOp()->erase();
 
   if (lb.use_empty())
     lb.erase();
 }
if (!iv.use_empty()) {
     if (forOp.hasConstantLowerBound()) {
       OpBuilder topBuilder(forOp->getParentOfType<FuncOp>().getBody());
       auto constOp = topBuilder.create<arith::ConstantIndexOp>(
           forOp.getLoc(), forOp.getConstantLowerBound());
       iv.replaceAllUsesWith(constOp);
     } else {
       auto lbOperands = forOp.getLowerBoundOperands();
       auto lbMap = forOp.getLowerBoundMap();
-      OpBuilder builder(parentBlock, Block::iterator(forOp));
+      OpBuilder builder(forOp);
       if (lbMap == builder.getDimIdentityMap()) {
         // No need of generating an affine.apply.
         iv.replaceAllUsesWith(lbOperands[0]);
       } else {
         auto affineApplyOp =
             builder.create<AffineApplyOp>(forOp.getLoc(), lbMap, lbOperands);
         iv.replaceAllUsesWith(affineApplyOp);
       }
for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++)
     annotateFn(0, &*it, builder);
 
   // Update operands of the yield statement.
   loopBodyBlock->getTerminator()->setOperands(lastYielded);
 }
 
 /// Helper to generate cleanup loop for unroll or unroll-and-jam when the trip
 /// count is not a multiple of `unrollFactor`.
-static void generateCleanupLoopForUnroll(AffineForOp forOp,
+static LogicalResult generateCleanupLoopForUnroll(AffineForOp forOp,
                                                   uint64_t unrollFactor) {
   // Insert the cleanup loop right after 'forOp'.
   OpBuilder builder(forOp->getBlock(), std::next(Block::iterator(forOp)));
   auto cleanupForOp = cast<AffineForOp>(builder.clone(*forOp));
 
   // Update uses of `forOp` results. `cleanupForOp` should use `forOp` result
   // and produce results for the original users of `forOp` results.
   auto results = forOp.getResults();
   auto cleanupResults = cleanupForOp.getResults();
   auto cleanupIterOperands = cleanupForOp.getIterOperands();
 
   for (auto e : llvm::zip(results, cleanupResults, cleanupIterOperands)) {
     std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
     cleanupForOp->replaceUsesOfWith(std::get<2>(e), std::get<0>(e));
   }
 
   AffineMap cleanupMap;
   SmallVector<Value, 4> cleanupOperands;
   getCleanupLoopLowerBound(forOp, unrollFactor, cleanupMap, cleanupOperands);
-  assert(cleanupMap &&
-         "cleanup loop lower bound map for single result lower bound maps "
-         "can always be determined");
+  if (!cleanupMap)
+    return failure();
+
   cleanupForOp.setLowerBound(cleanupOperands, cleanupMap);
   // Promote the loop body up if this has turned into a single iteration loop.
   (void)promoteIfSingleIteration(cleanupForOp);
 
   // Adjust upper bound of the original loop; this is the same as the lower
   // bound of the cleanup loop.
   forOp.setUpperBound(cleanupOperands, cleanupMap);
+  return success();
 }
 
 /// Unrolls this loop by the specified factor. Returns success if the loop
 /// is successfully unrolled.
 LogicalResult mlir::loopUnrollByFactor(
     AffineForOp forOp, uint64_t unrollFactor,
     function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn) {
   assert(unrollFactor > 0 && "unroll factor should be positive");
 
   if (unrollFactor == 1)
     return promoteIfSingleIteration(forOp);
 
   // Nothing in the loop body other than the terminator.
   if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
     return success();
 
-  // Loops where the lower bound is a max expression isn't supported for
-  // unrolling since the trip count can be expressed as an affine function when
-  // both the lower bound and the upper bound are multi-result maps. However,
-  // one meaningful way to do such unrolling would be to specialize the loop for
-  // the 'hotspot' case and unroll that hotspot.
-  if (forOp.getLowerBoundMap().getNumResults() != 1)
-    return failure();
-
   // If the trip count is lower than the unroll factor, no unrolled body.
   // TODO: option to specify cleanup loop unrolling.
   Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
   if (mayBeConstantTripCount.hasValue() &&
       mayBeConstantTripCount.getValue() < unrollFactor)
     return failure();
 
   // Generate the cleanup loop if trip count isn't a multiple of unrollFactor.
-  if (getLargestDivisorOfTripCount(forOp) % unrollFactor != 0)
-    generateCleanupLoopForUnroll(forOp, unrollFactor);
+  if (getLargestDivisorOfTripCount(forOp) % unrollFactor != 0) {
+    // Loops where the lower bound is a max expression or the upper bound is
+    // a min expression and the trip count doesn't divide the unroll factor
+    // can't be unrolled since the lower bound of the cleanup loop in such cases
+    // cannot be expressed as an affine function or a max over affine functions.
+    if (forOp.getLowerBoundMap().getNumResults() != 1 ||
+        forOp.getUpperBoundMap().getNumResults() != 1)
+      return failure();
+    if (failed(generateCleanupLoopForUnroll(forOp, unrollFactor)))
+      assert(false && "cleanup loop lower bound map for single result lower "
+                      "and upper bound maps can always be determined");
+  }
 
   ValueRange iterArgs(forOp.getRegionIterArgs());
   auto yieldedValues = forOp.getBody()->getTerminator()->getOperands();
 
   // Scale the step of loop being unrolled by unroll factor.
   int64_t step = forOp.getStep();
   forOp.setStep(step * unrollFactor);
   generateUnrolledLoop(
LogicalResult mlir::loopUnrollJamByFactor(AffineForOp forOp,
 
   if (unrollJamFactor == 1)
     return promoteIfSingleIteration(forOp);
 
   // Nothing in the loop body other than the terminator.
   if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
     return success();
 
-  // Loops where both lower and upper bounds are multi-result maps won't be
-  // unrolled (since the trip can't be expressed as an affine function in
-  // general).
-  // TODO: this may not be common, but we could support the case
-  // where the lower bound is a multi-result map and the ub is a single result
-  // one.
-  if (forOp.getLowerBoundMap().getNumResults() != 1)
-    return failure();
-
   Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
   // If the trip count is lower than the unroll jam factor, no unroll jam.
   if (mayBeConstantTripCount.hasValue() &&
       mayBeConstantTripCount.getValue() < unrollJamFactor) {
     LLVM_DEBUG(llvm::dbgs() << "[failed] trip count < unroll-jam factor\n");
     return failure();
   }
 
LogicalResult mlir::loopUnrollJamByFactor(AffineForOp forOp,
 
   // Get supported reductions to be used for creating reduction ops at the end.
   SmallVector<LoopReduction> reductions;
   if (forOp.getNumIterOperands() > 0)
     getSupportedReductions(forOp, reductions);
 
   // Generate the cleanup loop if trip count isn't a multiple of
   // unrollJamFactor.
-  if (getLargestDivisorOfTripCount(forOp) % unrollJamFactor != 0)
-    generateCleanupLoopForUnroll(forOp, unrollJamFactor);
+  if (getLargestDivisorOfTripCount(forOp) % unrollJamFactor != 0) {
+    // Loops where the lower bound is a max expression or the upper bound is
+    // a min expression and the trip count doesn't divide the unroll factor
+    // can't be unrolled since the lower bound of the cleanup loop in such cases
+    // cannot be expressed as an affine function or a max over affine functions.
+    if (forOp.getLowerBoundMap().getNumResults() != 1 ||
+        forOp.getUpperBoundMap().getNumResults() != 1)
+      return failure();
+    if (failed(generateCleanupLoopForUnroll(forOp, unrollJamFactor)))
+      assert(false && "cleanup loop lower bound map for single result lower "
+                      "and upper bound maps can always be determined");
+  }
 
   // `operandMaps[i - 1]` carries old->new operand mapping for the ith unrolled
   // iteration. There are (`unrollJamFactor` - 1) iterations.
   SmallVector<BlockAndValueMapping, 4> operandMaps(unrollJamFactor - 1);
 
   // For any loop with iter_args, replace it with a new loop that has
   // `unrollJamFactor` copies of its iterOperands, iter_args and yield
   // operands.

mlir/test/Dialect/Affine/unroll-jam.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll-jam="unroll-jam-factor=2" | FileCheck %s
 // RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll-jam="unroll-jam-factor=4" | FileCheck --check-prefix=UJAM-FOUR %s
 
 // CHECK-DAG: [[$MAP_PLUS_1:#map[0-9]+]] = affine_map<(d0) -> (d0 + 1)>
 // CHECK-DAG: [[$MAP_DIV_OFFSET:#map[0-9]+]] = affine_map<()[s0] -> (((s0 - 1) floordiv 2) * 2 + 1)>
-// CHECK-DAG: [[$MAP_MULTI_RES:#map[0-9]+]] = affine_map<()[s0, s1] -> ((s0 floordiv 2) * 2, (s1 floordiv 2) * 2, 1024)>
 // CHECK-DAG: [[$MAP_SYM_UB:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0, s1, 1024)>
 
 // UJAM-FOUR-DAG: [[$UBMAP:#map[0-9]+]] = affine_map<()[s0] -> (s0 + 8)>
 // UJAM-FOUR-DAG: [[$MAP_PLUS_1:#map[0-9]+]] = affine_map<(d0) -> (d0 + 1)>
 // UJAM-FOUR-DAG: [[$MAP_PLUS_2:#map[0-9]+]] = affine_map<(d0) -> (d0 + 2)>
 // UJAM-FOUR-DAG: [[$MAP_PLUS_3:#map[0-9]+]] = affine_map<(d0) -> (d0 + 3)>
 
 // CHECK-LABEL: func @unroll_jam_imperfect_nest() {

 
 // CHECK-LABEL: func @loop_nest_symbolic_and_min_upper_bound
 // CHECK-SAME: [[M:arg[0-9]+]]: index
 // CHECK-SAME: [[N:arg[0-9]+]]: index
 // CHECK-SAME: [[K:arg[0-9]+]]: index
 func @loop_nest_symbolic_and_min_upper_bound(%M : index, %N : index, %K : index) {
   affine.for %i = 0 to min affine_map<()[s0, s1] -> (s0, s1, 1024)>()[%M, %N] {
     affine.for %j = 0 to %K {
-      "foo"(%i, %j) : (index, index) -> ()
+      "test.foo"(%i, %j) : (index, index) -> ()
     }
   }
   return
 }
-// CHECK-NEXT:  affine.for [[IV0:%arg[0-9]+]] = 0 to min [[$MAP_MULTI_RES]]()[%[[M]], %[[N]]] step 2 {
+// No unroll-and-jam possible here as the lower bound for the cleanup loop won't
+// be representable.
+// CHECK-NEXT:  affine.for [[IV0:%arg[0-9]+]] = 0 to min #map{{.*}}()[%[[M]], %[[N]]] {
 // CHECK-NEXT:    affine.for [[IV1:%arg[0-9]+]] = 0 to %[[K]] {
-// CHECK-NEXT:      "foo"([[IV0]], [[IV1]])
-// CHECK-NEXT:      [[RES:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])
-// CHECK-NEXT:      "foo"([[RES]], [[IV1]])
-// CHECK-NEXT:    }
-// CHECK-NEXT:  }
-// CHECK-NEXT:  affine.for [[IV0]] = max [[$MAP_MULTI_RES]]()[%[[M]], %[[N]]] to min [[$MAP_SYM_UB]]()[%[[M]], %[[N]]] {
-// CHECK-NEXT:    affine.for [[IV1]] = 0 to %[[K]] {
-// CHECK-NEXT:      "foo"([[IV0]], [[IV1]])
+// CHECK-NEXT:      "test.foo"([[IV0]], [[IV1]])
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
 
 // The inner loop trip count changes each iteration of outer loop.
 // Do no unroll-and-jam.
 // CHECK-LABEL: func @no_unroll_jam_dependent_ubound
 func @no_unroll_jam_dependent_ubound(%in0: memref<?xf32, 1>) {

mlir/test/Dialect/Affine/unroll.mlir

 // UNROLL-BY-4-DAG: [[$MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0 + 1)>
 // UNROLL-BY-4-DAG: [[$MAP1:#map[0-9]+]] = affine_map<(d0) -> (d0 + 2)>
 // UNROLL-BY-4-DAG: [[$MAP2:#map[0-9]+]] = affine_map<(d0) -> (d0 + 3)>
 // UNROLL-BY-4-DAG: [[$MAP3:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 + 1)>
 // UNROLL-BY-4-DAG: [[$MAP4:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 + 3)>
 // UNROLL-BY-4-DAG: [[$MAP5:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 + s0 + 1)>
 // UNROLL-BY-4-DAG: [[$MAP6:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 * 16 + d1)>
 // UNROLL-BY-4-DAG: [[$MAP11:#map[0-9]+]] = affine_map<(d0) -> (d0)>
-// UNROLL-BY-4-DAG: [[$MAP_TRIP_COUNT_MULTIPLE_FOUR:#map[0-9]+]] = affine_map<()[s0, s1, s2] -> (s0 + ((-s0 + s1) floordiv 4) * 4, s0 + ((-s0 + s2) floordiv 4) * 4, s0 + ((-s0) floordiv 4) * 4 + 1024)>
 
 // UNROLL-FULL-LABEL: func @loop_nest_simplest() {
 func @loop_nest_simplest() {
   // UNROLL-FULL: affine.for %arg0 = 0 to 100 step 2 {
   affine.for %i = 0 to 100 step 2 {
     // UNROLL-FULL: %c1_i32 = arith.constant 1 : i32
     // UNROLL-FULL-NEXT: %c1_i32_0 = arith.constant 1 : i32
     // UNROLL-FULL-NEXT: %c1_i32_1 = arith.constant 1 : i32

 // UNROLL-BY-4-NEXT: }
   }
   return
 }
 
 // UNROLL-BY-4-LABEL: func @loop_nest_symbolic_and_min_upper_bound
 func @loop_nest_symbolic_and_min_upper_bound(%M : index, %N : index, %K : index) {
   affine.for %i = %M to min affine_map<()[s0, s1] -> (s0, s1, 1024)>()[%N, %K] {
-    "foo"() : () -> ()
+    "test.foo"() : () -> ()
   }
   return
 }
-// CHECK-NEXT:  affine.for %arg0 = %arg0 to min [[$MAP_TRIP_COUNT_MULTIPLE_FOUR]]()[%arg0, %arg1, %arg2] step 4 {
-// CHECK-NEXT:    "foo"() : () -> ()
-// CHECK-NEXT:    "foo"() : () -> ()
-// CHECK-NEXT:    "foo"() : () -> ()
-// CHECK-NEXT:    "foo"() : () -> ()
-// CHECK-NEXT:  }
-// CHECK-NEXT:  affine.for %arg1 = max [[$MAP_TRIP_COUNT_MULTIPLE_FOUR]]()[%arg0, %arg1, %arg2] to min #map28()[%arg1, %arg2] {
-// CHECK-NEXT:    "foo"() : () -> ()
-// CHECK-NEXT:  }
-// CHECK-NEXT:  return
+// No unrolling here.
+// UNROLL-BY-4:      affine.for %{{.*}} = %{{.*}} to min #map{{.*}}()[%{{.*}}, %{{.*}}] {
+// UNROLL-BY-4-NEXT:   "test.foo"() : () -> ()
+// UNROLL-BY-4-NEXT: }
+// UNROLL-BY-4-NEXT: return
 
 // The trip count here is a multiple of four, but this can be inferred only
 // through composition. Check for no cleanup scf.
 // UNROLL-BY-4-LABEL: func @loop_nest_non_trivial_multiple_upper_bound
 func @loop_nest_non_trivial_multiple_upper_bound(%M : index, %N : index) {
   %T = affine.apply affine_map<(d0) -> (4*d0 + 1)>(%M)
   %K = affine.apply affine_map<(d0) -> (d0 - 1)> (%T)
   affine.for %i = 0 to min affine_map<(d0, d1) -> (4 * d0, d1, 1024)>(%N, %K) {
     "foo"() : () -> ()
   }
   return
 }
 // UNROLL-BY-4: affine.for %arg2 = 0 to min
 // UNROLL-BY-4-NOT: for
 // UNROLL-BY-4: return
 
+// UNROLL-BY-4-LABEL: func @multi_upper_bound
+func @multi_upper_bound(%arg0: index) {
+  affine.for %i = 0 to min affine_map<()[s0] -> (8 * s0, 12 * s0)>()[%arg0] {
+    "test.foo"() : () -> ()
+  }
+  // No unrolling possible here.
+  // UNROLL-BY-4: affine.for %{{.*}} = 0 to min #map{{.*}}()[%{{.*}}]
+  return
+}
+
+// UNROLL-BY-4-LABEL: func @multi_lower_bound
+func @multi_lower_bound(%arg0: index) {
+  affine.for %i = max affine_map<()[s0] -> (8 * s0, 12 * s0)>()[%arg0] to 100 {
+    "test.foo"() : () -> ()
+  }
+  // TODO: Extend getTripCountMapAndOperands to handle multi-result lower bound
+  // maps.
+  // UNROLL-BY-4: affine.for %{{.*}} = max #map{{.*}}()[%{{.*}}] to 100
+  // UNROLL-BY-4-NOT: affine.for
+  return
+}
+
 // UNROLL-BY-4-LABEL: func @loop_nest_non_trivial_multiple_upper_bound_alt
 func @loop_nest_non_trivial_multiple_upper_bound_alt(%M : index, %N : index) {
   %K = affine.apply affine_map<(d0) -> (4*d0)> (%M)
   affine.for %i = 0 to min affine_map<()[s0, s1] -> (4 * s0, s1, 1024)>()[%N, %K] {
     "foo"() : () -> ()
   }
   // UNROLL-BY-4: affine.for %arg2 = 0 to min
   // UNROLL-BY-4-NEXT: "foo"