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

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D135927

Author: Peiming Liu

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

@@ -219,9 +219,12 @@ 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];
@@ -232,11 +235,38 @@ Operation *SparseTensorLoopEmitter::enterLoopOverTensorAtDim(
   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());
+    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);
 
-  scf::ForOp forOp = builder.create<scf::ForOp>(loc, lo, hi, step, reduc);
-  builder.setInsertionPointToStart(forOp.getBody());
-  Value iv = forOp.getInductionVar();
-  assert(iv);
   if (isSparseInput) {
     pidxs[tid][dim] = iv;
     // Generating a load on the indices array yields the coordinate.
@@ -253,16 +283,12 @@ Operation *SparseTensorLoopEmitter::enterLoopOverTensorAtDim(
 
   // 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(
@@ -434,17 +460,73 @@ void SparseTensorLoopEmitter::emitExtraLocalsForTensorsAtDenseDims(
   }
 }
 
-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);
-  if (!reduc.empty()) {
-    assert(reduc.size() == forOp.getNumResults());
-    builder.setInsertionPointToEnd(forOp.getBody());
-    builder.create<scf::YieldOp>(loc, reduc);
+  auto forOp = llvm::dyn_cast<scf::ForOp>(loopInfo.loop);
+  if (forOp) {
+    if (!reduc.empty()) {
+      assert(reduc.size() == forOp.getNumResults());
+      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
@@ -458,14 +540,10 @@ SparseTensorLoopEmitter::exitForLoop(OpBuilder &builder, Location loc,
     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;
@@ -499,10 +577,10 @@ SparseTensorLoopEmitter::exitCoiterationLoop(OpBuilder &builder, Location loc,
   }
 
   // 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.
@@ -517,26 +595,24 @@ 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,
-                                         ArrayRef<Value> reduc) {
+void SparseTensorLoopEmitter::exitCurrentLoop(RewriterBase &rewriter,
+                                              Location loc,
+                                              MutableArrayRef<Value> reduc) {
   // Clean up the values, it would help use to discover potential bug at a
   // earlier stage (instead of silently using a wrong value).
   LoopLevelInfo &loopInfo = loopStack.back();
   assert(loopInfo.tids.size() == loopInfo.dims.size());
   SmallVector<Value, 2> red;
   if (llvm::isa<scf::WhileOp>(loopInfo.loop)) {
-    red = exitCoiterationLoop(builder, loc, reduc);
+    exitCoIterationLoop(rewriter, loc, reduc);
   } else {
-    red = exitForLoop(builder, loc, reduc);
+    exitForLoop(rewriter, loc, reduc);
   }
 
   assert(loopStack.size() == loopSeqStack.size());
   loopStack.pop_back();
-  return red;
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -380,8 +380,8 @@ class SparseTensorLoopEmitter {
       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 {
@@ -452,17 +452,35 @@ class SparseTensorLoopEmitter {
                                             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.