[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

384

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

462

ditto

483

ditto

[Peiming]

rebase

[Peiming]

rebase

[Peiming]

minor changes..

[Peiming]

rebase

[Peiming]

minor fix

[Peiming]

rebase

[aartbik]

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

455

Can you make this

For sequential for loops:

(the For for reads a bit like a typo)

470

do we only support addition reduction?

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

413

Variables

mlir/test/Dialect/SparseTensor/sparse_parallel.mlir

154

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

470

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

[Peiming]

+ address comments
+ add a check test

[Peiming]

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

455

Haha, indeed!

470

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

249

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

460–461

period

461

end with :

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

912

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

mlir/test/Dialect/SparseTensor/sparse_parallel_reduce.mlir

43

yeah, nice!

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

6

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

249

You are right ;-(.

[aartbik]

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

268

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

414

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!

913

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

268

yeah, that's a good idea!

[Peiming]

address comments from Aart.

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

913

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 sequential for loops:
   /// %ret = for () {
   ///   ...
+  ///   %val = addi %args, %c
   ///   yield %val
   /// }
-  /// Return %ret to user, while %val is provided by users (`reduc`)
-  SmallVector<Value, 2> exitForLoop(OpBuilder &builder, Location loc,
-                                    ArrayRef<Value> reduc);
+  /// For parallel loops, the following generated code by users:
+  /// %ret = parallel () init(%args) {
+  ///   ...
+  ///   %val = op %args, %c
+  /// }
+  /// will be transformed into
+  /// %ret = parallel () init(%args) {
+  ///   ...
+  ///   scf.reduce(%c) bb0(%0, %1){
+  ///     %val = op %0, %1
+  ///     scf.reduce.return %val
+  ///   }
+  /// }
+  /// 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.
-  SmallVector<Value, 2> exitCoiterationLoop(OpBuilder &builder, Location loc,
-                                            ArrayRef<Value> reduc);
+  void exitCoIterationLoop(OpBuilder &builder, Location loc,
+                           MutableArrayRef<Value> reduc);
 
   // Whether the loop emitter needs to treat the last tensor as the output
   // tensor.
   bool hasOutput;
   bool isSparseOut;
   /// Input and (optional) output tensors.
   std::vector<Value> tensors;
   /// The dim type array for each tensor.

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 (temporarily) 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 && !isSparseOutput(tid))
       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 && !isSparseOutput(tid)) {
         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.
+      // NOTE: 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);
+      // Replaces arguments of the reduction expression by using the block
+      // arguments from scf.reduce.
+      rewriter.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));
+    }
+    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;
 }
 
 //===----------------------------------------------------------------------===//
 // ExecutionEngine/SparseTensorUtils helper functions.
 //===----------------------------------------------------------------------===//
 
 OverheadType mlir::sparse_tensor::overheadTypeEncoding(unsigned width) {
   switch (width) {

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 loop boundary statements (entering/exiting loops). The function
+/// passes and updates the reduction value.
+static Optional<Operation *> genLoopBoundary(
+    CodeGen &codegen, Merger &merger,
+    function_ref<Optional<Operation *>(MutableArrayRef<Value> reduc)>
+        callback) {
+  SmallVector<Value, 4> reduc;
+  if (codegen.redVal)
+    reduc.push_back(codegen.redVal);
+  if (codegen.expValues)
+    reduc.push_back(codegen.expCount);
+  if (codegen.insChain)
+    reduc.push_back(codegen.insChain);
+
+  auto r = callback(reduc);
+
+  // Callback should do in-place update on reduction value vector.
+  unsigned i = 0;
+  if (codegen.redVal)
+    updateReduc(merger, codegen, reduc[i++]);
+  if (codegen.expValues)
+    codegen.expCount = reduc[i++];
+  if (codegen.insChain)
+    codegen.insChain = reduc[i];
+
+  return r;
+}
+
 /// 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.getNumDpsInputs() + 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;
+  // Parallel loops on tensor expansion can cause data races.
+  if (codegen.expCount)
+    return false;
   // Inspect strategy.
   switch (codegen.options.parallelizationStrategy) {
   case SparseParallelizationStrategy::kNone:
     return false;
   case SparseParallelizationStrategy::kDenseOuterLoop:
-    return isOuter && !isSparse && !isReduction;
+    return isOuter && !isSparse;
   case SparseParallelizationStrategy::kAnyStorageOuterLoop:
-    return isOuter && !isReduction;
+    return isOuter;
   case SparseParallelizationStrategy::kDenseAnyLoop:
-    return !isSparse && !isReduction;
+    return !isSparse;
   case SparseParallelizationStrategy::kAnyStorageAnyLoop:
-    return !isReduction;
+    return true;
   }
   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 for loop.
-  SmallVector<Value, 4> operands;
-  if (codegen.redVal)
-    operands.push_back(codegen.redVal);
-  if (codegen.expValues)
-    operands.push_back(codegen.expCount);
-  if (codegen.insChain)
-    operands.push_back(codegen.insChain);
-
-  Operation *loop = codegen.loopEmitter.enterLoopOverTensorAtDim(
-      builder, loc, tid, dim, operands, isParallel, extraTids, extraDims);
-
-  unsigned o = 0;
-  if (codegen.redVal)
-    updateReduc(merger, codegen, operands[o++]);
-  if (codegen.expValues)
-    codegen.expCount = operands[o++];
-  if (codegen.insChain)
-    codegen.insChain = operands[o++];
-  assert(o == operands.size());
+  bool isParallel = isParallelFor(codegen, isOuter, isSparse);
 
+  Operation *loop =
+      genLoopBoundary(codegen, merger, [&](MutableArrayRef<Value> reduc) {
+        return codegen.loopEmitter.enterLoopOverTensorAtDim(
+            builder, loc, tid, dim, reduc, isParallel, extraTids, extraDims);
+      }).value();
+  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 =
+      genLoopBoundary(codegen, merger, [&](MutableArrayRef<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);
-  if (codegen.insChain)
-    operands.push_back(codegen.insChain);
-
-  Operation *loop = codegen.loopEmitter.enterCoIterationOverTensorsAtDims(
-      builder, op.getLoc(), condTids, condDims, needsUniv, operands, extraTids,
-      extraDims);
-
-  unsigned o = 0;
-  if (codegen.redVal)
-    updateReduc(merger, codegen, operands[o++]);
-  if (codegen.expValues)
-    codegen.expCount = operands[o++];
-  if (codegen.insChain)
-    codegen.insChain = operands[o++];
-  assert(o == operands.size());
-
+        return codegen.loopEmitter.enterCoIterationOverTensorsAtDims(
+            builder, op.getLoc(), condTids, condDims, needsUniv, reduc,
+            extraTids, extraDims);
+      }).value();
+  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,
translateBitsToTidDimPairs(merger, codegen, li, codegen.topSort[at], condTids,
                              condDims, extraTids, extraDims);
   // Emit the for/while-loop control.
   Operation *loop = genLoop(merger, codegen, builder, op, at, needsUniv,
                             condTids, condDims, extraTids, extraDims);
   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);
