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

[Peiming]

rebase

[Peiming]

rebase

[Peiming]

rebase.

[Peiming]

fix rebase errors..

[Peiming]

rebase.

[Peiming]

fix rebase error.

[Peiming]

rebase.

[Peiming]

rebase.

[Peiming]

code movement.

[Peiming]

rebase.

[Peiming]

add comments,

[Peiming]

code cleanup.

[Peiming]

fix rebase mistakes.

[aartbik]

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

411

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

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

31

"appear in the same order" -> sort?

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

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

1548

why is this 4 to the right suddenly?
lint?

[Peiming]

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

1548

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

[Peiming]

address comments.

[Peiming]

address comments.

[aartbik]

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

411

Oh, the split is fine, just add a comment

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

31

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

[Peiming]

fix comments.

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

public:
   /// operation [Bik96]. For instance: a(i) *= 2.0 or a(i) += a(i) for
   /// sparse vector a.
   bool isSingleCondition(TensorId t, ExprId e) const;
 
   /// Returns true if any `TensorLoopId` in the bitvector corresponds
   /// to sparse level-type.
   bool hasAnySparse(const BitVector &bits) const;
 
+  /// Returns true if bits contains a dependent index reduction condition on
+  /// sparse levels.
+  bool hasSparseIdxReduction(const BitVector &bits) const;
+
   /// Gets the level-type of the `t`th tensor on `i`th loop.
   DimLevelType getDimLevelType(TensorId t, LoopId i) const {
     assert(t < numTensors && i < numLoops);
     return lvlTypes[t][i];
   }
