[mlir][sparse] support dual sparse convolution.

[Peiming]

cleanup

[Peiming]

cleanup

[Peiming]

cleanup.

[Peiming]

cleanup.

[Peiming]

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

1518 ↗

(On Diff #531838)

How come that this does not causing any problem before?

[Peiming]

revert unintended changes.

[Peiming]

add code comments.

[Peiming]

revert unintended change.

[Peiming]

rebase

[Peiming]

revert unintended changes.

[K-Wu]

Hey Peiming, I noticed that you didn't commit this yet. Is it abandoned or are you waiting for other reviewers?

[Peiming]

In D152601#4479358, @K-Wu wrote:

Hey Peiming, I noticed that you didn't commit this yet. Is it abandoned or are you waiting for other reviewers?

Oh, I forget about this, but let me also wait for Aart's review @aartbik

[Peiming]

rebase.

[aartbik]

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

45

I would say, non-continuous

46

tuple -> triple ;-)

1464

Finds, or Prune below

1465

Since we forward (no s) an iterator (a -> an)

1775

Wraps

an extra counting ...

2133

triple

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

1489 ↗

(On Diff #538169)

typo: multiple

1518 ↗

(On Diff #531838)

That is strange yes, do we have no tests for this at all?!

[Peiming]

address comments.

[Peiming]

fix asan.

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

private:
   bool isSparseOutput(TensorId tid) const {
     return isOutputTensor(tid) && isSparseOut;
   }
 
   bool isValidLevel(TensorId tid, Level lvl) const {
     return tid < lvlTypes.size() && lvl < lvlTypes[tid].size();
   }
 
+  void forwardsReducedSliceLevelTreeIt(OpBuilder &builder, Location loc,
+                                       TensorId tid, Level lvl, Value fcnt);
+
   /// 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.

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

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
 // File local shorthand macros
 //===----------------------------------------------------------------------===//
 
 #define CMPI(p, l, r)                                                          \
   (builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::p, (l), (r))       \
        .getResult())
 
 #define C_IDX(v) (constantIndex(builder, loc, (v)))
 #define YIELD(vs) (builder.create<scf::YieldOp>(loc, (vs)))
 #define ADDI(lhs, rhs) (builder.create<arith::AddIOp>(loc, (lhs), (rhs)))
 #define ANDI(lhs, rhs) (builder.create<arith::AndIOp>(loc, (lhs), (rhs)))
 #define SUBI(lhs, rhs) (builder.create<arith::SubIOp>(loc, (lhs), (rhs)))
 #define MULI(lhs, rhs) (builder.create<arith::MulIOp>(loc, (lhs), (rhs)))
 #define SELECT(c, l, r) (builder.create<arith::SelectOp>(loc, (c), (l), (r)))
 
 //===----------------------------------------------------------------------===//
 // File local helper functions.
 //===----------------------------------------------------------------------===//
 
+// For index reduction loops, since the tensor are sliced into uncontinuous
+// fragments, we need a tuple [pLo, pHi, pPtr], in which the pair (pLo, pHi)
+// specifies the range of the fragment, and pPtr specifies the index of the
+// corresponding fragment in the child level (i.e., a pointer to the sliced
+// position array).
+static constexpr unsigned kSliceIterWidth = 3;
+
 static Value genSliceOffset(OpBuilder &builder, Location loc, Value tensor,
                             Level lvl) {
   auto enc = getSparseTensorEncoding(tensor.getType());
   // FIXME: `toOrigDim` is deprecated
   return createOrFoldSliceOffsetOp(builder, loc, tensor, toOrigDim(enc, lvl));
 }
 
 static Value genSliceStride(OpBuilder &builder, Location loc, Value tensor,
static Value genSparseReducedAffineCond(OpBuilder &builder, Location loc,
   //   yield false
   builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
   YIELD(constantI1(builder, loc, false));
 
   builder.setInsertionPointAfter(ifOp);
   return ifOp.getResult(0);
 }
 
+static void dumpIndexMemRef(OpBuilder &builder, Location loc, Value memref) {
+  memref = builder.create<memref::CastOp>(
+      loc, UnrankedMemRefType::get(builder.getIndexType(), 0), memref);
+  createFuncCall(builder, loc, "printMemrefInd", TypeRange{},
+                 ValueRange{memref}, EmitCInterface::On);
+}
+
+static Value loadSlicePosPtr(OpBuilder &builder, Location loc, Value sPosBuf) {
+  return genIndexLoad(builder, loc, sPosBuf, C_IDX(1));
+}
+
+static void updateSlicePosPtr(OpBuilder &builder, Location loc, Value sPosBuf,
+                              Value pPtr) {
+  builder.create<memref::StoreOp>(loc, pPtr, sPosBuf, C_IDX(1));
+}
+
+static Value loadSliceNextPosPtrStart(OpBuilder &builder, Location loc,
+                                      Value sPosBuf, Value pPtr) {
+  return genIndexLoad(builder, loc, sPosBuf, ADDI(pPtr, C_IDX(4)));
+}
+
 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());
void LoopEmitter::exitCurrentLoopSeq(OpBuilder &builder, Location loc) {
   // Depending on whether the slice is resolved or not at current loop sequence,
   // end them in different ways.
   for (auto [tid, lvl, res] : slicedTids) {
     if (!res) {
       // If this is a unresolved-slice-driven loop, pops out the slice.
       assert(sliceStack[tid].back().slicedOnLvl == lvl);
       sliceStack[tid].pop_back();
     } else {
+      continue;
+
       if (!isDenseDLT(lvlTypes[tid][lvl])) {
         // Else this is a resolved-slice, and advance posit similar to TACO.
-        Value c1 = C_IDX(1), c2 = C_IDX(2);
+        Value c1 = C_IDX(1);
         // pIdx += 2, we finished the current lvl, advance the pointer index of
         // the previous level by two to skip the [pLo, pHi] for current level.
         Value sPtrBuf = slicePosBuffer[tid][lvl].back();
         Value curP = genIndexLoad(builder, loc, sPtrBuf, c1);
         // TODO: we could probably use an SSA value for it.
-        Value nexP = ADDI(curP, c2);
+        Value nexP = ADDI(curP, C_IDX(kSliceIterWidth));
         builder.create<memref::StoreOp>(loc, nexP, sPtrBuf, c1);
       }
     }
   }
   loopSeqStack.pop_back();
 }
 
 Value LoopEmitter::genAffine(OpBuilder &builder, Location loc, AffineExpr a) {
void LoopEmitter::enterTensorsAtDenseLvls(
   for (auto [dnTidLvl, denseLoopCond] : dnConds) {
     auto [tid, lvl] = unpackTensorLevel(dnTidLvl);
     assert(isDenseDLT(lvlTypes[tid][lvl]));
 
     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.
+      sliceInfo.emplace_back(tid, lvl, /*fullyReduced=*/!unReduc);
+      if (unReduc) {
         assert(*info.slicedOnLvl == lvl);
+        // Update the slice information as we enter the new loop.
         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 {
if (isAffineIdxCond(denseLoopCond)) {
         }
         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,
+void LoopEmitter::exitForLoop(RewriterBase &builder, Location loc,
                               MutableArrayRef<Value> reduc) {
   const LoopInfo &loopInfo = loopStack.back();
   for (auto [tid, lvl, reduced] : loopInfo.sliceDrivenInfo) {
+    if (!reduced) {
       SliceInfo &info = sliceStack[tid].back();
       assert(isDenseDLT(lvlTypes[tid][lvl]));
-    assert(*info.slicedOnLvl == lvl && !reduced);
+      assert(*info.slicedOnLvl == lvl);
       (void)reduced;
       // Resets slices pointers as the resolved slices are invalidated after we
       // moves forward to the next slice.
-    invalidateSliceIterIdx(rewriter, loc, tid, lvl);
+      invalidateSliceIterIdx(builder, loc, tid, lvl);
       info.minCrd = info.offset = info.isNonEmpty = Value();
       levelReducedDep[tid][lvl]--;
+    } else {
+      forwardsReducedSliceLevelTreeIt(builder, loc, tid, lvl, C_IDX(1));
+    }
   }
   if (auto forOp = llvm::dyn_cast<scf::ForOp>(loopInfo.loop)) {
     if (!reduc.empty()) {
       assert(reduc.size() == forOp.getNumResults());
-      rewriter.create<scf::YieldOp>(loc, reduc);
+      builder.create<scf::YieldOp>(loc, reduc);
     }
     // Exit the loop.
-    rewriter.setInsertionPointAfter(forOp);
+    builder.setInsertionPointAfter(forOp);
     // In-place update reduction variables.
     for (unsigned i = 0, e = forOp.getResults().size(); i < e; i++)
       reduc[i] = forOp.getResult(i);
   } else {
     auto parOp = llvm::cast<scf::ParallelOp>(loopInfo.loop);
     if (!reduc.empty()) {
       assert(reduc.size() == parOp.getInitVals().size() && reduc.size() == 1);
       Operation *redExp = reduc.front().getDefiningOp();
#ifndef NDEBUG
       unsigned numUsers = 0;
       for (Operation *op : redVal.getUsers()) {
         if (op->getParentOp() == parOp)
           numUsers++;
       }
       assert(numUsers == 1);
 #endif // NDEBUG
 
-      rewriter.setInsertionPointAfter(redExp);
-      auto redOp = rewriter.create<scf::ReduceOp>(loc, curVal);
+      builder.setInsertionPointAfter(redExp);
+      auto redOp = builder.create<scf::ReduceOp>(loc, curVal);
       // Attach to the reduction op.
       Block *redBlock = &redOp.getRegion().getBlocks().front();
-      rewriter.setInsertionPointToEnd(redBlock);
-      Operation *newRed = rewriter.clone(*redExp);
+      builder.setInsertionPointToEnd(redBlock);
+      Operation *newRed = builder.clone(*redExp);
       // Replaces arguments of the reduction expression by using the block
       // arguments from scf.reduce.
-      rewriter.updateRootInPlace(
+      builder.updateRootInPlace(
           newRed, [&]() { newRed->setOperands(redBlock->getArguments()); });
       // Erases the out-dated reduction expression.
-      rewriter.eraseOp(redExp);
-      rewriter.setInsertionPointToEnd(redBlock);
-      rewriter.create<scf::ReduceReturnOp>(loc, newRed->getResult(0));
+      builder.eraseOp(redExp);
+      builder.setInsertionPointToEnd(redBlock);
+      builder.create<scf::ReduceReturnOp>(loc, newRed->getResult(0));
     }
-    rewriter.setInsertionPointAfter(parOp);
+    builder.setInsertionPointAfter(parOp);
     // 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.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::forwardsReducedSliceLevelTreeIt(OpBuilder &builder,
+                                                  Location loc, TensorId tid,
+                                                  Level rootLvl, Value fcnt) {
+  auto stt = getSparseTensorType(tensors[tid]);
+
+  // Find a [Lvl, leafLvl) range, and all level in between are fully reduced
+  // level (but not resolved). Since we forwards a iterator at higher level of
+  // the tree, the subtree need to be pruned.
+  Level leafLvl = rootLvl + 1;
+  while (leafLvl < stt.getLvlRank() && !dependentLvlMap[tid][leafLvl].empty()) {
+    assert(depFullyReduced(tid, leafLvl));
+    leafLvl++;
+  }
+
+  Level curLvl = rootLvl + 1;
+  // Prunes all denses subtree.
+  while (curLvl < leafLvl && isDenseDLT(lvlTypes[tid][curLvl])) {
+    fcnt = MULI(sliceSizes[tid][curLvl].back(), fcnt);
+    curLvl++;
+  }
+
+  Value nxPosPtr = nullptr;
+  if (curLvl < leafLvl) {
+    assert(!isDenseDLT(lvlTypes[tid][curLvl]));
+    Value sPosBuf = slicePosBuffer[tid][curLvl].back();
+    Value fPosPtr = MULI(fcnt, C_IDX(kSliceIterWidth));     // forward ptr
+    Value pPosPtr = loadSlicePosPtr(builder, loc, sPosBuf); // previous ptr
+    Value cPosPtr = ADDI(fPosPtr, pPosPtr);                 // current ptr
+    updateSlicePosPtr(builder, loc, sPosBuf, cPosPtr);
+    // dumpIndexMemRef(builder, loc, sPosBuf);
+    // Loads the position pointer start for next level.
+    nxPosPtr = genIndexLoad(builder, loc, sPosBuf, ADDI(cPosPtr, C_IDX(1)));
+    curLvl++;
+  }
+
+  for (; curLvl < leafLvl; curLvl++) {
+    assert(nxPosPtr);
+    if (!isDenseDLT(lvlTypes[tid][curLvl])) {
+      nxPosPtr = MULI(nxPosPtr, C_IDX(kSliceIterWidth));
+      Value sPosBuf = slicePosBuffer[tid][curLvl].back();
+      updateSlicePosPtr(builder, loc, sPosBuf, nxPosPtr);
+      nxPosPtr = genIndexLoad(builder, loc, sPosBuf, ADDI(nxPosPtr, C_IDX(1)));
+    }
+  }
+}
+
 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
for (auto [tid, lvl, resolved] : loopInfo.sliceDrivenInfo) {
     if (!resolved) {
       // TODO: support coiterating multiple slices
       assert(loopInfo.trivialTidLvls.empty() &&
              loopInfo.sliceDrivenInfo.size() == 1);
       genSliceNextInduction(builder, loc, whileOp, tid, lvl, operands, o);
       continue;
     }
 
+    Value forwarded = nullptr;
     if (loopInfo.trivialTidLvls.empty() &&
         loopInfo.sliceDrivenInfo.size() == 1) {
       // Forwards the position iterator.
       operands.push_back(ADDI(posits[tid][lvl], one));
+      forwarded = constantI1(builder, loc, true);
     } 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));
+      forwarded = CMPI(eq, coords[tid][lvl], iv);
+      operands.push_back(SELECT(forwarded, nxPos, pos));
+    }
+    {
+      OpBuilder::InsertionGuard guard(builder);
+      auto ifOp = builder.create<scf::IfOp>(loc, TypeRange{}, forwarded,
+                                            /*else=*/false);
+      builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+      forwardsReducedSliceLevelTreeIt(builder, loc, tid, lvl, one);
     }
-
     // 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++);
   };
 
   for (auto [tid, dstLvl] : unpackTensorLevelRange(loopInfo.trivialTidLvls)) {
scf::WhileOp whileOp = builder.create<scf::WhileOp>(
         YIELD(yields);
       });
 
   builder.setInsertionPointAfter(whileOp);
   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];
 //
 //   // Then the same loop generated by genSliceLvlTraverse above.
 //   while (loopLo < loopHI) {
 //     if (pos[loopLo] < sliceHi) {
ValueRange LoopEmitter::genUnResolvedSliceTreeTraverse(
     std::optional<std::pair<TensorId, Level>> firstResLvl, ValueRange userReduc,
     LoopBodyBuilder bodyBuilder) {
 
   Value c0 = C_IDX(0), c1 = C_IDX(1), c2 = C_IDX(2);
   Value pos = c0;
   OpBuilder::InsertPoint ip;
   SmallVector<Value> innerArgs(userReduc.begin(), userReduc.end());
   scf::ForOp outerMost = nullptr; // the outtermost loop.
+
+  // Wrap body builder and insert a extr counting instruction at the end.
+  auto wrapped = [bodyBuilder](OpBuilder &builder, Location loc, Value iv,
+                               MutableArrayRef<Value> reduc) {
+    bodyBuilder(builder, loc, iv, reduc.drop_back());
+    // Increments the counter.
+    reduc.back() = ADDI(reduc.back(), C_IDX(1));
+  };
+
   if (firstResLvl.has_value()) {
     // Overwrite position when the first level is fully resolved.
     pos = posits[firstResLvl->first][firstResLvl->second];
     ip = builder.saveInsertionPoint();
   } else {
     const SliceInfo &frontSlice = *unResLvls.back();
     Level firstLvl = *frontSlice.slicedOnLvl;
     if (!lvlFullyResolved(tid, firstLvl)) {
       if (isCompressedDLT(lvlTypes[tid][firstLvl])) {
+        // An extra counter that tracks how many segments are there in the child
+        // compressed level.
+        innerArgs.push_back(c0);
+        // Overrides the user-provided builder.
+        bodyBuilder = wrapped;
         unsigned depth = frontSlice.depth - 1;
         Value offset = frontSlice.offset;
         Value sPtrBuf = slicePosBuffer[tid][firstLvl][depth];
         Value mSz = genIndexLoad(builder, loc, sPtrBuf, c0); // memSize
         outerMost = builder.create<scf::ForOp>(
-            loc, c2, mSz, c2, innerArgs,
-            [this, c1, tid, firstLvl, offset, sPtrBuf, &ip, &pos,
+            loc, c2, mSz, C_IDX(kSliceIterWidth), innerArgs,
+            [this, c1, c2, tid, firstLvl, offset, sPtrBuf, &ip, &pos,
              &innerArgs](OpBuilder &builder, Location loc, Value iv,
                          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,
                       [&](OpBuilder &builder, Location, Value iv,
                           MutableArrayRef<Value> reduc) {
                         ip = builder.saveInsertionPoint();
                         pos = iv;
                         innerArgs.assign(reduc.begin(), reduc.end());
                       })
                       .second;
+              // Marks downs the pPtr for next level.
+              builder.create<memref::StoreOp>(loc, itArgs.back(), sPtrBuf,
+                                              ADDI(iv, c2).getResult());
               YIELD(itArgs);
             });
       } else if (isDenseDLT(lvlTypes[tid][firstLvl])) {
         assert(firstLvl == 0); // This must be the first level.
         Value lb = frontSlice.offset;
         Value sliceSz =
             sliceSizes[tid][*frontSlice.slicedOnLvl][frontSlice.depth - 1];
         Value ub = ADDI(lb, sliceSz);

 //     }
 //     loopLo ++;
 //   }
 // }
 void LoopEmitter::genUnResolvedSliceBegin(OpBuilder &builder, Location loc,
                                           TensorId tid, Level lvl) {
   Value c0 = C_IDX(0), c1 = C_IDX(1), c2 = C_IDX(2);
   unsigned depth = levelReducedDep[tid][lvl];
-  Value size = sliceSizes[tid][lvl][depth];
-  // Dense slice begin is trivial
+  Value size = sliceSizes[tid][lvl][depth]; // Dense slice begin is trivial
   if (isDenseDLT(lvlTypes[tid][lvl])) {
     sliceStack[tid].emplace_back(c0, c0, constantI1(builder, loc, false), lvl,
                                  depth + 1);
     return;
   }
 
   assert(isCompressedDLT(lvlTypes[tid][lvl]));
   // Unhandled Cases:
void LoopEmitter::genUnResolvedSliceBegin(OpBuilder &builder, Location loc,
   SmallVector<Value, 3> reduc = {
       constantI1(builder, loc, false), // isNonEmpty
       lvlSizes[tid][lvl],              // minCoord
       c2,                              // memSize
   };
 
   ValueRange result = genUnResolvedSliceTreeTraverse(
       builder, loc, tid, unResSlices, firstResLvl, reduc,
-      [this, c1, c2, tid, lvl, sPtrBuf](OpBuilder &builder, Location loc,
-                                        Value iv,
-                                        MutableArrayRef<Value> reduc) {
+      [this, c1, tid, lvl, sPtrBuf](OpBuilder &builder, Location loc, Value iv,
         Value &nonEmpty = reduc[0];
         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);
ValueRange result = genUnResolvedSliceTreeTraverse(
           YIELD(newMin);
           builder.setInsertionPointToStart(ifNonEmpty.elseBlock());
           YIELD(minCrd);
         }
         minCrd = ifNonEmpty.getResult(0);
         builder.create<memref::StoreOp>(loc, sPLo, sPtrBuf, curMemSz);
         Value nxtMemSize = ADDI(curMemSz, c1);
         builder.create<memref::StoreOp>(loc, sPHi, sPtrBuf, nxtMemSize);
-        // curMemSize += 2
-        curMemSz = ADDI(curMemSz, c2);
+        // curMemSize += kSliceIterWidth
+        curMemSz = ADDI(curMemSz, C_IDX(kSliceIterWidth));
       });
 
   Value isNonEmpty = result[0];
   Value minCrd = result[1];
   // Two metadata [memSize, idx].
   // TODO: Can use an SSA value for these two metadata
   builder.create<memref::StoreOp>(loc, result[2], sPtrBuf, c0);
   builder.create<memref::StoreOp>(loc, c0, sPtrBuf, c1);
   // 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);
 }
 
 bool LoopEmitter::genSliceBegin(OpBuilder &builder, Location loc, TensorId tid,
                                 Level lvl) {
   Value c1 = C_IDX(1), c2 = C_IDX(2);
 
   if (depFullyReduced(tid, lvl)) {
+
     // Do not need to prepare for slice driven loop on dense level after it is
     // fully reduced.
     if (isDenseDLT(lvlTypes[tid][lvl]))
       return true;
     // If constraints on the tensor is fully resolved. We do not need to
     // generates slice begin any more, instead we fall back to TACO-based
     // algorithm to (co)iterates over the slice.
+    // dumpIndexMemRef(builder, loc, slicePosBuffer[tid][lvl].back());
     Value pLoPtr =
-        genIndexLoad(builder, loc, slicePosBuffer[tid][lvl].back(), c1);
+        loadSlicePosPtr(builder, loc, slicePosBuffer[tid][lvl].back());
     pLoPtr = ADDI(pLoPtr, c2);
     Value pHiPtr = ADDI(pLoPtr, c1);
     posits[tid][lvl] =
         genIndexLoad(builder, loc, slicePosBuffer[tid][lvl].back(), pLoPtr);
     highs[tid][lvl] =
         genIndexLoad(builder, loc, slicePosBuffer[tid][lvl].back(), pHiPtr);
     return true;
   }
if (slicePosBuffer[tid][lvl][0] == nullptr && !isDenseDLT(lvlType)) {
     // 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
+    // For a tuple of [pLo, pHi, N]. 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);
+    bufSize = MULI(bufSize, C_IDX(kSliceIterWidth));
     // 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());
         });
   }
 
   if (sliceInfo.isInitialTensor() ||
       (lvl >= 1 && lvlFullyResolved(tid, lvl - 1))) {
     // First level or previous level has been full resolved.
     genResolvedSliceBegin(builder, loc, tid, lvl);
   } else {
     // The previous level has not been full resolved.
     genUnResolvedSliceBegin(builder, loc, tid, lvl);
   }
   return false;
 }
 
 void LoopEmitter::invalidateSliceIterIdx(OpBuilder &builder, Location loc,
                                          TensorId tid, Level lvl) {
   for (unsigned i = 0; i <= lvl; i++) {
     if (!isDenseDLT(lvlTypes[tid][i]) && !dependentLvlMap[tid][i].empty()) {
-      builder.create<memref::StoreOp>(loc, C_IDX(0),
-                                      slicePosBuffer[tid][i].back(), C_IDX(1));
+      updateSlicePosPtr(builder, loc, slicePosBuffer[tid][i].back(), C_IDX(0));
     }
   }
 }
 
 void LoopEmitter::genSliceNextInduction(OpBuilder &builder, Location loc,
                                         const Operation *op, TensorId tid,
                                         Level lvl,
                                         SmallVectorImpl<Value> &operands,
auto ifOp = builder.create<scf::IfOp>(loc, ValueRange(reduc).getTypes(),
     //   return offset += 1
     // }
     builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
     reduc[2] = ADDI(absOffset, c1);
     // Yield offset + 1.
     YIELD(reduc);
 
     // else /*minCrd == offset*/ {
-    //    for (i = 0; i < slicePos.size(); i+=2) {
+    //    for (i = 0; i < slicePos.size(); i+=kSliceIterWidth) {
     //       if (crd[pos[slicePos[i]]] == minCrd) {
     //          slicePos[i]++;
     //       }
     //       minCrd=min(minCrd, crd[pos[slicePos[i]]]);
     //    }
     //    offset = minCrd - size + 1;
     // }
     builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
     reduc[2] = absOffset; // restore value.
     Value pSt = c2;       // pointer starting index
     Value mSz = genIndexLoad(builder, loc, sPtrBuf, c0); // memSize
     reduc[0] = lvlSizes[tid][lvl];                       // next min coord
     reduc[1] = constantI1(builder, loc, false);          // isNonEmpty
     auto loopArgs = static_cast<ValueRange>(reduc).drop_back();
     auto forOp = scf::buildLoopNest(
-        builder, loc, pSt, mSz, c2, loopArgs,
+        builder, loc, pSt, mSz, C_IDX(kSliceIterWidth), loopArgs,
         [this, tid, lvl, c1, sPtrBuf,
          &info](OpBuilder &builder, Location loc, ValueRange ivs,
                 ValueRange iterArgs) -> scf::ValueVector {
           Value curMinCrd = iterArgs[0];
           Value isNonEmpty = iterArgs[1];
 
           Type idxTp = builder.getIndexType();
           Value pLo = genIndexLoad(builder, loc, sPtrBuf, ivs.front());

mlir/test/Integration/Dialect/SparseTensor/CPU/dual_sparse_conv_2d.mlir

 // DEFINE: %{option} = "enable-runtime-library=true enable-index-reduction=true"
 // DEFINE: %{compile} = mlir-opt %s --sparse-compiler=%{option}
 // DEFINE: %{run} = mlir-cpu-runner \
 // DEFINE:  -e entry -entry-point-result=void  \
-// DEFINE:  -shared-libs=%mlir_c_runner_utils | \
+// DEFINE:  -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils | \
 // DEFINE: FileCheck %s
 //
 // RUN: %{compile} | %{run}
 //
 // Do the same run, but now with direct IR generation.
 // REDEFINE: %{option} = "enable-runtime-library=false enable-index-reduction=true"
 // RUN: %{compile} | %{run}
 //
func.func @entry() {
     //
     %v = vector.transfer_read %0[%c0, %c0], %i0
       : tensor<6x6xi32>, vector<6x6xi32>
     vector.print %v : vector<6x6xi32>
 
     // FIXME: DCSR still wrong
     //
     // Should be the same as dense output
-    // C_HECK:    ( ( 0, 0, -1, -6, -1, 6 ),
-    // C_HECK-SAME: ( -1, 0, 1, 0, 1, 0 ),
-    // C_HECK-SAME: ( 0, -1, 1, 0, 0, 0 ),
-    // C_HECK-SAME: ( -1, 0, 0, 0, 0, 0 ),
-    // C_HECK-SAME: ( 0, 0, 3, 6, -3, -6 ),
-    // C_HECK-SAME: ( 2, -1, 3, 0, -3, 0 ) )
-    //
-    // %all_sparse_DCSR = sparse_tensor.convert %2
-    //   : tensor<6x6xi32, #DCSR> to tensor<6x6xi32>
-    // %v2 = vector.transfer_read %all_sparse_DCSR[%c0, %c0], %i0
-    //   : tensor<6x6xi32>, vector<6x6xi32>
-    // vector.print %v2 : vector<6x6xi32>
+    // CHECK:    ( ( 0, 0, -1, -6, -1, 6 ),
+    // CHECK-SAME: ( -1, 0, 1, 0, 1, 0 ),
+    // CHECK-SAME: ( 0, -1, 1, 0, 0, 0 ),
+    // CHECK-SAME: ( -1, 0, 0, 0, 0, 0 ),
+    // CHECK-SAME: ( 0, 0, 3, 6, -3, -6 ),
+    // CHECK-SAME: ( 2, -1, 3, 0, -3, 0 ) )
+    //
+    %all_sparse_DCSR = sparse_tensor.convert %2
+      : tensor<6x6xi32, #DCSR> to tensor<6x6xi32>
+    %v2 = vector.transfer_read %all_sparse_DCSR[%c0, %c0], %i0
+      : tensor<6x6xi32>, vector<6x6xi32>
+    vector.print %v2 : vector<6x6xi32>
 
     //
     // Should be the same as dense output
     // CHECK:    ( ( 0, 0, -1, -6, -1, 6 ),
     // CHECK-SAME: ( -1, 0, 1, 0, 1, 0 ),
     // CHECK-SAME: ( 0, -1, 1, 0, 0, 0 ),
     // CHECK-SAME: ( -1, 0, 0, 0, 0, 0 ),
     // CHECK-SAME: ( 0, 0, 3, 6, -3, -6 ),

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_3d.mlir

 // DEFINE: %{option} = "enable-index-reduction=true enable-runtime-library=true"
 // DEFINE: %{compile} = mlir-opt %s --sparse-compiler=%{option}
 // DEFINE: %{run} = mlir-cpu-runner \
 // DEFINE:  -e entry -entry-point-result=void  \
-// DEFINE:  -shared-libs=%mlir_c_runner_utils | \
+// DEFINE:  -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils | \
 // DEFINE: FileCheck %s
 //
 // RUN: %{compile} | %{run}
 //
 // Do the same run, but now with direct IR generation.
 // REDEFINE: %{option} = "enable-runtime-library=false enable-buffer-initialization=true enable-index-reduction=true"
 // RUN: %{compile} | %{run}
 //