[mlir][sparse] comment cleanup in iteration graph sorter (#75508)

Author: aartbik

mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp

@@ -29,22 +29,23 @@ class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> {
   explicit AffineDimFinder(ArrayRef<utils::IteratorType> itTypes)
       : iterTypes(itTypes) {}
 
-  // Override method from AffineExprVisitor.
+  /// Overrides the visit method from AffineExprVisitor.
   void visitDimExpr(AffineDimExpr expr) {
     if (pickedDim == nullptr || pickIterType == iterTypes[expr.getPosition()])
       pickedDim = expr;
   }
 
-  /// Set the desired iterator type that we want to pick.
+  /// Sets the desired iterator type that we want to pick.
   void setPickedIterType(utils::IteratorType iterType) {
     pickIterType = iterType;
   }
 
-  /// Get the desired AffineDimExpr.
+  /// Gets the desired AffineDimExpr.
   AffineDimExpr getDimExpr() const {
     return llvm::cast<AffineDimExpr>(pickedDim);
   }
 
+  /// Walks the graph in post order to find dim expr.
   void walkPostOrder(AffineExpr expr) {
     pickedDim = nullptr;
     AffineExprVisitor<AffineDimFinder>::walkPostOrder(expr);
@@ -55,11 +56,11 @@ class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> {
   AffineExpr pickedDim;
   /// The iterator type that we want.
   utils::IteratorType pickIterType;
-  /// The mapping between dim=>iterator type.
+  /// The mapping between levels and iterator types.
   ArrayRef<utils::IteratorType> iterTypes;
 };
 
-// Flattens an affine expression into a list of AffineDimExprs.
+/// Flattens an affine expression into a list of AffineDimExprs.
 struct AffineDimCollector : public AffineExprVisitor<AffineDimCollector> {
   // Overrides method from AffineExprVisitor.
   void visitDimExpr(AffineDimExpr expr) { dims.push_back(expr); }
@@ -97,8 +98,8 @@ AffineMap IterationGraphSorter::topoSort() {
 
   SmallVector<unsigned> loopOrder;
   while (!redIt.empty() || !parIt.empty()) {
-    // We always prefer parallel loop over reduction loop because putting
-    // reduction loop early might make the loop sequence inadmissible.
+    // We always prefer a parallel loop over a reduction loop because putting
+    // a reduction loop early might make the loop sequence inadmissible.
     auto &it = !parIt.empty() ? parIt : redIt;
     auto src = it.back();
     loopOrder.push_back(src);
@@ -114,6 +115,7 @@ AffineMap IterationGraphSorter::topoSort() {
     }
   }
 
+  // Return the topological sort on success.
   if (loopOrder.size() == numLoops)
     return AffineMap::getPermutationMap(loopOrder, out.getContext());
 
@@ -164,13 +166,14 @@ IterationGraphSorter::IterationGraphSorter(
 }
 
 AffineMap IterationGraphSorter::sort(SortMask mask, Value ignored) {
-  // Reset the interation graph.
+  // Reset the adjacency matrix that represents the iteration graph.
   for (auto &row : itGraph)
     std::fill(row.begin(), row.end(), false);
 
-  // Reset cached in-degree.
+  // Reset in-degree.
   std::fill(inDegree.begin(), inDegree.end(), 0);
 
+  // Add the constraints for the loop to level map.
   for (auto [in, map] : llvm::zip(ins, loop2InsLvl)) {
     // Get map and encoding.
     const auto enc = getSparseTensorEncoding(in.getType());
@@ -180,11 +183,12 @@ AffineMap IterationGraphSorter::sort(SortMask mask, Value ignored) {
     addConstraints(in, map);
   }
 
-  // Get map and encoding.
+  // Add the constraints for the output map.
   const auto enc = getSparseTensorEncoding(out.getType());
   if ((enc || includesDenseOutput(mask)) && out != ignored)
     addConstraints(out, loop2OutLvl);
 
+  // Return the topological sort (empty for cyclic).
   return topoSort();
 }
 
@@ -196,6 +200,7 @@ void IterationGraphSorter::addConstraints(Value t, AffineMap loop2LvlMap) {
     }
   };
 
+  // Set up a reduction finder.
   AffineDimFinder finder(iterTypes);
   finder.setPickedIterType(utils::IteratorType::reduction);
 

mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h

@@ -79,7 +79,7 @@ class IterationGraphSorter {
   // Loop itation types;
   SmallVector<utils::IteratorType> iterTypes;
 
-  // Adjacent matrix that represents the iteration graph.
+  // Adjacency matrix that represents the iteration graph.
   std::vector<std::vector<bool>> itGraph;
 
   // InDegree used for topo sort.