+
+  /// Gets the level-type of the TensorLoopId.
   DimLevelType getDimLevelType(TensorLoopId b) const {
     return getDimLevelType(tensor(b), loop(b));
   }
 
   /// Gets the loop identifier for the `lvl`th level of the `t`th tensor.
   std::optional<LoopId> getLoopId(TensorId t, Level lvl) const {
     assert(t < numTensors && lvl < lvlToLoop[t].size());
     return lvlToLoop[t][lvl];

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

public:
   LoopOrd topSortSize() const { return topSort.size(); }
   LoopId topSortAt(LoopOrd n) const { return topSort.at(n); }
   void topSortPushBack(LoopId i) { topSort.push_back(i); }
   void topSortClear(size_t capacity = 0) {
     topSort.clear();
     topSort.reserve(capacity);
   }
 
-  ArrayRef<LoopId> getTopSort() const { return topSort; };
   ArrayRef<LoopId> getTopSortSlice(LoopOrd n, LoopOrd m) const;
   ArrayRef<LoopId> getLoopStackUpTo(LoopOrd n) const;
   ArrayRef<LoopId> getCurrentLoopStack() const;
   /// Returns the induction-variable for the loop identified by the given
   /// `LoopId`.  This method handles application of the topological sort
   /// in order to convert the `LoopId` into the corresponding `LoopOrd`.
   Value getLoopVar(LoopId i) const;
 

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

 //===----------------------------------------------------------------------===//
 
 /// Returns true if tensor materializes uninitialized into the computation.
 static bool isMaterializing(Value val) {
   return val.getDefiningOp<tensor::EmptyOp>() ||
          val.getDefiningOp<bufferization::AllocTensorOp>();
 }
 
+/// Makes target array's elements sorted according to the `order` array.
+static void sortArrayBasedOnOrder(std::vector<LoopId> &target,
+                                  ArrayRef<LoopId> order) {
+  std::sort(target.begin(), target.end(), [&order](LoopId l, LoopId r) {
+    assert(l != r);
+    int idxL = -1, idxR = -1;
+    for (int i = 0, e = order.size(); i < e; i++) {
+      if (order[i] == l)
+        idxL = i;
+      if (order[i] == r)
+        idxR = i;
+    }
+    assert(idxL >= 0 && idxR >= 0);
+    return idxL < idxR;
+  });
+}
+
 //===----------------------------------------------------------------------===//
 // Code generation environment constructor and general methods
 //===----------------------------------------------------------------------===//
 
 CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
                        unsigned numTensors, unsigned numLoops,
                        unsigned numFilterLoops)
     : linalgOp(linop), sparseOptions(opts),

 }
 
 void CodegenEnv::startEmit() {
   assert(insChain == nullptr && "must only start emitting once");
   if (sparseOut) {
     insChain = sparseOut->get();
     latticeMerger.setHasSparseOut(true);
   }
+
+  // Sort the related loop array such that they are in the same order as they
+  // appears on the topoOrder.
+  // TODO: since we only handle affine addition for slice based codegen, and
+  // addition is assoicative, the order how we evaluate the expression does
+  // not matter. However, to support multiplication, the order of the loop
+  // index should match the evaluation order to the affine expression AST.
+
   // Initialize loop emitter.
-  SmallVector<Value> tensors;
-  for (OpOperand &t : linalgOp->getOpOperands())
+  SmallVector<Value> tensors; // input tensors passed to loop emitter
+  for (OpOperand &t : linalgOp->getOpOperands()) {
     tensors.push_back(t.get());
-  loopEmitter.initialize(tensors,
-                         StringAttr::get(linalgOp.getContext(),
-                                         linalg::GenericOp::getOperationName()),
-                         /*hasOutput=*/true,
-                         /*isSparseOut=*/sparseOut != nullptr, topSort);
+    Level rank = linalgOp.getMatchingIndexingMap(&t).getNumResults();
+    for (Level lvl = 0; lvl < rank; lvl++) {
+      sortArrayBasedOnOrder(
+          latticeMerger.getDependentLoops(t.getOperandNumber(), lvl), topSort);
+    }
+  }
+
+  loopEmitter.initialize(
+      tensors,
+      /*isSparseOut=*/sparseOut != nullptr, topSort,
+      // TODO: compute the map and pass it to loop emitter directly instead of
+      // passing in a callback.
+      [this](TensorId t, Level lvl) -> std::vector<std::pair<TensorId, Level>> {
+        // Translates from a list of loop index to a list of [tid, dim] pair.
+        std::vector<LoopId> rLoops = this->merger().getDependentLoops(t, lvl);
+        std::vector<std::pair<TensorId, Level>> ret;
+        ret.reserve(rLoops.size());
+        for (LoopId l : rLoops)
+          ret.emplace_back(this->merger().getLoopDefiningLvl(l));
+        return ret;
+      });
 }
 
 std::optional<Operation *> CodegenEnv::genLoopBoundary(
     function_ref<std::optional<Operation *>(MutableArrayRef<Value> parameters)>
         callback) {
   SmallVector<Value> params;
   if (isReduc()) {
     params.push_back(redVal);

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

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

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

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

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

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

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

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

BitVector Merger::simplifyCond(LatSetId s0, LatPointId p0) {
   for (const LatPointId p1 : latSets[s0]) {
     if (p0 != p1 && latGT(p0, p1)) {
       isSingleton = false;
       break;
     }
   }
 
   BitVector simple(latPoints[p0].bits);
-  bool reset = isSingleton && hasAnySparse(simple);
+  bool reset =
+      isSingleton && (hasAnySparse(simple) || hasSparseIdxReduction(simple));
   const TensorLoopId be = simple.size();
   TensorLoopId offset = 0; // relative to the end
   if (!reset)
     // Starts resetting from a dense level, so that the first bit (if kept)
     // is not undefined level-type.
     for (TensorLoopId b = 0; b < be; b++) {
       if (simple[b] && isDenseDLT(getDimLevelType(b))) {
         offset = be - b - 1; // relative to the end
         break;
       }
     }
 
   // Now apply the two basic rules. We also iterate the bits reversely to always
   // keep the rightmost bit (which could possibly be a synthetic tensor).
   for (TensorLoopId b = be - 1 - offset, i = 0; i < be;
        b = b == 0 ? be - 1 : b - 1, i++) {
-    if (simple[b]) {
+    // FIXME: better name? also slice on dense level has locate property as
+    // well. Handle it correctly!
+    if (simple[b] && !isLvlWithNonTrivialIdxExp(b)) {
       const auto dlt = getDimLevelType(b);
       if (!isCompressedDLT(dlt) && !isSingletonDLT(dlt)) {
         if (reset)
           simple.reset(b);
         reset = true;
       }
     }
   }
if (bitsi.count() > bitsj.count()) {
     return true;
   }
   return false;
 }
 
 bool Merger::onlyDenseDiff(LatPointId i, LatPointId j) const {
   BitVector tmp(latPoints[j].bits);
   tmp ^= latPoints[i].bits;
-  return !hasAnySparse(tmp);
+  return !hasAnySparse(tmp) && !hasSparseIdxReduction(tmp);
 }
 
 bool Merger::expContainsTensor(ExprId e, TensorId t) const {
   if (tensorExps[e].kind == kTensor)
     return tensorExps[e].tensor == t;
 
   switch (getExpArity(tensorExps[e].kind)) {
   case ExpArity::kNullary:
for (TensorLoopId b = 0, be = bits.size(); b < be; b++)
     if (bits[b]) {
       const auto dlt = getDimLevelType(b);
       if (isCompressedDLT(dlt) || isSingletonDLT(dlt))
         return true;
     }
   return false;
 }
 
+bool Merger::hasSparseIdxReduction(const BitVector &bits) const {
+  // TODO: return false on dense levels.
+  for (unsigned b = 0, be = bits.size(); b < be; b++)
+    if (bits[b] && isLvlWithNonTrivialIdxExp(b))
+      return true;
+  return false;
+}
+
 #ifndef NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Print methods (for debugging).
 //===----------------------------------------------------------------------===//
 
 static const char *kindToOpSymbol(Kind kind) {
   switch (kind) {