[mlir][sparse] extend loop emitter and optimize lattices with the awareness of slice based iteration

[Peiming]

rebase

[Peiming]

rebase

[Peiming]

rebase.

[Peiming]

fix rebase errors..

[Peiming]

rebase.

[Peiming]

fix rebase error.

[Peiming]

rebase.

[Peiming]

rebase.

[Peiming]

code movement.

[Peiming]

rebase.

[Peiming]

add comments,

[Peiming]

code cleanup.

[Peiming]

fix rebase mistakes.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

410

note that these two functions were "touching" since they had the same function doc comment initially.
so if you split them, give L412 version its own doc too

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

31

"appear in the same order" -> sort?

(and yes, I realize it is moved code, but I just noticed the comment ;-)

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

1574

why is this 4 to the right suddenly?
lint?

[Peiming]

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

1574

Idk :-(, done by clang-format automatically\.

[Peiming]

address comments.

[Peiming]

address comments.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

410

Oh, the split is fine, just add a comment

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

31

sorted as -> sorted according to... ? or something of the sort [pun intended ;-)]

[Peiming]

fix comments.

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

public:
   /// operation [Bik96]. For instance: a(i) *= 2.0 or a(i) += a(i) for
   /// sparse vector a.
   bool isSingleCondition(TensorId t, ExprId e) const;
 
   /// Returns true if any `TensorLoopId` in the bitvector corresponds
   /// to sparse level-type.
   bool hasAnySparse(const BitVector &bits) const;
 
-  /// Gets the level-type of the `t`th tensor on `i`th loop.
+  /// Returns true if any set bit corresponds to sparse dimension level type.
+  bool hasAnySlice(const BitVector &bits) const;
+
+  /// Gets the dimension level type of the `t`th tensor on `i`th loop.
   DimLevelType getDimLevelType(TensorId t, LoopId i) const {
     assert(t < numTensors && i < numLoops);
     return lvlTypes[t][i];
   }
