[mlir][sparse] unifying enterLoopOverTensorAtLvl and enterCoIterationOverTensorsAtLvls

[Peiming]

Summary

The tensor levels are now explicitly categorized into different LoopCondKind to instruct LoopEmitter generate different code for different kinds of condition (e.g., SparseCond, SparseSliceCond, SparseAffineIdxCond, etc)

The process of generating a while loop is now dissembled into three steps and they are dispatched to different LoopCondKind handler.

1. Generate LoopCondition (e.g., pos <= posHi for SparseCond, slice.isNonEmpty for SparseAffineIdxCond)
2. Generate LoopBody (e.g., compute the coordinates)
3. Generate ExtraChecks (e.g., if (onSlice(crd)) for SparseSliceCond)

[Peiming]

fix bugs.

[Peiming]

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

1462

How come that this did not cause any issue before?

[Peiming]

cleanup.

[Peiming]

cleanup.

[Peiming]

excludes unrelated changes.

[aartbik]

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

144

Add comment, since we document all methods pretty well in this header

148

same here, and below

270

.... that hold different

294

"that the used" seems to be missing something

299

same, can be used?

331

loop conditions

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

1308–1309

since this returns a bool, I would just say

isParallelLoop(....)

or so (since isParallelFor is already taken)

[Peiming]

address comments.

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

1308–1309

How about the current name?

[Peiming]

rebase.

[aartbik]

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

144

.. loop that tries to ...

270

I still see "holds" not "hold" (subject is plural). so please use hold

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

1308–1309

works too

1309

add

/// Comment

1325

I liked the original genFor and genWhile terminology here, since it was really what was generated
(the filter and coiteration live less for me)

[Peiming]

address comments.

[Peiming]

update comments.

[Peiming]

rebase.

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

public:
   ///   // loop sequence end.
   /// }
   void enterNewLoopSeq(OpBuilder &builder, Location loc,
                        ArrayRef<TensorLevel> tidLvls);
 
   /// Exits the current loop sequence, this will reset universal index to 0.
   void exitCurrentLoopSeq(OpBuilder &builder, Location loc);
 
-  // TODO: Get rid of `lvls` in the argument list? Track the level we
-  // are currently at internally. Then it would be enterNextLvlForTensor.
-  // Still need a way to specify the lvl for non-annotated tensors though,
-  // as those can be accessed out of order.
-  //
-  /// Emits loop over tensor_tid_lvl, it assumes that loops between
-  /// tensor_tid_[0, lvl - 1] have already been generated.
-  /// The function will also perform in-place update on the `reduc` vector to
-  /// return the reduction variable used inside the generated loop.
-  Operation *enterLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
-                                      ArrayRef<TensorLevel> tidLvls,
-                                      MutableArrayRef<Value> reduc = {},
-                                      bool isParallel = false);
-
+  /// Enters a loop that trying to locate a coordinates in a sparse level based
+  /// on the value evaluated by the provided affine expression.
+  /// DEPRECATED: affine index expression should be handled by index reduction
+  /// loop, filter loop-based solution is slow.
   Operation *enterFilterLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                             TensorId tid, Level lvl,
                                             AffineExpr affine,
                                             MutableArrayRef<Value> reduc = {});
 
+  /// Emits the address for a dense level based on the value evaluated by the
+  /// provided affine expression.
+  /// DEPRECATED: affine index expression should be handled by index reduction
+  /// loop, filter loop-based solution is slow.
   void genDenseAffineAddress(OpBuilder &builder, Location loc,
                              TensorLevel tidLvl, AffineExpr lvlExpr);
 
+  // TODO: Get rid of `lvls` in the argument list? Track the level we
+  // are currently at internally. Then it would be enterNextLvlForTensor.
+  // Still need a way to specify the lvl for non-annotated tensors though,
+  // as those can be accessed out of order.
+  //
   /// Emits a co-iteration loop over a set of tensors.
+  /// Emits loop over tensor_tid_lvl, it assumes that loops between
+  /// tensor_tid_[0, lvl - 1] have already been generated.
+  /// The function will also perform in-place update on the `reduc` vector to
+  /// return the reduction variable used inside the generated loop.
   Operation *enterCoIterationOverTensorsAtLvls(
       OpBuilder &builder, Location loc, ArrayRef<TensorLevel> tidLvls,
-      bool needsUniv, MutableArrayRef<Value> reduc = {});
+      MutableArrayRef<Value> reduc = {}, bool isParallel = false,
+      bool genDedup = false, bool needsUniv = false);
 
