[mlir][NFC] Fix 'gatherLoops' utility

[dcaballe]

Summary

It replaces DenseMap output with a SmallVector and it
removes empty loop levels from the output.

[mehdi_amini]

mlir/lib/Transforms/Utils/LoopUtils.cpp

1808

Why is this pop_back() needed now? Or do you just view this as an "optimization"?

[dcaballe]

Thanks!

mlir/lib/Transforms/Utils/LoopUtils.cpp

1808

This is not to keep an empty level around in the output. For example, let's say we have two nested loops in a function. The algorithm would proceed as follows:

• Visit function body, add level 0 to output, add the outer loop to level 0.
• Visit the outer loop body, add level 1 to output, add the inner loop to level 1.
• Visit inner loop body, add level 2 to output, no loops are added to level 2.
• Remove level 2 since it's empty.

I'm now thinking that I could move this removal to the public 'gatherLoops' function so that it's only executed once since the empty level will always be one and the last one in the output vector, no matter how many or the structure of the loop nests visited.

[dcaballe]

Move empty level removal to public function so that it's executed only once.

[andydavis1]

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

46

Typically we don't see SmallVector<SmallVector. These patterns are typically std::vector<SmallVector.

[dcaballe]

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

46

Thanks Andy! I'll change that before committing. For my understanding, I'm not sure I follow the rationale. Isn't this one of the ideal cases for SmallVector? I.e., the object to be contained has a "large" memory footprint (a SmallVector with two op elements, in this case) and we want to avoid allocating unnecessary slots?

mlir/include/mlir/Transforms/LoopUtils.h

 ///      ...
 ///    }
 /// ```
 void mapLoopToProcessorIds(loop::ForOp forOp, ArrayRef<Value> processorId,
                            ArrayRef<Value> numProcessors);
 
 /// Gathers all AffineForOps in 'func' grouped by loop depth.
 void gatherLoops(FuncOp func,
-                 DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops);
+                 std::vector<SmallVector<AffineForOp, 2>> &depthToLoops);
 
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_LOOP_UTILS_H

mlir/lib/Transforms/Utils/LoopUtils.cpp

StringRef str = "Total size of all copy buffers' for this block "
                     "exceeds fast memory capacity\n";
     block->getParentOp()->emitError(str);
   }
 
   return totalCopyBuffersSizeInBytes;
 }
 
 /// Gathers all AffineForOps in 'block' at 'currLoopDepth' in 'depthToLoops'.
-static void gatherLoopsInBlock(
-    Block *block, unsigned currLoopDepth,
-    DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops) {
-  auto &loopsAtDepth = depthToLoops[currLoopDepth];
+static void
+gatherLoopsInBlock(Block *block, unsigned currLoopDepth,
+                   std::vector<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>());
+
   for (auto &op : *block) {
     if (auto forOp = dyn_cast<AffineForOp>(op)) {
-      loopsAtDepth.push_back(forOp);
+      depthToLoops[currLoopDepth].push_back(forOp);
       gatherLoopsInBlock(forOp.getBody(), currLoopDepth + 1, depthToLoops);
     }
   }
 }
 
 /// Gathers all AffineForOps in 'func' grouped by loop depth.
-void mlir::gatherLoops(
-    FuncOp func,
-    DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops) {
+void mlir::gatherLoops(FuncOp func,
+                       std::vector<SmallVector<AffineForOp, 2>> &depthToLoops) {
   for (auto &block : func)
     gatherLoopsInBlock(&block, /*currLoopDepth=*/0, depthToLoops);
+
+  // Remove last loop level from output since it's empty.
+  if (!depthToLoops.empty()) {
+    assert(depthToLoops.back().empty() && "Last loop level is not empty?");
+    depthToLoops.pop_back();
+  }
 }

mlir/test/lib/Transforms/TestAffineDataCopy.cpp

Option<bool> clMemRefFilter{
           "Enable memref filter testing in affine data copy optimization"),
       llvm::cl::init(false)};
 };
 
 } // end anonymous namespace
 
 void TestAffineDataCopy::runOnFunction() {
   // Gather all AffineForOps by loop depth.
-  DenseMap<unsigned, SmallVector<AffineForOp, 2>> depthToLoops;
+  std::vector<SmallVector<AffineForOp, 2>> 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;
 
   auto loopNest = depthToLoops[0][0];
   auto innermostLoop = depthToLoops[innermostLoopIdx][0];
   Optional<Value> memrefFilter;
   if (clMemRefFilter) {
     // Gather MemRef filter. For simplicity, we use the first loaded memref

mlir/test/lib/Transforms/TestLoopFusion.cpp

 }
 
 using LoopFunc = function_ref<bool(AffineForOp, AffineForOp, unsigned, unsigned,
                                    unsigned, unsigned)>;
 
 // 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,
+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) {
         if (j != k)
           changed |=
               fn(loops[j], loops[k], j, k, loopDepth, depthToLoops.size());
         if (changed && return_on_change)
           return true;
       }
     }
   }
   return changed;
 }
 
 void TestLoopFusion::runOnFunction() {
-  DenseMap<unsigned, SmallVector<AffineForOp, 2>> depthToLoops;
+  std::vector<SmallVector<AffineForOp, 2>> depthToLoops;
   if (clTestLoopFusionTransformation) {
     // Run loop fusion until a fixed point is reached.
     do {
       depthToLoops.clear();
       // Gather all AffineForOps by loop depth.
       gatherLoops(getFunction(), depthToLoops);
 
       // Try to fuse all combinations of src/dst loop nests in 'depthToLoops'.