[mlir][sparse] Make getNumTensors() consistent between LoopEmitter and Merger.

[Peiming]

remove outdated TODOs

[Peiming]

update comments.

[Peiming]

fix typo.

mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h

public:
   DimLevelType dlt(TensorLoopId b) const { return latticeMerger.getLvlType(b); }
 
   //
   // LoopEmitter delegates.
   //
 
   TensorLevel makeTensorLevel(TensorId t, Level l) const {
     // Make sure LoopEmitter, GenericOp, and Merger agree on the number of
-    // tensors. Merger has one more synthetic tensor for loop invariants.
-    assert(loopEmitter.getNumTensors() == linalgOp->getNumOperands() &&
-           loopEmitter.getNumTensors() == latticeMerger.getNumTensors() - 1);
+    // tensors.
+    assert(loopEmitter.getNumManifestTensors() == linalgOp->getNumOperands() &&
+           loopEmitter.getNumTensors() == latticeMerger.getNumTensors() &&
+           loopEmitter.getOutTensorId() == latticeMerger.getOutTensorID() &&
+           loopEmitter.getSynTensorId() == latticeMerger.getSynTensorID());
     return loopEmitter.makeTensorLevel(t, l);
   }
   std::pair<TensorId, Level> unpackTensorLevel(TensorLevel tl) const {
     return loopEmitter.unpackTensorLevel(tl);
   }
   template <class ContainerTy>
   auto unpackTensorLevelRange(ContainerTy &&c) const {
     return loopEmitter.unpackTensorLevelRange(std::forward<ContainerTy>(c));

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

public:
   /// the type `LoopOrd` for the same reason as one-past-the-end iterators.
   LoopOrd getCurrentDepth() const { return loopStack.size(); }
 
   /// Gets loop induction variable for the given `LoopOrd`.
   Value getLoopIV(LoopOrd n) const {
     return n < getCurrentDepth() ? loopStack[n].iv : Value();
   }
 
+  /// Gets the total number of manifest tensors (excluding the synthetic
+  /// tensor).
+  unsigned getNumManifestTensors() const { return tensors.size(); }
+
   /// Gets the total number of tensors that loopEmitter is operating on.
-  unsigned getNumTensors() const { return tensors.size(); }
+  unsigned getNumTensors() const {
+    // Manifest tensors with one synthetic tensor at the end.
+    return getNumManifestTensors() + 1;
+  }
 
   /// Gets the TensorId for synthetic tensor.
   TensorId getSynTensorId() const { return tensors.size(); }
 
+  /// Gets the TensorId for output tensor.
+  TensorId getOutTensorId() const {
+    assert(hasOutput);
+    return getNumManifestTensors() - 1;
+  }
+
   /// Compresses a TensorId and Level into a TensorLevel.
   TensorLevel makeTensorLevel(TensorId t, Level l) const {
-    // TODO: getNumTensor() should include synthetic tensor.
-    return l * (getNumTensors() + 1) + t;
+    return l * getNumTensors() + t;
   }
 
   /// De-compresses a TensorLevel back to a pair of TensorId and Level.
   std::pair<TensorId, Level> unpackTensorLevel(TensorLevel tidLvl) const {
-    unsigned nt = getNumTensors() + 1;
+    unsigned nt = getNumTensors();
     return std::make_pair(tidLvl % nt, tidLvl / nt);
   }
 
   /// Converts a range of TensorLevel to a range of std::pair<TensorId, Level>
   template <class ContainerTy>
   auto unpackTensorLevelRange(ContainerTy &&c) const {
     using EltTy = decltype(*c.begin());
     static_assert(std::is_same_v<llvm::remove_cvref_t<EltTy>, TensorLevel>,
private:
 
   /// Generates a predicate to determine whether the tranformed coordinates are
   /// in the given slice.
   /// Returns std::pair<Transformed coordinates, Predicate>
   std::pair<Value, Value> genSliceLegitPredicate(OpBuilder &builder,
                                                  Location loc, Value crd,
                                                  TensorId tid, Level lvl);
 
-  bool isSynTensor(TensorId tid) const { return tid == getNumTensors(); }
+  bool isSynTensor(TensorId tid) const { return tid == getSynTensorId(); }
 
   bool isOutputTensor(TensorId tid) const {
-    return hasOutput && tid == getNumTensors() - 1;
+    return hasOutput && tid == getOutTensorId();
   }
 
   bool isSparseOutput(TensorId tid) const {
     return isOutputTensor(tid) && isSparseOut;
   }
 
   bool isValidLevel(TensorId tid, Level lvl) const {
     return tid < lvlTypes.size() && lvl < lvlTypes[tid].size();
private:
   ///
   /// NOTE: the result uses `Level` rather than the `int64_t` of
   /// `ReassociationIndices`, since the former gives clarity to what
   /// the values actually mean.
   ///
   /// TODO: why not do this computation when we first store the reassoc,
   /// instead of doing it every time we look it up?
   SmallVector<Level, 2> getCollapseReassociation(TensorId tid, Level dstLvl) {
-    assert(tid < getNumTensors() + 1 && "Invalid TensorId");
-    assert(collapseReassoc.size() == getNumTensors() + 1);
+    assert(tid < getNumTensors() && "Invalid TensorId");
+    assert(collapseReassoc.size() == getNumTensors());
     if (const auto reassoc = collapseReassoc[tid]) {
       assert(!isSynTensor(tid) && !isOutputTensor(tid) &&
              "Output/Synthetic tensor should not have reassociation");
       // TODO: store the dstLvlRank in the LoopEmitter so that we can
       // check `dstLvl < dstLvlRank` at the top; and only here need to
       // assert that `reassoc.size() == dstLvlRank`.
       assert(dstLvl < reassoc.size() && "Level is out-of-bounds");
       const auto srcLvls = cast<ArrayAttr>(reassoc[dstLvl]);

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

void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
   // First initialize the top-level type of the fields.
   this->loopTag = loopTag;
   this->hasOutput = hasOutput;
   this->isSparseOut = isSparseOut;
 
   const unsigned numManifestTensors = ts.size();
   const unsigned synTensorId = numManifestTensors;
   const unsigned numTensors = numManifestTensors + 1;
-
+  // tensors array (len == numManifestTensor).
   this->tensors.assign(ts.begin(), ts.end());
+  // Arrays with len == numTensor.
   this->lvlTypes.assign(numTensors, std::vector<DimLevelType>());
   this->lvlSizes.assign(numTensors, std::vector<Value>());
   this->highs.assign(numTensors, std::vector<Value>());
   this->segHi.assign(numTensors, std::vector<Value>());
   this->posits.assign(numTensors, std::vector<Value>());
   this->coords.assign(numTensors, std::vector<Value>());
   this->positionsBuffers.assign(numTensors, std::vector<Value>());
   this->coordinatesBuffers.assign(numTensors, std::vector<Value>());
void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
   // used in loop emitter.
   // FIXME: This map should be maintained outside loop emitter.
   for (LoopOrd n = 0; n < numLoops; n++)
     loopIdToOrd[topSort[n]] = n;
 }
 
 void LoopEmitter::initializeLoopEmit(OpBuilder &builder, Location loc,
                                      LoopEmitter::OutputUpdater updater) {
-  // For every tensor:
+  // For every manifest tensor:
   // * get the values buffer.
   // * For every level:
   //   * get the positions and coordinates buffers
   //   * get/compute the level-size, which is also used as the upper-bound
   //     on positions.
-  for (TensorId t = 0, numTensors = getNumTensors(); t < numTensors; t++) {
+  for (TensorId t = 0, numTensors = getNumManifestTensors(); t < numTensors;
+       t++) {
     const Value tensor = tensors[t];
     const auto rtp = dyn_cast<RankedTensorType>(tensor.getType());
     if (!rtp)
       // Skips only scalar, zero ranked tensor still need to be bufferized and
       // (probably) filled with zeros by users.
       continue;
     // FIXME: the definition of `lvlRank` looks more like a dim-rank;
     // but the variable is used as a level everywhere below, which