+  /// Generates code to exit the current loop (e.g., generates yields, forwards
+  /// loop induction variables, etc).
   void exitCurrentLoop(RewriterBase &rewriter, Location loc,
                        MutableArrayRef<Value> reduc = {});
 
   /// Fills the out-parameter with the loop induction variables for all
   /// loops in the current loop-stack.  The variables are given in the
   /// same order as the loop-stack, hence `ivs` should be indexed into
   /// by `LoopOrd` (not `LoopId`).
   void getLoopIVs(SmallVectorImpl<Value> &ivs) const {
auto unpackTensorLevelRange(ContainerTy &&c) const {
     using EltTy = decltype(*c.begin());
     static_assert(std::is_same_v<llvm::remove_cvref_t<EltTy>, TensorLevel>,
                   "Must be unpacking a TensorLevel range");
     return llvm::map_range(std::forward<ContainerTy>(c), [this](EltTy tl) {
       return this->unpackTensorLevel(tl);
     });
   }
 
+  template <class ContainerTy>
+  auto unpackTensorLevelFromCondRange(ContainerTy &&c) const {
+    using EltTy = decltype(*c.begin());
+    static_assert(std::is_same_v<llvm::remove_cvref_t<EltTy>, TensorLvlCond>,
+                  "Must be unpacking a TensorLvlCond range");
+    return unpackTensorLevelRange(
+        llvm::make_first_range(std::forward<ContainerTy>(c)));
+  }
+
   ///
   /// Getters.
   ///
   const std::vector<std::vector<Value>> &getPosits() const { return posits; };
   const std::vector<std::vector<Value>> &getCoords() const { return coords; };
   const std::vector<std::vector<Value>> &getHighs() const { return highs; };
   const std::vector<std::vector<Value>> &getPositionBuffers() const {
     return positionsBuffers;
   };
   const std::vector<std::vector<Value>> &getCoordinateBuffers() const {
     return coordinatesBuffers;
   };
   const std::vector<Value> &getValBuffer() const { return valBuffer; };
 
   constexpr static llvm::StringLiteral getLoopEmitterLoopAttrName() {
     return llvm::StringLiteral("Emitted from");
   }
 
 private:
+  ///
+  /// Structure definitions that holds different kinds of loops information.
+  ///
+
   // A tuple that stored the slice-driven loop information.
   struct SliceLoopInfo final {
     SliceLoopInfo(TensorId tid, Level lvl, bool reduced)
         : tid(tid), lvl(lvl), reduced(reduced) {}
     TensorId tid;
     Level lvl;
     bool reduced;
   };
   // LoopInfo stores information of a loop generated by LoopEmitter. E.g.,
   // the set of tensors levels that the loop is iterating over.
   struct LoopInfo final {
-    LoopInfo(ArrayRef<TensorLevel> tidLvls,
+    LoopInfo(ArrayRef<TensorLevel> trivialTidLvls,
              ArrayRef<SliceLoopInfo> sliceDrivenInfo, Operation *loop,
              Block *userBlock, Value iv, StringAttr loopTag)
-        : tidLvls(tidLvls), sliceDrivenInfo(sliceDrivenInfo), loop(loop),
-          userCodeBlock(userBlock), iv(iv) {
+        : trivialTidLvls(trivialTidLvls), sliceDrivenInfo(sliceDrivenInfo),
+          loop(loop), userCodeBlock(userBlock), iv(iv) {
       // Attached a special tag to loop emitter generated loop.
       if (loopTag)
         loop->setAttr(LoopEmitter::getLoopEmitterLoopAttrName(), loopTag);
     }
-    // The set of <tensor, lvl> that the loop is operating on
-    const llvm::SmallVector<TensorLevel> tidLvls;
-    // Slice-driven loop conditions.
+    // The set of <tensor, lvl>, with *only* trivial index expressions, that are
+    // used as the condition for the generated loop. Extra information is
+    // required for levels with non-tivial index expressions, which is
+    // maintained by the sliceDrivenInfo array below.
+    const llvm::SmallVector<TensorLevel> trivialTidLvls;
+    // The set of <tensor, lvl>, with *only* non-trivial index expressions, that
+    // are used as the condition for the generated loop.
     const llvm::SmallVector<SliceLoopInfo> sliceDrivenInfo;
     const Operation *loop;      // the loop operation
     Block *const userCodeBlock; // the block holding users' generated code.
     const Value iv;             // the induction variable for the loop
   };
 
   // SliceInfo stores information of an extracted slice for slice-driven loop.
   // E.g., the in-scope SSA values for the minimum coordinates and offset for
struct SliceInfo final {
 
     Value minCrd;                     // the minimum coordinate of the slice.
     Value offset;                     // the offset of the current slice.
     Value isNonEmpty;                 // whether the slice is empty.
     std::optional<Level> slicedOnLvl; // the level on which the slice is done
     unsigned depth; // the depth (relative to dependentDimMap[tid][lvl]).
   };
 
+  ///
+  /// Enums for different kinds of loop conditions.
+  ///
+
+  // The bit indicating whether the loop conditions is sparse.
+  static constexpr uint8_t kSparseCond = 1 << 3;
+  // The bit indicating whether the loop iterates over sparse tensor slices
+  // (i.e., with non-empty SliceDimAttr).
+  static constexpr uint8_t kSliceCond = 1 << 2;
+  // The bit indicating whether the loop iterates over tensor levels with
+  // non-trivial affine index reduction.
+  static constexpr uint8_t kAffineIdxCond = 1 << 1;
+  // The bit indicating whether the loop iterates over tensor levels with
+  // non-trivial affine index reduction, and it is not fully reduced.
+  static constexpr uint8_t kAffineIdxCondUnRed = 1 << 0;
+
+  enum class LoopCondKind : uint8_t {
+    // Dense conditions.
+    DenseCond = 0,
+    DenseSliceCond = kSliceCond,
+    DenseAffineCond = kAffineIdxCond,
+    DenseAffineUnRedCond = kAffineIdxCond | kAffineIdxCondUnRed,
+    // Sparse Conditions.
+    SparseCond = kSparseCond,
+    SparseSliceCond = kSparseCond | kSliceCond,
+    SparseAffineCond = kSparseCond | kAffineIdxCond,
+    SparseAffineUnRedCond = kSparseCond | kAffineIdxCond | kAffineIdxCondUnRed,
+  };
+  using TensorLvlCond = std::pair<TensorLevel, LoopCondKind>;
+
+  /// Sparse or dense loop condition.
+  static bool isSparseCond(LoopCondKind k) {
+    return static_cast<uint8_t>(k) & kSparseCond;
+  }
+  static bool isDenseCond(LoopCondKind k) { return !isSparseCond(k); }
+
+  /// Whether loops over sparse tensor slices or sparse tensors.
+  static bool isSliceCond(LoopCondKind k) {
+    return static_cast<uint8_t>(k) & kSliceCond;
+  }
+
+  /// Affine or trivial index expression loop condition.
+  static bool isAffineIdxCond(LoopCondKind k) {
+    return static_cast<uint8_t>(k) & kAffineIdxCond;
+  }
+  static bool isTrivalIdxCond(LoopCondKind k) { return !isAffineIdxCond(k); }
+
+  /// Whether the affine index expression is not fully reduced.
+  static bool isAffineIdxUnRedCond(LoopCondKind k) {
+    return isAffineIdxCond(k) && static_cast<uint8_t>(k) & kAffineIdxCondUnRed;
+  }
+  static bool isAffineIdxRedCond(LoopCondKind k) {
+    return isAffineIdxCond(k) && !isAffineIdxUnRedCond(k);
+  }
+
+  // Whether the loop condition kind requires extra check inside the loop body.
+  // E.g., to iterate over sparse tensor slice, we need to check whether the
+  // current cooridnate is on the slice (e.g., due to stride) or not.
+  static bool isCondWithExtraCheck(LoopCondKind k) {
+    return isSparseCond(k) && isSliceCond(k);
+  }
+
+  static LoopCondKind makeLoopCondKind(bool isSparse, bool isSlice,
+                                       bool isAffine, bool isUnRedu) {
+    assert(!isUnRedu || isAffine);
+    uint8_t bits = 0;
+    bits = isSparse ? bits | kSparseCond : bits;
+    bits = isSlice ? bits | kSliceCond : bits;
+    bits = isAffine ? bits | kAffineIdxCond : bits;
+    bits = isUnRedu ? bits | kAffineIdxCondUnRed : bits;
+    LoopCondKind kind = static_cast<LoopCondKind>(bits);
+
+    // Sanity checks.
+    assert(isSparse == isSparseCond(kind));
+    assert(isSlice == isSliceCond(kind));
+    assert(isAffine == isAffineIdxCond(kind));
+    assert(isUnRedu == isAffineIdxUnRedCond(kind));
+    return kind;
+  }
+
+  void categorizeLoopCondition(ArrayRef<TensorLevel> tidLvls,
+                               SmallVectorImpl<TensorLvlCond> &dnConds,
+                               SmallVectorImpl<TensorLvlCond> &spConds);
+
+  ///
+  /// LoopEmitter internal helper functions.
+  ///
+
   using LoopBodyBuilder = llvm::function_ref<void(OpBuilder &, Location, Value,
                                                   MutableArrayRef<Value>)>;
 
+  /// Whether the list of the sparse condition should be iterated by for loop.
+  bool shouldIteratedByForLoop(ArrayRef<TensorLvlCond> spConds, bool genDedup);
+
   /// Linearizes address for dense dimension (i.e., p = (i * d0) + j).
   Value genAddress(OpBuilder &builder, Location loc, TensorId tid, Level lvl,
                    Value iv);
 
   /// Generates the segment high for a non-unique level (to fast forward
   /// duplicated coordinates).  That is, it generates the code:
   ///
   ///   crd = coordinates_tid_lvl[pos]
bool isValidLevel(TensorId tid, Level lvl) const {
     return tid < lvlTypes.size() && lvl < lvlTypes[tid].size();
   }
 
   /// Prepares loop for iterating over `tensor[lvl]`, under the assumption
   /// that `tensor[0...lvl-1]` loops have already been set up.
   void prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                   TensorId tid, Level lvl);
 
-  /// Emits extra locals, since the locals might not be in simplified lattices
-  /// point used to generate the loops, but are still required to generate
-  /// expressions.
-  void emitExtraLocalsForTensorsAtDenseLvls(OpBuilder &builder, Location loc,
-                                            ArrayRef<TensorLevel> tidLvls);
-
-  /// Emits a for loop to iterate over a tensor level with the provided lower
-  /// bound `lo` and upper bound `hi`.
-  /// Apart from iterating just single tensor level, for loops can be used for
-  /// slice-driven loop on dense level too.
-  Operation *emitForLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
-                                        TensorId tid, Level lvl, Value lo,
-                                        Value hi, MutableArrayRef<Value> reduc,
-                                        bool isParallel);
-
-  /// Emits a while loop to iterate over a sparse level that has been sliced.
-  /// Inserts break statement when the coordinate exceeds the sliceSize;
-  /// The method sets the insertion point inside the generated while loop body
-  /// after the break statement before return (so that callers need to handle
-  /// only in-bound coordinates).
-  Operation *emitWhileLoopOverSliceAtSparseLvl(OpBuilder &builder, Location loc,
-                                               Value pLo, Value pHi,
-                                               Value offset, Value sliceSize,
-                                               TensorId tid, Level lvl,
-                                               MutableArrayRef<Value> reduc);
+  /// Enter dense tensor levels. Since the dense tensor condition could be
+  /// optimized from the loop condition, we need to compute the
+  /// positions/coordinates inside the loop body.
+  void enterTensorsAtDenseLvls(OpBuilder &builder, Location loc,
+                               ArrayRef<TensorLvlCond> dnConds, Value iv,
+                               SmallVectorImpl<SliceLoopInfo> &sliceInfo);
+
+  /// Emits a for loop to iterate over a tensor level with the provided
+  /// lower bound `lo` and upper bound `hi`. Apart from iterating just
+  /// single tensor level, for loops can be used for slice-driven loop on
+  /// dense level too.
+  /// Returns a pair: the loop generated and the value for the induction
+  /// variable.
+  std::pair<Operation *, Value>
+  emitForLoopOverTensorAtLvl(OpBuilder &builder, Location loc, TensorId tid,
+                             Level lvl, Value lo, Value hi,
+                             MutableArrayRef<Value> reduc, bool isParallel);
+
+  /// Emits a while loop to co-iterate over a list of sparse condition, or
+  /// (complex) single sparse condition that can not be handled by for loop
+  /// (e.g., index reduction loop).
+  /// Returns a pair: the loop generated and the value for the induction
+  /// variable (which is the minimum coordinate of all the tensor that being
+  /// iterated).
+  std::pair<Operation *, Value>
+  emitWhileLoopOverTensorsAtLvls(OpBuilder &builder, Location loc,
+                                 ArrayRef<TensorLvlCond> spConds,
+                                 MutableArrayRef<Value> reduc, bool needsUniv);
+
+  /// Generates the while loop condition for the given tensor level condition.
+  Value genWhileLoopConditions(OpBuilder &builder, Location loc, ValueRange ivs,
+                               TensorLvlCond cond);
+
+  /// Generates the while loop body for the given tensor level condition.
+  std::optional<Value> genWhileLoopBody(OpBuilder &builder, Location loc,
+                                        ValueRange ivs, TensorLvlCond cond);
+
+  /// Generates the values (to forward the loop) if the extra check failes.
+  /// E.g., to iterate over a sparse tensor slice, we need:
+  ///
+  /// pos = onSlice(curCrd) ? pos : pos + 1
+  ///
+  /// to skip invalid coordinate that is included in the slice.
+  ValueRange genCheckedValue(OpBuilder &builder, Location loc, Value pred,
+                             ValueRange curArg, TensorLvlCond cond);
 
   /// Exits a for loop, returns the reduction results, e.g.,
   /// For sequential for loops:
   /// %ret = for () {
   ///   ...
   ///   %val = addi %args, %c
   ///   yield %val
   /// }
private:
 
   /// Generates a whileOp to iterate over a subset of coordinates on tid on lvl
   /// using the pHi and pLo provided, the loop break on the first coordinate
   /// that exceeds the slice boundary (i.e., coord >= slice.offset +
   /// slice.size).
   std::pair<Operation *, ValueRange>
   genSliceLvlTraverseLoop(OpBuilder &builder, Location loc, Value pLo,
                           Value pHi, Value offset, Value size, TensorId tid,
-                          Level lvl, ValueRange userReduc, bool genYield,
+                          Level lvl, ValueRange userReduc,
                           LoopBodyBuilder bodyBuilder);
 
   /// Generates a nested loop that iterates over tid on all the coordinates on
   /// lvl.
   ValueRange genUnResolvedSliceTreeTraverse(
       OpBuilder &builder, Location loc, TensorId tid,
       ArrayRef<const SliceInfo *> unResLvls,
       std::optional<std::pair<TensorId, Level>> firstResLvl,
private:
   bool genSliceBegin(OpBuilder &builder, Location loc, TensorId tid, Level lvl);
 
   /// Generates code to get the next non-empty slices of tid on lvl.
   void genSliceNextInduction(OpBuilder &builder, Location loc,
                              const Operation *whileOp, TensorId tid, Level lvl,
                              SmallVectorImpl<Value> &operands,
                              unsigned &retIdx);
 
-  /// Generates a slice-driven while loop as follows.
-  ///
-  /// curSlice = getFirstNonEmptySlice(tensor).
-  ///
-  /// while(isNonEmpty) {
-  ///   ..user code..
-  ///   isNonEmpty, curSlice = getNextNonEmptySlice(curSlice)
-  /// }
-  Operation *emitSliceDrivenLoopOverTensorAtLvl(OpBuilder &builder,
-                                                Location loc, TensorId tid,
-                                                Level lvl,
-                                                MutableArrayRef<Value> reduc);
-
   /// A optional string attribute that should be attached to the loop
   /// generated by loop emitter, it might help following passes to identify
   /// loops that operates on sparse tensors more easily.
   StringAttr loopTag;
   /// Whether the loop emitter needs to treat the last tensor as the output
   /// tensor.
   bool hasOutput;
   bool isSparseOut;

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

static std::pair<Value, Value> fromSliceCrd(OpBuilder &builder, Location loc,
                                             Level lvl) {
   // sliceCrd = (tensorCrd - offset) / stride
   crd = SUBI(crd, offset);
   Value rem = builder.create<arith::RemUIOp>(loc, crd, stride);
   crd = builder.create<arith::DivUIOp>(loc, crd, stride);
   return std::make_pair(crd, rem);
 }
 
+// Generates a bool value for while loop condition that tries to iterate over a
+// fully reduced level with affine index expression.
+static Value genSparseReducedAffineCond(OpBuilder &builder, Location loc,
+                                        Value crdBuf, Value crdHi, Value posit,
+                                        Value posHi) {
+  Value inBound = CMPI(ult, posit, posHi);
+  auto ifOp =
+      builder.create<scf::IfOp>(loc, builder.getI1Type(), inBound, true);
+  // if (inbound)
+  //   yield coord < crdHi
+  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+  Value crd = genIndexLoad(builder, loc, crdBuf, posit);
+  YIELD(CMPI(ult, crd, crdHi));
+  // else
+  //   yield false
+  builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
+  YIELD(constantI1(builder, loc, false));
+
+  builder.setInsertionPointAfter(ifOp);
+  return ifOp.getResult(0);
+}
+
 std::pair<Value, Value>
 LoopEmitter::genSliceLegitPredicate(OpBuilder &builder, Location loc, Value crd,
                                     TensorId tid, Level lvl) {
   assert(isSparseSlices[tid]);
   Value slice = tensors[tid];
   Value offset = sliceOffsets[tid][lvl];
   Value stride = sliceStrides[tid][lvl];
   auto enc = getSparseTensorEncoding(slice.getType());
for (Level lvl = 0; lvl < lvlRank; lvl++) {
           sliceSizes[t][lvl][e - 1] = size;
         }
       }
     }
   }
   localInsertPos = builder.getInsertionPoint()->getPrevNode();
 }
 
+void LoopEmitter::categorizeLoopCondition(
+    ArrayRef<TensorLevel> tidLvls, SmallVectorImpl<TensorLvlCond> &dnConds,
+    SmallVectorImpl<TensorLvlCond> &spConds) {
+  // Finds out the tensor level that we should use to generate loops. Amongs all
+  // the tensor levels, there is at most one sparse tensor level.
+  for (auto [t, l] : unpackTensorLevelRange(tidLvls)) {
+    assert(lvlTypes[t].size() > l); // Must be a valid tid, dim pair
+    auto lvlType = lvlTypes[t][l];
+    // Must be a recognizable DLT.
+    assert(isDenseDLT(lvlType) || isCompressedDLT(lvlType) ||
+           isCompressedWithHiDLT(lvlType) || isSingletonDLT(lvlType));
+
+    bool isSparse = !isDenseDLT(lvlType);
+    bool isSlice = isSparseSlices[t];
+    bool isAffine = !dependentLvlMap[t][l].empty();
+    bool isUnRedu = false;
+    // TODO: Supports affine index expression on sparse tensor slices.
+    assert(!isSlice || !isAffine);
+
+    // Whether the affine index expression has been fully reduced or not.
+    if (!dependentLvlMap[t][l].empty())
+      isUnRedu = !depFullyReduced(t, l);
+
+    auto &dstVec = isSparse ? spConds : dnConds;
+    dstVec.emplace_back(
+        makeTensorLevel(t, l),
+        makeLoopCondKind(isSparse, isSlice, isAffine, isUnRedu));
+  }
+
+  std::sort(spConds.begin(), spConds.end(), [](auto lhs, auto rhs) {
+    // AffineUnRed > Affine > Slice > Trivial
+    return static_cast<uint8_t>(lhs.second) > static_cast<uint8_t>(rhs.second);
+  });
+}
+
 void LoopEmitter::enterNewLoopSeq(OpBuilder &builder, Location loc,
                                   ArrayRef<TensorLevel> tidLvls) {
   // TODO: sort
   assert(loopSeqStack.size() == loopStack.size());
   // Prepares for all the tensors used in the current loop sequence.
   std::vector<std::tuple<TensorId, Level, bool>> slicedTids;
 
   bool hasSynTensor = false;
case AffineExprKind::Constant: {
     int64_t c = a.cast<AffineConstantExpr>().getValue();
     return C_IDX(c);
   }
   default:
     llvm_unreachable("unexpected affine subscript");
   }
 }
 
