[Mlir-commits] [mlir] 376c685 - [mlir][NFC] Fix 'gatherLoops' utility

[Diego Caballero]

Author: Diego Caballero
Date: 2020-02-19T10:48:14-08:00
New Revision: 376c68539c5a4b8baf7dc8f34ffa71daadada28f

URL: https://github.com/llvm/llvm-project/commit/376c68539c5a4b8baf7dc8f34ffa71daadada28f
DIFF: https://github.com/llvm/llvm-project/commit/376c68539c5a4b8baf7dc8f34ffa71daadada28f.diff

LOG: [mlir][NFC] Fix 'gatherLoops' utility

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

Reviewed By: andydavis1, mehdi_amini

Differential Revision: https://reviews.llvm.org/D74658

Added: 
    

Modified: 
    mlir/include/mlir/Transforms/LoopUtils.h
    mlir/lib/Transforms/Utils/LoopUtils.cpp
    mlir/test/lib/Transforms/TestAffineDataCopy.cpp
    mlir/test/lib/Transforms/TestLoopFusion.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Transforms/LoopUtils.h b/mlir/include/mlir/Transforms/LoopUtils.h
index cf6316cae643..4748e13cf5eb 100644
--- a/mlir/include/mlir/Transforms/LoopUtils.h
+++ b/mlir/include/mlir/Transforms/LoopUtils.h
@@ -225,7 +225,7 @@ void mapLoopToProcessorIds(loop::ForOp forOp, ArrayRef<Value> processorId,
 
 /// 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
 

diff  --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index da3d819cbc3e..a8825c11bb12 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -1785,22 +1785,31 @@ uint64_t mlir::affineDataCopyGenerate(Block::iterator begin,
 }
 
 /// 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();
+  }
 }

diff  --git a/mlir/test/lib/Transforms/TestAffineDataCopy.cpp b/mlir/test/lib/Transforms/TestAffineDataCopy.cpp
index e03d45cb9dd4..f5caa6922710 100644
--- a/mlir/test/lib/Transforms/TestAffineDataCopy.cpp
+++ b/mlir/test/lib/Transforms/TestAffineDataCopy.cpp
@@ -43,13 +43,13 @@ struct TestAffineDataCopy : public FunctionPass<TestAffineDataCopy> {
 
 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;
 

diff  --git a/mlir/test/lib/Transforms/TestLoopFusion.cpp b/mlir/test/lib/Transforms/TestLoopFusion.cpp
index 9214fa9fc433..148eeaa3a60e 100644
--- a/mlir/test/lib/Transforms/TestLoopFusion.cpp
+++ b/mlir/test/lib/Transforms/TestLoopFusion.cpp
@@ -154,13 +154,12 @@ using LoopFunc = function_ref<bool(AffineForOp, AffineForOp, 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,
-             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 @@ iterateLoops(DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops,
 }
 
 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 {