[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

353

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

431

ditto

435

ditto

[Peiming]

rebase

[Peiming]

rebase

[Peiming]

minor changes..

[Peiming]

rebase

[Peiming]

minor fix

[Peiming]

rebase

[aartbik]

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

424

Can you make this

For sequential for loops:

(the For for reads a bit like a typo)

437

do we only support addition reduction?

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

476

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

437

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

[Peiming]

+ address comments
+ add a check test

[Peiming]

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

424

Haha, indeed!

437

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

199

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

463

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

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

428

end with :

428–429

period

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

1018–1019

"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

199

You are right ;-(.

[aartbik]

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

218

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

477

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!

1019

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

218

yeah, that's a good idea!

[Peiming]

address comments from Aart.

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

1019

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 coiteration loop over a set of tensors.
   Operation *enterCoiterationOverTensorsAtDims(
       OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
       ArrayRef<size_t> dims, SmallVectorImpl<Value> &reduc, bool needsUniv,
       ArrayRef<size_t> extraTids = {}, ArrayRef<size_t> extraDims = {});
 
-  SmallVector<Value, 2> exitCurrentLoop(OpBuilder &builder, Location loc,
+  SmallVector<Value, 2> exitCurrentLoop(RewriterBase &rewriter, Location loc,
                                         ArrayRef<Value> reduc = {});
 
   /// Return 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);
   }
 
   /// Get loop induction variable at the given level.
   Value getLoopIV(size_t level) const {
private:
   }
 
   bool isLastOutput; // Is the last tensor output tensor
   bool isOutputTensor(size_t tid) {
     return isLastOutput && tid == tensors.size() - 1;
   }
 
   /// Exits a for loop, returns the reduction results, e.g.,
   /// %ret = for () {
   ///   ...
   ///   yield %val
   /// }
   /// Return %ret to user, while %val is provided by users (`reduc`)
-  SmallVector<Value, 2> exitForLoop(OpBuilder &builder, Location loc,
+  SmallVector<Value, 2> exitForLoop(RewriterBase &rewriter, Location loc,
                                     ArrayRef<Value> reduc);
 
   /// Exits a while loop, returns the reduction results.
   SmallVector<Value, 2> exitCoiterationLoop(OpBuilder &builder, Location loc,
                                             ArrayRef<Value> reduc);
 
   struct LoopLevelInfo {
     LoopLevelInfo(ArrayRef<size_t> tids, ArrayRef<size_t> dims, Operation *loop,
                   Value iv)
         : tids(tids), dims(dims), loop(loop), iv(iv) {}
     const llvm::SmallVector<size_t, 4> tids;
     const llvm::SmallVector<size_t, 4> dims;
     const Operation *loop; // the loop operation
     const Value iv;        // the induction variable for the loop
   };
 

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,
     SmallVectorImpl<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];
   // Note that we generate loop based on dim bound for output tensor even on
   // sparse dimension.
   bool isSparseInput = isCompressedDim(dimType) || isSingletonDim(dimType);
 
   assert(isDenseDim(dimType) || isCompressedDim(dimType) ||
          isSingletonDim(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.
+    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, SmallVectorImpl<Value> &reduc, bool needsUniv,
     ArrayRef<size_t> extraTids, ArrayRef<size_t> extraDims) {
   assert(tids.size() == dims.size());
   SmallVector<Type, 4> types;
if (enc) {
       }
       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,
+SparseTensorLoopEmitter::exitForLoop(RewriterBase &rewriter, Location loc,
                                      ArrayRef<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);
+  SmallVector<Value, 2> results;
+  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);
+    results = forOp.getResults();
+  } 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();
+      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);
+      rewriter.setInsertionPointAfter(redExp);
+      auto redOp = rewriter.create<scf::ReduceOp>(loc, curVal);
+      redExp->remove(); // detach from previous parents
+      // Attach to the reduction op.
+      Block *redBlock = &redOp.getRegion().getBlocks().front();
+      redBlock->push_back(redExp);
+      // replaced arguments
+      rewriter.updateRootInPlace(
+          redExp, [&]() { redExp->setOperands(redBlock->getArguments()); });
+      rewriter.setInsertionPointToEnd(redBlock);
+      rewriter.create<scf::ReduceReturnOp>(loc, redExp->getResult(0));
+      rewriter.setInsertionPointToEnd(parOp.getBody());
+    }
+    rewriter.setInsertionPointAfter(parOp);
+    results = parOp.getResults();
   }
 
   // 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 (!isDenseDim(dimTypes[tid][dim]))
       highs[tid][dim] = Value();
   }
-  // exit the loop
-  builder.setInsertionPointAfter(forOp);
-  return forOp.getResults();
+  return results;
 }
 
 SmallVector<Value, 2>
 SparseTensorLoopEmitter::exitCoiterationLoop(OpBuilder &builder, Location loc,
                                              ArrayRef<Value> reduc) {
   auto whileOp = llvm::cast<scf::WhileOp>(loopStack.back().loop);
   auto &dims = loopStack.back().dims;
   auto &tids = loopStack.back().tids;
SparseTensorLoopEmitter::exitCoiterationLoop(OpBuilder &builder, Location loc,
 
   assert(o == operands.size());
   builder.create<scf::YieldOp>(loc, operands);
   builder.setInsertionPointAfter(whileOp);
   return ret;
 }
 
 SmallVector<Value, 2>
-SparseTensorLoopEmitter::exitCurrentLoop(OpBuilder &builder, Location loc,
+SparseTensorLoopEmitter::exitCurrentLoop(RewriterBase &rewriter, Location loc,
                                          ArrayRef<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);
+    red = exitCoiterationLoop(rewriter, loc, reduc);
   } else {
-    red = exitForLoop(builder, loc, reduc);
+    red = 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).
 //===----------------------------------------------------------------------===//
 
 /// 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.getNumInputsAndOutputs() == op.getNumInputs() + 1);
 
   codegen.loopEmitter.initializeLoopEmit(
       builder, loc,
       /// Generates buffer for the output tensor. Note that all sparse kernels
       /// assume that when all elements are written to (viz. x(i) = y(i) *
static bool isParallelFor(CodeGen &codegen, bool isOuter, bool isReduction,
   // 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 = merger.isDimLevelType(tid, idx, DimLvlType::kCompressed) ||
                   merger.isDimLevelType(tid, idx, DimLvlType::kSingleton);
   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, false, extraTids, extraDims);
+      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();
 
   return loop;
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) {
   codegen.curVecLength = 1;
 
   // 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);
+      codegen.loopEmitter.exitCurrentLoop(rewriter, op.getLoc(), reduc);
   assert(reduc.size() == loopRet.size());
 
   if (codegen.redVal)
     updateReduc(merger, codegen, loopRet.front());
   if (codegen.expValues)
     codegen.expCount = loopRet.back();
 
   return needsUniv;

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
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 // FIXME: support vector in sparse loop emitter.
 // R_U_N: mlir-opt %s \
 // R_U_N:   --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=16 enable-simd-index32" | \
 // R_U_N: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
 // R_U_N: mlir-cpu-runner \
 // R_U_N:  -e entry -entry-point-result=void  \
 // R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.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 SIMDization as well. This should not change the outcome.
-// FIXME: we do not support vectorization/parallel loops in loop emitter right now
-// R_U_N: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
-// R_U_N: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
-// R_U_N: mlir-cpu-runner \
-// R_U_N:  -e entry -entry-point-result=void  \
-// R_U_N:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// R_U_N: FileCheck %s
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
+// 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" ]
 }>
 
 #spmm = {