-Operation *LoopEmitter::emitForLoopOverTensorAtLvl(
+std::pair<Operation *, Value> LoopEmitter::emitForLoopOverTensorAtLvl(
     OpBuilder &builder, Location loc, TensorId tid, Level dstLvl, Value lo,
     Value hi, MutableArrayRef<Value> reduc, bool isParallel) {
   bool isSparseCond = isCompressedDLT(lvlTypes[tid][dstLvl]) ||
                       isCompressedWithHiDLT(lvlTypes[tid][dstLvl]) ||
                       isSingletonDLT(lvlTypes[tid][dstLvl]);
 
   const auto reassoc = getCollapseReassociation(tid, dstLvl);
   // TODO: support dynamic slices.
if (isSparseSlices[tid] && isSparseCond) {
     }
     // Set the insertion point to matched branch.
     builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
     crd = trans;
   }
 
   assert(crd);
   coords[tid][dstLvl] = crd;
-  return loop;
+  return {loop, crd};
 }
 
-Operation *LoopEmitter::emitWhileLoopOverSliceAtSparseLvl(
-    OpBuilder &builder, Location loc, Value pLo, Value pHi, Value offset,
-    Value sliceSize, TensorId tid, Level lvl, MutableArrayRef<Value> reduc) {
-  // TODO: we should generalize the method to support iteration over for
-  // normal slices as well to allow early break.
-  Operation *insertPoint = nullptr;
-  Operation *loop =
-      genSliceLvlTraverseLoop(
-          builder, loc, pLo, pHi, offset, sliceSize, tid, lvl, reduc,
-          /*genYield=*/false, // unaware of the yield values from user yet
-          [this, tid, lvl, reduc, offset,
-           &insertPoint](OpBuilder &builder, Location loc, Value iv,
-                         MutableArrayRef<Value> innerReduc) {
-            assert(innerReduc.size() == reduc.size());
-            // Updates users' reduction variable inplace
-            for (unsigned i = 0, e = reduc.size(); i < e; i++)
-              reduc[i] = innerReduc[i];
-            // Loads the coordinates.
-            Value absC =
-                genIndexLoad(builder, loc, coordinatesBuffers[tid][lvl], iv);
+Value LoopEmitter::genWhileLoopConditions(OpBuilder &builder, Location loc,
+                                          ValueRange ivs, TensorLvlCond cond) {
+  auto [tid, lvl] = unpackTensorLevel(cond.first);
+
+  switch (cond.second) {
+  case LoopCondKind::SparseCond: {
+    const auto reassoc = getCollapseReassociation(tid, lvl);
+    assert(reassoc.size() == ivs.size());
+    assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
+    // We used the first level bound as the bound the collapsed set of levels.
+    return CMPI(ult, ivs.back(), highs[tid][reassoc.front()]);
+  }
+  case LoopCondKind::SparseSliceCond: {
+    assert(ivs.size() == 1);
+    return CMPI(ult, ivs.back(), highs[tid][lvl]);
+  }
+  case LoopCondKind::SparseAffineCond: {
+    assert(ivs.size() == 1);
+    Value crdHi; // loop upper bound
+    {
+      OpBuilder::InsertionGuard guard(builder);
+      Operation *loop = builder.getInsertionBlock()->getParentOp();
+      // crdHi is a loop invariant, hosit the computation outside the loop.
+      if (llvm::isa_and_nonnull<scf::WhileOp>(loop))
+        builder.setInsertionPoint(loop);
+      crdHi = ADDI(getMostRecentSliceOnLvl(tid, lvl).offset,
+                   sliceSizes[tid][lvl].back());
+    }
+    assert(crdHi);
+    return genSparseReducedAffineCond(builder, loc,
+                                      coordinatesBuffers[tid][lvl], crdHi,
+                                      ivs[0], highs[tid][lvl]);
+  }
+  case LoopCondKind::SparseAffineUnRedCond: {
+    assert(ivs.size() == 3);
+    return ivs.front(); // isNonEmpty
+  }
+  default:
+    llvm_unreachable("Unhandled LoopCondKind");
+  }
+  llvm_unreachable("Unhandled LoopCondKind");
+}
+
+std::optional<Value> LoopEmitter::genWhileLoopBody(OpBuilder &builder,
+                                                   Location loc, ValueRange ivs,
+                                                   TensorLvlCond cond) {
+  auto [tid, lvl] = unpackTensorLevel(cond.first);
 
-            // We need to substract the offset to get relative coordinates.
-            // TODO: how to assert relC >=0 during runtime?
-            insertPoint = builder.create<arith::SubIOp>(loc, absC, offset);
-            posits[tid][lvl] = iv;
-            coords[tid][lvl] = insertPoint->getResult(0);
-          })
-          .first;
-  // Sets the insertionn pointer inside loop body.
-  builder.setInsertionPointAfter(insertPoint);
-  return loop;
+  switch (cond.second) {
+  case LoopCondKind::SparseCond: {
+    const auto reassoc = getCollapseReassociation(tid, lvl);
+    assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
+    // Links the SSA chain for segHi.
+    for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++)
+      if (!isUniqueDLT(lvlTypes[tid][reassoc[i]]))
+        segHi[tid][reassoc[i]] = ivs[i];
+
+    // Updates position. For collapsed COO, the position is the same across
+    // consecutive levels.
+    for (auto srcLvl : reassoc)
+      posits[tid][srcLvl] = ivs.back();
+
+    // Update coordinates.
+    coords[tid][lvl] = genSparseCrd(builder, loc, tid, lvl);
+    return std::nullopt;
+  }
+  case LoopCondKind::SparseSliceCond: {
+    assert(ivs.size() == 1);
+    posits[tid][lvl] = ivs.front();
+    Value sCrd = genSparseCrd(builder, loc, tid, lvl);
+    // Converts the coordinate loaded from the actual sparse tensor to the
+    // coordinates in the sparse slice.
+    auto [dCrd, pred] = genSliceLegitPredicate(builder, loc, sCrd, tid, lvl);
+    coords[tid][lvl] = dCrd;
+    return pred;
+  }
+  case LoopCondKind::SparseAffineCond: {
+    assert(ivs.size() == 1);
+    // Coord is the relative offset related to its parents.
+    assert(sliceStack[tid].back().depth == 1 && "TODO: not yet implement");
+    // Update c = absOffset[lvl][depth] - absOffset[lvl][depth - 1]
+    Value posit = ivs[0];
+    Value crdBuf = coordinatesBuffers[tid][lvl];
+    // TODO: Maybe assert relC >=0 during runtime in debug build?
+    Value absC = genIndexLoad(builder, loc, crdBuf, posit);
+    auto relC = SUBI(absC, getFinalSliceOnLvl(tid, lvl).offset);
+    posits[tid][lvl] = posit;
+    coords[tid][lvl] = relC;
+    return std::nullopt;
+  }
+  case LoopCondKind::SparseAffineUnRedCond: {
+    assert(ivs.size() == 3);
+    // Coord is the relative offset related to its parents.
+    // Update c = absOffset[lvl][depth] - absOffset[lvl][depth - 1]
+    assert(sliceStack[tid].back().depth == 1 && "TODO: not yet implement");
+    // Updates the current slice info
+    SliceInfo &sliceInfo = sliceStack[tid].back();
+    sliceInfo.isNonEmpty = ivs[0];
+    sliceInfo.minCrd = ivs[1];
+    sliceInfo.offset = ivs[2];
+    coords[tid][lvl] = sliceInfo.offset;
+    // No extra check is needed before accessing the tensor level.
+    return std::nullopt;
+  }
+  default:
+    llvm_unreachable("Unhandled LoopCondKind");
+  }
+  llvm_unreachable("Unhandled LoopCondKind");
 }
 
