[MLIR][NFC] loop tiling - improve comments / naming

[bondhugula]

Summary

Improve comments, naming, and other cleanup

Signed-off-by: Uday Bondhugula <[e-mail]>

[bondhugula]

This is in preparation for another major feature addition.

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp

constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
   }
 }
 
 /// Tiles the specified band of perfectly nested loops creating tile-space loops
 /// and intra-tile loops. A band is a contiguous set of loops.
 //  TODO(bondhugula): handle non hyper-rectangular spaces.
 LogicalResult mlir::tileCodeGen(MutableArrayRef<AffineForOp> band,
                                 ArrayRef<unsigned> tileSizes) {
-  assert(!band.empty());
-  assert(band.size() == tileSizes.size() && "Incorrect number of tile sizes");
+  assert(!band.empty() && "no loops in band");
+  assert(band.size() == tileSizes.size() && "Too few/many tile sizes");
 
   // Check if the supplied for op's are all successively nested.
-  for (unsigned i = 1, e = band.size(); i < e; i++) {
-    assert(band[i].getParentOp() == band[i - 1].getOperation());
-  }
+  for (unsigned i = 1, e = band.size(); i < e; i++)
+    assert(band[i].getParentOp() == band[i - 1] && "not a perfect nest / band");
 
   auto origLoops = band;
 
   AffineForOp rootAffineForOp = origLoops[0];
   auto loc = rootAffineForOp.getLoc();
   // Note that width is at least one since band isn't empty.
   unsigned width = band.size();
 
-  SmallVector<AffineForOp, 12> newLoops(2 * width);
-  AffineForOp innermostPointLoop;
+  SmallVector<AffineForOp, 6> tiledLoops(2 * width);
 
   // The outermost among the loops as we add more..
   auto *topLoop = rootAffineForOp.getOperation();
+  AffineForOp innermostPointLoop;
 
   // Add intra-tile (or point) loops.
   for (unsigned i = 0; i < width; i++) {
     OpBuilder b(topLoop);
     // Loop bounds will be set later.
     auto pointLoop = b.create<AffineForOp>(loc, 0, 0);
     pointLoop.getBody()->getOperations().splice(
         pointLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
         topLoop);
-    newLoops[2 * width - 1 - i] = pointLoop;
+    tiledLoops[2 * width - 1 - i] = pointLoop;
     topLoop = pointLoop.getOperation();
     if (i == 0)
       innermostPointLoop = pointLoop;
   }
 
   // Add tile space loops;
   for (unsigned i = width; i < 2 * width; i++) {
     OpBuilder b(topLoop);
     // Loop bounds will be set later.
     auto tileSpaceLoop = b.create<AffineForOp>(loc, 0, 0);
     tileSpaceLoop.getBody()->getOperations().splice(
         tileSpaceLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
         topLoop);
-    newLoops[2 * width - i - 1] = tileSpaceLoop;
+    tiledLoops[2 * width - i - 1] = tileSpaceLoop;
     topLoop = tileSpaceLoop.getOperation();
   }
 
   // Move the loop body of the original nest to the new one.
   moveLoopBody(origLoops[origLoops.size() - 1], innermostPointLoop);
 
   SmallVector<Value, 8> origLoopIVs;
   extractForInductionVars(band, &origLoopIVs);
   SmallVector<Optional<Value>, 6> ids(origLoopIVs.begin(), origLoopIVs.end());
   FlatAffineConstraints cst;
   getIndexSet(band, &cst);
 
   if (!cst.isHyperRectangular(0, width)) {
-    rootAffineForOp.emitError("tiled code generation unimplemented for the "
-                              "non-hyperrectangular case");
+    llvm::dbgs() << "tiled code generation unimplemented for the "
+                    "non-hyperrectangular case, op:"
+                 << *rootAffineForOp << "\n";
     return failure();
   }
 
-  constructTiledIndexSetHyperRect(origLoops, newLoops, tileSizes);
-  // In this case, the point loop IVs just replace the original ones.
-  for (unsigned i = 0; i < width; i++) {
-    origLoopIVs[i].replaceAllUsesWith(newLoops[i + width].getInductionVar());
-  }
+  constructTiledIndexSetHyperRect(origLoops, tiledLoops, tileSizes);
+
+  // Replace original IVs with intra-tile loop IVs.
+  for (unsigned i = 0; i < width; i++)
+    origLoopIVs[i].replaceAllUsesWith(tiledLoops[i + width].getInductionVar());
 
   // Erase the old loop nest.
   rootAffineForOp.erase();
 
   return success();
 }
 
 // Identify valid and profitable bands of loops to tile. This is currently just
void LoopTiling::runOnFunction() {
   // Override cache size if provided on command line.
   if (clCacheSizeKiB.getNumOccurrences() > 0)
     cacheSizeBytes = clCacheSizeKiB * 1024;
 
   // Bands of loops to tile.
   std::vector<SmallVector<AffineForOp, 6>> bands;
   getTileableBands(getFunction(), &bands);
 
+  // Tile each band.
   for (auto &band : bands) {
     // Set up tile sizes; fill missing tile sizes at the end with default tile
     // size or clTileSize if one was provided.
     SmallVector<unsigned, 6> tileSizes;
     getTileSizes(band, &tileSizes);
     if (llvm::DebugFlag) {
       auto diag = band[0].emitRemark("using tile sizes [");
       for (auto tSize : tileSizes)
-        diag << tSize << " ";
+        diag << tSize << ' ';
       diag << "]\n";
     }
     if (failed(tileCodeGen(band, tileSizes)))
       return signalPassFailure();
   }
 }
 
 constexpr unsigned LoopTiling::kDefaultTileSize;
 constexpr uint64_t LoopTiling::kDefaultCacheMemCapacity;
 
 static PassRegistration<LoopTiling> pass("affine-loop-tile", "Tile loop nests");