[mlir][sparse] support Parallel for/reduction.

[Peiming]

rebase

[Peiming]

update comments

[Peiming]

update comments..

[Peiming]

add more assertions

[Peiming]

minor fix

[Peiming]

code cleanup

[Peiming]

minor updates.

[Peiming]

cleanup.

[Peiming]

rebase

[wrengr]

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

383

Like I commented on D136185, do you actually need SmallVectorImpl here or can you use MutableArrayRef instead?

473

ditto

477

ditto

[Peiming]

rebase

[Peiming]

rebase

[Peiming]

minor changes..

[Peiming]

rebase

[Peiming]

minor fix

[Peiming]

rebase

[aartbik]

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

452

Can you make this

For sequential for loops:

(the For for reads a bit like a typo)

467

do we only support addition reduction?

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

442

Variables

mlir/test/Dialect/SparseTensor/sparse_parallel.mlir

155

this is good for *this* test, since it looks for structure

but can you also add a more strict CHECK test for all details of scf.reduce?
that all values in and out are right and such, preferably not in this test but a new one

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

8

parallelizaiton

[Peiming]

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

467

No, but we only support "one-line reduction".

[Peiming]

+ address comments
+ add a check test

[Peiming]

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

452

Haha, indeed!

467

I added some comment here to make it more clear.

[Peiming]

rebase.

[aartbik]

some super nits to wrap this up ;-)

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

248

although correct english, do you really need temporally or temporarily?

516

Here and below, follow the "full sentence" style we usually follow

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

458

end with :

471–472

period

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

910–911

"reduction*s*" or "a parallel reduction"

mlir/test/Dialect/SparseTensor/sparse_parallel_reduce.mlir

42 ↗