-Operation *LoopEmitter::enterLoopOverTensorAtLvl(OpBuilder &builder,
-                                                 Location loc,
-                                                 ArrayRef<TensorLevel> tidLvls,
-                                                 MutableArrayRef<Value> reduc,
-                                                 bool isParallel) {
-  // TODO: support multiple return on parallel for?
-  assert(!isParallel || reduc.size() <= 1);
-  bool isSparseCond = false, isSparseSliceCond = false;
-  auto [tid, lvl] = unpackTensorLevel(tidLvls.front());
+ValueRange LoopEmitter::genCheckedValue(OpBuilder &builder, Location loc,
+                                        Value pred, ValueRange curArgs,
+                                        TensorLvlCond cond) {
+  // Currently only sparse slice condition need extra check.
+  assert(isSliceCond(cond.second) && isSparseCond(cond.second));
+  assert(curArgs.size() == 1);
+  Value nextPos = ADDI(curArgs.front(), C_IDX(1));
+  return SELECT(pred, curArgs.front(), nextPos)->getResults();
+}
 
-  // Finds out the tensor level that we should use to generate loops. Amongs all
-  // the tensor levels, there is at most one sparse tensor level.
+std::pair<Operation *, Value> LoopEmitter::emitWhileLoopOverTensorsAtLvls(
+    OpBuilder &builder, Location loc, ArrayRef<TensorLvlCond> spConds,
+    MutableArrayRef<Value> reduc, bool needsUniv) {
+  // NOTE: the slice driven tensor-related reduction variable must
+  // appear before normal tensors.
+  assert(!spConds.empty());
+
+  // The set of induction variables for the while loop.
+  SmallVector<Value> ivs;
+  // Segement sizes for induction variables used for different kinds of loop
+  // conditions.
+  SmallVector<unsigned> opSegSize;
+
+  // Construct the while-loop with a parameter for each coordinate.
+  for (auto [tl, cKind] : spConds) {
+    auto [tid, lvl] = unpackTensorLevel(tl);
+    const auto lvlTp = lvlTypes[tid][lvl];
+    // Dense level are handled by the shared univeral index.
+    assert(!isDenseCond(cKind));
+    // Must be a recognizable sparse level.
+    assert(isCompressedDLT(lvlTp) || isCompressedWithHiDLT(lvlTp) ||
+           isSingletonDLT(lvlTp));
+
+    unsigned prevSz = ivs.size();
+    const auto reassoc = getCollapseReassociation(tid, lvl);
+    if (isAffineIdxCond(cKind)) {
+      // TODO: Support view-based reshape on sparse levels with affine index
+      // expressions.
+      assert(reassoc.size() == 1);
+      if (isAffineIdxUnRedCond(cKind)) {
+        SliceInfo &sliceInfo = sliceStack[tid].back();
+        // The order matters!
+        ivs.push_back(sliceInfo.isNonEmpty);
+        ivs.push_back(sliceInfo.minCrd);
+        ivs.push_back(sliceInfo.offset);
+      } else {
+        ivs.push_back(posits[tid][lvl]); // loop lower bound (pos low).
+      }
+      // We reduced one more dependency after entering the loop.
+      levelReducedDep[tid][lvl]++;
+    } else {
+      assert(dependentLvlMap[tid][lvl].empty());
+      for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
+        // This is the segment high for each non-unique levels.
+        if (!isUniqueDLT(lvlTypes[tid][reassoc[i]]))
+          ivs.push_back(C_IDX(0));
+      }
+      const Value pos = posits[tid][reassoc.front()];
+      ivs.push_back(pos);
+    }
+    opSegSize.push_back(ivs.size() - prevSz);
+  }
+
+  // The position where user-supplied reduction variable starts.
+  ivs.append(reduc.begin(), reduc.end());
+  // Update universal index.
+  if (needsUniv)
+    ivs.push_back(loopSeqStack.back().first);
+
+  // Ensures all operands are valid.
+  assert(llvm::all_of(ivs, [](Value v) { return v != nullptr; }));
+  TypeRange types = ValueRange(ivs).getTypes();
+  auto whileOp = builder.create<scf::WhileOp>(loc, types, ivs);
+
+  SmallVector<Location> locs(types.size(), loc);
+  Block *before = builder.createBlock(&whileOp.getBefore(), {}, types, locs);
+  Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);
+
+  // Generates loop conditions.
+  builder.setInsertionPointToStart(before);
+  ValueRange bArgs = before->getArguments();
+  Value whileCond = nullptr; // bool values for loop condition.
+  for (auto [c, segSz] : llvm::zip_equal(spConds, opSegSize)) {
+    Value cv = genWhileLoopConditions(builder, loc, bArgs.take_front(segSz), c);
+    bArgs = bArgs.drop_front(segSz);
+    whileCond = !whileCond ? cv : ANDI(whileCond, cv);
+  }
+  // The remaining block arguments are user-provided reduction values and an
+  // optional universal index. Make sure their sizes match.
+  assert(bArgs.size() == reduc.size() + needsUniv ? 1 : 0);
+  builder.create<scf::ConditionOp>(loc, whileCond, before->getArguments());
+
+  // Generates loop body.
+  builder.setInsertionPointToStart(after);
+  ValueRange aArgs = after->getArguments();
+  // Since some LoopCondKind might need extra checks to filter out invalid
+  // iterations, we maintains another array to hold the iteration arguments to
+  // yield if the checks fails.
+  SmallVector<Value> nextArgs(aArgs.begin(), aArgs.end());
+  // A mutable alias for convenient slicing.
+  MutableArrayRef<Value> nextArgsRef = nextArgs;
+  Value extraPred = nullptr;
+  for (auto [c, segSz] : llvm::zip_equal(spConds, opSegSize)) {
+    ValueRange condArgs = aArgs.take_front(segSz);
+    auto pred = genWhileLoopBody(builder, loc, condArgs, c);
+    assert(pred.has_value() == isCondWithExtraCheck(c.second));
+    if (pred.has_value()) {
+      // We need all extra checks to pass.
+      extraPred = extraPred == nullptr ? *pred : ANDI(*pred, extraPred);
+      ValueRange nxArgs = genCheckedValue(builder, loc, *pred, condArgs, c);
+      assert(nxArgs.size() == segSz);
+      // Update the value for cases when some check fails.
+      for (unsigned i = 0; i < segSz; i++) {
+        nextArgsRef[i] = nxArgs[i];
+      }
+    }
+    aArgs = aArgs.drop_front(segSz);
+    nextArgsRef = nextArgsRef.drop_front(segSz);
+  }
+
+  if (extraPred) {
+    auto ifOp = builder.create<scf::IfOp>(loc, types, extraPred, /*else*/ true);
+    // Marks this special IfOp so that Sparsification does not finalizing it.
+    ifOp->setAttr(getLoopEmitterLoopAttrName(),
+                  StringAttr::get(builder.getContext(), "slice"));
+    // Links the SSA chain outside the if statement.
+    YIELD(ifOp->getResults());
+
+    // If not all slices are legit, yield the updated value.
+    builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
+    YIELD(nextArgs);
+
+    // If all slices are legit, start the user generated code.
+    builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+  }
+
+  for (auto [tid, dstLvl] : unpackTensorLevelFromCondRange(spConds)) {
+    const auto reassoc = getCollapseReassociation(tid, dstLvl);
+    assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
+    // TODO: Refactors this into smaller functions.
+    // NOTE: For all the collapsed level (except for the last one, that is why
+    // the loop ends with `reassoc.size() - 1`), as each iteration is advanced
+    // by the segment size of the last level, which does not always invalidate
+    // the segment size for the previous levels, thus we need to propagate the
+    // segment sizes across loop iterations and only forward if needed.
+    //
+    // E.g., for a COO tensor with the following coordinates array.
+    // (0, 0, 1),
+    // (0, 0, 2),
+    // (1, 1, 1),
+    // segHi[lvl=0] = segHi[lvl=1] = 2
+    // segHi[lvl=2] = 1,
+    // the first iteration does not invalidate segHi[0] and segHi[1]
+    for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
+      const Level srcLvl = reassoc[i];
+      if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
+        const Value pos = posits[tid][srcLvl];
+        const auto oldSegHi = segHi[tid][srcLvl];
+        assert(oldSegHi);
+        Value newSegHi = builder.create<arith::CmpIOp>(
+            loc, arith::CmpIPredicate::uge, pos, oldSegHi);
+        auto ifNewSegHi = builder.create<scf::IfOp>(loc, builder.getIndexType(),
+                                                    newSegHi, true);
+        {
+          OpBuilder::InsertionGuard guard(builder);
+          builder.setInsertionPointToStart(ifNewSegHi.thenBlock());
+          YIELD(genSegmentHigh(builder, loc, tid, srcLvl, pos,
+                               highs[tid][srcLvl]));
+          // Else, resues the same segment high.
+          builder.setInsertionPointToStart(ifNewSegHi.elseBlock());
+          YIELD(oldSegHi);
+        }
+        highs[tid][srcLvl + 1] = segHi[tid][srcLvl] = ifNewSegHi.getResult(0);
+      }
+    };
+    const auto srcLvl = reassoc.back();
+    if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
+      segHi[tid][srcLvl] = genSegmentHigh(
+          builder, loc, tid, srcLvl, posits[tid][srcLvl], highs[tid][srcLvl]);
+    }
+  }
+
+  // In-place update on reduction variable.
+  assert(aArgs.size() == reduc.size() + needsUniv ? 1 : 0);
+  for (unsigned i = 0, e = reduc.size(); i < e; i++)
+    reduc[i] = aArgs[i];
+
+  Value min;
+  // Finds the minimum coordinate
+  if (!needsUniv) {
+    for (auto [tid, lvl] : unpackTensorLevelFromCondRange(spConds)) {
+      const auto lvlTp = lvlTypes[tid][lvl];
+      if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
+          isCompressedWithHiDLT(lvlTp)) {
+        const auto crd = coords[tid][lvl];
+        if (min) {
+          Value cmp = CMPI(ult, coords[tid][lvl], min);
+          min = SELECT(cmp, coords[tid][lvl], min);
+        } else {
+          min = crd;
+        }
+      }
+    }
+  } else {
+    assert(!min);
+    // Otherwise, universal index is the minimal pos.
+    min = whileOp.getAfterArguments().back();
+  }
+
+  return {whileOp, min};
+}
+
+bool LoopEmitter::shouldIteratedByForLoop(ArrayRef<TensorLvlCond> sparseConds,
+                                          bool genDedup) {
+  assert(llvm::all_of(sparseConds,
+                      [](TensorLvlCond c) { return isSparseCond(c.second); }));
+
+  // If we need to co-iterate over two sparse tensors, we need a while loop
+  if (sparseConds.size() > 1)
+    return false;
+
+  // We also need a while loop for levels with affine index expression for
+  // non-unique levels when deduplication is required.
+  if (sparseConds.size() == 1) {
+    auto [tid, lvl] = unpackTensorLevel(sparseConds.back().first);
+    auto reassoc = getCollapseReassociation(tid, lvl);
+    return !isAffineIdxCond(sparseConds.back().second) &&
+           !(genDedup && !isUniqueDLT(lvlTypes[tid][reassoc.back()]));
+  }
+
+  return true;
+}
+
+Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
+    OpBuilder &builder, Location loc, ArrayRef<TensorLevel> tidLvls,
+    MutableArrayRef<Value> reduc, bool tryParallel, bool genDedup,
+    bool needsUniv) {
+  // Sanity checks.
+  assert(!tidLvls.empty());
   for (auto [t, l] : unpackTensorLevelRange(tidLvls)) {
-    assert(lvlTypes[t].size() > l);         // Must be a valid tid, dim pair
     assert(!coords[t][l] ||                 // We cannot re-enter the same level
            !dependentLvlMap[t][l].empty()); // unless it is a slice-driver loop
-    auto lvlType = lvlTypes[t][l];
-    // Must be a recognizable DLT.
-    assert(isDenseDLT(lvlType) || isCompressedDLT(lvlType) ||
-           isCompressedWithHiDLT(lvlType) || isSingletonDLT(lvlType));
-
-    // This is a slice-driven loop on sparse level.
-    if (!dependentLvlMap[t][l].empty() && !isDenseDLT(lvlType)) {
-      assert(!isSparseSliceCond && !isSparseCond);
-      isSparseSliceCond = true;
-      tid = t;
-      lvl = l;
-      continue;
-    }
+  // TODO: support multiple return on parallel for?
+  tryParallel = tryParallel && reduc.size() <= 1;
 
-    bool isSparse = isCompressedDLT(lvlType) || isSingletonDLT(lvlType) ||
-                    isCompressedWithHiDLT(lvlType);
-    // We can at most have one sparse input, otherwise, a while loop is
-    // required to co-iterate multiple sparse tensors.
-    assert(!isSparseCond || !isSparse);
-    assert(!isSparseSliceCond || !isSparseCond);
-    if (isSparse) {
-      tid = t;
-      lvl = l;
-    }
-    isSparseCond = isSparseCond || isSparse;
-  }
-
-  DimLevelType lvlType = lvlTypes[tid][lvl];
-  // TODO: Dense slice driven loop can be generated using for loop as well.
-  assert(!isSparseSliceCond || !isDenseDLT(lvlType));
-  bool isDenseSliceCond =
-      isDenseDLT(lvlType) && !dependentLvlMap[tid][lvl].empty();
-  // if the slice is fully reduced, we can now use TACO-based algorithm to
-  // iterate it.
+  SmallVector<TensorLvlCond> spConds;
+  SmallVector<TensorLvlCond> dnConds;
+  categorizeLoopCondition(tidLvls, dnConds, spConds);
 
+  // Only when there is at least one sparse conditions, do we really need the
+  // universal index.
+  // TODO: Maybe we should instead requires merger to pass in a valid value at
+  // the first place instead of adjusting it in LoopEmitter?
+  needsUniv = !spConds.empty() && needsUniv;
+  // The TensorLevel used for loop conditions.
+  // If there is any sparse level, we need to use the sparse condition.
+  // If all levels are dense, we can pick arbitary one (dense slice-driven loop
+  // can be generated using a simple ForOp as well).
   Operation *l = nullptr;
-
-  // At most one tensor used as condition in for loop;
-  SmallVector<TensorLevel, 1> condTidLvl;
-  // There might be multiple dense slice driven tensor.
+  Value iv = nullptr;
   SmallVector<SliceLoopInfo> sliceDrivenInfo;
+  SmallVector<TensorLevel> trivialLvls;
 
   // Generates loops differently depending on whether we need a slice-driven
   // loop or a simple level traversal loop.
-  if (isSparseSliceCond) {
-    bool fullyReduced = depFullyReduced(tid, lvl);
-    if (!fullyReduced) {
-      l = emitSliceDrivenLoopOverTensorAtLvl(builder, loc, tid, lvl, reduc);
-    } else {
-      // If the slice is fully reduced, we can now use TACO-based algorithm to
-      // iterate it.
-      l = emitWhileLoopOverSliceAtSparseLvl(
-          builder, loc, posits[tid][lvl], highs[tid][lvl],
-          getFinalSliceOnLvl(tid, lvl).offset, sliceSizes[tid][lvl].back(), tid,
-          lvl, reduc);
-    }
-    levelReducedDep[tid][lvl]++;
-    sliceDrivenInfo.emplace_back(tid, lvl, fullyReduced);
-  } else {
-    Value lo = isSparseCond ? posits[tid][lvl]           // current offset
-                            : loopSeqStack.back().first; // universal index
+  if (shouldIteratedByForLoop(spConds, genDedup) && !needsUniv) {
+    assert(spConds.size() <= 1);
+    TensorLvlCond tlCond = spConds.empty() ? dnConds.front() : spConds.front();
+    auto loopCondKind = tlCond.second;
+    auto [tid, lvl] = unpackTensorLevel(tlCond.first);
+    Value lo = isSparseCond(loopCondKind)
+                   ? posits[tid][lvl]           // current offset
     Value hi = highs[tid][lvl];
-    if (isDenseSliceCond) {
-      bool fullyReduced = depFullyReduced(tid, lvl);
-      Value sliceSz = sliceSizes[tid][lvl][sliceStack[tid].back().depth - 1];
+    if (isDenseCond(loopCondKind) && isAffineIdxCond(loopCondKind)) {
+      bool unReduc = isAffineIdxUnRedCond(loopCondKind);
+      assert(unReduc == !depFullyReduced(tid, lvl));
+      hi = sliceSizes[tid][lvl][sliceStack[tid].back().depth - 1];
+      if (unReduc) {
         // Adjust for loop hi for dense slice-driven loop.
-      if (fullyReduced) {
-        hi = sliceSz;
-        condTidLvl.push_back(makeTensorLevel(tid, lvl));
-      } else {
-        hi = SUBI(lvlSizes[tid][lvl], sliceSz);
+        hi = SUBI(lvlSizes[tid][lvl], hi);
         hi = ADDI(hi, C_IDX(1));
       }
-    } else {
-      condTidLvl.push_back(makeTensorLevel(tid, lvl));
-    }
-    l = emitForLoopOverTensorAtLvl(builder, loc, tid, lvl, lo, hi, reduc,
-                                   isParallel);
     }
-  Value iv = coords[tid][lvl];
-  for (auto [t, l] : unpackTensorLevelRange(tidLvls)) {
-    // We only need to handle slice-driven loops on dense level here.
-    // If it is a slice-driven loop on sparse level, it needs a while loop to
-    // insert break statements, and it must have been handled correctly in L692.
-    if (!dependentLvlMap[t][l].empty() && isDenseDLT(lvlTypes[t][l])) {
-      // Pushes sliced levels to build correct LoopInfo.
-      bool fullyReduc = depFullyReduced(t, l);
-      SliceInfo &info = sliceStack[t].back();
-      if (fullyReduc) {
-        posits[t][l] = genAddress(builder, loc, t, l, ADDI(info.offset, iv));
-      } else {
-        // Puts sliced dense loop into LoopInfo so that LoopEmitter knows how to
-        // exit it.
-        sliceDrivenInfo.emplace_back(t, l, fullyReduc);
-        // Update the slice information as we enter the new loop.
-        assert(*info.slicedOnLvl == l);
-        info.minCrd = info.offset = iv;
-        info.isNonEmpty = constantI1(builder, loc, true);
-        levelReducedDep[t][l]++;
+    std::tie(l, iv) = emitForLoopOverTensorAtLvl(builder, loc, tid, lvl, lo, hi,
+                                                 reduc, tryParallel);
+    // For loop condition must be a trivial condition (levels without affine
+    // index expression).
+    trivialLvls.push_back(tlCond.first);
+  } else {
+    for (auto [tl, cKind] : spConds) {
+      if (isAffineIdxCond(cKind)) {
+        auto [tid, lvl] = unpackTensorLevel(tl);
+        bool unReduc = isAffineIdxUnRedCond(cKind);
+        assert(unReduc == !depFullyReduced(tid, lvl));
+        sliceDrivenInfo.emplace_back(tid, lvl, /*fullyReduced=*/!unReduc);
+      } else {
+        trivialLvls.push_back(tl);
       }
     }
+
+    std::tie(l, iv) =
+        emitWhileLoopOverTensorsAtLvls(builder, loc, spConds, reduc, needsUniv);
   }
+
+  // Enter dense tensor levels.
+  enterTensorsAtDenseLvls(builder, loc, dnConds, iv, sliceDrivenInfo);
   // NOTE: we can also prepare for next dim here in advance
+
   // Pushes the loop into stack.
-  loopStack.emplace_back(condTidLvl, sliceDrivenInfo, l,
+  loopStack.emplace_back(trivialLvls, sliceDrivenInfo, l,
                          builder.getInsertionBlock(), iv, loopTag);
-  // Emit extra locals.
-  emitExtraLocalsForTensorsAtDenseLvls(builder, loc, tidLvls);
   return l;
 }
 
 Operation *LoopEmitter::enterFilterLoopOverTensorAtLvl(
     OpBuilder &builder, Location loc, TensorId tid, Level lvl,
     AffineExpr affine, MutableArrayRef<Value> reduc) {
   assert(isValidLevel(tid, lvl));
   assert(!affine.isa<AffineDimExpr>() && !isDenseDLT(lvlTypes[tid][lvl]));
Operation *LoopEmitter::enterFilterLoopOverTensorAtLvl(
   return forOp;
 }
 
 void LoopEmitter::genDenseAffineAddress(OpBuilder &builder, Location loc,
                                         TensorLevel tidLvl,
                                         AffineExpr lvlExpr) {
   auto [tid, lvl] = unpackTensorLevel(tidLvl);
   assert(isDenseDLT(lvlTypes[tid][lvl]));
-  // For dense levels, the level-coordinate also serves as the position.
+  // For dense levels, the vel-coordinate also serves as the position.
   Value lvlCrd = genAffine(builder, loc, lvlExpr);
   posits[tid][lvl] = genAddress(builder, loc, tid, lvl, lvlCrd);
 }
 
-Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
-    OpBuilder &builder, Location loc, ArrayRef<TensorLevel> tidLvls,
-    bool needsUniv, MutableArrayRef<Value> reduc) {
-  // NOTE: the slice driven tensor-related reduction variable must
-  // appear before normal tensors.
-  SmallVector<Type> types;
-  SmallVector<Value> operands;
-  // Construct the while-loop with a parameter for each coordinate.
-  const Type indexType = builder.getIndexType();
-  for (auto [tid, lvl] : unpackTensorLevelRange(tidLvls)) {
-    // TODO: support coiteration with slice driven tensors.
-    const auto lvlTp = lvlTypes[tid][lvl];
-    assert(dependentLvlMap[tid][lvl].empty() && "TODO: not yet implemented");
-    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
-        isCompressedWithHiDLT(lvlTp)) {
-      const auto reassoc = getCollapseReassociation(tid, lvl);
-      for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
-        if (!isUniqueDLT(lvlTypes[tid][reassoc[i]])) {
-          // This is the segment high for each non-unique levels.
-          types.push_back(indexType);
-          operands.push_back(C_IDX(0));
-        }
-      }
-      const auto pos = posits[tid][reassoc.front()];
-      assert(pos);
-      types.push_back(indexType);
-      operands.push_back(pos);
-    }
-  }
-  // The position where user-supplied reduction variable starts.
-  for (Value rec : reduc) {
-    types.push_back(rec.getType());
-    operands.push_back(rec);
-  }
-  if (needsUniv) {
-    types.push_back(indexType);
-    // Update universal index.
-    operands.push_back(loopSeqStack.back().first);
-  }
-  assert(types.size() == operands.size());
-  scf::WhileOp whileOp = builder.create<scf::WhileOp>(loc, types, operands);
-
-  SmallVector<Location> locs(types.size(), loc);
-  Block *before = builder.createBlock(&whileOp.getBefore(), {}, types, locs);
-  Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);
-
-  // Build the "before" region, which effectively consists
-  // of a conjunction of "i < upper" tests on all induction.
-  builder.setInsertionPointToStart(&whileOp.getBefore().front());
-  Value cond;
-  unsigned o = 0;
-  for (auto [t, lvl] : unpackTensorLevelRange(tidLvls)) {
-    const TensorId tid = t; // Why `t` can not be captured by lambda?
-    const auto lvlTp = lvlTypes[tid][lvl];
-    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
-        isCompressedWithHiDLT(lvlTp)) {
-      const auto reassoc = getCollapseReassociation(tid, lvl);
-      assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
-      for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
-        if (!isUniqueDLT(lvlTypes[tid][reassoc[i]])) {
-          // Links the SSA chain for segHi.
-          segHi[tid][reassoc[i]] = after->getArgument(o++);
-        }
-      }
-      Value op1 = before->getArgument(o);
-      // We used the first level bound as the bound the collapsed set of levels.
-      Value op2 = highs[tid][reassoc.front()];
-      Value opc = CMPI(ult, op1, op2);
-      cond = cond ? ANDI(cond, opc) : opc;
-      // Update positions
-      Value pos = after->getArgument(o++);
-      // For COO, the position is the same across consecutive levels.
-      /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
-      llvm::for_each(reassoc, [this, tid, pos](Level srcLvl) {
-        posits[tid][srcLvl] = pos;
-      });
-    }
-  }
-  builder.create<scf::ConditionOp>(loc, cond, before->getArguments());
-
-  // Generates while body.
-  builder.setInsertionPointToStart(&whileOp.getAfter().front());
-
-  SmallVector<std::pair<Value, unsigned>> slicesPreds;
-  unsigned i = 0;
-  for (auto [tid, lvl] : unpackTensorLevelRange(tidLvls)) {
-    // Prepares for next level.
-    const auto lvlTp = lvlTypes[tid][lvl];
-    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
-        isCompressedWithHiDLT(lvlTp)) {
-      coords[tid][lvl] = genSparseCrd(builder, loc, tid, lvl);
-      if (isSparseSlices[tid]) {
-        auto [trans, pred] =
-            genSliceLegitPredicate(builder, loc, coords[tid][lvl], tid, lvl);
-        slicesPreds.emplace_back(pred, i);
-        // Updates to the relative coordinate to the slice.
-        coords[tid][lvl] = trans;
-      }
-      i++;
-    }
-  }
-
-  if (!slicesPreds.empty()) {
-    // Skips invalid loop iteration when slice coordinate is inapplicable.
-    SmallVector<Value> yields(after->getArguments());
-    // Generates a list of if statments
-    //  pos = in_slice ? pos : pos + 1
-    // TODO: instead of always picking pos + 1, we should set pos = high to
-    // break to loop if the coordinates are larger than the slice size.
-    //
-    // This "idx" is the index into `llvm::zip(tids, lvls)`
-    for (auto [pred, idx] : slicesPreds) {
-      Value nextPos = ADDI(yields[idx], C_IDX(1));
-      yields[idx] = SELECT(pred, yields[idx], nextPos);
-    }
-
-    Value pred = slicesPreds.front().first;
-    for (int i = 1, e = slicesPreds.size(); i < e; i++) {
-      pred = ANDI(pred, slicesPreds[i].first);
-    }
-    auto ifOp = builder.create<scf::IfOp>(loc, types, pred, /*else*/ true);
-    ifOp->setAttr(getLoopEmitterLoopAttrName(),
-                  StringAttr::get(builder.getContext(), "slice"));
-    YIELD(ifOp->getResults());
-    assert(types.size() == yields.size());
-    // If not all slices are legit
-    builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
-    YIELD(yields);
-
-    // If all slices are legit, start the user generated code.
-    builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
-  }
-
-  Value min;
-  // Finds the minimum coordinate
-  if (!needsUniv) {
-    for (auto [tid, lvl] : unpackTensorLevelRange(tidLvls)) {
-      const auto lvlTp = lvlTypes[tid][lvl];
-      if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
-          isCompressedWithHiDLT(lvlTp)) {
-        const auto crd = coords[tid][lvl];
-        if (min) {
-          Value cmp = CMPI(ult, coords[tid][lvl], min);
-          min = SELECT(cmp, coords[tid][lvl], min);
-        } else {
-          min = crd;
-        }
-      }
-    }
-  } else {
-    assert(!min);
-    // Otherwise, universal index is the minimal pos.
-    min = after->getArguments().back();
-  }
-
-  // Sets up the loop stack.
-  loopStack.emplace_back(tidLvls, ArrayRef<SliceLoopInfo>(), whileOp,
-                         builder.getInsertionBlock(), min, loopTag);
-  assert(loopStack.size() == loopSeqStack.size());
-
-  for (auto [tid, dstLvl] : unpackTensorLevelRange(tidLvls)) {
-    const auto reassoc = getCollapseReassociation(tid, dstLvl);
-    assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
-    // TODO: Refactors this into smaller functions.
-    // NOTE: For all the collapsed level (except for the last one, that is why
-    // the loop ends with `reassoc.size() - 1`), as each iteration is advanced
-    // by the segment size of the last level, which does not always invalidate
-    // the segment size for the previous levels, thus we need to propagate the
-    // segment sizes across loop iterations and only forward if needed.
-    //
-    // E.g., for a COO tensor with the following coordinates array.
-    // (0, 0, 1),
-    // (0, 0, 2),
-    // (1, 1, 1),
-    // segHi[lvl=0] = segHi[lvl=1] = 2
-    // segHi[lvl=2] = 1,
-    // the first iteration does not invalidate segHi[0] and segHi[1]
-    for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
-      const Level srcLvl = reassoc[i];
-      if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
-        const Value pos = posits[tid][srcLvl];
-        const auto oldSegHi = segHi[tid][srcLvl];
-        assert(oldSegHi);
-        Value newSegHi = builder.create<arith::CmpIOp>(
-            loc, arith::CmpIPredicate::uge, pos, oldSegHi);
-        auto ifNewSegHi = builder.create<scf::IfOp>(loc, builder.getIndexType(),
-                                                    newSegHi, true);
-        {
-          OpBuilder::InsertionGuard guard(builder);
-          builder.setInsertionPointToStart(ifNewSegHi.thenBlock());
-          YIELD(genSegmentHigh(builder, loc, tid, srcLvl, pos,
-                               highs[tid][srcLvl]));
-          // Else, resues the same segment high.
-          builder.setInsertionPointToStart(ifNewSegHi.elseBlock());
-          YIELD(oldSegHi);
-        }
-        highs[tid][srcLvl + 1] = segHi[tid][srcLvl] = ifNewSegHi.getResult(0);
-      }
-    };
-    const auto srcLvl = reassoc.back();
-    if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
-      segHi[tid][srcLvl] = genSegmentHigh(
-          builder, loc, tid, srcLvl, posits[tid][srcLvl], highs[tid][srcLvl]);
-    }
-  }
-
-  // Emits extra locals
-  emitExtraLocalsForTensorsAtDenseLvls(builder, loc, tidLvls);
-
-  // Updates reduction variables
-  assert(after->getNumArguments() == o + reduc.size() + (needsUniv ? 1 : 0));
-  // In-place update on reduction variable.
-  for (unsigned i = 0, e = reduc.size(); i < e; i++)
-    reduc[i] = after->getArgument(o + i);
-
-  return whileOp;
-}
-
 void LoopEmitter::prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                              TensorId tid, Level dstLvl) {
   assert(isValidLevel(tid, dstLvl));
   const auto lvlTp = lvlTypes[tid][dstLvl];
 
   if (isDenseDLT(lvlTp))
     return;
 
if (isSingletonDLT(lvlTp)) {
               : ADDI(pLo, c1);
       return;
     }
   }
 
   llvm_unreachable("Unrecognized level-type!");
 }
 
