[mlir][sparse] add parallelization strategies to sparse compiler

[aartbik]

Summary

This CL adds the ability to request different parallelization strategies
for the generate code. Every "parallel" loop is a candidate, and converted
to a parallel op if it is an actual for-loop (not a while) and the strategy
allows dense/sparse outer/inner parallelization.

This will connect directly with the work of @ezhulenev on parallel loops.

Still TBD: vectorization strategy

[aartbik]

For ease of reviewing, I split out the invariant generalization in its own review: https://reviews.llvm.org/D91985

[penpornk]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

782

Typo: s/stategy/strategy/

784

Maybe explain the types of loop first just so it's clear for people who are not familiar with TACO?

• For loops that only iterates one tensor. (Can be called dense/sparse loop based on the compression of that tensor dimension.)
• For loops that co-iterates tensors (at most one tensor can have sparse storage for that dimension).
• While loops co-iterating tensors (more than one tensor have sparse storage for that dimension -- parallelization not supported).

785

Nit: Comma after e.g.

786

Nit: Missing a close parenthesis.

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

394

Typo: s/it it/it is/

579

Typo: s/loop loop/loop/

617–621

Please bear with my noob questions:
It looks like highs and sizes store the inferred dimension sizes. What exactly is lo (and pidxs and loops)? If it's the starting point of a dimension, isn't it 0? Why are we updating pidxs and loops through parOp.getInductionVars()?

Also, why do we need highs to be per tensor for sparse but sizes can be just per dimension? (If multiple tensors co-iterate on a dimension, the size should match anyway.)

It'd be great if you could help add more explanations about them (or examples) to struct CodeGen. (Or I can do it later once I'm more familiar with the code.)

mlir/test/Dialect/Linalg/sparse_2d.mlir

1133

Does alloca fill the allocated memory with 0s? If not, X has not been initialized to 0 before the loop, and the loop only set its values where A is nonzero.

mlir/test/Dialect/Linalg/sparse_parallel.mlir

148

If in the SCALE case (e.g., trait_ss) the inner loop can be parallelized, then I think in this case it could be parallelized too, because the amount of work solely depends on A.

[aartbik]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

784

Added some explanation at the top level.

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

617–621

Yes, I can add some more comments to the code, since it has become less self-explanatory then when I started this. In a nutshell, yes some of these are fixed before hand (like sizes), others are updated as we go (so that loops after loops iterate from where the previous loop left the induction).

mlir/test/Dialect/Linalg/sparse_2d.mlir

1133

None of these buffers are properly initialized. This is the temporary "local bufferization" solution we talked about earlier. Leaving these stubs here allows me to quickly hand covert this into something working. But eventually this will be replaced by proper initialization code (reading from file for example).

mlir/test/Dialect/Linalg/sparse_parallel.mlir

148

Well in this case, the inner loop is a "reduction". But you are right this could be parallelized too with a parallel tree reduction. And in fact, the "scf.parallel" operation supports the "scf.reduce" we may start using in the future....

[aartbik]

rebased + addressed comments

[penpornk]

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

617–621

Thank you very much!

mlir/test/Dialect/Linalg/sparse_parallel.mlir

148

Oh, right. I forgot that it is doing reduction so it's out of scope for now. Thank you for the clarification! :)

[aartbik]

Thanks for your review, Penporn! Keep those useful comments coming!

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 LogicalResult applyStagedPatterns(
     Operation *op, ArrayRef<FrozenRewritePatternList> stage1Patterns,
     const FrozenRewritePatternList &stage2Patterns,
     function_ref<LogicalResult(Operation *)> stage3Lambda = nullptr);
 
 //===----------------------------------------------------------------------===//
 // Support for sparse tensor code generation.
 //===----------------------------------------------------------------------===//