(On Diff #471245)

yeah, nice!

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

6 ↗

(On Diff #471245)

period at end

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

8

same

[Peiming]

rebase + address comments from Aart.

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

248

You are right ;-(.

[aartbik]

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

267

not for this revision, but the loop emitter methods are becoming a bit big (for example this one, but also the exit loop below!)

perhaps we can later refactor the public methods into small methods that call a number of private methods with a very specific function, e.g. prepare induction, do reduction etc.

WDYT?

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

443

The name genReducUpdateStmt is a bit misnomer, right?
You basically do all the logic needed for starting/ending a loop, including handling reductions?
Can we come up with something better?

But moving the two similar parts into one method is good refactoring!

911

testing for expCount is very hard to read in these expression

can we make this something like

isExpanion = codegenExpCount

and then use it in the expressions?

[Peiming]

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

267

yeah, that's a good idea!

[Peiming]

address comments from Aart.

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

911

I instead followed the way how you did with codegen.SparseOut, which seems even more clear!

[aartbik]

Nice work, Peiming!

This finishes your starter ;-)

[NathanielMcVicar]

This change appears to have broken the Windows buildbot. https://lab.llvm.org/buildbot/#/builders/13/builds/28002/steps/6/logs/stdio

C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\Transforms\CodegenUtils.cpp(227): error C2220: the following warning is treated as an error
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\Transforms\CodegenUtils.cpp(227): warning C4804: '<=': unsafe use of type 'bool' in operation

[Peiming]

In D135927#3912455, @NathanielMcVicar wrote:

This change appears to have broken the Windows buildbot. https://lab.llvm.org/buildbot/#/builders/13/builds/28002/steps/6/logs/stdio

C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\Transforms\CodegenUtils.cpp(227): error C2220: the following warning is treated as an error
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\Transforms\CodegenUtils.cpp(227): warning C4804: '<=': unsafe use of type 'bool' in operation

Thanks, I will take a look!

[Peiming]

In D135927#3912455, @NathanielMcVicar wrote:

This change appears to have broken the Windows buildbot. https://lab.llvm.org/buildbot/#/builders/13/builds/28002/steps/6/logs/stdio

C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\Transforms\CodegenUtils.cpp(227): error C2220: the following warning is treated as an error
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\Transforms\CodegenUtils.cpp(227): warning C4804: '<=': unsafe use of type 'bool' in operation

Please see https://reviews.llvm.org/D137560

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

Operation *enterLoopOverTensorAtDim(OpBuilder &builder, Location loc,
                                       ArrayRef<size_t> extraDims = {});
 
   /// Emits a co-iteration loop over a set of tensors.
   Operation *enterCoIterationOverTensorsAtDims(
       OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
       ArrayRef<size_t> dims, bool needsUniv, MutableArrayRef<Value> reduc = {},
       ArrayRef<size_t> extraTids = {}, ArrayRef<size_t> extraDims = {});
 
-  SmallVector<Value, 2> exitCurrentLoop(OpBuilder &builder, Location loc,
-                                        ArrayRef<Value> reduc = {});
+  void exitCurrentLoop(RewriterBase &rewriter, Location loc,
+                       MutableArrayRef<Value> reduc = {});
 
   /// Returns the array of coordinate for all the loop generated till now.
   void getCoordinateArray(SmallVectorImpl<Value> &coords) const {
     for (auto &l : loopStack)
       coords.push_back(l.iv);
   }
 
   /// Gets loop induction variable at the given level.
private:
   /// Emits extra locals, since the locals might not be in simplified lattices
   /// point used to generate the loops, but are still required to generates
   /// expressions.
   void emitExtraLocalsForTensorsAtDenseDims(OpBuilder &builder, Location loc,
                                             ArrayRef<size_t> tids,
                                             ArrayRef<size_t> dims);
 
   /// Exits a for loop, returns the reduction results, e.g.,
+  /// For for loops:
   /// %ret = for () {
   ///   ...
+  ///   %val = addi %args, %c
   ///   yield %val
   /// }
+  /// For parallel loops, the following generated code by users
+  /// %ret = parallel () init(%args) {
+  ///   ...
+  ///   %val = addi %args, %c
+  /// }
+  /// will be transformed into
+  /// %ret = parallel () init(%args) {
+  ///   ...
+  ///   scf.reduce(%c) bb0(%0, %1){
+  ///     %val = addi %0, %1
+  ///     scf.reduce.return %val
+  ///   }
+  /// }
   /// Return %ret to user, while %val is provided by users (`reduc`)
-  SmallVector<Value, 2> exitForLoop(OpBuilder &builder, Location loc,
-                                    ArrayRef<Value> reduc);
+  void exitForLoop(RewriterBase &rewriter, Location loc,
+                   MutableArrayRef<Value> reduc);
 
   /// Exits a while loop, returns the reduction results.
-  SmallVector<Value, 2> exitCoiterationLoop(OpBuilder &builder, Location loc,
-                                            ArrayRef<Value> reduc);
+  void exitCoiterationLoop(OpBuilder &builder, Location loc,
+                           MutableArrayRef<Value> reduc);
 
   // Is the last tensor output tensor
   bool isLastOutput;
   /// Input and (optional) output tensors.
   std::vector<Value> tensors;
   /// The dim type array for each tensor.
   std::vector<std::vector<DimLevelType>> dimTypes;
   /// Sparse iteration information (by tensor and dim). These arrays

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

void SparseTensorLoopEmitter::enterNewLoopSeq(OpBuilder &builder, Location loc,
   for (auto [tid, dim] : llvm::zip(tids, dims))
     prepareLoopOverTensorAtDim(builder, loc, tid, dim);
 }
 
 Operation *SparseTensorLoopEmitter::enterLoopOverTensorAtDim(
     OpBuilder &builder, Location loc, size_t tid, size_t dim,
     MutableArrayRef<Value> reduc, bool isParallel, ArrayRef<size_t> extraTids,
     ArrayRef<size_t> extraDims) {
+
   assert(dimTypes[tid].size() > dim);
   // We can not re-enter the same level.
   assert(!coord[tid][dim]);
+  // TODO: support multiple return on parallel for?
+  assert(!isParallel || reduc.empty() <= 1);
 
   Value step = constantIndex(builder, loc, 1);
   auto dimType = dimTypes[tid][dim];
   bool isSparseInput = isCompressedDLT(dimType) || isSingletonDLT(dimType);
   assert(isDenseDLT(dimType) || isCompressedDLT(dimType) ||
          isSingletonDLT(dimType));
 
   Value lo = isSparseInput ? pidxs[tid][dim]      // current offset
                            : loopSeqStack.back(); // univeral tid
   Value hi = highs[tid][dim];
+  Operation *loop = nullptr;
+  Value iv;
+  if (isParallel) {
+    scf::ParallelOp parOp =
+        builder.create<scf::ParallelOp>(loc, lo, hi, step, reduc);
+    builder.setInsertionPointToStart(parOp.getBody());
+    assert(parOp.getNumReductions() == reduc.size());
+    iv = parOp.getInductionVars()[0];
 
+    // In-place update on the reduction variable vector.
+    // Note that the init vals is not the actual reduction variables but instead
+    // used as a `special handle` to (temporally) represent them. The
+    // expression on init vals will be moved into scf.reduce and replaced with
+    // the block arguments when exiting the loop (see exitForLoop). This is
+    // needed as we can not build the actual reduction block and get the actual
+    // reduction varaible before users fill parallel loop body.
+    for (int i = 0, e = reduc.size(); i < e; i++)
+      reduc[i] = parOp.getInitVals()[i];
+    loop = parOp;
+  } else {
     scf::ForOp forOp = builder.create<scf::ForOp>(loc, lo, hi, step, reduc);
     builder.setInsertionPointToStart(forOp.getBody());
-  Value iv = forOp.getInductionVar();
-  assert(iv);
+    iv = forOp.getInductionVar();
+
+    // In-place update on the reduction variable vector.
+    assert(forOp.getNumRegionIterArgs() == reduc.size());
+    for (int i = 0, e = reduc.size(); i < e; i++)
+      reduc[i] = forOp.getRegionIterArg(i);
+    loop = forOp;
+  }
+  assert(loop && iv);
+
   if (isSparseInput) {
     pidxs[tid][dim] = iv;
     // Generating a load on the indices array yields the coordinate.
     Value ptr = idxBuffer[tid][dim];
     coord[tid][dim] = genIndexLoad(builder, loc, ptr, iv);
   } else {
     // Dense tensor, the coordinates is the inducation variable.
     coord[tid][dim] = iv;
     // generate pidx for dense dim (pidx = i * sz + j)
     auto enc = getSparseTensorEncoding(tensors[tid].getType());
     if (enc)
       pidxs[tid][dim] = genAddress(builder, loc, tid, dim, iv);
   }
 
   // NOTE: we can also prepares for next dim here in advance
   // Push the loop into stack
-  loopStack.emplace_back(ArrayRef<size_t>(tid), ArrayRef<size_t>(dim), forOp,
+  loopStack.emplace_back(ArrayRef<size_t>(tid), ArrayRef<size_t>(dim), loop,
                          coord[tid][dim]);
   // Emit extra locals.
   emitExtraLocalsForTensorsAtDenseDims(builder, loc, extraTids, extraDims);
 
-  // In-place update on the reduction variable vector.
-  assert(forOp.getNumRegionIterArgs() == reduc.size());
-  for (int i = 0, e = reduc.size(); i < e; i++)
-    reduc[i] = forOp.getRegionIterArg(i);
-  return forOp;
+  return loop;
 }
 
 Operation *SparseTensorLoopEmitter::enterCoIterationOverTensorsAtDims(
     OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
     ArrayRef<size_t> dims, bool needsUniv, MutableArrayRef<Value> reduc,
     ArrayRef<size_t> extraTids, ArrayRef<size_t> extraDims) {
   assert(tids.size() == dims.size());
   SmallVector<Type, 4> types;
if (enc) {
         continue;
       }
       pidxs[tid][dim] = genAddress(builder, loc, tid, dim, loopStack.back().iv);
       // NOTE: we can also prepares for next dim here in advance
     }
   }
 }
 
-SmallVector<Value, 2>
-SparseTensorLoopEmitter::exitForLoop(OpBuilder &builder, Location loc,
-                                     ArrayRef<Value> reduc) {
+void SparseTensorLoopEmitter::exitForLoop(RewriterBase &rewriter, Location loc,
+                                          MutableArrayRef<Value> reduc) {
   LoopLevelInfo &loopInfo = loopStack.back();
   auto &dims = loopStack.back().dims;
   auto &tids = loopStack.back().tids;
-  auto forOp = llvm::cast<scf::ForOp>(loopInfo.loop);
+  auto forOp = llvm::dyn_cast<scf::ForOp>(loopInfo.loop);
+  if (forOp) {
     if (!reduc.empty()) {
       assert(reduc.size() == forOp.getNumResults());
-    builder.setInsertionPointToEnd(forOp.getBody());
-    builder.create<scf::YieldOp>(loc, reduc);
+      rewriter.setInsertionPointToEnd(forOp.getBody());
+      rewriter.create<scf::YieldOp>(loc, reduc);
+    }
+    // Exit the loop.
+    rewriter.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();
+      // reduction expression should have no use.
+      assert(redExp->getUses().empty());
+      // this must be a binary operation.
+      // This is users' responsibilty to ensure the operation are commutative.
+      assert(redExp->getNumOperands() == 2 && redExp->getNumResults() == 1);
+
+      Value redVal = parOp.getInitVals().front();
+      Value curVal;
+      if (redExp->getOperand(0) == redVal)
+        curVal = redExp->getOperand(1);
+      else if (redExp->getOperand(1) == redVal)
+        curVal = redExp->getOperand(0);
+      // one of the operands must be the init value (which is also the
+      // previous reduction value).
+      assert(curVal);
+      // the reduction expression should be the only user of the reduction val
+      // inside the parallel for.
+      unsigned numUsers = 0;
+      for (Operation *op : redVal.getUsers()) {
+        if (op->getParentOp() == parOp)
+          numUsers++;
+      }
+      assert(numUsers == 1);
+      (void)numUsers; // to silence unused variable warning in release build
+
+      rewriter.setInsertionPointAfter(redExp);
+      auto redOp = rewriter.create<scf::ReduceOp>(loc, curVal);
+      // Attach to the reduction op.
+      Block *redBlock = &redOp.getRegion().getBlocks().front();
+      rewriter.setInsertionPointToEnd(redBlock);
+      Operation *newRed = rewriter.clone(*redExp);
+      // replaced arguments
+      rewriter.updateRootInPlace(
+          newRed, [&]() { newRed->setOperands(redBlock->getArguments()); });
+      // erase the previous reduction expression.
+      rewriter.eraseOp(redExp);
+      rewriter.setInsertionPointToEnd(redBlock);
+      rewriter.create<scf::ReduceReturnOp>(loc, newRed->getResult(0));
+    }
+    rewriter.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, dim] : llvm::zip(tids, dims)) {
     // Reset to null.
     coord[tid][dim] = Value();
     pidxs[tid][dim] = Value();
     // Dense dimension, high is fixed.
     if (!isDenseDLT(dimTypes[tid][dim]))
       highs[tid][dim] = Value();
   }
-  // exit the loop
-  builder.setInsertionPointAfter(forOp);
-  return forOp.getResults();
 }
 