-void LoopEmitter::emitExtraLocalsForTensorsAtDenseLvls(
-    OpBuilder &builder, Location loc, ArrayRef<TensorLevel> tidLvls) {
-  // Initialize dense positions. Note that we generate dense coordinates of the
-  // output tensor unconditionally, since they may not appear in the lattice,
-  // but may be needed for linearized codegen.
-  for (auto [tid, lvl] : unpackTensorLevelRange(tidLvls)) {
-    if (isSynTensor(tid))
-      continue;
+void LoopEmitter::enterTensorsAtDenseLvls(
+    OpBuilder &builder, Location loc, ArrayRef<TensorLvlCond> dnConds, Value iv,
+    SmallVectorImpl<SliceLoopInfo> &sliceInfo) {
+  for (auto [dnTidLvl, denseLoopCond] : dnConds) {
+    auto [tid, lvl] = unpackTensorLevel(dnTidLvl);
+    assert(isDenseDLT(lvlTypes[tid][lvl]));
 
-    if (isDenseDLT(lvlTypes[tid][lvl])) {
-      // Slice-driven dense level should have be handled already.
-      if (!dependentLvlMap[tid][lvl].empty())
+    if (isAffineIdxCond(denseLoopCond)) {
+      // Pushes sliced levels to build correct LoopInfo.
+      bool unReduc = isAffineIdxUnRedCond(denseLoopCond);
+      SliceInfo &info = sliceStack[tid].back();
+      if (unReduc) {
+        // Pushes sliced dense loop info to tell LoopEmitter how to exit it.
+        sliceInfo.emplace_back(tid, lvl, /*fullyReduced=*/false);
+        // Update the slice information as we enter the new loop.
+        assert(*info.slicedOnLvl == lvl);
+        info.minCrd = info.offset = iv;
+        info.isNonEmpty = constantI1(builder, loc, true);
+        levelReducedDep[tid][lvl]++;
+      } else {
+        posits[tid][lvl] =
+            genAddress(builder, loc, tid, lvl, ADDI(info.offset, iv));
+      }
+    } else {
+      // Skips the synthetic tensor
+      if (isSynTensor(tid))
         continue;
-
+      // A dense level with trivial index expression.
+      assert(dependentLvlMap[tid][lvl].empty());
       auto enc = getSparseTensorEncoding(tensors[tid].getType());
       if (enc && !isSparseOutput(tid)) {
         bool validPos = lvl == 0 || posits[tid][lvl - 1];
         if (!validPos) {
           // We might not find the pos for the sparse output tensor as it is
           // unconditionally required by the sparsification.
           assert(isOutputTensor(tid));
           continue;
         }
-        posits[tid][lvl] =
-            genAddress(builder, loc, tid, lvl, loopStack.back().iv);
+        posits[tid][lvl] = genAddress(builder, loc, tid, lvl, iv);
         // NOTE: we can also prepare for next lvl here in advance
       }
     }
   }
 }
 
 void LoopEmitter::exitForLoop(RewriterBase &rewriter, Location loc,
                               MutableArrayRef<Value> reduc) {
#endif // NDEBUG
     // In-place update reduction variables.
     for (unsigned i = 0, e = parOp.getResults().size(); i < e; i++)
       reduc[i] = parOp.getResult(i);
   }
 
   // Finished iterating a tensor, clean up
   // We only do the clean up on for loop as while loops do not necessarily
   // finish the iteration on a sparse tensor
-  for (auto [tid, lvl] : unpackTensorLevelRange(loopInfo.tidLvls)) {
+  for (auto [tid, lvl] : unpackTensorLevelRange(loopInfo.trivialTidLvls)) {
     // Reset to null.
     coords[tid][lvl] = Value();
     posits[tid][lvl] = Value();
     // Dense level, high is fixed.
     if (!isDenseDLT(lvlTypes[tid][lvl]))
       highs[tid][lvl] = Value();
   }
 }
 
 void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
                                 MutableArrayRef<Value> reduc) {
   const LoopInfo &loopInfo = loopStack.back();
   auto whileOp = llvm::cast<scf::WhileOp>(loopInfo.loop);
   Value iv = loopInfo.iv;
+  Value one = C_IDX(1);
 
   // 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
   // after the if-statements more closely resembles code generated by TACO.
   unsigned o = 0;
   SmallVector<Value> operands;
   unsigned delta = 0;
   for (auto [tid, lvl, resolved] : loopInfo.sliceDrivenInfo) {
     // TODO: handle dense.
     assert(isCompressedDLT(lvlTypes[tid][lvl]));
     levelReducedDep[tid][lvl]--;
     if (!resolved) {
+      // TODO: support coiterating multiple slices
+      assert(loopInfo.trivialTidLvls.empty() &&
+             loopInfo.sliceDrivenInfo.size() == 1);
       genSliceNextInduction(builder, loc, whileOp, tid, lvl, operands, o);
       continue;
     }
-    // TODO: We need to distinguish coiterate loop with slice-driven loop and
-    // fully reduced while op for iterating one slices.
-    // FIXME: since we didn't implement coiteration, this must be iteration
-    // just on fully resolved slice.
-    assert(loopInfo.sliceDrivenInfo.size() == 1 && loopInfo.tidLvls.empty());
-    // The if guard to filter out out-range coordinates.
-    assert(llvm::isa<scf::IfOp>(builder.getInsertionBlock()->getParentOp()));
+
+    if (loopInfo.trivialTidLvls.empty() &&
+        loopInfo.sliceDrivenInfo.size() == 1) {
+      // Forwards the position iterator.
+      operands.push_back(ADDI(posits[tid][lvl], one));
+    } else {
+      const Value pos = posits[tid][lvl];
+      const Value nxPos = ADDI(posits[tid][lvl], one);
+      Value cmp = CMPI(eq, coords[tid][lvl], iv);
+      operands.push_back(SELECT(cmp, nxPos, pos));
+    }
+
+    // The coordinate is invalid now.
+    coords[tid][lvl] = nullptr;
+
+    // Update the position iterator as we exit the while loop.
     posits[tid][lvl] = whileOp->getResult(o++);
-    // FIXME: we are not using continue here since we do not support
-    // coiteration on slices. But it need to be treated similarly as the
-    // universal index.
-    o++; // skip continue flag.
-    // Since we did not push two results from whileOp. The size of the
-    // operands vector is smaller than the actual number of return values from
-    // the whileOp.
-    // It is because we are actually generating yield in the IfOp inside the
-    // whileOp to only iterates over inbound coordinates within the slices.
-    delta += 2;
   };
 
-  Value one = C_IDX(1);
-  for (auto [tid, dstLvl] : unpackTensorLevelRange(loopInfo.tidLvls)) {
+  for (auto [tid, dstLvl] : unpackTensorLevelRange(loopInfo.trivialTidLvls)) {
     const auto lvlTp = lvlTypes[tid][dstLvl];
     if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
         isCompressedWithHiDLT(lvlTp)) {
       const auto reassoc = getCollapseReassociation(tid, dstLvl);
       assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
         const Level srcLvl = reassoc[i];
         if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
       const Value newPos = whileOp->getResult(o++);
       // We need to define a new local variable for `tid` to avoid
       // warnings about "captured structured bindings are a C++20 extension".
       // FIXME(wrengr): define a helper function to capture this idiom!
       const TensorId newTid = tid;
       llvm::for_each(reassoc, [this, newTid, newPos](Level srcLvl) {
         posits[newTid][srcLvl] = newPos;
       });
+
       // The coordinate is invalid now.
       coords[tid][dstLvl] = nullptr;
       // The segment high is invalid now.
       segHi[tid][dstLvl] = nullptr;
       // highs remains unchanged.
     }
   }
 
for (auto it = sliceStack[tid].rbegin(), ie = sliceStack[tid].rend(); it < ie;
        it++) {
     if (it->slicedOnLvl == lvl) { // the level matched
       return *it;
     }
   }
   llvm_unreachable("Failed to find sliceInfo");
 }
 
-static Value genSparseReducedAffineCond(OpBuilder &builder, Location loc,
-                                        Value crdBuf, Value crdHi, Value posit,
-                                        Value posHi, Value cont) {
-  Value inBound = CMPI(ult, posit, posHi);
-  auto ifOp = builder.create<scf::IfOp>(loc, cont.getType(), inBound, true);
-  // if (inbound)
-  //   yield coord < crdHi
-  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
-  Value crd = genIndexLoad(builder, loc, crdBuf, posit);
-  YIELD(CMPI(ult, crd, crdHi));
-  // else
-  //   yield false
-  builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
-  YIELD(constantI1(builder, loc, false));
-
-  builder.setInsertionPointAfter(ifOp);
-  return ifOp.getResult(0);
-}
-
 // Generates a while loop to iterate over a slice sparse level as follows.
 //
 // while(coords[loopLo] < offset + size) {
 //   body_builder
 //   loopLo ++;
 // }
 std::pair<Operation *, ValueRange> LoopEmitter::genSliceLvlTraverseLoop(
     OpBuilder &builder, Location loc, Value posLo, Value posHi, Value offset,
-    Value size, TensorId tid, Level lvl, ValueRange userReduc, bool genYield,
+    Value size, TensorId tid, Level lvl, ValueRange userReduc,
     LoopBodyBuilder bodyBuilder) {
   Value c1 = C_IDX(1);
   Value sliceHi = ADDI(offset, sliceSizes[tid][lvl].back());
+  SmallVector<Value> reduc{posLo}; // loop lower bounds
+  const unsigned numMetaReduc = reduc.size();
 
-  SmallVector<Value> reduc = {
-      posLo,                          // loop lower bounds
-      constantI1(builder, loc, true), // continue
-  };
   // Append user required reduction value.
   reduc.append(userReduc.begin(), userReduc.end());
   scf::WhileOp whileOp = builder.create<scf::WhileOp>(
       loc, ValueRange(reduc).getTypes(), reduc,
       /*beforeBuilder=*/
       [this, posHi, sliceHi, tid, lvl](OpBuilder &builder, Location loc,
                                        ValueRange args) {
-        Value cond = genSparseReducedAffineCond(
-            builder, loc, coordinatesBuffers[tid][lvl], sliceHi, args[0], posHi,
-            args[1]);
+        Value cond = genSparseReducedAffineCond(builder, loc,
+                                                coordinatesBuffers[tid][lvl],
+                                                sliceHi, args[0], posHi);
         // continue if not yet break nor out of bound.
         builder.create<scf::ConditionOp>(loc, cond, args);
       },
       /*afterBuilder=*/
-      [c1, genYield, bodyBuilder](OpBuilder &builder, Location loc,
+      [c1, numMetaReduc, bodyBuilder](OpBuilder &builder, Location loc,
                                       ValueRange args) {
         Value iv = args[0];
-        TypeRange types = args.drop_front(2).getTypes();
-        // The coordinate must be in bound as guaranteed by the loop condition.
-        // We generate a fake if operation here only to hide the two extra loop
-        // induction variable maintained by us from user, and relies on later
-        // optimization pass to remove it.
-        Value cont = constantI1(builder, loc, true);
-        auto ifOp = builder.create<scf::IfOp>(loc, types, cont,
+        TypeRange types = args.drop_front(numMetaReduc).getTypes();
+        // The coordinate must be in bound as guaranteed by the loop
+        // condition. We generate a fake if operation here only to hide the
+        // extra loop induction variables maintained by us from users, which
+        // will be removed by later optimization pass.
+        auto ifOp = builder.create<scf::IfOp>(loc, types,
+                                              constantI1(builder, loc, true),
                                               /*withElseBlock=*/!types.empty());
         {
           // 2 reduction variable maintained by us.
-          SmallVector<Value> ifRet = args.drop_front(2);
-          assert(ifRet.size() == args.size() - 2);
+          SmallVector<Value> ifRet = args.drop_front(numMetaReduc);
+          assert(ifRet.size() == args.size() - 1);
 
           OpBuilder::InsertionGuard guard(builder);
           // If coord >= sliceHi.
           if (!ifRet.empty()) {
             builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
             YIELD(ifRet);
           }
 
           // If coord < sliceHi.
           builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
           // Delegates to users' callback.
           bodyBuilder(builder, loc, iv, ifRet);
-          if (genYield) {
-            builder.setInsertionPointToEnd(&ifOp.getThenRegion().front());
-            YIELD(ifRet);
-          }
         }
         // Marks this speical ifOp to avoid sparisification finalizing it.
         ifOp->setAttr(getLoopEmitterLoopAttrName(),
                       StringAttr::get(builder.getContext(), "slice"));
         // Insertion point restored to after ifOp.
         SmallVector<Value> yields;
         // Increase induction variable.
         yields.push_back(ADDI(iv, c1));
-        yields.push_back(cont);
         yields.append(ifOp.getResults().begin(), ifOp.getResults().end());
         YIELD(yields);
       });
 
   builder.setInsertionPointAfter(whileOp);
-  return std::make_pair(whileOp, whileOp.getResults().drop_front(2));
+  return std::make_pair(whileOp, whileOp.getResults().drop_front(numMetaReduc));
 }
 
 // Generates a loop nest that traverse all the unresolved levels in between.
 // TODO: it can only handle all compressed tensors.
 //
 // for(int i = 0; i < slicePos.size(); i+=2) {
 //   loopLo = slicePos[i];
 //   loopHi = slicePos[i + 1];
if (!lvlFullyResolved(tid, firstLvl)) {
                          ValueRange iterArgs) {
               // generate traversal for each level.
               Value loopLo = genIndexLoad(builder, loc, sPtrBuf, iv);
               Value loopHi = genIndexLoad(builder, loc, sPtrBuf, ADDI(iv, c1));
               ValueRange itArgs =
                   genSliceLvlTraverseLoop(
                       builder, loc, loopLo, loopHi, offset,
                       sliceSizes[tid][firstLvl].back(), tid, firstLvl, iterArgs,
-                      false,
                       [&](OpBuilder &builder, Location, Value iv,
                           MutableArrayRef<Value> reduc) {
                         ip = builder.saveInsertionPoint();
                         pos = iv;
                         innerArgs.assign(reduc.begin(), reduc.end());
                       })
                       .second;
               YIELD(itArgs);
void LoopEmitter::genResolvedSliceBegin(OpBuilder &builder, Location loc,
   // FIXME: Technically we should load the coord only when the slice is
   // nonempty. though we assume that even on empty sparse tensors, a non-empty
   // ptr/idx buffer is allocated for each level so it would not cause OOB to
   // avoid generating a ifOp here.
   Value minCrd = genIndexLoad(builder, loc, coordinatesBuffers[tid][0], c0);
 
   // FIXME: We need the relative offset related to the base slice.
   Value absOffset = offsetFromMinCoord(builder, loc, minCrd, size, isNonEmpty);
-  sliceStack[tid].emplace_back(minCrd, absOffset, isNonEmpty, lvl, /*depth=*/1);
+  sliceStack[tid].emplace_back(minCrd, absOffset, isNonEmpty, lvl,
+                               /*depth=*/1);
 }
 
 // Fills in the slicePosBuffer before slice-driven loop begin.
 // TODO: it can only handle all compressed tensors.
 //
 // // Loop generated by `genUnResolvedSliceTreeTraverse`
 // for(int i = 0; i < slicePos.size(); i+=2) {
 //   loopLo = slicePos[i];
ValueRange result = genUnResolvedSliceTreeTraverse(
         Value &minCrd = reduc[1];
         Value &curMemSz = reduc[2];
 
         Value pHi = ADDI(iv, c1);
         Value sPLo = genIndexLoad(builder, loc, positionsBuffers[tid][lvl], iv);
         Value sPHi =
             genIndexLoad(builder, loc, positionsBuffers[tid][lvl], pHi);
 
-        // isNonEmpty = isNonEmpty || lvlNonEmpty, i.e., as long as there is one
-        // non-empty lvl, the slice is non-empty.
+        // isNonEmpty = isNonEmpty || lvlNonEmpty, i.e., as long as there is
+        // one non-empty lvl, the slice is non-empty.
         Value lvlNonEmpty = CMPI(ult, sPLo, sPHi);
         nonEmpty = builder.create<arith::OrIOp>(loc, lvlNonEmpty, nonEmpty);
 
         // Update the minimum coordinate.
         auto ifNonEmpty = builder.create<scf::IfOp>(loc, builder.getIndexType(),
                                                     lvlNonEmpty, true);
         {
           // Generate Code as follows.
bool LoopEmitter::genSliceBegin(OpBuilder &builder, Location loc, TensorId tid,
   if (baseEnc.isSlice())
     llvm_unreachable("TODO: not yet implemented");
 
   // Generate caches required to fast compute next-non-empty slices with
   // increasing offset for slice-base loop.
   // We do not need cache for dense levels.
   if (slicePosBuffer[tid][lvl][0] == nullptr && !isDenseDLT(lvlType)) {
     OpBuilder::InsertionGuard guard(builder);
-    // The buffer can be reused, and the size is loop invariant: it only depends
-    // on the iteration graph's toposort.
+    // The buffer can be reused, and the size is loop invariant: it only
+    // depends on the iteration graph's toposort.
     builder.setInsertionPointAfter(localInsertPos);
     Value bufSize = C_IDX(1);
     Value c2 = C_IDX(2);
     // Accumlates the size required to cache the pLo for the slice.
     // E.g., if we want to cache the pIdx for slice<d0xd1xf64> on the second
     // level. We at most need to a memref<d0xindex>.
     // NOTE: this is apperantly an over-approximation when the previous
     // level is compressed, and we can compute a precise memory size
     // inside the loops. But that would also requires us to allocate/free
     // memorys in loops.
     // TODO: Maybe using allocaScopeOp inside the loop to resolve the issue?
     for (Level curLevel = lvl;
          curLevel >= 1 && !lvlFullyResolved(tid, curLevel - 1); curLevel--) {
       auto depth = remDepOnLevel(tid, curLevel - 1);
       assert(sliceSizes[tid][lvl].size() >= depth);
       Value sz = *(sliceSizes[tid][lvl].rbegin() + depth - 1);
       bufSize = MULI(bufSize, sz);
     }
-    // For a pair of [pLo, pHi]. Note that we can not compress pHi because slice
-    // creates segments in the index buffer so that the pHi for the current
-    // level is no longer the pLo for the next level.
+    // For a pair of [pLo, pHi]. Note that we can not compress pHi because
+    // slice creates segments in the index buffer so that the pHi for the
+    // current level is no longer the pLo for the next level.
     bufSize = MULI(bufSize, c2);
     // Additional two metadata {memSize, idx} at head.
     bufSize = ADDI(bufSize, c2);
     llvm::for_each(
         slicePosBuffer[tid][lvl], [bufSize, loc, &builder](Value &cache) {
           cache = genAlloca(builder, loc, bufSize, builder.getIndexType());
         });
   }
void LoopEmitter::genSliceNextInduction(OpBuilder &builder, Location loc,
   operands.push_back(nextAbsOffset); // we push the absolute offset.
 
   // Update the slice stack.
   info.isNonEmpty = whileOp.getResult(retIdx++);
   info.minCrd = whileOp.getResult(retIdx++);
   info.offset = whileOp.getResult(retIdx++);
 }
 
-Operation *LoopEmitter::emitSliceDrivenLoopOverTensorAtLvl(
-    OpBuilder &builder, Location loc, TensorId tid, Level lvl,
-    MutableArrayRef<Value> reduc) {
-  assert(!depFullyReduced(tid, lvl));
-  SliceInfo &sliceInfo = sliceStack[tid].back();
-  assert(sliceInfo.slicedOnLvl == lvl);
-
-  // The order matters!
-  SmallVector<Value, 3> operands{sliceInfo.isNonEmpty, sliceInfo.minCrd,
-                                 sliceInfo.offset};
-  // number of reduction maintained by us.
-  size_t numMetaReduc = operands.size();
-
-  // Append user-required reduction values.
-  operands.append(reduc.begin(), reduc.end());
-  assert(operands.size() == numMetaReduc + reduc.size());
-
-  // while (slice.nonEmpty()) {
-  //   bodyBuilder();
-  //   SliceNext();
-  // }
-  auto whileOp = builder.create<scf::WhileOp>(
-      loc, ValueRange(operands).getTypes(), operands,
-      /*beforeBuilder=*/
-      [](OpBuilder &builder, Location loc, ValueRange args) {
-        builder.create<scf::ConditionOp>(loc, /*isNonEmpty*/ args[0], args);
-      },
-      /*afterBuilder=*/
-      [this, tid, lvl, reduc, numMetaReduc,
-       &sliceInfo](OpBuilder &builder, Location loc, ValueRange args) {
-        assert(args.size() == reduc.size() + numMetaReduc);
-        sliceInfo.isNonEmpty = args[0];
-        sliceInfo.minCrd = args[1];
-        sliceInfo.offset = args[2];
-        // The slice offset is used to coiterate with other tensors'
-        // coordinates.
-        Value c = sliceInfo.offset;
-        if (sliceInfo.depth > 1) {
-          // Coord is the relative offset related to its parents.
-          // Update c = absOffset[lvl][depth] - absOffset[lvl][depth - 1]
-          llvm_unreachable("TODO: not yet implement");
-        }
-        coords[tid][lvl] = c;
-
-        for (unsigned i = 0, e = reduc.size(); i < e; i++)
-          reduc[i] = args[i + numMetaReduc];
-      });
-
-  // Set the insertion point to while loop body.
-  builder.setInsertionPointToEnd(&whileOp.getAfter().front());
-  return whileOp;
-}
-
 #undef CMPI
 #undef C_IDX
 #undef YIELD
 #undef ADDI
 #undef ANDI
 #undef SUBI
 #undef MULI
 #undef SELECT

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

LogicalResult matchAndRewrite(ForeachOp op,
     for (Level l = 0; l < lvlRank; l++) {
       // TODO: provide utility function for loop sequences that only contains
       // one for loop?
       const SmallVector<TensorLevel, 1> tidLvls{
           loopEmitter.makeTensorLevel(0, l)};
       loopEmitter.enterNewLoopSeq(rewriter, loc, tidLvls);
       // Note that reduc will be taken care of by loop emitter and get updated
       // in place.
-      loopEmitter.enterLoopOverTensorAtLvl(rewriter, loc, tidLvls, reduc);
+      loopEmitter.enterCoIterationOverTensorsAtLvls(rewriter, loc, tidLvls,
+                                                    reduc);
     }
 
     SmallVector<Value> lcvs;
     lcvs.reserve(lvlRank);
     loopEmitter.getLoopIVs(lcvs);
 
     if (op.getOrder()) {
       // FIXME: There is some dim/lvl confusion here since `dimRank != lvlRank`

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

case SparseParallelizationStrategy::kAnyStorageOuterLoop:
     return isOuter;
   case SparseParallelizationStrategy::kDenseAnyLoop:
     return !isSparse;
   case SparseParallelizationStrategy::kAnyStorageAnyLoop:
     return true;
   }
   llvm_unreachable("unexpected parallelization strategy");
 }
 
-/// Generates a for-loop on a single index.
-static Operation *genFor(CodegenEnv &env, OpBuilder &builder, bool isOuter,
-                         bool isInner, LoopId ldx,
+static bool shouldTryParallize(CodegenEnv &env, LoopId ldx, bool isOuter,
                                ArrayRef<TensorLevel> tidLvls) {
   linalg::GenericOp op = env.op();
-  Location loc = op.getLoc();
   auto iteratorTypes = op.getIteratorTypesArray();
   bool isSparse = llvm::any_of(tidLvls, [ldx, &env](TensorLevel tidLvl) {
     // Queries the DLT based on the tensor id and loop idx, as requested by
     // `CodegenEnv::dlt(TensorId, LoopIdx)`. The returned DLT from CodegenEnv
     // should be consistent with the DLT indexed by <TensorId, Level>.
     const auto dlt = env.dlt(env.unpackTensorLevel(tidLvl).first, ldx);
     return isCompressedDLT(dlt) || isSingletonDLT(dlt);
   });
 
-  bool isParallel = isParallelFor(env, isOuter, isSparse);
+  return isParallelFor(env, isOuter, isSparse);
+}
 
+/// Generates a for-loop on a single index.
+static Operation *genFilterLoop(CodegenEnv &env, OpBuilder &builder, LoopId ldx,
+                                ArrayRef<TensorLevel> tidLvls) {
+  linalg::GenericOp op = env.op();
+  Location loc = op.getLoc();
   Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
-    if (env.merger().isFilterLoop(ldx)) {
-      const auto [tid, lvl] = env.unpackTensorLevel(tidLvls.front());
-      // tids/lvls must only have one value because filter loops only
-      // corresponding to the one and only sparse tensor level.
-      assert(isSparse && tidLvls.size() == 1);
-      OpOperand *t = &op->getOpOperand(tid);
-      auto enc = getSparseTensorEncoding(t->get().getType());
-      // Retrieves the affine expression for the filter loop.
-      // FIXME: `toOrigDim` is deprecated.
-      AffineExpr a =
-          op.getMatchingIndexingMap(t).getResult(toOrigDim(enc, lvl));
-      return env.emitter().enterFilterLoopOverTensorAtLvl(builder, loc, tid,
-                                                          lvl, a, reduc);
-    }
-    return env.emitter().enterLoopOverTensorAtLvl(builder, loc, tidLvls, reduc,
-                                                  isParallel);
+    assert(env.merger().isFilterLoop(ldx));
+    AffineExpr a = op.getMatchingIndexingMap(t).getResult(toOrigDim(enc, lvl));
+    return env.emitter().enterFilterLoopOverTensorAtLvl(builder, loc, tid, lvl,
+                                                        a, reduc);
   });
-  assert(loop);
   return loop;
 }
 
 /// Emit a while-loop for co-iteration over multiple indices.
-static Operation *genWhile(CodegenEnv &env, OpBuilder &builder, LoopId idx,
-                           bool needsUniv, ArrayRef<TensorLevel> tidLvls) {
+static Operation *genCoIteration(CodegenEnv &env, OpBuilder &builder,
+                                 LoopId idx, ArrayRef<TensorLevel> tidLvls,
+                                 bool tryParallel, bool needsUniv) {
   Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
     // Construct the while-loop with a parameter for each
     // index.
     return env.emitter().enterCoIterationOverTensorsAtLvls(
-        builder, env.op().getLoc(), tidLvls, needsUniv, reduc);
+        builder, env.op().getLoc(), tidLvls, reduc, tryParallel,
+        /*genDedup=*/true, needsUniv);
   });
   assert(loop);
   return loop;
 }
 
 /// Generates a for-loop or a while-loop, depending on whether it implements
 /// singleton iteration or co-iteration over the given conjunction.
 static Operation *genLoop(CodegenEnv &env, OpBuilder &builder, LoopOrd at,
-                          bool needsUniv, ArrayRef<TensorLevel> tidLvls,
-                          bool isFor) {
-  const LoopId idx = env.topSortAt(at);
-  if (isFor) {
-    bool isOuter = at == 0;
-    bool isInner = at == env.topSortSize() - 1;
-    return genFor(env, builder, isOuter, isInner, idx, tidLvls);
+                          bool needsUniv, ArrayRef<TensorLevel> tidLvls) {
+  const LoopId ldx = env.topSortAt(at);
+  if (env.merger().isFilterLoop(ldx)) {
+    return genFilterLoop(env, builder, ldx, tidLvls);
   }
-  return genWhile(env, builder, idx, needsUniv, tidLvls);
+
+  bool tryParallel = shouldTryParallize(env, ldx, at == 0, tidLvls);
+  return genCoIteration(env, builder, ldx, tidLvls, tryParallel, needsUniv);
 }
 
 /// Generates the induction structure for a while-loop.
 static void finalizeWhileOp(CodegenEnv &env, OpBuilder &builder, LoopId idx,
                             bool needsUniv, scf::WhileOp whileOp) {
   Location loc = env.op().getLoc();
   // Finalize each else branch of all if statements.
   if (env.isReduc() || env.isExpand() || env.getInsertionChain()) {
static void endIf(CodegenEnv &env, OpBuilder &builder, scf::IfOp ifOp,
                   Value insInput) {
   SmallVector<Value> operands;
   if (env.isReduc()) {
     operands.push_back(env.getReduc());
     env.updateReduc(redInput);
     if (env.getValidLexInsert())
       // Any overlapping indices during a reduction creates a valid lex insert.
       operands.push_back(constantI1(builder, env.op().getLoc(), true));
   }
   if (env.isExpand()) {
     operands.push_back(env.getExpandCount());
     env.updateExpandCount(cntInput);
   }
   if (env.getInsertionChain()) {
     operands.push_back(env.getInsertionChain());
     env.updateInsertionChain(insInput);
   }
static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
   // 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<std::pair<TensorLevel, AffineExpr>> affineTidLvls;
   bool isSingleCond = translateBitsToTidLvlPairs(env, li, env.topSortAt(at),
                                                  tidLvls, affineTidLvls);
 
   // Emit the for/while-loop control.
-  Operation *loop = genLoop(env, builder, at, needsUniv, tidLvls, isSingleCond);
+  Operation *loop = genLoop(env, builder, at, needsUniv, tidLvls);
   Location loc = env.op().getLoc();
   for (auto [tidLvl, exp] : affineTidLvls) {
     env.emitter().genDenseAffineAddress(builder, loc, tidLvl, 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.

mlir/test/Dialect/SparseTensor/sorted_coo.mlir

 // CHECK:             %[[VAL_23:.*]] = arith.cmpi ult, %[[VAL_21]], %[[VAL_17]] : index
 // CHECK:             %[[VAL_24:.*]] = arith.cmpi ult, %[[VAL_22]], %[[VAL_19]] : index
 // CHECK:             %[[VAL_25:.*]] = arith.andi %[[VAL_23]], %[[VAL_24]] : i1
 // CHECK:             scf.condition(%[[VAL_25]]) %[[VAL_21]], %[[VAL_22]] : index, index
 // CHECK:           } do {
 // CHECK:           ^bb0(%[[VAL_26:.*]]: index, %[[VAL_27:.*]]: index):
 // CHECK:             %[[VAL_28:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_26]]] : memref<?xindex, strided<[?], offset: ?>>
 // CHECK:             %[[VAL_29:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_27]]] : memref<?xindex, strided<[?], offset: ?>>
-// CHECK:             %[[VAL_30:.*]] = arith.cmpi ult, %[[VAL_29]], %[[VAL_28]] : index
-// CHECK:             %[[VAL_31:.*]] = arith.select %[[VAL_30]], %[[VAL_29]], %[[VAL_28]] : index
 // CHECK:             %[[VAL_32:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_26]]] : memref<?xindex, strided<[?], offset: ?>>
 // CHECK:             %[[VAL_33:.*]] = scf.while (%[[VAL_34:.*]] = %[[VAL_26]]) : (index) -> index {
 // CHECK:               %[[VAL_35:.*]] = arith.cmpi ult, %[[VAL_34]], %[[VAL_17]] : index
 // CHECK:               %[[VAL_36:.*]] = scf.if %[[VAL_35]] -> (i1) {
 // CHECK:                 %[[VAL_37:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_34]]] : memref<?xindex, strided<[?], offset: ?>>
 // CHECK:                 %[[VAL_38:.*]] = arith.cmpi eq, %[[VAL_37]], %[[VAL_32]] : index
 // CHECK:                 scf.yield %[[VAL_38]] : i1
 // CHECK:               } else {

 // CHECK:                 scf.yield %[[VAL_3]] : i1
 // CHECK:               }
 // CHECK:               scf.condition(%[[VAL_49:.*]]) %[[VAL_44]] : index
 // CHECK:             } do {
 // CHECK:             ^bb0(%[[VAL_50:.*]]: index):
 // CHECK:               %[[VAL_51:.*]] = arith.addi %[[VAL_50]], %[[VAL_6]] : index
 // CHECK:               scf.yield %[[VAL_51]] : index
 // CHECK:             }
+// CHECK:             %[[VAL_30:.*]] = arith.cmpi ult, %[[VAL_29]], %[[VAL_28]] : index
+// CHECK:             %[[VAL_31:.*]] = arith.select %[[VAL_30]], %[[VAL_29]], %[[VAL_28]] : index
 // CHECK:             %[[VAL_52:.*]] = arith.cmpi eq, %[[VAL_28]], %[[VAL_31]] : index
 // CHECK:             %[[VAL_53:.*]] = arith.cmpi eq, %[[VAL_29]], %[[VAL_31]] : index
 // CHECK:             %[[VAL_54:.*]] = arith.andi %[[VAL_52]], %[[VAL_53]] : i1
 // CHECK:             scf.if %[[VAL_54]] {
 // CHECK:               %[[VAL_55:.*]]:2 = scf.while (%[[VAL_56:.*]] = %[[VAL_26]], %[[VAL_57:.*]] = %[[VAL_27]]) : (index, index) -> (index, index) {
 // CHECK:                 %[[VAL_58:.*]] = arith.cmpi ult, %[[VAL_56]], %[[VAL_59:.*]] : index
 // CHECK:                 %[[VAL_60:.*]] = arith.cmpi ult, %[[VAL_57]], %[[VAL_61:.*]] : index
 // CHECK:                 %[[VAL_62:.*]] = arith.andi %[[VAL_58]], %[[VAL_60]] : i1