+
   DimLevelType getDimLevelType(TensorLoopId b) const {
     return getDimLevelType(tensor(b), loop(b));
   }
 
   /// Gets the loop identifier for the `lvl`th level of the `t`th tensor.
   std::optional<LoopId> getLoopId(TensorId t, Level lvl) const {
     assert(t < numTensors && lvl < lvlToLoop[t].size());
     return lvlToLoop[t][lvl];
public:
   void setLevelAndType(TensorId t, LoopId i, Level lvl, DimLevelType dlt) {
     assert(t < numTensors && i < numLoops && lvl < lvlToLoop[t].size() &&
            isValidDLT(dlt));
     lvlTypes[t][i] = dlt;
     loopToLvl[t][i] = lvl;
     lvlToLoop[t][lvl] = i;
   }
 
-  /// Iterates over a set of `TensorLoopId`s, invoking the callback
-  /// for each `TensorLoopId` and passing it the corresponding tensor
-  /// identifier, level, and level-type.
-  void
-  foreachTensorLoopId(const BitVector &bits,
-                      function_ref<void(TensorLoopId, TensorId,
-                                        std::optional<Level>, DimLevelType)>
-                          callback) const {
-    for (const TensorLoopId b : bits.set_bits())
-      callback(b, tensor(b), getLvl(b), getDimLevelType(b));
-  }
-
   /// Sets whether the output tensor is sparse or not.
   void setHasSparseOut(bool s) { hasSparseOut = s; }
 
   // FIXME: Dim or Level?
   void setLoopDependentSliceDim(unsigned l, unsigned t, unsigned dim) {
     assert(dim < numLoops);
     ldxToDependentSlice[t][l] = dim;
     sliceToRelatedldx[t][dim].push_back(l);
   }
 
   // Whether the ldx has dependent slice.
   bool hasDependentSliceDim(unsigned l) {
     return llvm::any_of(ldxToDependentSlice,
                         [l](const std::vector<std::optional<unsigned>> &m) {
                           return m[l].has_value();
                         });
   }
 
   // Returns the list of loop index appeared in the affine indexing expression
   // on t[d].
   std::vector<unsigned> &getRelatedLoops(unsigned t, unsigned d) {
     return sliceToRelatedldx[t][d];
   }
 
+  // Return the defining [tid, dim] for the loop.
+  std::pair<unsigned, unsigned> getLoopDefiningDim(unsigned l) const {
+    return loopBounds[l];
+  }
+
+  /// Gets the dimension level type of `b`.
+  bool isSliceDim(unsigned b) const {
+    return ldxToDependentSlice[tensor(b)][loop(b)].has_value();
+  }
+
+  /// Iterates over a set of `TensorLoopId`s, invoking the callback
+  /// for each `TensorLoopId` and passing it the corresponding tensor
+  /// identifier, level, and level-type.
+  void foreachTensorLoopId(
+      const BitVector &bits,
+      function_ref<void(TensorLoopId, TensorId, std::optional<Level>,
+                        DimLevelType, bool)>
+          callback) {
+    for (unsigned b : bits.set_bits()) {
+      unsigned t = tensor(b), l = loop(b);
+      bool isSlice = false;
+      std::optional<unsigned> lvl = getLvl(b);
+      if (isSliceDim(b)) {
+        // This must be an undefined dim.
+        assert(!lvl.has_value());
+        isSlice = true;
+        // Slice the tid along the dependent dim to iterate current loop.
+        lvl = ldxToDependentSlice[t][l];
+      }
+      callback(b, t, lvl, getDimLevelType(b), isSlice);
+    }
+  }
+
   /// Convenience getters to immediately access the stored nodes.
   /// Typically it is inadvisible to keep the reference around, as in
   /// `TensorExpr &te = merger.exp(e)`, since insertions into the merger
   /// may cause data movement and invalidate the underlying memory address.
   TensorExp &exp(ExprId e) { return tensorExps[e]; }
   LatPoint &lat(LatPointId p) { return latPoints[p]; }
   SmallVector<LatPointId> &set(LatSetId s) { return latSets[s]; }
 
private:
 
   // Map that converts pair<TensorId, Level> to the corresponding LoopId.
   std::vector<std::vector<std::optional<LoopId>>> lvlToLoop;
 
   // Map from loop idx to the dependent slice [tid, dim] pair (if any).
   std::vector<std::vector<std::optional<unsigned>>> ldxToDependentSlice;
   // Map from the dependent slice [tid, dim] pair to a list of loop idx.
   std::vector<std::vector<std::vector<unsigned>>> sliceToRelatedldx;
+  // Map from loop index to the [tid, dim] pair that defines the loop boundary.
+  std::vector<std::pair<unsigned, unsigned>> loopBounds;
 
   llvm::SmallVector<TensorExp> tensorExps;
   llvm::SmallVector<LatPoint> latPoints;
   llvm::SmallVector<SmallVector<LatPointId>> latSets;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_

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

 //===----------------------------------------------------------------------===//
 
 /// Returns true if tensor materializes uninitialized into the computation.
 static bool isMaterializing(Value val) {
   return val.getDefiningOp<tensor::EmptyOp>() ||
          val.getDefiningOp<bufferization::AllocTensorOp>();
 }
 
 //===----------------------------------------------------------------------===//
 // Code generation environment constructor and general methods
 //===----------------------------------------------------------------------===//
 
 CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
                        unsigned numTensors, unsigned numLoops,
                        unsigned numFilterLoops)
     : linalgOp(linop), sparseOptions(opts),
       latticeMerger(numTensors, numLoops, numFilterLoops), loopEmitter(),