-void populateSparsificationPatterns(MLIRContext *context,
-                                    OwningRewritePatternList &patterns);
+
+/// Defines a parallelization stategy for the sparse compiler. Any implicit loop
+/// in the Linalg operation that is marked "parallel" (viz. no "reduction") is
+/// a candidate for parallelization. The loop is made parallel if (1) allowed by
+/// the strategy (e.g. AnyStorageOuterLoop considers either a dense or sparse
+/// outermost loop only, and (2) the emitted code is an actual for-loop (and
+/// not a co-iterating while-loop).
+enum class SparseParallelizationStrategy {
+  kNone,
+  kDenseOuterLoop,
+  kAnyStorageOuterLoop,
+  kDenseAnyLoop,
+  kAnyStorageAnyLoop
+};
+
+/// Defines a vectorization strategy for the sparse compiler. Any implicit inner
+/// in the Linalg operation is a candidate (full SIMD for "parallel" loops and
+/// horizontal SIMD for "reduction" loops). A loop is actually vectorized if
+/// (1) allowed by the strategy, and (2) the emitted code is an actual
+/// for-loop (and not a co-iterating while-loop).
+enum class SparseVectorizationStrategy {
+  kNone,
+  kDenseInnerLoop,
+  kAnyStorageInnerLoop
+};
+
+/// Sparsification options.
+struct SparsificationOptions {
+  SparsificationOptions()
+      : parallelizationStrategy(SparseParallelizationStrategy::kNone),
+        vectorizationStrategy(SparseVectorizationStrategy::kNone) {}
+  SparseParallelizationStrategy parallelizationStrategy;
+  SparseVectorizationStrategy vectorizationStrategy;
+};
+
+/// Set up sparsification rewriting rules with the given options.
+void populateSparsificationPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns,
+    const SparsificationOptions &options = SparsificationOptions());
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_

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

 
 using namespace mlir;
 
 namespace {
 
 enum class Kind { kTensor, kInvariant, kMulF, kMulI, kAddF, kAddI };
 
 /// Tensor expression. Represents a MLIR expression in tensor index notation.
-/// For tensors and invariants, e0 denotes the tensor index. For all binary
-/// operations, e0 and e1 denote the index of the children tensor expressions.
+/// For tensors, e0 denotes the tensor index. For invariants, the IR value is
+/// stored directly. For binary operations, e0 and e1 denote the index of the
+/// children tensor expressions.
 struct TensorExp {
-  TensorExp(Kind k, unsigned x, unsigned y) : kind(k), e0(x), e1(y) {}
+  TensorExp(Kind k, unsigned x, unsigned y, Value v)
+      : kind(k), e0(x), e1(y), val(v) {}
   Kind kind;
   unsigned e0;
   unsigned e1;
+  Value val;
 };
 
 /// Lattice point. Each lattice point consist of a conjunction of tensor
 /// loop indices (encoded in a bitvector) and the index of the corresponding
 /// tensor expression.
 struct LatPoint {
   LatPoint(unsigned n, unsigned e, unsigned b) : bits(n, false), exp(e) {
     bits.set(b);
   }
   LatPoint(const llvm::BitVector &b, unsigned e) : bits(b), exp(e) {}
   llvm::BitVector bits;
   unsigned exp;
 };
 
 /// A class to handle all iteration lattice operations. This class abstracts
 /// away from some implementation details of storing iteration lattices and
 /// tensor expressions. This allows for fine-tuning performance characteristics
 /// independently from the basic algorithm if bottlenecks are identified.
 class Merger {
 public:
   Merger(unsigned t, unsigned l)
       : numTensors(t), numLoops(l), isSparse(t, std::vector<bool>(l, false)) {}
 
   /// Adds a tensor expression. Returns its index.
-  unsigned addExp(Kind k, unsigned e0, unsigned e1 = -1u) {
+  unsigned addExp(Kind k, unsigned e0, unsigned e1 = -1u, Value v = Value()) {
     unsigned e = tensorExps.size();
-    tensorExps.push_back(TensorExp(k, e0, e1));
+    tensorExps.push_back(TensorExp(k, e0, e1, v));
     return e;
   }
+  unsigned addExp(Kind k, Value v) { return addExp(k, -1u, -1u, v); }
 
   /// Adds an iteration lattice point. Returns its index.
   unsigned addLat(unsigned t, unsigned i, unsigned e) {
     assert(t < numTensors && i < numLoops);
     unsigned p = latPoints.size();
     latPoints.push_back(LatPoint(numLoops * numTensors, e, numTensors * i + t));
     return p;
   }
private:
   std::vector<std::vector<bool>> isSparse;
   llvm::SmallVector<TensorExp, 32> tensorExps;
   llvm::SmallVector<LatPoint, 16> latPoints;
   llvm::SmallVector<SmallVector<unsigned, 16>, 8> latSets;
 };
 
 // Code generation.
 struct CodeGen {
-  CodeGen(unsigned numTensors, unsigned numLoops)
-      : loops(numLoops), sizes(numLoops), buffers(numTensors),
+  CodeGen(linalg::SparsificationOptions o, unsigned numTensors,
+          unsigned numLoops)
+      : options(o), loops(numLoops), sizes(numLoops), buffers(numTensors),
         pointers(numTensors, std::vector<Value>(numLoops)),
         indices(numTensors, std::vector<Value>(numLoops)),
         highs(numTensors, std::vector<Value>(numLoops)),
         pidxs(numTensors, std::vector<Value>(numLoops)),
         idxs(numTensors, std::vector<Value>(numLoops)) {}
+  // Sparsification options.
+  linalg::SparsificationOptions options;
   // Universal dense indices and upper bounds (by index).
   std::vector<Value> loops;
   std::vector<Value> sizes;
   // Buffers for storing dense and sparse numerical values (by tensor).
   std::vector<Value> buffers;
   // Sparse storage schemes (1-D): pointers and indices (by tensor and index).
   std::vector<std::vector<Value>> pointers;
   std::vector<std::vector<Value>> indices;
static bool computeIterationGraph(linalg::GenericOp op,
   return true;
 }
 
 /// Traverses the SSA tree (possibly a DAG) to build a tensor expression.
 /// This simplifies constructing (sub)expressions during iteration lattice
 /// building (compared to using the SSA representation everywhere).
 static Optional<unsigned> buildTensorExp(Merger &merger, linalg::GenericOp op,
                                          Value val) {
-  Operation *def = val.getDefiningOp();
   if (auto arg = val.dyn_cast<BlockArgument>()) {
     unsigned argN = arg.getArgNumber();
     if (arg.getOwner()->getParentOp() == op) {
       // Any parameter of the generic op is considered a tensor,
       // indexed by the implicit loop bounds.
       auto map = op.getIndexingMap(argN);
       if (map.isProjectedPermutation())
         return merger.addExp(Kind::kTensor, argN);
-    } else {
-      // Any parameter of a higher op is invariant in the tensor expression.
-      return merger.addExp(Kind::kInvariant, argN);
+      // Cannot handle (yet).
+      return None;
     }
+    // Any parameter of a higher op is invariant.
+    return merger.addExp(Kind::kInvariant, val);
+  }
+  Operation *def = val.getDefiningOp();
+  if (def->getBlock() != &op.region().front()) {
+    // Something defined outside is invariant.
+    return merger.addExp(Kind::kInvariant, val);
   } else if (def->getNumOperands() == 2) {
     // Construct binary operations if subexpressions could be built.
     auto x = buildTensorExp(merger, op, def->getOperand(0));
     auto y = buildTensorExp(merger, op, def->getOperand(1));
     if (x.hasValue() && y.hasValue()) {
       unsigned e0 = x.getValue();
       unsigned e1 = y.getValue();
       if (isa<MulFOp>(def))

 }
 
 /// Builds the iteration lattices in a bottom-up traversal given the remaining
 /// tensor (sub)expression and the next loop index in the iteration graph.
 static unsigned buildLattices(Merger &merger, linalg::GenericOp op,
                               unsigned exp, unsigned idx) {
   Kind kind = merger.exp(exp).kind;
   if (kind == Kind::kTensor || kind == Kind::kInvariant) {
     // Either the index is really used in the tensor expression, or it it
-    // set to the "non-existing dense index" in that dimension.
+    // set to the "non-existing dense index" in that dimension. Invariant
+    // expressions borrow the output tensor indices.
     unsigned s = merger.addSet();
-    merger.set(s).push_back(merger.addLat(merger.exp(exp).e0, idx, exp));
+    unsigned t = kind == Kind::kTensor ? merger.exp(exp).e0
+                                       : op.getNumInputsAndOutputs() - 1;
+    merger.set(s).push_back(merger.addLat(t, idx, exp));
     return s;
   }
   unsigned s0 = buildLattices(merger, op, merger.exp(exp).e0, idx);
   unsigned s1 = buildLattices(merger, op, merger.exp(exp).e1, idx);
   switch (kind) {
   case Kind::kTensor:
   case Kind::kInvariant:
     llvm_unreachable("handled above");

 }
 
 /// Recursively generates tensor expression.
 static Value genExp(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
                     linalg::GenericOp op, unsigned exp) {
   if (merger.exp(exp).kind == Kind::kTensor)
     return genTensorLoad(merger, codegen, rewriter, op, merger.exp(exp).e0);
   else if (merger.exp(exp).kind == Kind::kInvariant)
-    return op.getParentRegion()->front().getArgument(merger.exp(exp).e0);
+    return merger.exp(exp).val;
   Value v0 = genExp(merger, codegen, rewriter, op, merger.exp(exp).e0);
   Value v1 = genExp(merger, codegen, rewriter, op, merger.exp(exp).e1);
   switch (merger.exp(exp).kind) {
   case Kind::kTensor:
   case Kind::kInvariant:
     llvm_unreachable("handled above");
   case Kind::kMulF:
     return rewriter.create<MulFOp>(op.getLoc(), v0, v1);
for (unsigned b = 0, be = inits.size(); b < be; b++) {
     }
   }
 
   // Initialize the universal dense index.
   codegen.loops[idx] = rewriter.create<ConstantIndexOp>(loc, 0);
   return needsUniv;
 }
 
-/// Generates a for-loop or a while-loop, depending on whether it implements
-/// singleton iteration or co-iteration over the given conjunction.
-static void genLoop(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
-                    linalg::GenericOp op, unsigned idx, bool needsUniv,
-                    llvm::BitVector &indices, scf::ForOp &forOp,
-                    scf::WhileOp &whileOp) {
-  Location loc = op.getLoc();
-
-  // Emit a for-loop for a single index.
-  if (indices.count() == 1) {
-    unsigned fb = indices.find_first();
-    unsigned tensor = merger.tensor(fb);
-    assert(idx == merger.index(fb));
-    // Emit a sparse for-loop or a dense for-loop.
-    Value one = rewriter.create<ConstantIndexOp>(loc, 1);
-    if (merger.isSparseBit(fb)) {
-      forOp = rewriter.create<scf::ForOp>(loc, codegen.pidxs[tensor][idx],
-                                          codegen.highs[tensor][idx], one);
-      codegen.pidxs[tensor][idx] = forOp.getInductionVar();
-    } else {
-      forOp = rewriter.create<scf::ForOp>(loc, codegen.loops[idx],
-                                          codegen.sizes[idx], one);
-      codegen.loops[idx] = forOp.getInductionVar();
-    }
-    rewriter.setInsertionPointToStart(forOp.getBody());
-    return;
-  }
+/// Generates a for-loop loop on a single index.
+static Operation *genFor(Merger &merger, CodeGen &codegen,
+                         PatternRewriter &rewriter, linalg::GenericOp op,
+                         bool isOuter, unsigned idx, llvm::BitVector &indices) {
+
+  // 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.
+  auto iteratorTypes = op.iterator_types().getValue();
+  bool isSparse = merger.isSparseBit(fb);
+  bool isParallel = linalg::isParallelIteratorType(iteratorTypes[idx]);
+  switch (codegen.options.parallelizationStrategy) {
+  case linalg::SparseParallelizationStrategy::kNone:
+    isParallel = false;
+    break;
+  case linalg::SparseParallelizationStrategy::kDenseOuterLoop:
+    isParallel &= isOuter && !isSparse;
+    break;
+  case linalg::SparseParallelizationStrategy::kAnyStorageOuterLoop:
+    isParallel &= isOuter;
+    break;
+  case linalg::SparseParallelizationStrategy::kDenseAnyLoop:
+    isParallel &= !isSparse;
+    break;
+  case linalg::SparseParallelizationStrategy::kAnyStorageAnyLoop:
+    break;
+  }
+
+  // Loop bounds and increment.
+  Location loc = op.getLoc();
+  Value lo;
+  Value hi;
+  Value step = rewriter.create<ConstantIndexOp>(loc, 1);
+  Value index;
+  if (isSparse) {
+    lo = codegen.pidxs[tensor][idx];
+    hi = codegen.highs[tensor][idx];
+    lo = codegen.loops[idx];
+    hi = codegen.sizes[idx];
+
+  // Emit a parallel loop.
+  if (isParallel) {
+    scf::ParallelOp parOp = rewriter.create<scf::ParallelOp>(loc, lo, hi, step);
+    if (isSparse)
+      codegen.pidxs[tensor][idx] = parOp.getInductionVars()[0];
+    else
+      codegen.loops[idx] = parOp.getInductionVars()[0];
+    rewriter.setInsertionPointToStart(parOp.getBody());
+    return parOp;
+  }
+
+  // Emit a sequential loop.
+  scf::ForOp forOp = rewriter.create<scf::ForOp>(loc, lo, hi, step);
+  if (isSparse)
+    codegen.pidxs[tensor][idx] = forOp.getInductionVar();
+  else
+    codegen.loops[idx] = forOp.getInductionVar();
+  return forOp;
 
-  // Otherwise, emit a while-loop for co-iteration.
+/// Emit a while-loop for co-iteration over multiple indices.
+static Operation *genWhile(Merger &merger, CodeGen &codegen,
+                           PatternRewriter &rewriter, linalg::GenericOp op,
+                           unsigned idx, bool needsUniv,
+                           llvm::BitVector &indices) {
   Type indexType = rewriter.getIndexType();
   SmallVector<Type, 4> types;
   SmallVector<Value, 4> operands;
+  // Construct the while-loop with a parameter for each index.
   for (unsigned b = 0, be = indices.size(); b < be; b++) {
     if (indices[b] && merger.isSparseBit(b)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       types.push_back(indexType);
       operands.push_back(codegen.pidxs[tensor][idx]);
     }
   }
   if (needsUniv) {
     types.push_back(indexType);
     operands.push_back(codegen.loops[idx]);
   }
-  whileOp = rewriter.create<scf::WhileOp>(loc, types, operands);
+  Location loc = op.getLoc();
+  scf::WhileOp whileOp = rewriter.create<scf::WhileOp>(loc, types, operands);
   Block *before = rewriter.createBlock(&whileOp.before(), {}, types);
   Block *after = rewriter.createBlock(&whileOp.after(), {}, types);
+
   // Build the "before" region, which effectively consists
   // of a conjunction of "i < upper" tests on all induction.
   rewriter.setInsertionPointToStart(&whileOp.before().front());
   Value cond;
   unsigned o = 0;
   for (unsigned b = 0, be = indices.size(); b < be; b++) {
     if (indices[b] && merger.isSparseBit(b)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value op1 = before->getArgument(o);
       Value op2 = codegen.highs[tensor][idx];
       Value opc = rewriter.create<CmpIOp>(loc, CmpIPredicate::ult, op1, op2);
       cond = cond ? rewriter.create<AndOp>(loc, cond, opc) : opc;
       codegen.pidxs[tensor][idx] = after->getArgument(o++);
     }
   }
   if (needsUniv)
     codegen.loops[idx] = after->getArgument(o++);
   assert(o == operands.size());
   rewriter.create<scf::ConditionOp>(loc, cond, before->getArguments());
   rewriter.setInsertionPointToStart(&whileOp.after().front());
+  return whileOp;
+}
+
+/// 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,
+                          PatternRewriter &rewriter, linalg::GenericOp op,
+                          bool isOuter, unsigned idx, bool needsUniv,
+                          llvm::BitVector &indices) {
+  if (indices.count() == 1)
+    return genFor(merger, codegen, rewriter, op, isOuter, idx, indices);
+  return genWhile(merger, codegen, rewriter, op, idx, needsUniv, indices);
 }
 
 /// Generates the local variables for this loop, consisting of the sparse
 /// indices, restored universal dense index, and dense positions.
 static void genLocals(Merger &merger, CodeGen &codegen,
                       PatternRewriter &rewriter, linalg::GenericOp op,
                       std::vector<unsigned> &topSort, unsigned at,
                       bool needsUniv, llvm::BitVector &locals) {
bool needsUniv =
       genInit(merger, codegen, rewriter, op, topSort, at, lat0.bits) &&
       lsize > 1;
 
   // Emit a loop for every lattice point L0 >= Li.
   for (unsigned li : merger.set(lts)) {
     LatPoint lati = merger.lat(li);
 
     // Emit loop.
-    scf::ForOp forOp;
-    scf::WhileOp whileOp;
     llvm::BitVector indices = lati.bits;
     optimizeIndices(merger, lsize, indices);
-    genLoop(merger, codegen, rewriter, op, idx, needsUniv, indices, forOp,
-            whileOp);
+    bool isOuter = at == 0;
+    Operation *loop = genLoop(merger, codegen, rewriter, op, isOuter, idx,
+                              needsUniv, indices);
     genLocals(merger, codegen, rewriter, op, topSort, at, needsUniv, lati.bits);
 
     // Visit all lattices points with Li >= Lj to generate the
     // loop-body, possibly with if statements for coiteration.
+    bool isWhile = dyn_cast<scf::WhileOp>(loop) != nullptr;
     scf::IfOp ifOp;
     for (unsigned lj : merger.set(lts)) {
       if (li == lj || merger.latGT(li, lj)) {
         LatPoint latj = merger.lat(lj);
         llvm::BitVector tmp = latj.bits;
         tmp ^= lati.bits;
         if (merger.hasAnyOf(tmp, false))
           continue; // dense exhausted within if/else
         // Recurse into body of each branch.
-        if (whileOp)
+        if (isWhile)
           genIf(merger, codegen, rewriter, op, idx, latj.bits, ifOp);
         genStmt(merger, codegen, rewriter, op, topSort, latj.exp, at + 1);
       }
     }
 
     // Wrap-up induction and restore insertion point.
-    if (forOp) {
-      needsUniv = false;
-      rewriter.setInsertionPointAfter(forOp);
-    } else {
+    if (isWhile) {
+      scf::WhileOp whileOp = cast<scf::WhileOp>(loop);
       rewriter.setInsertionPointToEnd(&whileOp.after().front());
       genWhileInduction(merger, codegen, rewriter, op, idx, needsUniv,
                         lati.bits, whileOp.results());
-      rewriter.setInsertionPointAfter(whileOp);
+    } else {
+      needsUniv = false;
     }
+    rewriter.setInsertionPointAfter(loop);
   }
 }
 
 namespace {
 
 /// Sparse rewriting rule for generic Lingalg operation.
 struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
-  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+public:
+  GenericOpSparsifier(MLIRContext *context, linalg::SparsificationOptions o)
+      : OpRewritePattern<linalg::GenericOp>(context), options(o) {}
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
     // Detects sparse annotations and translate the per-dimension sparsity
     // information for all tensors to loop indices in the kernel.
     if (!op.hasSparseSemantics())
       return failure();
     assert(op.getNumOutputs() == 1);
LogicalResult matchAndRewrite(linalg::GenericOp op,
     // Finds the terminating yield statement and builds the tensor
     // expression for the Linalg operation in SSA form.
     Operation *yield = op.region().front().getTerminator();
     Optional<unsigned> exp = buildTensorExp(merger, op, yield->getOperand(0));
     if (!exp.hasValue())
       return failure(); // build failure
 
     // Recursively generates code.
-    CodeGen codegen(numTensors, numLoops);
+    CodeGen codegen(options, numTensors, numLoops);
     genBuffers(merger, codegen, rewriter, op);
     genStmt(merger, codegen, rewriter, op, topSort, exp.getValue(), 0);
     Value result =
         rewriter.create<TensorLoadOp>(op.getLoc(), codegen.buffers.back());
     rewriter.replaceOp(op, result);
     return success();
   }
+
+private:
+  /// Options to control sparse code generation.
+  linalg::SparsificationOptions options;
 };
 
 } // namespace
 
 /// Populates the given patterns list with rewriting rules required for
 /// the sparsification of linear algebra operations.
-void mlir::linalg::populateSparsificationPatterns(
-    MLIRContext *context, OwningRewritePatternList &patterns) {
-  patterns.insert<GenericOpSparsifier>(context);
+void linalg::populateSparsificationPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns,
+    const SparsificationOptions &options) {
+  patterns.insert<GenericOpSparsifier>(context, options);
 }

mlir/test/Dialect/Linalg/sparse_2d.mlir

%0 = linalg.generic #trait_sum_reduction
     init(%argx : tensor<f32>) {
       ^bb(%a : f32, %x : f32):
         %0 = addf %x, %a  : f32
         linalg.yield %0: f32
   } -> tensor<f32>
   return %0 : tensor<f32>
 }
 
+#trait_scale = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // A
+    affine_map<(i,j) -> (i,j)>   // X (out)
+  ],
+  sparse = [
+    [ "D", "S" ],  // A
+    [ "D", "D" ]   // X
+  ],
+  iterator_types = ["parallel", "parallel"],
+  doc = "X(i,j) = A(i,j) * SCALE"
+}
+
+// CHECK-LABEL:   func @scale(
+// CHECK-SAME:                %[[VAL_0:.*]]: tensor<?x?xf64>) -> tensor<?x?xf64> {
+// CHECK:           %[[VAL_1:.*]] = constant 2.000000e+00 : f64
+// CHECK:           %[[VAL_2:.*]] = constant 999 : index
+// CHECK:           %[[VAL_3:.*]] = constant 0 : index
+// CHECK:           %[[VAL_4:.*]] = constant 1 : index
+// CHECK:           %[[VAL_5:.*]] = dim %[[VAL_0]], %[[VAL_3]] : tensor<?x?xf64>
+// CHECK:           %[[VAL_6:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
+// CHECK:           %[[VAL_7:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
+// CHECK:           %[[VAL_8:.*]] = dim %[[VAL_0]], %[[VAL_4]] : tensor<?x?xf64>
+// CHECK:           %[[VAL_9:.*]] = alloca(%[[VAL_2]]) : memref<?xf64>
+// CHECK:           %[[VAL_10:.*]] = alloca(%[[VAL_5]], %[[VAL_8]]) : memref<?x?xf64>
+// CHECK:           scf.for %[[VAL_11:.*]] = %[[VAL_3]] to %[[VAL_5]] step %[[VAL_4]] {
+// CHECK:             %[[VAL_12:.*]] = load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
+// CHECK:             %[[VAL_13:.*]] = addi %[[VAL_11]], %[[VAL_4]] : index
+// CHECK:             %[[VAL_14:.*]] = load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xindex>
+// CHECK:             scf.for %[[VAL_15:.*]] = %[[VAL_12]] to %[[VAL_14]] step %[[VAL_4]] {
+// CHECK:               %[[VAL_16:.*]] = load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex>
+// CHECK:               %[[VAL_17:.*]] = load %[[VAL_9]]{{\[}}%[[VAL_15]]] : memref<?xf64>
+// CHECK:               %[[VAL_18:.*]] = mulf %[[VAL_17]], %[[VAL_1]] : f64
+// CHECK:               store %[[VAL_18]], %[[VAL_10]]{{\[}}%[[VAL_11]], %[[VAL_16]]] : memref<?x?xf64>
+// CHECK:             }
+// CHECK:           }
+// CHECK:           %[[VAL_19:.*]] = tensor_load %[[VAL_10]] : memref<?x?xf64>
+// CHECK:           return %[[VAL_19]] : tensor<?x?xf64>
+// CHECK:         }
+func @scale(%arga: tensor<?x?xf64>) -> tensor<?x?xf64> {
+  %0 = constant 2.0 : f64
+  %1 = linalg.generic #trait_scale
+    ins(%arga: tensor<?x?xf64>) {
+      ^bb(%a: f64):
+        %2 = mulf %a, %0  : f64
+        linalg.yield %2 : f64
+  } -> tensor<?x?xf64>
+  return %1 : tensor<?x?xf64>
+}
+
 #trait_sampled_dense_dense = {
   indexing_maps = [
     affine_map<(i,j,k) -> (i,j)>,  // S
     affine_map<(i,j,k) -> (i,k)>,  // A
     affine_map<(i,j,k) -> (k,j)>,  // B
     affine_map<(i,j,k) -> (i,j)>   // X (out)
   ],
   sparse = [

mlir/test/Dialect/Linalg/sparse_parallel.mlir

+// RUN: mlir-opt %s -test-sparsification="parallelization-strategy=0" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-PAR0
+// RUN: mlir-opt %s -test-sparsification="parallelization-strategy=1" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-PAR1
+// RUN: mlir-opt %s -test-sparsification="parallelization-strategy=2" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-PAR2
+// RUN: mlir-opt %s -test-sparsification="parallelization-strategy=3" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-PAR3
+// RUN: mlir-opt %s -test-sparsification="parallelization-strategy=4" | \
+// RUN:   FileCheck %s --check-prefix=CHECK-PAR4
+
+#trait_dd = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // A
+    affine_map<(i,j) -> (i,j)>   // X (out)
+  ],
+  sparse = [
+    [ "D", "D" ],  // A
+    [ "D", "D" ]   // X
+  ],
+  iterator_types = ["parallel", "parallel"],
+  doc = "X(i,j) = A(i,j) * SCALE"
+}
+
+//
+// CHECK-PAR0-LABEL: func @scale_dd
+// CHECK-PAR0:         scf.for
+// CHECK-PAR0:           scf.for
+// CHECK-PAR0:         return
+//
+// CHECK-PAR1-LABEL: func @scale_dd
+// CHECK-PAR1:         scf.parallel
+// CHECK-PAR1:           scf.for
+// CHECK-PAR1:         return
+//
+// CHECK-PAR2-LABEL: func @scale_dd
+// CHECK-PAR2:         scf.parallel
+// CHECK-PAR2:           scf.for
+// CHECK-PAR2:         return
+//
+// CHECK-PAR3-LABEL: func @scale_dd
+// CHECK-PAR3:         scf.parallel
+// CHECK-PAR3:           scf.parallel
+// CHECK-PAR3:         return
+//
+// CHECK-PAR4-LABEL: func @scale_dd
+// CHECK-PAR4:         scf.parallel
+// CHECK-PAR4:           scf.parallel
+// CHECK-PAR4:         return
+//
+func @scale_dd(%scale: f32, %arga: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = linalg.generic #trait_dd
+    ins(%arga: tensor<?x?xf32>) {
+      ^bb(%a: f32):
+        %0 = mulf %a, %scale  : f32
+        linalg.yield %0 : f32
+  } -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+#trait_ss = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // A
+    affine_map<(i,j) -> (i,j)>   // X (out)
+  ],
+  sparse = [
+    [ "S", "S" ],  // A
+    [ "D", "D" ]   // X
+  ],
+  iterator_types = ["parallel", "parallel"],
+  doc = "X(i,j) = A(i,j) * SCALE"
+}
+
+//
+// CHECK-PAR0-LABEL: func @scale_ss
+// CHECK-PAR0:         scf.for
+// CHECK-PAR0:           scf.for
+// CHECK-PAR0:         return
+//
+// CHECK-PAR1-LABEL: func @scale_ss
+// CHECK-PAR1:         scf.for
+// CHECK-PAR1:           scf.for
+// CHECK-PAR1:         return
+//
+// CHECK-PAR2-LABEL: func @scale_ss
+// CHECK-PAR2:         scf.parallel
+// CHECK-PAR2:           scf.for
+// CHECK-PAR2:         return
+//
+// CHECK-PAR3-LABEL: func @scale_ss
+// CHECK-PAR3:         scf.for
+// CHECK-PAR3:           scf.for
+// CHECK-PAR3:         return
+//
+// CHECK-PAR4-LABEL: func @scale_ss
+// CHECK-PAR4:         scf.parallel
+// CHECK-PAR4:           scf.parallel
+// CHECK-PAR4:         return
+//
+func @scale_ss(%scale: f32, %arga: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = linalg.generic #trait_ss
+    ins(%arga: tensor<?x?xf32>) {
+      ^bb(%a: f32):
+        %0 = mulf %a, %scale  : f32
+        linalg.yield %0 : f32
+  } -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+#trait_matvec = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // A
+    affine_map<(i,j) -> (j)>,    // b
+    affine_map<(i,j) -> (i)>     // x (out)
+  ],
+  sparse = [
+    [ "D", "S" ],  // A
+    [ "D" ],       // b
+    [ "D" ]        // x
+  ],
+  iterator_types = ["parallel", "reduction"],
+  doc = "x(i) += A(i,j) * b(j)"
+}
+
+//
+// CHECK-PAR0-LABEL: func @matvec
+// CHECK-PAR0:         scf.for
+// CHECK-PAR0:           scf.for
+// CHECK-PAR0:         return
+//
+// CHECK-PAR1-LABEL: func @matvec
+// CHECK-PAR1:         scf.parallel
+// CHECK-PAR1:           scf.for
+// CHECK-PAR1:         return
+//
+// CHECK-PAR2-LABEL: func @matvec
+// CHECK-PAR2:         scf.parallel
+// CHECK-PAR2:           scf.for
+// CHECK-PAR2:         return
+//
+// 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:         return
+//
+func @matvec(%argA: tensor<16x32xf32>, %argb: tensor<32xf32>, %argx: tensor<16xf32>) -> tensor<16xf32> {
+  %0 = linalg.generic #trait_matvec
+      ins(%argA, %argb : tensor<16x32xf32>, tensor<32xf32>)
+      init(%argx : tensor<16xf32>) {
+    ^bb(%A: f32, %b: f32, %x: f32):
+      %0 = mulf %A, %b : f32
+      %1 = addf %0, %x : f32
+      linalg.yield %1 : f32
+  } -> tensor<16xf32>
+  return %0 : tensor<16xf32>
+}

mlir/test/lib/Transforms/TestSparsification.cpp

 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 
 struct TestSparsification
     : public PassWrapper<TestSparsification, FunctionPass> {
+
+  TestSparsification() = default;
+  TestSparsification(const TestSparsification &pass) {}
+
+  Option<int32_t> parallelization{
+      *this, "parallelization-strategy",
+      llvm::cl::desc("Set the parallelization strategy"), llvm::cl::init(0)};
+
+  Option<int32_t> vectorization{
+      *this, "vectorization-strategy",
+      llvm::cl::desc("Set the vectorization strategy"), llvm::cl::init(0)};
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<scf::SCFDialect>();
   }
+
   void runOnFunction() override {
     auto *ctx = &getContext();
     OwningRewritePatternList patterns;
-    linalg::populateSparsificationPatterns(ctx, patterns);
+    // Translate strategy flag to strategy options.
+    linalg::SparsificationOptions options;
+    switch (parallelization) {
+    case 0:
+      options.parallelizationStrategy =
+          linalg::SparseParallelizationStrategy::kNone;
+      break;
+    case 1:
+      options.parallelizationStrategy =
+          linalg::SparseParallelizationStrategy::kDenseOuterLoop;
+      break;
+    case 2:
+      options.parallelizationStrategy =
+          linalg::SparseParallelizationStrategy::kAnyStorageOuterLoop;
+      break;
+    case 3:
+      options.parallelizationStrategy =
+          linalg::SparseParallelizationStrategy::kDenseAnyLoop;
+      break;
+    case 4:
+      options.parallelizationStrategy =
+          linalg::SparseParallelizationStrategy::kAnyStorageAnyLoop;
+      break;
+    }
+    switch (vectorization) {
+    case 0:
+      options.vectorizationStrategy =
+          linalg::SparseVectorizationStrategy::kNone;
+      break;
+    case 1:
+      options.vectorizationStrategy =
+          linalg::SparseVectorizationStrategy::kDenseInnerLoop;
+      break;
+    case 2:
+      options.vectorizationStrategy =
+          linalg::SparseVectorizationStrategy::kAnyStorageInnerLoop;
+      break;
+    }
+    // Apply rewriting.
+    linalg::populateSparsificationPatterns(ctx, patterns, options);
     applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 };
 
 } // end anonymous namespace
 
 namespace mlir {
 namespace test {