[MLIR] Add loop coalesce utility for affine.for

[arnab-oss]

Summary

Add loop coalesce utility for affine.for. This expects loops to have
been normalized a-priori. This works for both constant as well non
constant upper bounds having single/multiple result upper bound affine
map.

With contributions from Arnab Dutta and Uday Bondhugula.

[arnab-oss]

This comment was removed by arnab-oss.

[arnab-oss]

This comment was removed by arnab-oss.

[arnab-oss]

This comment was removed by arnab-oss.

[bondhugula]

@ftynse Could you please review this? I can review too but I also contributed a part of the code!

[bondhugula]

mlir/include/mlir/Transforms/LoopUtils.h

246–249

It would good to mention in what order these loops appear - "outermost to innermost".

a list of ... loops outermost to innermost ...

mlir/lib/Transforms/LoopCoalescing.cpp

86–87

Can we handle both in one walk/pass?

mlir/lib/Transforms/Utils/LoopUtils.cpp

2116

More natural: orgiUbMap.getNumResults() == 1

2127–2128

1 can be on the first line? clang-format

2166–2185

This needs better code comments.

2167

Wrong tense here. Rephrase comment please.

2193

Avoid name like second.

[bondhugula]

This already crashes on an input like this:

#map0 = affine_map<(d0) -> (d0 * 110)>
#map1 = affine_map<(d0) -> (696, d0 * 110 + 110)>
func @test(%arg0: memref<2088x2048xf64>, %arg1: memref<2048x2048xf64>, %arg2: memref<2088x2048xf64>) { 
  affine.for %arg3 = 0 to 4 { 
    affine.for %arg4 = 0 to 7 { 
      affine.for %arg5 = 0 to 128 { 
        affine.for %arg6 = #map0(%arg4) to min #map1(%arg4) {
        }
      }
    }
  } 
  return
}