void CodegenEnv::startEmit() {
   if (sparseOut) {
     insChain = sparseOut->get();
     latticeMerger.setHasSparseOut(true);
   }
   // Initialize loop emitter.
   SmallVector<Value> tensors;
   for (OpOperand &t : linalgOp->getOpOperands())
     tensors.push_back(t.get());
-  loopEmitter.initialize(tensors,
-                         StringAttr::get(linalgOp.getContext(),
-                                         linalg::GenericOp::getOperationName()),
-                         /*hasOutput=*/true,
-                         /*isSparseOut=*/sparseOut != nullptr, topSort);
+  loopEmitter.initialize(
+      tensors,
+      /*isSparseOut=*/sparseOut != nullptr, topSort,
+      [this](unsigned t,
+             unsigned d) -> std::vector<std::pair<unsigned, unsigned>> {
+        // Translates from a list of loop index to a list of [tid, dim] pair.
+        std::vector<unsigned> rLoops = this->merger().getRelatedLoops(t, d);
+        std::vector<std::pair<unsigned, unsigned>> ret;
+        ret.reserve(rLoops.size());
+        for (unsigned l : rLoops)
+          ret.push_back(this->merger().getLoopDefiningDim(l));
+        return ret;
+      });
 }
 
 std::optional<Operation *> CodegenEnv::genLoopBoundary(
     function_ref<std::optional<Operation *>(MutableArrayRef<Value> parameters)>
         callback) {
   SmallVector<Value> params;
   if (isReduc()) {
     params.push_back(redVal);

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

 //===----------------------------------------------------------------------===//
 
 class LoopEmitter {
 public:
   /// Optional callback function to setup dense output tensors when
   /// initializing the loop emitter (e.g., to fill a dense output with zeros).
   using OutputUpdater = function_ref<Value(OpBuilder &builder, Location loc,
                                            Value memref, Value tensor)>;
+  // Map from [tid, dim] to a list of dependent [tid, dim] for affine expression
+  // index on sparse tensors.
+  // E.g., for affine index (d0 + d1), it depends on two [tid, dim] that defines
+  // d0 and d1 (for affine expression reduction).
+  // If the list is empty, it means that there is no affine expression on the
+  // input [tid, dim].
+  using DependentDimGetter =
+      function_ref<std::vector<std::pair<unsigned, unsigned>>(unsigned,
+                                                              unsigned)>;
 
   LoopEmitter() = default;
 
   /// Takes an array of input tensors, which the generated loops will
   /// iterate over.  Each tensor is given a `TensorId` (numerically equal
   /// to the position of that tensor `Value` in the array).  Setting
   /// `isSparseOut` indicates that the sparse output tensor is empty,
   /// so the loop emitter will generate loops over it according to the
   /// level-sizes.  The `topSort` array specifies the actual order in
   /// which loops are generated, thus providing a mapping from `LoopOrd`
   /// to `LoopId`.
   void initialize(ValueRange tensors, StringAttr loopTag = nullptr,
                   bool hasOutput = false, bool isSparseOut = false,
-                  ArrayRef<LoopId> topSort = {});
+                  ArrayRef<LoopId> topSort = {},
+                  DependentDimGetter getter = nullptr);
 
   explicit LoopEmitter(ValueRange tensors, StringAttr loopTag = nullptr,
                        bool hasOutput = false, bool isSparseOut = false,
-                       ArrayRef<LoopId> topSort = {});
+                       ArrayRef<LoopId> topSort = {},
+                       DependentDimGetter getter = nullptr);
 
   /// Starts a loop emitting session by generating all the buffers needed
   /// for iterating over the tensors.
   void initializeLoopEmit(OpBuilder &builder, Location loc,
                           OutputUpdater updater = nullptr);
 
   /// Generates code to compute an affine expression whose variables are
   /// `LoopId`s (i.e., `a.cast<AffineDimExpr>().getPosition()` is a valid
private:
   /// NOTE: only one instruction will be moved into reduce block,
   /// transformation will fail if multiple instructions are used to compute
   /// the reduction value. Return %ret to user, while %val is provided by
   /// users (`reduc`).
   void exitForLoop(RewriterBase &rewriter, Location loc,
                    MutableArrayRef<Value> reduc);
 
   /// Exits a while loop, returns the reduction results.
-  void exitCoIterationLoop(OpBuilder &builder, Location loc,
-                           MutableArrayRef<Value> reduc);
+  void exitWhileLoop(OpBuilder &builder, Location loc,
 
   //
   // View-based-reshape methods.
   //
 
   /// Get the collapse reassociation for `tensors[tid][dstLvl]`.
   /// For unreshaped operands, the reassociation is simply an identity
   /// transformation.
private:
   std::vector<Value> valBuffer;                       // to_value
 
   /// Whether the sparse input is a slice.
   std::vector<bool> isSparseSlices;
   /// Values related to slices.
   std::vector<std::vector<Value>> sliceOffsets;
   std::vector<std::vector<Value>> sliceStrides;
 
+  // Map from [tid, dim] to a list of dependent [tid, dim].
+  // See comments for `DependentDimGetter`.
+  std::vector<std::vector<std::vector<std::pair<unsigned, unsigned>>>>
+      dependentDimMap;
+
+  //
+  // View based reshape related-fields and methods
+  //
+
   /// Collapse Reassociations related to a specific tensor
   // TODO: support expand.
   std::vector<ArrayAttr> collapseReassoc;
 
   /// TODO: not yet used, it should track the current level for each tensor
   /// to help eliminate `lvls` paramters from above APIs.
   /// std::vector<Level> curLvl;
 

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

if (i != reassocSize - 1) {
       crd = builder.create<arith::MulIOp>(loc, crd,
                                           this->lvlSizes[tid][reassoc[i + 1]]);
     }
   }
   return crd;
 }
 
 LoopEmitter::LoopEmitter(ValueRange tensors, StringAttr loopTag, bool hasOutput,
-                         bool isSparseOut, ArrayRef<LoopId> topSort) {
-  initialize(tensors, loopTag, hasOutput, isSparseOut, topSort);
+                         bool isSparseOut, ArrayRef<LoopId> topSort,
+                         DependentDimGetter dimGetter) {
+  initialize(tensors, loopTag, hasOutput, isSparseOut, topSort, dimGetter);
 }
 
 void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
