address comments.

Author: Peiming Liu

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

@@ -1150,9 +1150,9 @@ struct ForeachRewriter : public OpRewritePattern<ForeachOp> {
 
     Operation &last = rewriter.getBlock()->back();
     if (llvm::isa<scf::YieldOp>(last)) {
-      // scf.for inserts a implicit yield op when there is no reduction
-      // variable upon creation, in this case we need to merge the block
-      // *before* the yield op.
+      // Because `scf.for` inserts an implicit yield op when there is no
+      // reduction variable upon creation, we reset the insertion point such
+      // that the block is inlined before *before* the yield op.
       rewriter.setInsertionPoint(&last);
     }
 

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

@@ -1032,9 +1032,9 @@ static bool getAllTidLvlsInLatPoints(
 
   if (isDenseLT(env.lt(outTid, curr))) {
     auto stt = getSparseTensorType(env.op().getOutputs().front());
-    // Note that we generate dense indices of the output tensor
-    // unconditionally, since they may not appear in the lattice, but may be
-    // needed for linearized env.
+    // Note that we generate dense indices of the output tensor unconditionally,
+    // since they may not appear in the lattice, but may be needed for
+    // linearized env.
     // TODO: we should avoid introducing corner cases for all-dense sparse
     // tensors.
     if (stt.hasEncoding() && stt.isAllDense())
@@ -1067,8 +1067,9 @@ static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
 
   SmallVector<TensorLevel> tidLvls;
   getAllTidLvlsInLatPoints(env, l0, curr, [&](TensorLevel tl, AffineExpr) {
-    // TODO: remove this! Duplication can be introduced due to the speical
-    // handling for all-dense "sparse" output tensor.
+    // TODO: remove this! The same tensor level might be added for multiple
+    // times due to the special handling for all-dense "sparse" output tensor
+    // (see L1038).
     if (llvm::find(tidLvls, tl) != tidLvls.end())
       return;
     tidLvls.emplace_back(tl);

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

@@ -408,7 +408,7 @@ class LoopEmitter {
   /// alive.
   std::vector<LoopInfo> loopStack;
 
-  // Loop Sequence Stack, stores the unversial index for the current loop
+  // Loop Sequence Stack, stores the universal index for the current loop
   // sequence. and a list of tid level that the loop sequence traverse.
   std::vector<std::pair<Value, std::vector<TensorLevel>>> loopSeqStack;
 };

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

@@ -164,17 +164,6 @@ static scf::ValueVector genWhenInBound(
     OpBuilder &b, Location l, SparseIterator &it, ValueRange elseRet,
     llvm::function_ref<scf::ValueVector(OpBuilder &, Location, Value)>
         builder) {
-  // Value isNotEnd = it.genNotEnd(b, l);
-  // Value crd = it.deref(b, l);
-  // scf::ValueVector ret = builder(b, l, crd);
-
-  // scf::ValueVector res;
-  // for (auto [notEnd, end] : llvm::zip_equal(ret, elseRet)) {
-  //   res.push_back(SELECT(isNotEnd, notEnd, end));
-  // };
-  // return res;
-
-  // !it.end() ? callback(*crd) : resOOB;
   TypeRange ifRetTypes = elseRet.getTypes();
   auto ifOp = b.create<scf::IfOp>(l, ifRetTypes, it.genNotEnd(b, l), true);
 
@@ -204,7 +193,7 @@ static scf::ValueVector genWhenInBound(
 static Value offsetFromMinCrd(OpBuilder &b, Location l, Value minCrd,
                               Value size) {
   Value geSize = CMPI(uge, minCrd, size);
-  // Computes minCrd - size + 1
+  // Compute minCrd - size + 1.
   Value mms = SUBI(ADDI(minCrd, C_IDX(1)), size);
   // This is the absolute offset related to the actual tensor.
   return SELECT(geSize, mms, C_IDX(0));
@@ -627,7 +616,7 @@ class NonEmptySubSectIterator : public SparseIterator {
 
 class SubSectIterator;
 
-// A simple helper that helps generating code to traverse a subsection, used
+// A wrapper that helps generating code to traverse a subsection, used
 // by both `NonEmptySubSectIterator`and `SubSectIterator`.
 struct SubSectIterHelper {
   explicit SubSectIterHelper(const SubSectIterator &iter);
@@ -778,7 +767,7 @@ class SubSectIterator : public SparseIterator {
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// Complex SparseIterator derived classes impl.
+// SparseIterator derived classes implementation.
 //===----------------------------------------------------------------------===//
 
 ValueRange SparseIterator::forwardIf(OpBuilder &b, Location l, Value cond) {
@@ -819,7 +808,6 @@ Value DedupIterator::genSegmentHigh(OpBuilder &b, Location l, Value pos) {
       },
       /*afterBuilder=*/
       [](OpBuilder &b, Location l, ValueRange ivs) {
-        // pos ++
         Value nxPos = ADDI(ivs[0], C_IDX(1));
         YIELD(nxPos);
       });
@@ -830,11 +818,11 @@ Value DedupIterator::genSegmentHigh(OpBuilder &b, Location l, Value pos) {
 Value FilterIterator::genCrdNotLegitPredicate(OpBuilder &b, Location l,
                                               Value wrapCrd) {
   Value crd = fromWrapCrd(b, l, wrapCrd);
-  // not on stride
+  // Test whether the coordinate is on stride.
   Value notlegit = CMPI(ne, toWrapCrd(b, l, crd), wrapCrd);
-  // wrapCrd < offset
+  // Test wrapCrd < offset
   notlegit = ORI(CMPI(ult, wrapCrd, offset), notlegit);
-  //  crd >= length
+  // Test crd >= length
   notlegit = ORI(CMPI(uge, crd, size), notlegit);
   return notlegit;
 }

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

@@ -29,12 +29,12 @@ class SparseTensorLevel {
   /// the given position `p` that the immediate parent level is current at.
   /// Returns a pair of values for *posLo* and *loopHi* respectively.
   ///
-  /// For dense level, the *posLo* is the linearized position at beginning,
+  /// For a dense level, the *posLo* is the linearized position at beginning,
   /// while *loopHi* is the largest *coordinate*, it also implies that the
   /// smallest *coordinate* to start the loop is 0.
   ///
-  /// For sparse level, [posLo, loopHi) specifies the range of index pointer to
-  /// load coordinate from the coordinate buffer.
+  /// For a sparse level, [posLo, loopHi) specifies the range of index pointer
+  /// to load coordinate from the coordinate buffer.
   ///
   /// `bound` is only used when the level is `non-unique` and deduplication is
   /// required. It specifies the max upper bound of the non-unique segment.
@@ -68,7 +68,7 @@ enum class IterKind : uint8_t {
   kFilter,
 };
 
-/// Helper class that helps generating loop conditions, etc, to traverse a
+/// Helper class that generates loop conditions, etc, to traverse a
 /// sparse tensor level.
 class SparseIterator {
   SparseIterator(SparseIterator &&) = delete;
@@ -103,17 +103,18 @@ class SparseIterator {
   //
 
   // Whether the iterator support random access (i.e., support look up by
-  // *coordinate*).
-  // A random access iterator also traverses a dense space.
+  // *coordinate*). A random access iterator must also traverses a dense space.
   virtual bool randomAccessible() const = 0;
+
   // Whether the iterator can simply traversed by a for loop.
   virtual bool iteratableByFor() const { return false; };
+
   // Get the upper bound of the sparse space that the iterator might visited. A
   // sparse space is a subset of a dense space [0, bound), this function returns
   // *bound*.
   virtual Value upperBound(OpBuilder &b, Location l) const = 0;
 
-  // Serialize and deserialize the current status to/from a set of values. The
+  // Serializes and deserializes the current status to/from a set of values. The
   // ValueRange should contain values that specifies the current postion and
   // loop bound.
   //
@@ -131,7 +132,7 @@ class SparseIterator {
   // Core functions.
   //
 
-  // Get the current position and the optional *position high* (for non-unique
+  // Gets the current position and the optional *position high* (for non-unique
   // iterators), the value is essentially the number of sparse coordinate that
   // the iterator is current visiting. It should be able to uniquely identify
   // the sparse range for the next level. See SparseTensorLevel::peekRangeAt();
@@ -143,16 +144,17 @@ class SparseIterator {
     llvm_unreachable("unsupported");
   };
 
-  // Initialize the iterator according to the parent iterator's state.
+  // Initializes the iterator according to the parent iterator's state.
   virtual void genInit(OpBuilder &, Location, const SparseIterator *) = 0;
 
-  // Return a pair of values for *upper*, *lower* bound respectively.
+  // Returns a pair of values for *upper*, *lower* bound respectively.
   virtual std::pair<Value, Value> genForCond(OpBuilder &b, Location l) {
     assert(randomAccessible());
     // Random-access iterator is traversed by coordinate, i.e., [curCrd, UB).
     return {getCrd(), upperBound(b, l)};
   }
 
+  // Returns a boolean value that equals `!it.end()`
   virtual Value genNotEnd(OpBuilder &b, Location l) = 0;
   std::pair<Value, ValueRange> genWhileCond(OpBuilder &b, Location l,
                                             ValueRange vs) {
@@ -221,21 +223,30 @@ std::unique_ptr<SparseTensorLevel> makeSparseTensorLevel(OpBuilder &builder,
                                                          Location loc, Value t,
                                                          unsigned tid, Level l);
 
-/// Helper function to create a SparseIterator object.
+/// Helper function to create a simple SparseIterator object that iterate over
+/// the SparseTensorLevel.
 std::unique_ptr<SparseIterator>
 makeSimpleIterator(const SparseTensorLevel &stl);
 
+/// Helper function to create a synthetic SparseIterator object that iterate
+/// over a dense space specified by [0,`sz`).
 std::pair<std::unique_ptr<SparseTensorLevel>, std::unique_ptr<SparseIterator>>
 makeSynLevelAndIterator(Value sz, unsigned tid, unsigned lvl);
 
+/// Helper function to create a SparseIterator object that iterate over a
+/// sliced space, the orignal space (before slicing) is traversed by `sit`.
 std::unique_ptr<SparseIterator>
 makeSlicedLevelIterator(std::unique_ptr<SparseIterator> &&sit, Value offset,
                         Value stride, Value size);
 
+/// Helper function to create a SparseIterator object that iterate over the
+/// non-empty subsections set.
 std::unique_ptr<SparseIterator> makeNonEmptySubSectIterator(
     OpBuilder &b, Location l, const SparseIterator *parent,
     std::unique_ptr<SparseIterator> &&delegate, Value size, unsigned stride);
 
+/// Helper function to create a SparseIterator object that iterate over a
+/// non-empty subsection created by NonEmptySubSectIterator.
 std::unique_ptr<SparseIterator> makeTraverseSubSectIterator(
     const SparseIterator &subsectIter, const SparseIterator &parent,
     std::unique_ptr<SparseIterator> &&delegate, Value size, unsigned stride);