-SmallVector<Value, 2>
-SparseTensorLoopEmitter::exitCoiterationLoop(OpBuilder &builder, Location loc,
-                                             ArrayRef<Value> reduc) {
+void SparseTensorLoopEmitter::exitCoiterationLoop(
+    OpBuilder &builder, Location loc, MutableArrayRef<Value> reduc) {
   auto whileOp = llvm::cast<scf::WhileOp>(loopStack.back().loop);
   auto &dims = loopStack.back().dims;
   auto &tids = loopStack.back().tids;
   Value iv = loopStack.back().iv;
   // Generation while loop induction at the end.
   builder.setInsertionPointToEnd(&whileOp.getAfter().front());
   // Finalize the induction. Note that the induction could be performed
   // in the individual if-branches to avoid re-evaluating the conditions.
if (isCompressedDLT(dimTypes[tid][dim]) ||
       pidxs[tid][dim] = whileOp->getResult(o++);
       // The coordinates are invalid now.
       coord[tid][dim] = nullptr;
       // highs remains unchanged.
     }
   }
 
   // Reduction value from users.
-  SmallVector<Value, 2> ret;
-  for (auto red : reduc) {
-    operands.push_back(red);
-    ret.push_back(whileOp->getResult(o++));
+  for (unsigned i = 0, e = reduc.size(); i < e; i++) {
+    operands.push_back(reduc[i]);
+    // In place update reduction variable.
+    reduc[i] = whileOp->getResult(o++);
   }
 
   // An (optional) universal index.
   if (operands.size() < whileOp.getNumResults()) {
     assert(operands.size() + 1 == whileOp.getNumResults());
     // The last one is the universial index.
     operands.push_back(builder.create<arith::AddIOp>(loc, iv, one));
     // update the loop starting point of current loop sequence
     loopSeqStack.back() = whileOp->getResult(o++);
   }
 
   assert(o == operands.size());
   builder.create<scf::YieldOp>(loc, operands);
   builder.setInsertionPointAfter(whileOp);
-  return ret;
 }
 
-SmallVector<Value, 2>
-SparseTensorLoopEmitter::exitCurrentLoop(OpBuilder &builder, Location loc,
-                                         ArrayRef<Value> reduc) {
+void SparseTensorLoopEmitter::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).
   LoopLevelInfo &loopInfo = loopStack.back();
   assert(loopInfo.tids.size() == loopInfo.dims.size());
   SmallVector<Value, 2> red;
   if (llvm::isa<scf::WhileOp>(loopInfo.loop)) {
-    red = exitCoiterationLoop(builder, loc, reduc);
+    exitCoiterationLoop(rewriter, loc, reduc);
   } else {
-    red = exitForLoop(builder, loc, reduc);
+    exitForLoop(rewriter, loc, reduc);
   }
 
   assert(loopStack.size() == loopSeqStack.size());
   loopStack.pop_back();
   return red;
 }
 
 //===----------------------------------------------------------------------===//

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

 static Value getCustomRedId(Operation *op) {
   return dyn_cast<sparse_tensor::ReduceOp>(op).getIdentity();
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (statements and expressions).
 //===----------------------------------------------------------------------===//
 
+/// Generates stmts that updates reduction varaibles.
+static void
+genReducUpdateStmt(CodeGen &codegen, Merger &merger,
+                   function_ref<void(SmallVectorImpl<Value> &reduc)> callback) {
+  SmallVector<Value, 4> reduc;
+  if (codegen.redVal)
+    reduc.push_back(codegen.redVal);
+  if (codegen.expValues)
+    reduc.push_back(codegen.expCount);
+
+  callback(reduc);
+
+  // Callback should do in-place update on reduction value vector.
+  if (codegen.redVal)
+    updateReduc(merger, codegen, reduc.front());
+  if (codegen.expValues)
+    codegen.expCount = reduc.back();
+}
+
 /// Local bufferization of all dense and sparse data structures.
 static void genBuffers(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                        linalg::GenericOp op) {
   Location loc = op.getLoc();
   assert(op.getNumOperands() == op.getNumInputs() + 1);
 
   codegen.loopEmitter.initializeLoopEmit(
       builder, loc,
if (atStart) {
     codegen.expValues = codegen.expFilled = codegen.expAdded =
         codegen.expCount = Value();
   }
 }
 
 /// Returns parallelization strategy. Any implicit loop in the Linalg
 /// operation that is marked "parallel" is a candidate. Whether it is actually
 /// converted to a parallel operation depends on the requested strategy.
-static bool isParallelFor(CodeGen &codegen, bool isOuter, bool isReduction,
-                          bool isSparse) {
+static bool isParallelFor(CodeGen &codegen, bool isOuter, bool isSparse) {
   // Reject parallelization of sparse output.
   if (codegen.sparseOut)
     return false;
   // Inspect strategy.
   switch (codegen.options.parallelizationStrategy) {
   case SparseParallelizationStrategy::kNone:
     return false;
   case SparseParallelizationStrategy::kDenseOuterLoop:
-    return isOuter && !isSparse && !isReduction;
+    // We do not support parallel reduction on tensor expansion.
+    return isOuter && !isSparse && !codegen.expCount;
   case SparseParallelizationStrategy::kAnyStorageOuterLoop:
-    return isOuter && !isReduction;
+    return isOuter && !codegen.expCount;
   case SparseParallelizationStrategy::kDenseAnyLoop:
-    return !isSparse && !isReduction;
+    return !isSparse && !codegen.expCount;
   case SparseParallelizationStrategy::kAnyStorageAnyLoop:
-    return !isReduction;
+    return !codegen.expCount;
   }
   llvm_unreachable("unexpected parallelization strategy");
 }
 
 /// Generates a for-loop on a single index.
 static Operation *genFor(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, bool isOuter, bool isInner,
                          unsigned idx, size_t tid, size_t dim,
                          ArrayRef<size_t> extraTids,
                          ArrayRef<size_t> extraDims) {
   Location loc = op.getLoc();
   auto iteratorTypes = op.getIteratorTypesArray();
-  bool isReduction = linalg::isReductionIterator(iteratorTypes[idx]);
   bool isSparse = isCompressedDLT(merger.getDimLevelType(tid, idx)) ||
                   isSingletonDLT(merger.getDimLevelType(tid, idx));
-  bool isParallel = isParallelFor(codegen, isOuter, isReduction, isSparse);
-  assert(!isParallel);
-
-  // Emit a sequential or vector loop.
-  SmallVector<Value, 4> operands;
-  if (codegen.redVal)
-    operands.push_back(codegen.redVal);
-  if (codegen.expValues)
-    operands.push_back(codegen.expCount);
-
-  Operation *loop = codegen.loopEmitter.enterLoopOverTensorAtDim(
-      builder, loc, tid, dim, operands, isParallel, extraTids, extraDims);
-
-  // The operands should be updated by loop emitter already.
-  if (codegen.redVal)
-    updateReduc(merger, codegen, operands.front());
-  if (codegen.expValues)
-    codegen.expCount = operands.back();
+  bool isParallel = isParallelFor(codegen, isOuter, isSparse);
 
+  Operation *loop = nullptr;
+  genReducUpdateStmt(codegen, merger, [&](SmallVectorImpl<Value> &reduc) {
+    loop = codegen.loopEmitter.enterLoopOverTensorAtDim(
+        builder, loc, tid, dim, reduc, isParallel, extraTids, extraDims);
+  });
+  assert(loop);
   return loop;
 }
 
 /// Emit a while-loop for co-iteration over multiple indices.
 static Operation *genWhile(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                            linalg::GenericOp op, unsigned idx, bool needsUniv,
                            ArrayRef<size_t> condTids, ArrayRef<size_t> condDims,
                            ArrayRef<size_t> extraTids,
                            ArrayRef<size_t> extraDims) {
 
-  SmallVector<Value, 4> operands;
-
+  Operation *loop = nullptr;
+  genReducUpdateStmt(codegen, merger, [&](SmallVectorImpl<Value> &reduc) {
     // Construct the while-loop with a parameter for each index.
-  if (codegen.redVal)
-    operands.push_back(codegen.redVal);
-  if (codegen.expValues)
-    operands.push_back(codegen.expCount);
-
-  Operation *loop = codegen.loopEmitter.enterCoIterationOverTensorsAtDims(
-      builder, op.getLoc(), condTids, condDims, needsUniv, operands, extraTids,
+    loop = codegen.loopEmitter.enterCoIterationOverTensorsAtDims(
+        builder, op.getLoc(), condTids, condDims, needsUniv, reduc, extraTids,
         extraDims);
-
-  if (codegen.redVal)
-    updateReduc(merger, codegen, operands.front());
-  if (codegen.expValues)
-    codegen.expCount = operands.back();
-
+  });
+  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(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                           linalg::GenericOp op, unsigned at, bool needsUniv,
                           ArrayRef<size_t> condTids, ArrayRef<size_t> condDims,
Operation *loop = genLoop(merger, codegen, builder, op, at, needsUniv,
                             condTids, condDims, extraTids, extraDims);
   // Emit the locals for this loop.
   // genLocals(merger, codegen, builder, op, at, needsUniv,
   // merger.lat(li).bits);
   return loop;
 }
 
 /// Ends a single loop in current sequence. Returns new values for needsUniv.
-static bool endLoop(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+static bool endLoop(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
                     linalg::GenericOp op, Operation *loop, unsigned idx,
                     unsigned li, bool needsUniv) {
   // End a while-loop.
   if (auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
-    finalizeWhileOp(merger, codegen, builder, op, idx, needsUniv,
+    finalizeWhileOp(merger, codegen, rewriter, op, idx, needsUniv,
                     merger.lat(li).bits, whileOp);
   } else {
     needsUniv = false;
   }
 
-  SmallVector<Value, 2> reduc;
-  if (codegen.redVal)
-    reduc.push_back(codegen.redVal);
-  if (codegen.expValues)
-    reduc.push_back(codegen.expCount);
-
-  auto loopRet =
-      codegen.loopEmitter.exitCurrentLoop(builder, op.getLoc(), reduc);
-  assert(reduc.size() == loopRet.size());
-
-  if (codegen.redVal)
-    updateReduc(merger, codegen, loopRet.front());
-  if (codegen.expValues)
-    codegen.expCount = loopRet.back();
+  genReducUpdateStmt(codegen, merger, [&](SmallVectorImpl<Value> &reduc) {
+    codegen.loopEmitter.exitCurrentLoop(rewriter, op.getLoc(), reduc);
+  });
 
   return needsUniv;
 }
 
 /// Ends a loop sequence at given level.
 static void endLoopSeq(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                        linalg::GenericOp op, unsigned exp, unsigned at,
                        unsigned idx, unsigned ldx) {

mlir/test/Dialect/SparseTensor/sparse_parallel.mlir

 // RUN: mlir-opt %s -sparsification="parallelization-strategy=none" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR0
 // FIXME: we do not support vectorization/parallel loops in loop emitter right now
-// R_U_N: mlir-opt %s -sparsification="parallelization-strategy=dense-outer-loop" | \
-// R_U_N:   FileCheck %s --check-prefix=CHECK-PAR1
-// R_U_N: mlir-opt %s -sparsification="parallelization-strategy=any-storage-outer-loop" | \
-// R_U_N:   FileCheck %s --check-prefix=CHECK-PAR2
-// R_U_N: mlir-opt %s -sparsification="parallelization-strategy=dense-any-loop" | \
-// R_U_N:   FileCheck %s --check-prefix=CHECK-PAR3
-// R_U_N: mlir-opt %s -sparsification="parallelization-strategy=any-storage-any-loop" | \
-// R_U_N:   FileCheck %s --check-prefix=CHECK-PAR4
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=dense-outer-loop" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-PAR1
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=any-storage-outer-loop" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-PAR2
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=dense-any-loop" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-PAR3
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=any-storage-any-loop" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-PAR4
 
 #DenseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "dense" ]
 }>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ]
 }>

 //
 // CHECK-PAR3-LABEL: func @matvec
 // CHECK-PAR3:         scf.parallel
 // CHECK-PAR3:           scf.for
 // CHECK-PAR3:         return
 //
 // CHECK-PAR4-LABEL: func @matvec
 // CHECK-PAR4:         scf.parallel
-// CHECK-PAR4:           scf.for
+// CHECK-PAR4:           scf.parallel
+// CHECK-PAR4:             scf.reduce
 // CHECK-PAR4:         return
 //
 func.func @matvec(%arga: tensor<16x32xf32, #CSR>,
              %argb: tensor<32xf32>,
        %argx: tensor<16xf32>) -> tensor<16xf32> {
   %0 = linalg.generic #trait_matvec
       ins(%arga, %argb : tensor<16x32xf32, #CSR>, tensor<32xf32>)
      outs(%argx: tensor<16xf32>) {
     ^bb(%A: f32, %b: f32, %x: f32):
       %0 = arith.mulf %A, %b : f32
       %1 = arith.addf %0, %x : f32
       linalg.yield %1 : f32
   } -> tensor<16xf32>
   return %0 : tensor<16xf32>
 }

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

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// Do the same run, but now with parallization
+//
+// RUN: mlir-opt %s \
+// RUN:   --sparse-compiler="parallelization-strategy=any-storage-any-loop" | \
+// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ],
   pointerBitWidth = 8,
   indexBitWidth = 8
 }>