... mlir/lib/IR/AffineMap.cpp:297: int64_t mlir::AffineMap::getSingleConstantResult() const: Assertion `isSingleConstant() && "map must have a single constant result"' failed.
PLEASE submit a bug report to https://bugs.llvm.org/ and include the crash backtrace.
Stack dump:
0.	Program arguments: /home/uday/llvm-project-upstream/build/bin/mlir-opt -loop-coalescing coalesce_crash.mlir
 #0 0x00007f7fa4887c54 PrintStackTraceSignalHandler(void*) Signals.cpp:0:0
 #1 0x00007f7fa488536e SignalHandler(int) Signals.cpp:0:0
 #2 0x00007f7fa42a6210 (/lib/x86_64-linux-gnu/libc.so.6+0x46210)
 #3 0x00007f7fa42a618b raise /build/glibc-eX1tMB/glibc-2.31/signal/../sysdeps/unix/sysv/linux/raise.c:51:1
 #4 0x00007f7fa4285859 abort /build/glibc-eX1tMB/glibc-2.31/stdlib/abort.c:81:7
 #5 0x00007f7fa4285729 get_sysdep_segment_value /build/glibc-eX1tMB/glibc-2.31/intl/loadmsgcat.c:509:8
 #6 0x00007f7fa4285729 _nl_load_domain /build/glibc-eX1tMB/glibc-2.31/intl/loadmsgcat.c:970:34
 #7 0x00007f7fa4296f36 (/lib/x86_64-linux-gnu/libc.so.6+0x36f36)
 #8 0x00007f7fa495e0fb mlir::AffineMap::getSingleConstantResult() const (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRIR.so.14git+0xba0fb)
 #9 0x00007f7fa4d39248 mlir::AffineForOp::getConstantLowerBound() (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRAffine.so.14git+0x4d248)
#10 0x00007f7fa4230048 mlir::coalesceLoops(llvm::MutableArrayRef<mlir::AffineForOp>) (/home/uday/llvm-project-upstream/build/bin/../lib/../lib/libMLIRTransformUtils.so.14git+0x97048)
#11 0x00007f7fa5435045 _ZN4llvm12function_refIFvPN4mlir9OperationEEE11callback_fnIZNS1_6detail4walkILNS1_9WalkOrderE1EZN12_GLOBAL__N_118LoopCoalescingPass13runOnFunctionEvEUlNS1_11AffineForOpEE0_SC_vEENSt9enable_ifIXaantsrNS_11disjunctionIJSt7is_sameIT1_S3_ESG_ISH_PNS1_6RegionEESG_ISH_PNS1_5BlockEEEEE5valuesrSG_IT2_vE5valueESQ_E4typeES3_OT0_EUlS3_E_EEvlS3_ LoopCoalescing.cpp:0:0
#12 0x00007f7fa4a3b2bb mlir::detail::walk(mlir::Operation*, llvm::function_ref<void (mlir::Operation*)>, mlir::WalkOrder) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRIR.so.14git+0x1972bb)
#13 0x00007f7fa5435405 (anonymous namespace)::LoopCoalescingPass::runOnFunction() LoopCoalescing.cpp:0:0
#14 0x00007f7fa54009c5 mlir::FunctionPass::runOnOperation() BufferDeallocation.cpp:0:0
#15 0x00007f7fa512fcdd mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x29cdd)
#16 0x00007f7fa5130020 mlir::detail::OpToOpPassAdaptor::runPipeline(llvm::iterator_range<llvm::pointee_iterator<std::unique_ptr<mlir::Pass, std::default_delete<mlir::Pass> >*, mlir::Pass> >, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x2a020)
#17 0x00007f7fa512e971 mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x28971)
#18 0x00007f7fa512fac7 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x29ac7)
#19 0x00007f7fa5130020 mlir::detail::OpToOpPassAdaptor::runPipeline(llvm::iterator_range<llvm::pointee_iterator<std::unique_ptr<mlir::Pass, std::default_delete<mlir::Pass> >*, mlir::Pass> >, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x2a020)
#20 0x00007f7fa51313c1 mlir::PassManager::run(mlir::Operation*) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIRPass.so.14git+0x2b3c1)
#21 0x00007f7fa68291eb performActions(llvm::raw_ostream&, bool, bool, llvm::SourceMgr&, mlir::MLIRContext*, mlir::PassPipelineCLParser const&) (.constprop.0) MlirOptMain.cpp:0:0
#22 0x00007f7fa68296c1 processBuffer(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer> >, bool, bool, bool, bool, mlir::PassPipelineCLParser const&, mlir::DialectRegistry&) MlirOptMain.cpp:0:0
#23 0x00007f7fa6829b4e mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer> >, mlir::PassPipelineCLParser const&, mlir::DialectRegistry&, bool, bool, bool, bool, bool) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIROptLib.so.14git+0x7b4e)
#24 0x00007f7fa682b62b mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&, bool) (/home/uday/llvm-project-upstream/build/bin/../lib/libMLIROptLib.so.14git+0x962b)
#25 0x0000556e63a101c8 main (/home/uday/llvm-project-upstream/build/bin/mlir-opt+0x191c8)
#26 0x00007f7fa42870b3 __libc_start_main /build/glibc-eX1tMB/glibc-2.31/csu/../csu/libc-start.c:342:3
#27 0x0000556e63a0834e _start (/home/uday/llvm-project-upstream/build/bin/mlir-opt+0x1134e)
Aborted (core dumped)

The test coverage has to be improved as well.

mlir/test/Transforms/loop-coalescing.mlir

229

Instead of MAP, use IDENTITY. Easier to read.

301

Add a comment here as to which scenario it's testing or rename the function to make it descriptive.

337

Also add a negative test case - one that is not supported for coalescing to make sure it doesn't do the wrong thing.

[ayzhuang]

mlir/lib/Transforms/Utils/LoopUtils.cpp

2092

Return failure if !loop.hasConstantLowerBound()

2115

is this check needed? The upper bound map should have at least 1 result.

[arnab-oss]

Updating D108126: [MLIR] Add loop coalesce utility for affine.for

Addressed reviews

*Add extra checks so that code does not crash on running the loop
coalescing pass on unnormalized loop.

*Refactored the code and improved documentation.

*Check for and coalesce affine.for and scf.for loops in a single pass.

[arnab-oss]

This comment was removed by arnab-oss.

[arnab-oss]

This comment was removed by arnab-oss.

[bondhugula]

LGTM - thanks.

[bondhugula]

@ayzhuang Do you have any more comments here?

[ayzhuang]

LGTM, thanks.

[bondhugula]

@arnab-oss Your revision adds five whitespace errors. You can check these with a $ git diff --check HEAD~. I'll fix these before committing this time.

$ git diff --check HEAD~
mlir/test/Transforms/loop-coalescing.mlir:230: trailing whitespace.
+// Check coalescing of affine.for loops when all the loops have non constant upper bounds. 
mlir/test/Transforms/loop-coalescing.mlir:349: trailing whitespace.
+func @test_loops_do_not_get_coalesced() { 
mlir/test/Transforms/loop-coalescing.mlir:350: trailing whitespace.
+  affine.for %i = 0 to 7 { 
mlir/test/Transforms/loop-coalescing.mlir:351: trailing whitespace.
+    affine.for %j = #map0(%i) to min #map1(%i) { 
mlir/test/Transforms/loop-coalescing.mlir:353: trailing whitespace.
+    } 

mlir/include/mlir/Transforms/LoopUtils.h

 /// as defined in `sizes`.
 TileLoops extractFixedOuterLoops(scf::ForOp rootFOrOp, ArrayRef<int64_t> sizes);
 
 /// Replace a perfect nest of "for" loops with a single linearized loop. Assumes
 /// `loops` contains a list of perfectly nested loops with bounds and steps
 /// independent of any loop induction variable involved in the nest.
 void coalesceLoops(MutableArrayRef<scf::ForOp> loops);
 
+/// Replace a perfect nest of "for" loops with a single linearized loop. Assumes
+/// `loops` contains a list of perfectly nested loops that are normalized (step
+/// one and lower bound of zero) and with bounds and steps independent of any
+/// loop induction variable involved in the nest. Coalescing affine.for loops is
+/// not always possible, i.e., the result may not be representable using
+/// affine.for.
+LogicalResult coalesceLoops(MutableArrayRef<AffineForOp> loops);
+
+/// Vectorizes a loop (either outer or inner, with a perfect or imperfectly
+/// nested body). `simdWidth` is the bit width of the vectors on target. `forOp`
+/// is never invalidated. This method does not check for contiguity of accesses
+/// or validity of vectorization. A memref.vector_cast op is used to create a
+/// vector element type version of the memref. If the memref involved with the
+/// vectorization is a function argument attribute, its `memref.alignment`
+/// argument attribute is checked to determine if it has the desired alignment.
+/// If not, the loop isn't vectorize and this returns failure.
+LogicalResult loopVectorize(AffineForOp forOp, unsigned simdWidth,
+                            DenseMap<Value, Value> *vecMemRefs = nullptr);
+
 /// Take the ParallelLoop and for each set of dimension indices, combine them
 /// into a single dimension. combinedDimensions must contain each index into
 /// loops exactly once.
 void collapseParallelLoops(scf::ParallelOp loops,
                            ArrayRef<std::vector<unsigned>> combinedDimensions);
 
 /// Maps `forOp` for execution on a parallel grid of virtual `processorIds` of
 /// size given by `numProcessors`. This is achieved by embedding the SSA values

mlir/lib/Transforms/LoopCoalescing.cpp

 //===- LoopCoalescing.cpp - Pass transforming loop nests into single loops-===//
 //
 // 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 "PassDetail.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Transforms/LoopUtils.h"
 #include "mlir/Transforms/Passes.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/Support/Debug.h"
 
 #define PASS_NAME "loop-coalescing"
 #define DEBUG_TYPE PASS_NAME
 
 using namespace mlir;
 
 namespace {
 struct LoopCoalescingPass : public LoopCoalescingBase<LoopCoalescingPass> {
-  void runOnFunction() override {
-    FuncOp func = getFunction();
 
-    func.walk([](scf::ForOp op) {
-      // Ignore nested loops.
-      if (op->getParentOfType<scf::ForOp>())
-        return;
+  /// Walk either an scf.for or an affine.for to find a band to coalesce.
+  template <typename LoopOpTy>
+  static void walkLoop(LoopOpTy op) {
+    if (op->template getParentOfType<LoopOpTy>())
 
-      SmallVector<scf::ForOp, 4> loops;
-      getPerfectlyNestedLoops(loops, op);
-      LLVM_DEBUG(llvm::dbgs()
-                 << "found a perfect nest of depth " << loops.size() << '\n');
+    SmallVector<LoopOpTy, 4> loops;
 
-      // Look for a band of loops that can be coalesced, i.e. perfectly nested
-      // loops with bounds defined above some loop.
-      // 1. For each loop, find above which parent loop its operands are
-      // defined.
-      SmallVector<unsigned, 4> operandsDefinedAbove(loops.size());
-      for (unsigned i = 0, e = loops.size(); i < e; ++i) {
-        operandsDefinedAbove[i] = i;
-        for (unsigned j = 0; j < i; ++j) {
-          if (areValuesDefinedAbove(loops[i].getOperands(),
-                                    loops[j].region())) {
-            operandsDefinedAbove[i] = j;
-            break;
-          }
-        }
-        LLVM_DEBUG(llvm::dbgs()
-                   << "  bounds of loop " << i << " are known above depth "
-                   << operandsDefinedAbove[i] << '\n');
-      }
+        if (areValuesDefinedAbove(loops[i].getOperands(), loops[j].region())) {
 
-      // 2. Identify bands of loops such that the operands of all of them are
-      // defined above the first loop in the band.  Traverse the nest bottom-up
-      // so that modifications don't invalidate the inner loops.
-      for (unsigned end = loops.size(); end > 0; --end) {
-        unsigned start = 0;
-        for (; start < end - 1; ++start) {
-          auto maxPos =
-              *std::max_element(std::next(operandsDefinedAbove.begin(), start),
-                                std::next(operandsDefinedAbove.begin(), end));
-          if (maxPos > start)
-            continue;
 
-          assert(maxPos == start &&
-                 "expected loop bounds to be known at the start of the band");
-          LLVM_DEBUG(llvm::dbgs() << "  found coalesceable band from " << start
-                                  << " to " << end << '\n');
 
-          auto band =
-              llvm::makeMutableArrayRef(loops.data() + start, end - start);
-          coalesceLoops(band);
-          break;
-        }
-        // If a band was found and transformed, keep looking at the loops above
-        // the outermost transformed loop.
-        if (start != end - 1)
-          end = start + 1;
-      }
-    });
+        (void)coalesceLoops(band);
+  }
+
+  void runOnFunction() override {
+    FuncOp func = getFunction();
+    func.walk([&](scf::ForOp op) { walkLoop(op); });
+    func.walk([&](AffineForOp op) { walkLoop(op); });
   }
 };
 
 } // namespace
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLoopCoalescingPass() {
   return std::make_unique<LoopCoalescingPass>();
 }

mlir/lib/Transforms/Utils/LoopUtils.cpp

void mlir::coalesceLoops(MutableArrayRef<scf::ForOp> loops) {
   Location loc = outermost.getLoc();
   Value upperBound = outermost.upperBound();
   for (auto loop : loops.drop_front())
     upperBound = builder.create<MulIOp>(loc, upperBound, loop.upperBound());
   outermost.setUpperBound(upperBound);
 
   builder.setInsertionPointToStart(outermost.getBody());
 
-  // 3. Remap induction variables.  For each original loop, the value of the
+  // 3. Remap induction variables. For each original loop, the value of the
   // induction variable can be obtained by dividing the induction variable of
   // the linearized loop by the total number of iterations of the loops nested
   // in it modulo the number of iterations in this loop (remove the values
   // related to the outer loops):
   //   iv_i = floordiv(iv_linear, product-of-loop-ranges-until-i) mod range_i.
   // Compute these iteratively from the innermost loop by creating a "running
   // quotient" of division by the range.
   Value previous = outermost.getInductionVar();
void mlir::coalesceLoops(MutableArrayRef<scf::ForOp> loops) {
   scf::ForOp second = loops[1];
   innermost.getBody()->back().erase();
   outermost.getBody()->getOperations().splice(
       Block::iterator(second.getOperation()),
       innermost.getBody()->getOperations());
   second.erase();
 }
 
+LogicalResult mlir::coalesceLoops(MutableArrayRef<AffineForOp> loops) {
+  if (loops.size() < 2)
+    return success();
+
+  AffineForOp innermost = loops.back();
+  AffineForOp outermost = loops.front();
+  AffineBound ub = outermost.getUpperBound();
+  AffineMap origUbMap = ub.getMap();
+  Location loc = outermost.getLoc();
+  OpBuilder builder(outermost);
+  for (AffineForOp loop : loops) {
+    // We only work on normalized loops.
+    if (loop.getStep() != 1 || loop.getConstantLowerBound() != 0)
+      return failure();
+  }
+  SmallVector<Value, 4> upperBoundSymbols;
+  SmallVector<Value, 4> ubOperands(ub.getOperands().begin(),
+                                   ub.getOperands().end());
+
+  // 1. Store the upper bound of the outermost loop in a variable.
+  Value prev;
+  if (!llvm::hasSingleElement(origUbMap.getResults()))
+    prev = builder.create<AffineMinOp>(loc, origUbMap, ubOperands);
+  else
+    prev = builder.create<AffineApplyOp>(loc, origUbMap, ubOperands);
+  upperBoundSymbols.push_back(prev);
+
+  // 2. Emit code computing the upper bound of the coalesced loop as product of
+  // the number of iterations of all loops.
+  for (AffineForOp loop : loops.drop_front()) {
+    ub = loop.getUpperBound();
+    origUbMap = ub.getMap();
+    ubOperands = ub.getOperands();
+    Value upperBound;
+    // If upper bound map has more than one result, take their minimum.
+    if (!(origUbMap.getResults().empty() ||
+          llvm::hasSingleElement(origUbMap.getResults())))
+      upperBound = builder.create<AffineMinOp>(loc, origUbMap, ubOperands);
+    else
+      upperBound = builder.create<AffineApplyOp>(loc, origUbMap, ubOperands);
+    upperBoundSymbols.push_back(upperBound);
+    SmallVector<Value, 4> operands;
+    operands.push_back(prev);
+    operands.push_back(upperBound);
+    // Maintain running product of loop upper bounds.
+    prev = builder.create<AffineApplyOp>(
+        loc,
+        AffineMap::get(/*numDims=*/
+                       1, /*numSymbols=*/1,
+                       builder.getAffineDimExpr(0) *
+                           builder.getAffineSymbolExpr(0)),
+        operands);
+  }
+  // Set upper bound of the coalesced loop.
+  AffineMap newUbMap = AffineMap::get(
+      /*numDims=*/0,
+      /*numSymbols=*/1, builder.getAffineSymbolExpr(0), builder.getContext());
+  outermost.setUpperBound(prev, newUbMap);
+
+  builder.setInsertionPointToStart(outermost.getBody());
+
+  // 3. Remap induction variables. For each original loop, the value of the
+  // induction variable can be obtained by dividing the induction variable of
+  // the linearized loop by the total number of iterations of the loops nested
+  // in it modulo the number of iterations in this loop (remove the values
+  // related to the outer loops):
+  //   iv_i = floordiv(iv_linear, product-of-loop-ranges-until-i) mod range_i.
+  // Compute these iteratively from the innermost loop by creating a "running
+  // quotient" of division by the range.
+  Value previous = outermost.getInductionVar();
+  for (unsigned i = 0, e = loops.size(); i < e; ++i) {
+    unsigned idx = loops.size() - i - 1;
+    if (i != 0) {
+      ub = loops[idx + 1].getUpperBound();
+      origUbMap = ub.getMap();
+      ubOperands = ub.getOperands();
+      SmallVector<Value, 4> operands;
+      operands.push_back(previous);
+      operands.push_back(upperBoundSymbols[idx + 1]);
+      previous = builder.create<AffineApplyOp>(
+          loc,
+          AffineMap::get(
+              /*numDims=*/1, /*numSymbols=*/1,
+              builder.getAffineDimExpr(0).floorDiv(
+                  builder.getAffineSymbolExpr(0))),
+          operands);
+    }
+    // Modified value of the induction variables of the nested loops after
+    // coalescing.
+    Value inductionVariable;
+    if (i == e - 1) {
+      inductionVariable = previous;
+    } else {
+      ub = loops[idx].getUpperBound();
+      origUbMap = ub.getMap();
+      ubOperands = ub.getOperands();
+      SmallVector<Value, 4> applyOperands;
+      applyOperands.push_back(previous);
+      applyOperands.push_back(upperBoundSymbols[idx]);
+      inductionVariable = builder.create<AffineApplyOp>(
+          loc,
+          AffineMap::get(
+              /*numDims=*/1, /*numSymbols=*/1,
+              builder.getAffineDimExpr(0) % builder.getAffineSymbolExpr(0)),
+          applyOperands);
+    }
+
+    replaceAllUsesInRegionWith(loops[idx].getInductionVar(), inductionVariable,
+                               loops.back().region());
+  }
+
+  // 4. Move the operations from the innermost just above the second-outermost
+  // loop, delete the extra terminator and the second-outermost loop.
+  AffineForOp second = loops[1];
+  innermost.getBody()->back().erase();
+  outermost.getBody()->getOperations().splice(
+      Block::iterator(second.getOperation()),
+      innermost.getBody()->getOperations());
+  second.erase();
+  return success();
+}
+
 void mlir::collapseParallelLoops(
     scf::ParallelOp loops, ArrayRef<std::vector<unsigned>> combinedDimensions) {
   OpBuilder outsideBuilder(loops);
   Location loc = loops.getLoc();
 
   // Presort combined dimensions.
   auto sortedDimensions = llvm::to_vector<3>(combinedDimensions);
   for (auto &dims : sortedDimensions)

mlir/test/Transforms/loop-coalescing.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -loop-coalescing %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -loop-coalescing %s | FileCheck %s
 
 // CHECK-LABEL: @one_3d_nest
 func @one_3d_nest() {
   // Capture original bounds.  Note that for zero-based step-one loops, the
   // upper bound is also the number of iterations.
   // CHECK: %[[orig_lb:.*]] = constant 0
   // CHECK: %[[orig_step:.*]] = constant 1
   // CHECK: %[[orig_ub_k:.*]] = constant 3
scf.for %j = %c0 to %c10 step %c1 {
         scf.for %l = %i to %j step %c1 {
           "foo"() : () -> ()
         }
       }
     }
   }
   return
 }
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<() -> (16)>
+// CHECK-DAG: #[[MAP1:.*]] = affine_map<() -> (64)>
+// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[MAP3:.*]] = affine_map<() -> (8)>
+// CHECK-DAG: #[[MAP4:.*]] = affine_map<(d0)[s0] -> (d0 mod s0)>
+// CHECK-DAG: #[[MAP5:.*]] = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+func @coalesce_affine_for() {
+  affine.for %i = 0 to 16 {
+    affine.for %j = 0 to 64 {
+      affine.for %k = 0 to 8 {
+        "test.foo"(%i, %j, %k) : (index, index, index) -> ()
+      }
+    }
+  }
+  return
+}
+// CHECK-DAG: %[[T0:.*]] = affine.apply #[[MAP0]]()
+// CHECK-DAG: %[[T1:.*]] = affine.apply #[[MAP1]]()
+// CHECK-DAG: %[[T2:.*]] = affine.apply #[[MAP2]](%[[T0]])[%[[T1]]]
+// CHECK-DAG: %[[T3:.*]] = affine.apply #[[MAP3]]()
+// CHECK-DAG: %[[T4:.*]] = affine.apply #[[MAP2]](%[[T2]])[%[[T3]]]
+// CHECK:       affine.for %[[IV:.*]] = 0 to %[[T4]]
+// CHECK-DAG:    %[[K:.*]] =  affine.apply #[[MAP4]](%[[IV]])[%[[T3]]]
+// CHECK-DAG:    %[[T6:.*]] = affine.apply #[[MAP5]](%[[IV]])[%[[T3]]]
+// CHECK-DAG:    %[[J:.*]] =  affine.apply #[[MAP4]](%[[T6]])[%[[T1]]]
+// CHECK-DAG:    %[[I:.*]] =  affine.apply #[[MAP5]](%[[T6]])[%[[T1]]]
+// CHECK-NEXT:    "test.foo"(%[[I]], %[[J]], %[[K]])
+// CHECK-NEXT:  }
+// CHECK-NEXT:  return
+
+// -----
+
+// CHECK-DAG: #[[MAP:.*]] = affine_map<()[s0] -> (s0)>
+// CHECK-DAG: #[[PRODUCT:.*]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[MOD:.*]] = affine_map<(d0)[s0] -> (d0 mod s0)>
+// CHECK-DAG: #[[FLOOR:.*]] = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+func @coalesce_affine_for(%arg0: memref<?x?xf32>) {
+  %c0 = constant 0 : index
+  %M = memref.dim %arg0, %c0 : memref<?x?xf32>
+  %N = memref.dim %arg0, %c0 : memref<?x?xf32>
+  %K = memref.dim %arg0, %c0 : memref<?x?xf32>
+  affine.for %i = 0 to %M {
+    affine.for %j = 0 to %N {
+      affine.for %k = 0 to %K {
+      "test.foo"(%i, %j, %k) : (index, index, index) -> ()
+      }
+    }
+  }
+  return
+}
+// CHECK: %[[T0:.*]] = memref.dim %arg{{.*}}, %c{{.*}} : memref<?x?xf32>
+// CHECK: %[[T1:.*]] = memref.dim %arg{{.*}}, %c{{.*}} : memref<?x?xf32>
+// CHECK: %[[T2:.*]] = memref.dim %arg{{.*}}, %c{{.*}} : memref<?x?xf32>
+// CHECK-DAG: %[[T3:.*]] = affine.apply #[[MAP]]()[%[[T0]]]
+// CHECK-DAG: %[[T4:.*]] = affine.apply #[[MAP]]()[%[[T1]]]
+// CHECK-DAG: %[[T5:.*]] = affine.apply #[[PRODUCT]](%[[T3]])[%[[T4]]]
+// CHECK-DAG: %[[T6:.*]] = affine.apply #[[MAP]]()[%[[T2]]]
+// CHECK-DAG: %[[T7:.*]] = affine.apply #[[PRODUCT]](%[[T5]])[%[[T6]]]
+// CHECK: affine.for %[[IV:.*]] = 0 to %[[T7]]
+// CHECK-DAG:    %[[K:.*]] = affine.apply #[[MOD]](%[[IV]])[%[[T6]]]
+// CHECK-DAG:    %[[T9:.*]] = affine.apply #[[FLOOR]](%[[IV]])[%[[T6]]]
+// CHECK-DAG:    %[[J:.*]] = affine.apply #[[MOD]](%[[T9]])[%[[T4]]]
+// CHECK-DAG:    %[[I:.*]] = affine.apply #[[FLOOR]](%[[T9]])[%[[T4]]]
+// CHECK-NEXT:    "test.foo"(%[[I]], %[[J]], %[[K]])
+// CHECK-NEXT:  }
+// CHECK-NEXT:  return
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0)>
+// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[MAP2:.*]] = affine_map<() -> (64)>
+// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0)[s0] -> (d0 mod s0)>
+// CHECK-DAG: #[[MAP4:.*]] = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+func @coalesce_affine_for(%arg0: memref<?x?xf32>) {
+  %c0 = constant 0 : index
+  %M = memref.dim %arg0, %c0 : memref<?x?xf32>
+  %N = memref.dim %arg0, %c0 : memref<?x?xf32>
+  affine.for %i = 0 to %M {
+    affine.for %j = 0 to %N {
+      affine.for %k = 0 to 64 {
+      "test.foo"(%i, %j, %k) : (index, index, index) -> ()
+      }
+    }
+  }
+  return
+}
+// CHECK: %[[T0:.*]] = memref.dim %arg{{.*}}, %c{{.*}} : memref<?x?xf32>
+// CHECK: %[[T1:.*]] = memref.dim %arg{{.*}}, %c{{.*}} : memref<?x?xf32>
+// CHECK-DAG: %[[T2:.*]] = affine.apply #[[MAP0]]()[%[[T0]]]
+// CHECK-DAG: %[[T3:.*]] = affine.apply #[[MAP0]]()[%[[T1]]]
+// CHECK-DAG: %[[T4:.*]] = affine.apply #[[MAP1]](%[[T2]])[%[[T3]]]
+// CHECK-DAG: %[[T5:.*]] = affine.apply #[[MAP2]]()
+// CHECK-DAG: %[[T6:.*]] = affine.apply #[[MAP1]](%[[T4]])[%[[T5]]]
+// CHECK: affine.for %[[IV:.*]] = 0 to %[[T6]]
+// CHECK-DAG:    %[[K:.*]] = affine.apply #[[MAP3]](%[[IV]])[%[[T5]]]
+// CHECK-DAG:    %[[T8:.*]] = affine.apply #[[MAP4]](%[[IV]])[%[[T5]]]
+// CHECK-DAG:    %[[J:.*]] = affine.apply #[[MAP3]](%[[T8]])[%[[T3]]]
+// CHECK-DAG:    %[[I:.*]] = affine.apply #[[MAP4]](%[[T8]])[%[[T3]]]
+// CHECK-NEXT:    "test.foo"(%[[I]], %[[J]], %[[K]])
+// CHECK-NEXT:  }
+// CHECK-NEXT:  return
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0, -s0)>
+// CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0)>
+// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0)[s0] -> (d0 mod s0)>
+// CHECK-DAG: #[[MAP4:.*]] = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+#myMap = affine_map<()[s1] -> (s1, -s1)>
+func @coalesce_affine_for(%arg0: memref<?x?xf32>) {
+ %c0 = constant 0 : index
+ %M = memref.dim %arg0, %c0 : memref<?x?xf32>
+ %N = memref.dim %arg0, %c0 : memref<?x?xf32>
+ %K = memref.dim %arg0, %c0 : memref<?x?xf32>
+ affine.for %i = 0 to min #myMap()[%M] {
+   affine.for %j = 0 to %N {
+     affine.for %k = 0 to %K {
+     "test.foo"(%i, %j, %k) : (index, index, index) -> ()
+     }
+   }
+ }
+ return
+}
+// CHECK: %[[T0:.*]] = memref.dim %arg{{.*}}, %c{{.*}} : memref<?x?xf32>
+// CHECK: %[[T1:.*]] = memref.dim %arg{{.*}}, %c{{.*}} : memref<?x?xf32>
+// CHECK: %[[T2:.*]] = memref.dim %arg{{.*}}, %c{{.*}} : memref<?x?xf32>
+// CHECK-DAG: %[[T3:.*]] = affine.min #[[MAP0]]()[%[[T0]]]
+// CHECK-DAG: %[[T4:.*]] = affine.apply #[[MAP1]]()[%[[T1]]]
+// CHECK-DAG: %[[T5:.*]] = affine.apply #[[MAP2]](%[[T3]])[%[[T4]]]
+// CHECK-DAG: %[[T6:.*]] = affine.apply #[[MAP1]]()[%[[T2]]]
+// CHECK-DAG: %[[T7:.*]] = affine.apply #[[MAP2]](%[[T5]])[%[[T6]]]
+// CHECK: affine.for %[[IV:.*]] = 0 to %[[T7]]
+// CHECK-DAG:    %[[K:.*]] = affine.apply #[[MAP3]](%[[IV]])[%[[T6]]]
+// CHECK-DAG:    %[[T9:.*]] = affine.apply #[[MAP4]](%[[IV]])[%[[T6]]]
+// CHECK-DAG:    %[[J:.*]] = affine.apply #[[MAP3]](%[[T9]])[%[[T4]]]
+// CHECK-DAG:    %[[I:.*]] = affine.apply #[[MAP4]](%[[T9]])[%[[T4]]]
+// CHECK-NEXT:    "test.foo"(%[[I]], %[[J]], %[[K]])
+// CHECK-NEXT:  }
+// CHECK-NEXT:  return