diff --git a/mlir/include/mlir/Transforms/LoopUtils.h b/mlir/include/mlir/Transforms/LoopUtils.h
--- a/mlir/include/mlir/Transforms/LoopUtils.h
+++ b/mlir/include/mlir/Transforms/LoopUtils.h
@@ -225,7 +225,7 @@
 
 /// Gathers all AffineForOps in 'func' grouped by loop depth.
 void gatherLoops(FuncOp func,
-                 DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops);
+                 SmallVectorImpl<SmallVector<AffineForOp, 2>> &depthToLoops);
 
 } // end namespace mlir
 
diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -1785,9 +1785,14 @@
 }
 
 /// Gathers all AffineForOps in 'block' at 'currLoopDepth' in 'depthToLoops'.
-static void gatherLoopsInBlock(
-    Block *block, unsigned currLoopDepth,
-    DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops) {
+static void
+gatherLoopsInBlock(Block *block, unsigned currLoopDepth,
+                   SmallVectorImpl<SmallVector<AffineForOp, 2>> &depthToLoops) {
+  // Add a new empty level to output if it doesn't exist level already.
+  assert(currLoopDepth <= depthToLoops.size() && "Unexpected currLoopDepth");
+  if (currLoopDepth == depthToLoops.size())
+    depthToLoops.push_back(SmallVector<AffineForOp, 2>());
+
   auto &loopsAtDepth = depthToLoops[currLoopDepth];
   for (auto &op : *block) {
     if (auto forOp = dyn_cast<AffineForOp>(op)) {
@@ -1795,12 +1800,18 @@
       gatherLoopsInBlock(forOp.getBody(), currLoopDepth + 1, depthToLoops);
     }
   }
+
+  // Remove level from output if no loops were found.
+  if (loopsAtDepth.empty()) {
+    assert(currLoopDepth == depthToLoops.size() - 1 &&
+           "Removing an intermediate loop depth?");
+    depthToLoops.pop_back();
+  }
 }
 
 /// Gathers all AffineForOps in 'func' grouped by loop depth.
 void mlir::gatherLoops(
-    FuncOp func,
-    DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops) {
+    FuncOp func, SmallVectorImpl<SmallVector<AffineForOp, 2>> &depthToLoops) {
   for (auto &block : func)
     gatherLoopsInBlock(&block, /*currLoopDepth=*/0, depthToLoops);
 }
diff --git a/mlir/test/lib/Transforms/TestAffineDataCopy.cpp b/mlir/test/lib/Transforms/TestAffineDataCopy.cpp
--- a/mlir/test/lib/Transforms/TestAffineDataCopy.cpp
+++ b/mlir/test/lib/Transforms/TestAffineDataCopy.cpp
@@ -43,13 +43,13 @@
 
 void TestAffineDataCopy::runOnFunction() {
   // Gather all AffineForOps by loop depth.
-  DenseMap<unsigned, SmallVector<AffineForOp, 2>> depthToLoops;
+  SmallVector<SmallVector<AffineForOp, 2>, 8> depthToLoops;
   gatherLoops(getFunction(), depthToLoops);
   assert(depthToLoops.size() && "Loop nest not found");
 
   // Only support tests with a single loop nest and a single innermost loop
   // for now.
-  unsigned innermostLoopIdx = depthToLoops.size() - 2;
+  unsigned innermostLoopIdx = depthToLoops.size() - 1;
   if (depthToLoops[0].size() != 1 || depthToLoops[innermostLoopIdx].size() != 1)
     return;
 
diff --git a/mlir/test/lib/Transforms/TestLoopFusion.cpp b/mlir/test/lib/Transforms/TestLoopFusion.cpp
--- a/mlir/test/lib/Transforms/TestLoopFusion.cpp
+++ b/mlir/test/lib/Transforms/TestLoopFusion.cpp
@@ -154,13 +154,12 @@
 // Run tests on all combinations of src/dst loop nests in 'depthToLoops'.
 // If 'return_on_change' is true, returns on first invocation of 'fn' which
 // returns true.
-static bool
-iterateLoops(DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops,
-             LoopFunc fn, bool return_on_change = false) {
+static bool iterateLoops(ArrayRef<SmallVector<AffineForOp, 2>> depthToLoops,
+                         LoopFunc fn, bool return_on_change = false) {
   bool changed = false;
-  for (auto &depthAndLoops : depthToLoops) {
-    unsigned loopDepth = depthAndLoops.first;
-    auto &loops = depthAndLoops.second;
+  for (unsigned loopDepth = 0, end = depthToLoops.size(); loopDepth < end;
+       ++loopDepth) {
+    auto &loops = depthToLoops[loopDepth];
     unsigned numLoops = loops.size();
     for (unsigned j = 0; j < numLoops; ++j) {
       for (unsigned k = 0; k < numLoops; ++k) {
@@ -176,7 +175,7 @@
 }
 
 void TestLoopFusion::runOnFunction() {
-  DenseMap<unsigned, SmallVector<AffineForOp, 2>> depthToLoops;
+  SmallVector<SmallVector<AffineForOp, 2>, 8> depthToLoops;
   if (clTestLoopFusionTransformation) {
     // Run loop fusion until a fixed point is reached.
     do {