-                             bool isSparseOut, ArrayRef<LoopId> topSort) {
+                             bool isSparseOut, ArrayRef<LoopId> topSort,
+                             DependentDimGetter dimGetter) {
   // First initialize the top-level type of the fields.
   this->loopTag = loopTag;
   this->hasOutput = hasOutput;
   this->isSparseOut = isSparseOut;
 
   const TensorId numTensors = ts.size();
   this->tensors.assign(ts.begin(), ts.end());
   this->lvlTypes.assign(numTensors, std::vector<DimLevelType>());
void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
 
   const LoopOrd numLoops = topSort.size();
   // These zeros will be overwritten below, but we need to initialize
   // them to something since we'll need random-access assignment.
   this->loopIdToOrd.assign(numLoops, 0);
   this->loopStack.reserve(numLoops);
   this->loopSeqStack.reserve(numLoops);
 
+  // TODO: dim or level?
+  this->dependentDimMap.assign(
+      numTensors, std::vector<std::vector<std::pair<unsigned, unsigned>>>());
+
   // Initialize nested types of `TensorId`-indexed fields.
   for (TensorId tid = 0; tid < numTensors; tid++) {
     const Value t = tensors[tid];
     // a scalar or 0-dimension tensors
     if (isZeroRankedTensorOrScalar(t.getType()))
       continue;
 
     auto rtp = getRankedTensorType(t);
for (TensorId tid = 0; tid < numTensors; tid++) {
     highs[tid].assign(lvlRank, Value());
     segHi[tid].assign(lvlRank, Value());
     posits[tid].assign(lvlRank, Value());
     coords[tid].assign(lvlRank, Value());
     positionsBuffers[tid].assign(lvlRank, Value());
     coordinatesBuffers[tid].assign(lvlRank, Value());
     sliceOffsets[tid].assign(lvlRank, Value());
     sliceStrides[tid].assign(lvlRank, Value());
+
+    dependentDimMap[tid].assign(lvlRank,
+                                std::vector<std::pair<unsigned, unsigned>>());
+    if (dimGetter)
+      for (Level i = 0; i < lvlRank; i++)
+        dependentDimMap[tid][i] = dimGetter(tid, i);
   }
 
   // Construct the inverse of the `topSort` from the sparsifier.
   // This is needed to map `AffineDimExpr`s back to the `LoopOrd`
   // used in loop emitter.
   // FIXME: This map should be maintained outside loop emitter.
   for (LoopOrd n = 0; n < numLoops; n++)
     loopIdToOrd[topSort[n]] = n;
for (auto [tid, lvl] : llvm::zip(loopInfo.tids, loopInfo.lvls)) {
     coords[tid][lvl] = Value();
     posits[tid][lvl] = Value();
     // Dense level, high is fixed.
     if (!isDenseDLT(lvlTypes[tid][lvl]))
       highs[tid][lvl] = Value();
   }
 }
 
-void LoopEmitter::exitCoIterationLoop(OpBuilder &builder, Location loc,
-                                      MutableArrayRef<Value> reduc) {
+void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
   const LoopInfo &loopInfo = loopStack.back();
   auto whileOp = llvm::cast<scf::WhileOp>(loopInfo.loop);
   builder.setInsertionPointToEnd(loopInfo.userCodeBlock);
   Value iv = loopInfo.iv;
   // Finalize the induction. Note that the induction could be performed
   // in the individual if-branches to avoid re-evaluating the conditions.
   // However, that would result in a rather elaborate forest of yield
   // instructions during code generation. Moreover, performing the induction

 void LoopEmitter::exitCurrentLoop(RewriterBase &rewriter, Location loc,
                                   MutableArrayRef<Value> reduc) {
   // Clean up the values, it would help use to discover potential bug at a
   // earlier stage (instead of silently using a wrong value).
   const LoopInfo &loopInfo = loopStack.back();
   assert(loopInfo.tids.size() == loopInfo.lvls.size());
   SmallVector<Value> red;
   if (llvm::isa<scf::WhileOp>(loopInfo.loop)) {
-    exitCoIterationLoop(rewriter, loc, reduc);
+    exitWhileLoop(rewriter, loc, reduc);
   } else {
     exitForLoop(rewriter, loc, reduc);
   }
 
   assert(loopStack.size() == loopSeqStack.size());
   loopStack.pop_back();
 }

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

static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
   genExpand(env, builder, at, /*atStart=*/true);
   // Emit further intitialization at this loop sequence level.
   const LatPointId l0 = env.set(lts)[0];
   bool needsUniv = false;
 
   SmallVector<TensorId> tids;
   SmallVector<Level> lvls;
   env.merger().foreachTensorLoopId(
-      env.lat(l0).bits, [&](TensorLoopId b, TensorId tid,
-                            std::optional<Level> lvl, DimLevelType dlt) {
+      env.lat(l0).bits,
+      [&](TensorLoopId b, TensorId tid, std::optional<Level> lvl,
+          DimLevelType dlt, bool isSlice) {
         assert(env.merger().loop(b) == idx);
-        if (isDenseDLT(dlt) || isUndefDLT(dlt)) {
+        if (!isSlice && (isDenseDLT(dlt) || isUndefDLT(dlt))) {
           needsUniv = true;
         } else {
           // sparse/singleton levels.
           tids.push_back(tid);
           lvls.push_back(*lvl);
         }
       });
 
   env.emitter().enterNewLoopSeq(builder, env.op().getLoc(), tids, lvls);
 
   // Maintain the universal index only if it is actually
   // consumed by a subsequent lattice point.
   if (needsUniv) {
     unsigned lsize = env.set(lts).size();
     for (unsigned i = 1; i < lsize; i++) {
       const LatPointId li = env.set(lts)[i];
-      if (!env.merger().hasAnySparse(env.lat(li).simple))
+      if (!env.merger().hasAnySparse(env.lat(li).simple) &&
+          !env.merger().hasAnySlice(env.lat(li).simple))
         return true;
     }
   }
   return false;
 }
 
 static void genConstantDenseAddressFromLevel(CodegenEnv &env,
                                              OpBuilder &builder, TensorId tid,
static bool translateBitsToTidLvlPairs(
     SmallVectorImpl<Level> &affineLvls, SmallVectorImpl<AffineExpr> &exps) {
   const BitVector &all = env.lat(li).bits;
   const BitVector &simple = env.lat(li).simple;
   const TensorId outTid = env.merger().getOutTensorID();
   const std::optional<Level> outLvl = env.merger().getLvl(outTid, ldx);
 
   unsigned numloopCond = 0;
   bool hasNonUnique = false;
+  // Converts bits to array + dim pair
   env.merger().foreachTensorLoopId(all, [&, ldx](TensorLoopId b, TensorId tid,
-                                                 std::optional<Level> lvl,
-                                                 DimLevelType dlt) {
+                                                 std::optional<Level> dim,
+                                                 DimLevelType dlt,
+                                                 bool isSlice) {
+    if (isSlice) {
+      // We need to coiterate using slice unconditinally.
+      tids.push_back(tid);
+      lvls.push_back(*dim);
+      numloopCond++;
+      return;
+    }
     if (simple.test(b)) {
       if (isUndefDLT(dlt)) {
-        // An undefined dlt in the lattices, we probably mean to
-        // iterate based on the level of output tensor.  E.g., this
-        // could be a synthetic tensor (for invariants and sparse
-        // output tensor).
-        // out[i][j] = invariant; or a broadcast
+        // An undefined dlt in the lattices, we probably mean to iterate
+        // based on the dim of output tensor.
+        // E.g., this could be a synthetic tensor (for invariants and
+        // sparse output tensor). out[i][j] = invariant; or a broadcast
         // out[i][j] = in[i] (j is undef for input)
-        tid = outTid;
-        lvl = outLvl;
-        // Skips invalid lvl (e.g., when this is a zero ranked tensor).
-        if (!lvl)
+        tid = env.merger().getOutTensorID();
+        dim = env.merger().getLvl(tid, ldx);
+        // Skips invalid dim (e.g., when this is a zero ranked tensor).
+        if (!dim)
           return;
       }
       hasNonUnique = !isUniqueDLT(dlt) || hasNonUnique;
       tids.push_back(tid);
-      lvls.push_back(*lvl);
+      lvls.push_back(*dim);
       numloopCond++;
     } else if (isDenseDLT(dlt)) {
       tids.push_back(tid);
-      lvls.push_back(*lvl);
+      lvls.push_back(*dim);
     } else {
       assert(isUndefDLT(dlt));
       linalg::GenericOp op = env.op();
       if (tid >= op.getNumDpsInputs())
         // We only handle affine expression on input tensors (for now).
         return;
       OpOperand *operand = &op->getOpOperand(tid);
       const auto stt = getSparseTensorType(operand->get());
       // Non-annotated dense tensors requires no special handling.
       if (!stt.hasEncoding())
         return;
 
+      auto enc = stt.getEncoding();
       ArrayRef<AffineExpr> affines =
           op.getMatchingIndexingMap(operand).getResults();
-      const Level lvlRank = stt.getLvlRank();
-      assert(affines.size() == static_cast<size_t>(lvlRank));
-      for (Level l = 0; l < lvlRank; l++) {
-        // FIXME: `toOrigDim` is deprecated.
-        AffineExpr exp = affines[toOrigDim(stt.getEncoding(), l)];
-        // Skip simple affine expression and non-dense levels (which
-        // have their own filter loop).
-        if (exp.isa<AffineDimExpr>() || !stt.isDenseLvl(l))
+      assert(affines.size() == stt.getLvlRank());
+      for (unsigned l = 0, e = stt.getLvlRank(); l < e; l++) {
+        AffineExpr exp = affines[toOrigDim(enc, l)];
+        // Skip simple affine expression and non dense dimensions (which
+        // has it own filter loop).
+        if (exp.isa<AffineDimExpr>() || !isDenseDLT(enc.getDimLevelType()[l]))
           continue;
 
         // Constant affine expression are handled in genLoop
         if (!exp.isa<AffineConstantExpr>()) {
-          bool isAtLoop = false;
-          if (isInvariantAffine(env, exp, ldx, isAtLoop) && isAtLoop) {
+          bool atLevel = false;
+          if (isInvariantAffine(env, exp, ldx, atLevel) && atLevel) {
             // If the compound affine is invariant and we are right at the
             // level. We need to generate the address according to the
             // affine expression. This is also the best place we can do it
             // to avoid putting it inside inner loops.
             // NOTE: It assumes that the levels of the input tensor are
-            // initialized in order (and it is also currently guaranteed by
-            // computeIterationGraph), another more admissible approach
+            // initialized in order (and it is also currently guaranteed
+            // by computeIterationGraph), another more admissible approach
             // might be accepting out-of-order access between consecutive
             // dense levels.
             affineTids.push_back(tid);
             affineLvls.push_back(l);
             exps.push_back(exp);
           }
         }
       }
static bool translateBitsToTidLvlPairs(
 
   assert(numloopCond > 0);
   // If we just need to one loop conditions and the conditions is not imposed on
   // non-unique level, the loop can be generated by a for loop.
   return numloopCond == 1 && !hasNonUnique;
 }
 
 /// Starts a single loop in current sequence.
-static Operation *startLoop(CodegenEnv &env, OpBuilder &builder, LoopOrd at,
-                            LatPointId li, bool needsUniv) {
+static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
+                                              OpBuilder &builder, unsigned at,
+                                              unsigned li, bool needsUniv) {
   // The set of tensors + lvls to generate loops on
   SmallVector<TensorId> tids, affineTids;
   SmallVector<Level> lvls, affineLvls;
   // The set of dense tensors with non-trivial affine expression that just
   // becomes invariant and the address shall now be generated at the current
   // level.
   SmallVector<AffineExpr> affines;
-  bool isFor = translateBitsToTidLvlPairs(
+  bool isSingleCond = translateBitsToTidLvlPairs(
       env, li, env.topSortAt(at), tids, lvls, affineTids, affineLvls, affines);
 
   // Emit the for/while-loop control.
-  Operation *loop = genLoop(env, builder, at, needsUniv, tids, lvls, isFor);
+  Operation *loop =
+      genLoop(env, builder, at, needsUniv, tids, lvls, isSingleCond);
   Location loc = env.op().getLoc();
   for (auto [tid, lvl, exp] : llvm::zip(affineTids, affineLvls, affines)) {
     env.emitter().genDenseAffineAddress(builder, loc, tid, lvl, exp);
   }
 
   // Until now, we have entered every <tid, lvl> pair in {cond, extra,
   // affine}Tids/Lvls. The addresses of the upcoming levels which are dependent
   // on constant affines expression may now be determined.
   auto allTids = llvm::concat<TensorId>(tids, affineTids);
   auto allLvls = llvm::concat<Level>(lvls, affineLvls);
   for (auto [tid, lvl] : llvm::zip(allTids, allLvls)) {
     if (tid != env.merger().getOutTensorID())
       genConstantDenseAddressFromLevel(env, builder, tid, lvl + 1);
   }
 
-  return loop;
+  return std::make_pair(loop, isSingleCond);
 }
 
 /// Ends a single loop in current sequence. Returns new values for needsUniv.
 static bool endLoop(CodegenEnv &env, RewriterBase &rewriter, Operation *loop,
                     LoopId idx, LatPointId li, bool needsUniv) {
   // End a while-loop.
   if (auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
     finalizeWhileOp(env, rewriter, idx, needsUniv, env.lat(li).bits, whileOp);
static void genStmt(CodegenEnv &env, RewriterBase &rewriter, ExprId exp,
   // Start a loop sequence.
   bool needsUniv = startLoopSeq(env, rewriter, exp, at, idx, ldx, lts);
 
   // Emit a loop for every lattice point L0 >= Li in this loop sequence.
   unsigned lsize = env.set(lts).size();
   for (unsigned i = 0; i < lsize; i++) {
     // Start a loop.
     const LatPointId li = env.set(lts)[i];
-    Operation *loop = startLoop(env, rewriter, at, li, needsUniv);
+    auto [loop, isSingleCond] = startLoop(env, rewriter, at, li, needsUniv);
 
     // Visit all lattices points with Li >= Lj to generate the
     // loop-body, possibly with if statements for coiteration.
     Value redInput = env.getReduc();
     Value cntInput = env.getExpandCount();
     Value insInput = env.getInsertionChain();
-    bool isWhile = dyn_cast<scf::WhileOp>(loop) != nullptr;
     for (unsigned j = 0; j < lsize; j++) {
       const LatPointId lj = env.set(lts)[j];
       const ExprId ej = env.lat(lj).exp;
       if (li == lj || env.merger().latGT(li, lj)) {
         // Recurse into body of each branch.
-        if (isWhile) {
+        if (!isSingleCond) {
           scf::IfOp ifOp = genIf(env, rewriter, idx, env.lat(lj).simple);
           genStmt(env, rewriter, ej, at + 1);
           endIf(env, rewriter, ifOp, loop, redInput, cntInput, insInput);
         } else {
           genStmt(env, rewriter, ej, at + 1);
         }
       }
     }
if (hasCycle) {
       return sliceBased
                  ? failure() // TODO: should cycle be resolved differently?
                  : resolveCycle(env, rewriter); // one last shot
     }
 
     if (!isAdmissible)
       return failure(); // inadmissible expression, reject
 
+    // Sort the related loop array such that they are in the same order as they
+    // appears on the topoOrder.
+    // TODO: since we only handle affine addition for slice based codegen, and
+    // addition is assoicative, the order how we evaluate the expression does
+    // not matter. However, to support multiplication, the order of the loop
+    // index should match the evaluation order to the affine expression AST.
     for (OpOperand &t : env.op()->getOpOperands()) {
       unsigned rank = env.op().getMatchingIndexingMap(&t).getNumResults();
       for (unsigned i = 0; i < rank; i++) {
         sortArrayBasedOnOrder(
             env.merger().getRelatedLoops(t.getOperandNumber(), i),
             env.getTopSort());
       }
     }
 
     // Recursively generates code if admissible.
     env.startEmit();
     genBuffers(env, rewriter);
+    // TODO: Constant affine expression should be handled differently when using
+    // slice-based codegen, it does not matter now becasue we already reject the
+    // constant expression at a earlier stage.
     genInitConstantDenseAddress(env, rewriter);
     genStmt(env, rewriter, env.getExprId(), 0);
     genResult(env, rewriter);
     return success();
   }
 
 private:
   // Last resort cycle resolution.

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

: outTensor(numInputOutputTensors - 1),
                std::vector<DimLevelType>(numLoops, DimLevelType::Undef)),
       loopToLvl(numTensors,
                 std::vector<std::optional<Level>>(numLoops, std::nullopt)),
       lvlToLoop(numTensors,
                 std::vector<std::optional<LoopId>>(numLoops, std::nullopt)),
       ldxToDependentSlice(numTensors, std::vector<std::optional<unsigned>>(
                                           numLoops, std::nullopt)),
       sliceToRelatedldx(numTensors, std::vector<std::vector<unsigned>>(
-                                        numLoops, std::vector<unsigned>())) {}
+                                        numLoops, std::vector<unsigned>())),
+      loopBounds(numLoops, std::make_pair(numTensors, numLoops)) {}
 
 //===----------------------------------------------------------------------===//
 // Lattice methods.
 //===----------------------------------------------------------------------===//
 
 ExprId Merger::addExp(Kind k, unsigned x, ExprId y, Value v, Operation *op) {
   const ExprId e = tensorExps.size();
   assert((k != kTensor || x < numTensors) && (k != kLoopVar || x < numLoops));
BitVector Merger::simplifyCond(LatSetId s0, LatPointId p0) {
   for (const LatPointId p1 : latSets[s0]) {
     if (p0 != p1 && latGT(p0, p1)) {
       isSingleton = false;
       break;
     }
   }
 
   BitVector simple(latPoints[p0].bits);
-  bool reset = isSingleton && hasAnySparse(simple);
+  bool reset = isSingleton && (hasAnySparse(simple) || hasAnySlice(simple));
   const TensorLoopId be = simple.size();
   TensorLoopId offset = 0; // relative to the end
   if (!reset)
     // Starts resetting from a dense level, so that the first bit (if kept)
     // is not undefined level-type.
     for (TensorLoopId b = 0; b < be; b++) {
       if (simple[b] && isDenseDLT(getDimLevelType(b))) {
         offset = be - b - 1; // relative to the end
         break;
       }
     }
 
   // Now apply the two basic rules. We also iterate the bits reversely to always
   // keep the rightmost bit (which could possibly be a synthetic tensor).
   for (TensorLoopId b = be - 1 - offset, i = 0; i < be;
        b = b == 0 ? be - 1 : b - 1, i++) {
-    if (simple[b]) {
+    // TODO: better name? also slice on dense level has locate property as well.
+    // Handle it correctly!
+    if (simple[b] && !isSliceDim(b)) {
       const auto dlt = getDimLevelType(b);
       if (!isCompressedDLT(dlt) && !isSingletonDLT(dlt)) {
         if (reset)
           simple.reset(b);
         reset = true;
       }
     }
   }
if (bitsi.count() > bitsj.count()) {
     return true;
   }
   return false;
 }
 
 bool Merger::onlyDenseDiff(LatPointId i, LatPointId j) const {
   BitVector tmp(latPoints[j].bits);
   tmp ^= latPoints[i].bits;
-  return !hasAnySparse(tmp);
+  return !hasAnySparse(tmp) && !hasAnySlice(tmp);
 }
 
 bool Merger::expContainsTensor(ExprId e, TensorId t) const {
   if (tensorExps[e].kind == kTensor)
     return tensorExps[e].tensor == t;
 
   switch (getExpArity(tensorExps[e].kind)) {
   case ExpArity::kNullary:
for (TensorLoopId b = 0, be = bits.size(); b < be; b++)
     if (bits[b]) {
       const auto dlt = getDimLevelType(b);
       if (isCompressedDLT(dlt) || isSingletonDLT(dlt))
         return true;
     }
   return false;
 }
 
+bool Merger::hasAnySlice(const BitVector &bits) const {
+  for (unsigned b = 0, be = bits.size(); b < be; b++)
+    if (bits[b] && isSliceDim(b))
+      return true;
+  return false;
+}
+
 #ifndef NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Print methods (for debugging).
 //===----------------------------------------------------------------------===//
 
 static const char *kindToOpSymbol(Kind kind) {
   switch (kind) {

 }
 
 void Merger::dumpBits(const BitVector &bits) const {
   for (TensorLoopId b = 0, be = bits.size(); b < be; b++) {
     if (bits[b]) {
       const TensorId t = tensor(b);
       const LoopId i = loop(b);
       const auto dlt = lvlTypes[t][i];
+      if (isSliceDim(b))
+        llvm::dbgs() << " DEP_" << t << "_" << i;
+      else
         llvm::dbgs() << " i_" << t << "_" << i << "_" << toMLIRString(dlt);
-      if (ldxToDependentSlice[t][i])
-        llvm::dbgs() << "_D_" << *ldxToDependentSlice[t][i];
     }
   }
 }
 
 #endif // NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Builder methods.