-  if (codegen.insChain)
-    reduc.push_back(codegen.insChain);
-
-  auto loopRet =
-      codegen.loopEmitter.exitCurrentLoop(builder, op.getLoc(), reduc);
-  assert(reduc.size() == loopRet.size());
-
-  unsigned o = 0;
-  if (codegen.redVal)
-    updateReduc(merger, codegen, loopRet[o++]);
-  if (codegen.expValues)
-    codegen.expCount = loopRet[o++];
-  if (codegen.insChain)
-    codegen.insChain = loopRet[o++];
-  assert(o == loopRet.size());
+  genLoopBoundary(codegen, merger, [&](MutableArrayRef<Value> reduc) {
+    codegen.loopEmitter.exitCurrentLoop(rewriter, op.getLoc(), reduc);
+    return llvm::None;
+  });
 
   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/Dialect/SparseTensor/sparse_parallel_reduce.mlir

+// RUN: mlir-opt %s -sparsification="parallelization-strategy=any-storage-any-loop" | \
+// RUN:   FileCheck %s
+
+#CSR = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ]
+}>
+
+#trait_matvec = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // A
+    affine_map<(i,j) -> (j)>,    // b
+    affine_map<(i,j) -> (i)>     // x (out)
+  ],
+  iterator_types = ["parallel", "reduction"],
+  doc = "x(i) += A(i,j) * b(j)"
+}
+// CHECK-LABEL:  func.func @matvec(
+//  CHECK-SAME:    %[[TMP_arg0:.*]]: tensor<16x32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>,
+//  CHECK-SAME:    %[[TMP_arg1:.*]]: tensor<32xf32>,
+//  CHECK-SAME:    %[[TMP_arg2:.*]]: tensor<16xf32>) -> tensor<16xf32> {
+//   CHECK-DAG:  %[[TMP_c16:.*]] = arith.constant 16 : index
+//   CHECK-DAG:  %[[TMP_c0:.*]] = arith.constant 0 : index
+//   CHECK-DAG:  %[[TMP_c1:.*]] = arith.constant 1 : index
+//       CHECK:  %[[TMP_0:.*]] = sparse_tensor.pointers %[[TMP_arg0]] {dimension = 1 : index}
+//       CHECK:  %[[TMP_1:.*]] = sparse_tensor.indices %[[TMP_arg0]] {dimension = 1 : index}
+//       CHECK:  %[[TMP_2:.*]] = sparse_tensor.values %[[TMP_arg0]]
+//       CHECK:  %[[TMP_3:.*]] = bufferization.to_memref %[[TMP_arg1]] : memref<32xf32>
+//       CHECK:  %[[TMP_4:.*]] = bufferization.to_memref %[[TMP_arg2]] : memref<16xf32>
+//       CHECK:  scf.parallel (%[[TMP_arg3:.*]]) = (%[[TMP_c0]]) to (%[[TMP_c16]]) step (%[[TMP_c1]]) {
+//       CHECK:    %[[TMP_6:.*]] = memref.load %[[TMP_4]][%[[TMP_arg3]]] : memref<16xf32>
+//       CHECK:    %[[TMP_7:.*]] = memref.load %[[TMP_0]][%[[TMP_arg3]]] : memref<?xindex>
+//       CHECK:    %[[TMP_8:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
+//       CHECK:    %[[TMP_9:.*]] = memref.load %[[TMP_0]][%[[TMP_8]]] : memref<?xindex>
+//       CHECK:    %[[TMP_10:.*]] = scf.parallel (%[[TMP_arg4:.*]]) = (%[[TMP_7]]) to (%[[TMP_9]]) step (%[[TMP_c1]]) init (%[[TMP_6]]) -> f32 {
+//       CHECK:      %[[TMP_11:.*]] = memref.load %[[TMP_1]][%[[TMP_arg4]]] : memref<?xindex>
+//       CHECK:      %[[TMP_12:.*]] = memref.load %[[TMP_2]][%[[TMP_arg4]]] : memref<?xf32>
+//       CHECK:      %[[TMP_13:.*]] = memref.load %[[TMP_3]][%[[TMP_11]]] : memref<32xf32>
+//       CHECK:      %[[TMP_14:.*]] = arith.mulf %[[TMP_12]], %[[TMP_13]] : f32
+//       CHECK:      scf.reduce(%[[TMP_14]])  : f32 {
+//       CHECK:      ^bb0(%[[TMP_arg5:.*]]: f32, %[[TMP_arg6:.*]]: f32):
+//       CHECK:        %[[TMP_15:.*]] = arith.addf %[[TMP_arg5]], %[[TMP_arg6]] : f32
+//       CHECK:        scf.reduce.return %[[TMP_15]] : f32
+//       CHECK:      }
+//       CHECK:      scf.yield
+//       CHECK:    }
+//       CHECK:    memref.store %[[TMP_10]], %[[TMP_4]][%[[TMP_arg3]]] : memref<16xf32>
+//       CHECK:    scf.yield
+//       CHECK:  }
+//       CHECK:  %[[TMP_5:.*]] = bufferization.to_tensor %[[TMP_4]] : memref<16xf32>
+//       CHECK:  return %[[TMP_5]] : tensor<16xf32>
+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_matmul.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner -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 parallelization.
+//
+// RUN: mlir-opt %s --sparse-compiler="parallelization-strategy=any-storage-any-loop" | \
+// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
+// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ],
   dimOrdering = affine_map<(i,j) -> (i,j)>
 }>
 
 #DCSR = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ],

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 parallelization.
+//
+// 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
 }>