fix bugs

Author: Peiming Liu

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

@@ -1103,7 +1103,7 @@ static void genConstantDenseAddressFromLevel(CodegenEnv &env,
     for (Level l = startLvl; l < lvlRank; l++) {
       AffineExpr lvlExpr = lvlExprs[l];
       if (enc.isDenseLvl(l) && isa<AffineConstantExpr>(lvlExpr))
-        env.emitter().genDenseAffineAddress(
+        env.emitter().locateLvlAtAffineAddress(
             builder, loc, env.makeTensorLevel(tid, l), lvlExpr);
       else
         return; // break on first non-dense non-constant level
@@ -1152,7 +1152,7 @@ static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
   Operation *loop = genLoop(env, builder, curr, needsUniv, tidLvls);
   Location loc = env.op().getLoc();
   for (auto [tidLvl, exp] : affineTidLvls) {
-    env.emitter().genDenseAffineAddress(builder, loc, tidLvl, exp);
+    env.emitter().locateLvlAtAffineAddress(builder, loc, tidLvl, exp);
   }
 
   // Until now, we have entered every <tid, lvl> pair in {cond, extra,

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

@@ -603,11 +603,16 @@ Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
   return l;
 }
 
-void LoopEmitter::genDenseAffineAddress(OpBuilder &builder, Location loc,
-                                        TensorLevel tidLvl,
-                                        AffineExpr lvlExpr) {
+void LoopEmitter::locateLvlAtAffineAddress(OpBuilder &builder, Location loc,
+                                           TensorLevel tidLvl,
+                                           AffineExpr lvlExpr) {
   auto [tid, lvl] = unpackTensorLevel(tidLvl);
+
+  const SparseIterator *parent =
+      lvl == 0 ? nullptr : iters[tid][lvl - 1].back().get();
   auto &it = getCurIterator(tid, lvl);
+  it.genInit(builder, loc, parent);
+
   assert(it.kind == IterKind::kTrivial && it.randomAccessible());
   Value lvlCrd = genAffine(builder, loc, lvlExpr);
   it.locate(builder, loc, lvlCrd);
@@ -710,9 +715,7 @@ void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
   // 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;
   ValueRange whileRes = whileOp.getResults();
 
   for (auto [tid, lvl] : unpackTensorLevelRange(loopInfo.tidLvls)) {
@@ -722,7 +725,6 @@ void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
       Value cmp = CMPI(eq, it.getCrd(), iv);
       it.forwardIf(builder, loc, cmp);
       operands.append(it.getItVals().begin(), it.getItVals().end());
-      o += it.getItVals().size();
       // const Value newPos = whileOp->getResult(o++);
       // Following loops continue iteration from the break point of the
       // current while loop.
@@ -738,20 +740,20 @@ void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
   // Reduction value from users.
   for (auto &i : reduc) {
     operands.push_back(i);
-    // In place update reduction variable.
-    i = whileOp->getResult(o++);
+    // Update user reduction variables.
+    i = whileRes.front();
+    whileRes = whileRes.drop_front();
   }
 
   // An (optional) universal index.
-  if (operands.size() + delta < whileOp.getNumResults()) {
-    assert(operands.size() + delta + 1 == whileOp.getNumResults());
+  if (operands.size() < whileOp.getNumResults()) {
+    assert(operands.size() + 1 == whileOp.getNumResults());
     // The last one is the universial index.
     operands.push_back(ADDI(iv, one));
     // update the loop starting point of current loop sequence
-    loopSeqStack.back().first = whileOp->getResult(o++);
+    loopSeqStack.back().first = whileOp->getResults().back();
   }
 
-  assert(o == operands.size() + delta);
   if (!operands.empty())
     YIELD(operands);
 

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

@@ -126,8 +126,8 @@ class LoopEmitter {
 
   /// Emits the address for a dense level based on the value evaluated by the
   /// provided affine expression.
-  void genDenseAffineAddress(OpBuilder &builder, Location loc,
-                             TensorLevel tidLvl, AffineExpr lvlExpr);
+  void locateLvlAtAffineAddress(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.