MLIR  14.0.0git
AsyncParallelFor.cpp
Go to the documentation of this file.
1 //===- AsyncParallelFor.cpp - Implementation of Async Parallel For --------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements scf.parallel to scf.for + async.execute conversion pass.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include <utility>
14 
15 #include "PassDetail.h"
20 #include "mlir/Dialect/SCF/SCF.h"
24 #include "mlir/IR/Matchers.h"
25 #include "mlir/IR/PatternMatch.h"
28 
29 using namespace mlir;
30 using namespace mlir::async;
31 
32 #define DEBUG_TYPE "async-parallel-for"
33 
34 namespace {
35 
36 // Rewrite scf.parallel operation into multiple concurrent async.execute
37 // operations over non overlapping subranges of the original loop.
38 //
39 // Example:
40 //
41 // scf.parallel (%i, %j) = (%lbi, %lbj) to (%ubi, %ubj) step (%si, %sj) {
42 // "do_some_compute"(%i, %j): () -> ()
43 // }
44 //
45 // Converted to:
46 //
47 // // Parallel compute function that executes the parallel body region for
48 // // a subset of the parallel iteration space defined by the one-dimensional
49 // // compute block index.
50 // func parallel_compute_function(%block_index : index, %block_size : index,
51 // <parallel operation properties>, ...) {
52 // // Compute multi-dimensional loop bounds for %block_index.
53 // %block_lbi, %block_lbj = ...
54 // %block_ubi, %block_ubj = ...
55 //
56 // // Clone parallel operation body into the scf.for loop nest.
57 // scf.for %i = %blockLbi to %blockUbi {
58 // scf.for %j = block_lbj to %block_ubj {
59 // "do_some_compute"(%i, %j): () -> ()
60 // }
61 // }
62 // }
63 //
64 // And a dispatch function depending on the `asyncDispatch` option.
65 //
66 // When async dispatch is on: (pseudocode)
67 //
68 // %block_size = ... compute parallel compute block size
69 // %block_count = ... compute the number of compute blocks
70 //
71 // func @async_dispatch(%block_start : index, %block_end : index, ...) {
72 // // Keep splitting block range until we reached a range of size 1.
73 // while (%block_end - %block_start > 1) {
74 // %mid_index = block_start + (block_end - block_start) / 2;
75 // async.execute { call @async_dispatch(%mid_index, %block_end); }
76 // %block_end = %mid_index
77 // }
78 //
79 // // Call parallel compute function for a single block.
80 // call @parallel_compute_fn(%block_start, %block_size, ...);
81 // }
82 //
83 // // Launch async dispatch for [0, block_count) range.
84 // call @async_dispatch(%c0, %block_count);
85 //
86 // When async dispatch is off:
87 //
88 // %block_size = ... compute parallel compute block size
89 // %block_count = ... compute the number of compute blocks
90 //
91 // scf.for %block_index = %c0 to %block_count {
92 // call @parallel_compute_fn(%block_index, %block_size, ...)
93 // }
94 //
95 struct AsyncParallelForPass
96  : public AsyncParallelForBase<AsyncParallelForPass> {
97  AsyncParallelForPass() = default;
98 
99  AsyncParallelForPass(bool asyncDispatch, int32_t numWorkerThreads,
100  int32_t minTaskSize) {
101  this->asyncDispatch = asyncDispatch;
102  this->numWorkerThreads = numWorkerThreads;
103  this->minTaskSize = minTaskSize;
104  }
105 
106  void runOnOperation() override;
107 };
108 
109 struct AsyncParallelForRewrite : public OpRewritePattern<scf::ParallelOp> {
110 public:
111  AsyncParallelForRewrite(
112  MLIRContext *ctx, bool asyncDispatch, int32_t numWorkerThreads,
113  AsyncMinTaskSizeComputationFunction computeMinTaskSize)
114  : OpRewritePattern(ctx), asyncDispatch(asyncDispatch),
115  numWorkerThreads(numWorkerThreads),
116  computeMinTaskSize(std::move(computeMinTaskSize)) {}
117 
118  LogicalResult matchAndRewrite(scf::ParallelOp op,
119  PatternRewriter &rewriter) const override;
120 
121 private:
122  bool asyncDispatch;
123  int32_t numWorkerThreads;
124  AsyncMinTaskSizeComputationFunction computeMinTaskSize;
125 };
126 
127 struct ParallelComputeFunctionType {
128  FunctionType type;
129  SmallVector<Value> captures;
130 };
131 
132 // Helper struct to parse parallel compute function argument list.
133 struct ParallelComputeFunctionArgs {
134  BlockArgument blockIndex();
135  BlockArgument blockSize();
136  ArrayRef<BlockArgument> tripCounts();
137  ArrayRef<BlockArgument> lowerBounds();
138  ArrayRef<BlockArgument> upperBounds();
139  ArrayRef<BlockArgument> steps();
140  ArrayRef<BlockArgument> captures();
141 
142  unsigned numLoops;
143  ArrayRef<BlockArgument> args;
144 };
145 
146 struct ParallelComputeFunctionBounds {
147  SmallVector<IntegerAttr> tripCounts;
148  SmallVector<IntegerAttr> lowerBounds;
149  SmallVector<IntegerAttr> upperBounds;
151 };
152 
153 struct ParallelComputeFunction {
154  unsigned numLoops;
155  FuncOp func;
156  llvm::SmallVector<Value> captures;
157 };
158 
159 } // namespace
160 
161 BlockArgument ParallelComputeFunctionArgs::blockIndex() { return args[0]; }
162 BlockArgument ParallelComputeFunctionArgs::blockSize() { return args[1]; }
163 
164 ArrayRef<BlockArgument> ParallelComputeFunctionArgs::tripCounts() {
165  return args.drop_front(2).take_front(numLoops);
166 }
167 
168 ArrayRef<BlockArgument> ParallelComputeFunctionArgs::lowerBounds() {
169  return args.drop_front(2 + 1 * numLoops).take_front(numLoops);
170 }
171 
172 ArrayRef<BlockArgument> ParallelComputeFunctionArgs::upperBounds() {
173  return args.drop_front(2 + 2 * numLoops).take_front(numLoops);
174 }
175 
176 ArrayRef<BlockArgument> ParallelComputeFunctionArgs::steps() {
177  return args.drop_front(2 + 3 * numLoops).take_front(numLoops);
178 }
179 
180 ArrayRef<BlockArgument> ParallelComputeFunctionArgs::captures() {
181  return args.drop_front(2 + 4 * numLoops);
182 }
183 
184 template <typename ValueRange>
186  SmallVector<IntegerAttr> attrs(values.size());
187  for (unsigned i = 0; i < values.size(); ++i)
188  matchPattern(values[i], m_Constant(&attrs[i]));
189  return attrs;
190 }
191 
192 // Converts one-dimensional iteration index in the [0, tripCount) interval
193 // into multidimensional iteration coordinate.
195  ArrayRef<Value> tripCounts) {
196  SmallVector<Value> coords(tripCounts.size());
197  assert(!tripCounts.empty() && "tripCounts must be not empty");
198 
199  for (ssize_t i = tripCounts.size() - 1; i >= 0; --i) {
200  coords[i] = b.create<arith::RemSIOp>(index, tripCounts[i]);
201  index = b.create<arith::DivSIOp>(index, tripCounts[i]);
202  }
203 
204  return coords;
205 }
206 
207 // Returns a function type and implicit captures for a parallel compute
208 // function. We'll need a list of implicit captures to setup block and value
209 // mapping when we'll clone the body of the parallel operation.
210 static ParallelComputeFunctionType
211 getParallelComputeFunctionType(scf::ParallelOp op, PatternRewriter &rewriter) {
212  // Values implicitly captured by the parallel operation.
213  llvm::SetVector<Value> captures;
214  getUsedValuesDefinedAbove(op.getRegion(), op.getRegion(), captures);
215 
216  SmallVector<Type> inputs;
217  inputs.reserve(2 + 4 * op.getNumLoops() + captures.size());
218 
219  Type indexTy = rewriter.getIndexType();
220 
221  // One-dimensional iteration space defined by the block index and size.
222  inputs.push_back(indexTy); // blockIndex
223  inputs.push_back(indexTy); // blockSize
224 
225  // Multi-dimensional parallel iteration space defined by the loop trip counts.
226  for (unsigned i = 0; i < op.getNumLoops(); ++i)
227  inputs.push_back(indexTy); // loop tripCount
228 
229  // Parallel operation lower bound, upper bound and step. Lower bound, upper
230  // bound and step passed as contiguous arguments:
231  // call @compute(%lb0, %lb1, ..., %ub0, %ub1, ..., %step0, %step1, ...)
232  for (unsigned i = 0; i < op.getNumLoops(); ++i) {
233  inputs.push_back(indexTy); // lower bound
234  inputs.push_back(indexTy); // upper bound
235  inputs.push_back(indexTy); // step
236  }
237 
238  // Types of the implicit captures.
239  for (Value capture : captures)
240  inputs.push_back(capture.getType());
241 
242  // Convert captures to vector for later convenience.
243  SmallVector<Value> capturesVector(captures.begin(), captures.end());
244  return {rewriter.getFunctionType(inputs, TypeRange()), capturesVector};
245 }
246 
247 // Create a parallel compute fuction from the parallel operation.
248 static ParallelComputeFunction createParallelComputeFunction(
249  scf::ParallelOp op, const ParallelComputeFunctionBounds &bounds,
250  unsigned numBlockAlignedInnerLoops, PatternRewriter &rewriter) {
251  OpBuilder::InsertionGuard guard(rewriter);
252  ImplicitLocOpBuilder b(op.getLoc(), rewriter);
253 
254  ModuleOp module = op->getParentOfType<ModuleOp>();
255 
256  ParallelComputeFunctionType computeFuncType =
257  getParallelComputeFunctionType(op, rewriter);
258 
259  FunctionType type = computeFuncType.type;
260  FuncOp func = FuncOp::create(op.getLoc(),
261  numBlockAlignedInnerLoops > 0
262  ? "parallel_compute_fn_with_aligned_loops"
263  : "parallel_compute_fn",
264  type);
265  func.setPrivate();
266 
267  // Insert function into the module symbol table and assign it unique name.
268  SymbolTable symbolTable(module);
269  symbolTable.insert(func);
270  rewriter.getListener()->notifyOperationInserted(func);
271 
272  // Create function entry block.
273  Block *block =
274  b.createBlock(&func.getBody(), func.begin(), type.getInputs(),
275  SmallVector<Location>(type.getNumInputs(), op.getLoc()));
276  b.setInsertionPointToEnd(block);
277 
278  ParallelComputeFunctionArgs args = {op.getNumLoops(), func.getArguments()};
279 
280  // Block iteration position defined by the block index and size.
281  BlockArgument blockIndex = args.blockIndex();
282  BlockArgument blockSize = args.blockSize();
283 
284  // Constants used below.
285  Value c0 = b.create<arith::ConstantIndexOp>(0);
286  Value c1 = b.create<arith::ConstantIndexOp>(1);
287 
288  // Materialize known constants as constant operation in the function body.
289  auto values = [&](ArrayRef<BlockArgument> args, ArrayRef<IntegerAttr> attrs) {
290  return llvm::to_vector(
291  llvm::map_range(llvm::zip(args, attrs), [&](auto tuple) -> Value {
292  if (IntegerAttr attr = std::get<1>(tuple))
293  return b.create<ConstantOp>(attr);
294  return std::get<0>(tuple);
295  }));
296  };
297 
298  // Multi-dimensional parallel iteration space defined by the loop trip counts.
299  auto tripCounts = values(args.tripCounts(), bounds.tripCounts);
300 
301  // Parallel operation lower bound and step.
302  auto lowerBounds = values(args.lowerBounds(), bounds.lowerBounds);
303  auto steps = values(args.steps(), bounds.steps);
304 
305  // Remaining arguments are implicit captures of the parallel operation.
306  ArrayRef<BlockArgument> captures = args.captures();
307 
308  // Compute a product of trip counts to get the size of the flattened
309  // one-dimensional iteration space.
310  Value tripCount = tripCounts[0];
311  for (unsigned i = 1; i < tripCounts.size(); ++i)
312  tripCount = b.create<arith::MulIOp>(tripCount, tripCounts[i]);
313 
314  // Find one-dimensional iteration bounds: [blockFirstIndex, blockLastIndex]:
315  // blockFirstIndex = blockIndex * blockSize
316  Value blockFirstIndex = b.create<arith::MulIOp>(blockIndex, blockSize);
317 
318  // The last one-dimensional index in the block defined by the `blockIndex`:
319  // blockLastIndex = min(blockFirstIndex + blockSize, tripCount) - 1
320  Value blockEnd0 = b.create<arith::AddIOp>(blockFirstIndex, blockSize);
321  Value blockEnd1 = b.create<arith::MinSIOp>(blockEnd0, tripCount);
322  Value blockLastIndex = b.create<arith::SubIOp>(blockEnd1, c1);
323 
324  // Convert one-dimensional indices to multi-dimensional coordinates.
325  auto blockFirstCoord = delinearize(b, blockFirstIndex, tripCounts);
326  auto blockLastCoord = delinearize(b, blockLastIndex, tripCounts);
327 
328  // Compute loops upper bounds derived from the block last coordinates:
329  // blockEndCoord[i] = blockLastCoord[i] + 1
330  //
331  // Block first and last coordinates can be the same along the outer compute
332  // dimension when inner compute dimension contains multiple blocks.
333  SmallVector<Value> blockEndCoord(op.getNumLoops());
334  for (size_t i = 0; i < blockLastCoord.size(); ++i)
335  blockEndCoord[i] = b.create<arith::AddIOp>(blockLastCoord[i], c1);
336 
337  // Construct a loop nest out of scf.for operations that will iterate over
338  // all coordinates in [blockFirstCoord, blockLastCoord] range.
339  using LoopBodyBuilder =
340  std::function<void(OpBuilder &, Location, Value, ValueRange)>;
341  using LoopNestBuilder = std::function<LoopBodyBuilder(size_t loopIdx)>;
342 
343  // Parallel region induction variables computed from the multi-dimensional
344  // iteration coordinate using parallel operation bounds and step:
345  //
346  // computeBlockInductionVars[loopIdx] =
347  // lowerBound[loopIdx] + blockCoord[loopIdx] * step[loopIdx]
348  SmallVector<Value> computeBlockInductionVars(op.getNumLoops());
349 
350  // We need to know if we are in the first or last iteration of the
351  // multi-dimensional loop for each loop in the nest, so we can decide what
352  // loop bounds should we use for the nested loops: bounds defined by compute
353  // block interval, or bounds defined by the parallel operation.
354  //
355  // Example: 2d parallel operation
356  // i j
357  // loop sizes: [50, 50]
358  // first coord: [25, 25]
359  // last coord: [30, 30]
360  //
361  // If `i` is equal to 25 then iteration over `j` should start at 25, when `i`
362  // is between 25 and 30 it should start at 0. The upper bound for `j` should
363  // be 50, except when `i` is equal to 30, then it should also be 30.
364  //
365  // Value at ith position specifies if all loops in [0, i) range of the loop
366  // nest are in the first/last iteration.
367  SmallVector<Value> isBlockFirstCoord(op.getNumLoops());
368  SmallVector<Value> isBlockLastCoord(op.getNumLoops());
369 
370  // Builds inner loop nest inside async.execute operation that does all the
371  // work concurrently.
372  LoopNestBuilder workLoopBuilder = [&](size_t loopIdx) -> LoopBodyBuilder {
373  return [&, loopIdx](OpBuilder &nestedBuilder, Location loc, Value iv,
374  ValueRange args) {
375  ImplicitLocOpBuilder nb(loc, nestedBuilder);
376 
377  // Compute induction variable for `loopIdx`.
378  computeBlockInductionVars[loopIdx] = nb.create<arith::AddIOp>(
379  lowerBounds[loopIdx], nb.create<arith::MulIOp>(iv, steps[loopIdx]));
380 
381  // Check if we are inside first or last iteration of the loop.
382  isBlockFirstCoord[loopIdx] = nb.create<arith::CmpIOp>(
383  arith::CmpIPredicate::eq, iv, blockFirstCoord[loopIdx]);
384  isBlockLastCoord[loopIdx] = nb.create<arith::CmpIOp>(
385  arith::CmpIPredicate::eq, iv, blockLastCoord[loopIdx]);
386 
387  // Check if the previous loop is in its first or last iteration.
388  if (loopIdx > 0) {
389  isBlockFirstCoord[loopIdx] = nb.create<arith::AndIOp>(
390  isBlockFirstCoord[loopIdx], isBlockFirstCoord[loopIdx - 1]);
391  isBlockLastCoord[loopIdx] = nb.create<arith::AndIOp>(
392  isBlockLastCoord[loopIdx], isBlockLastCoord[loopIdx - 1]);
393  }
394 
395  // Keep building loop nest.
396  if (loopIdx < op.getNumLoops() - 1) {
397  if (loopIdx + 1 >= op.getNumLoops() - numBlockAlignedInnerLoops) {
398  // For block aligned loops we always iterate starting from 0 up to
399  // the loop trip counts.
400  nb.create<scf::ForOp>(c0, tripCounts[loopIdx + 1], c1, ValueRange(),
401  workLoopBuilder(loopIdx + 1));
402 
403  } else {
404  // Select nested loop lower/upper bounds depending on our position in
405  // the multi-dimensional iteration space.
406  auto lb = nb.create<SelectOp>(isBlockFirstCoord[loopIdx],
407  blockFirstCoord[loopIdx + 1], c0);
408 
409  auto ub = nb.create<SelectOp>(isBlockLastCoord[loopIdx],
410  blockEndCoord[loopIdx + 1],
411  tripCounts[loopIdx + 1]);
412 
413  nb.create<scf::ForOp>(lb, ub, c1, ValueRange(),
414  workLoopBuilder(loopIdx + 1));
415  }
416 
417  nb.create<scf::YieldOp>(loc);
418  return;
419  }
420 
421  // Copy the body of the parallel op into the inner-most loop.
422  BlockAndValueMapping mapping;
423  mapping.map(op.getInductionVars(), computeBlockInductionVars);
424  mapping.map(computeFuncType.captures, captures);
425 
426  for (auto &bodyOp : op.getLoopBody().getOps())
427  nb.clone(bodyOp, mapping);
428  };
429  };
430 
431  b.create<scf::ForOp>(blockFirstCoord[0], blockEndCoord[0], c1, ValueRange(),
432  workLoopBuilder(0));
433  b.create<ReturnOp>(ValueRange());
434 
435  return {op.getNumLoops(), func, std::move(computeFuncType.captures)};
436 }
437 
438 // Creates recursive async dispatch function for the given parallel compute
439 // function. Dispatch function keeps splitting block range into halves until it
440 // reaches a single block, and then excecutes it inline.
441 //
442 // Function pseudocode (mix of C++ and MLIR):
443 //
444 // func @async_dispatch(%block_start : index, %block_end : index, ...) {
445 //
446 // // Keep splitting block range until we reached a range of size 1.
447 // while (%block_end - %block_start > 1) {
448 // %mid_index = block_start + (block_end - block_start) / 2;
449 // async.execute { call @async_dispatch(%mid_index, %block_end); }
450 // %block_end = %mid_index
451 // }
452 //
453 // // Call parallel compute function for a single block.
454 // call @parallel_compute_fn(%block_start, %block_size, ...);
455 // }
456 //
457 static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc,
458  PatternRewriter &rewriter) {
459  OpBuilder::InsertionGuard guard(rewriter);
460  Location loc = computeFunc.func.getLoc();
461  ImplicitLocOpBuilder b(loc, rewriter);
462 
463  ModuleOp module = computeFunc.func->getParentOfType<ModuleOp>();
464 
465  ArrayRef<Type> computeFuncInputTypes =
466  computeFunc.func.type().cast<FunctionType>().getInputs();
467 
468  // Compared to the parallel compute function async dispatch function takes
469  // additional !async.group argument. Also instead of a single `blockIndex` it
470  // takes `blockStart` and `blockEnd` arguments to define the range of
471  // dispatched blocks.
472  SmallVector<Type> inputTypes;
473  inputTypes.push_back(async::GroupType::get(rewriter.getContext()));
474  inputTypes.push_back(rewriter.getIndexType()); // add blockStart argument
475  inputTypes.append(computeFuncInputTypes.begin(), computeFuncInputTypes.end());
476 
477  FunctionType type = rewriter.getFunctionType(inputTypes, TypeRange());
478  FuncOp func = FuncOp::create(loc, "async_dispatch_fn", type);
479  func.setPrivate();
480 
481  // Insert function into the module symbol table and assign it unique name.
482  SymbolTable symbolTable(module);
483  symbolTable.insert(func);
484  rewriter.getListener()->notifyOperationInserted(func);
485 
486  // Create function entry block.
487  Block *block = b.createBlock(&func.getBody(), func.begin(), type.getInputs(),
488  SmallVector<Location>(type.getNumInputs(), loc));
489  b.setInsertionPointToEnd(block);
490 
491  Type indexTy = b.getIndexType();
494 
495  // Get the async group that will track async dispatch completion.
496  Value group = block->getArgument(0);
497 
498  // Get the block iteration range: [blockStart, blockEnd)
499  Value blockStart = block->getArgument(1);
500  Value blockEnd = block->getArgument(2);
501 
502  // Create a work splitting while loop for the [blockStart, blockEnd) range.
503  SmallVector<Type> types = {indexTy, indexTy};
504  SmallVector<Value> operands = {blockStart, blockEnd};
505  SmallVector<Location> locations = {loc, loc};
506 
507  // Create a recursive dispatch loop.
508  scf::WhileOp whileOp = b.create<scf::WhileOp>(types, operands);
509  Block *before = b.createBlock(&whileOp.getBefore(), {}, types, locations);
510  Block *after = b.createBlock(&whileOp.getAfter(), {}, types, locations);
511 
512  // Setup dispatch loop condition block: decide if we need to go into the
513  // `after` block and launch one more async dispatch.
514  {
515  b.setInsertionPointToEnd(before);
516  Value start = before->getArgument(0);
517  Value end = before->getArgument(1);
518  Value distance = b.create<arith::SubIOp>(end, start);
519  Value dispatch =
520  b.create<arith::CmpIOp>(arith::CmpIPredicate::sgt, distance, c1);
521  b.create<scf::ConditionOp>(dispatch, before->getArguments());
522  }
523 
524  // Setup the async dispatch loop body: recursively call dispatch function
525  // for the seconds half of the original range and go to the next iteration.
526  {
527  b.setInsertionPointToEnd(after);
528  Value start = after->getArgument(0);
529  Value end = after->getArgument(1);
530  Value distance = b.create<arith::SubIOp>(end, start);
531  Value halfDistance = b.create<arith::DivSIOp>(distance, c2);
532  Value midIndex = b.create<arith::AddIOp>(start, halfDistance);
533 
534  // Call parallel compute function inside the async.execute region.
535  auto executeBodyBuilder = [&](OpBuilder &executeBuilder,
536  Location executeLoc, ValueRange executeArgs) {
537  // Update the original `blockStart` and `blockEnd` with new range.
538  SmallVector<Value> operands{block->getArguments().begin(),
539  block->getArguments().end()};
540  operands[1] = midIndex;
541  operands[2] = end;
542 
543  executeBuilder.create<CallOp>(executeLoc, func.sym_name(),
544  func.getCallableResults(), operands);
545  executeBuilder.create<async::YieldOp>(executeLoc, ValueRange());
546  };
547 
548  // Create async.execute operation to dispatch half of the block range.
549  auto execute = b.create<ExecuteOp>(TypeRange(), ValueRange(), ValueRange(),
550  executeBodyBuilder);
551  b.create<AddToGroupOp>(indexTy, execute.token(), group);
552  b.create<scf::YieldOp>(ValueRange({start, midIndex}));
553  }
554 
555  // After dispatching async operations to process the tail of the block range
556  // call the parallel compute function for the first block of the range.
557  b.setInsertionPointAfter(whileOp);
558 
559  // Drop async dispatch specific arguments: async group, block start and end.
560  auto forwardedInputs = block->getArguments().drop_front(3);
561  SmallVector<Value> computeFuncOperands = {blockStart};
562  computeFuncOperands.append(forwardedInputs.begin(), forwardedInputs.end());
563 
564  b.create<CallOp>(computeFunc.func.sym_name(),
565  computeFunc.func.getCallableResults(), computeFuncOperands);
566  b.create<ReturnOp>(ValueRange());
567 
568  return func;
569 }
570 
571 // Launch async dispatch of the parallel compute function.
573  ParallelComputeFunction &parallelComputeFunction,
574  scf::ParallelOp op, Value blockSize,
575  Value blockCount,
576  const SmallVector<Value> &tripCounts) {
577  MLIRContext *ctx = op->getContext();
578 
579  // Add one more level of indirection to dispatch parallel compute functions
580  // using async operations and recursive work splitting.
581  FuncOp asyncDispatchFunction =
582  createAsyncDispatchFunction(parallelComputeFunction, rewriter);
583 
586 
587  // Appends operands shared by async dispatch and parallel compute functions to
588  // the given operands vector.
589  auto appendBlockComputeOperands = [&](SmallVector<Value> &operands) {
590  operands.append(tripCounts);
591  operands.append(op.getLowerBound().begin(), op.getLowerBound().end());
592  operands.append(op.getUpperBound().begin(), op.getUpperBound().end());
593  operands.append(op.getStep().begin(), op.getStep().end());
594  operands.append(parallelComputeFunction.captures);
595  };
596 
597  // Check if the block size is one, in this case we can skip the async dispatch
598  // completely. If this will be known statically, then canonicalization will
599  // erase async group operations.
600  Value isSingleBlock =
601  b.create<arith::CmpIOp>(arith::CmpIPredicate::eq, blockCount, c1);
602 
603  auto syncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
604  ImplicitLocOpBuilder nb(loc, nestedBuilder);
605 
606  // Call parallel compute function for the single block.
607  SmallVector<Value> operands = {c0, blockSize};
608  appendBlockComputeOperands(operands);
609 
610  nb.create<CallOp>(parallelComputeFunction.func.sym_name(),
611  parallelComputeFunction.func.getCallableResults(),
612  operands);
613  nb.create<scf::YieldOp>();
614  };
615 
616  auto asyncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
617  // Create an async.group to wait on all async tokens from the concurrent
618  // execution of multiple parallel compute function. First block will be
619  // executed synchronously in the caller thread.
620  Value groupSize = b.create<arith::SubIOp>(blockCount, c1);
621  Value group = b.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
622 
623  ImplicitLocOpBuilder nb(loc, nestedBuilder);
624 
625  // Launch async dispatch function for [0, blockCount) range.
626  SmallVector<Value> operands = {group, c0, blockCount, blockSize};
627  appendBlockComputeOperands(operands);
628 
629  nb.create<CallOp>(asyncDispatchFunction.sym_name(),
630  asyncDispatchFunction.getCallableResults(), operands);
631 
632  // Wait for the completion of all parallel compute operations.
633  b.create<AwaitAllOp>(group);
634 
635  nb.create<scf::YieldOp>();
636  };
637 
638  // Dispatch either single block compute function, or launch async dispatch.
639  b.create<scf::IfOp>(TypeRange(), isSingleBlock, syncDispatch, asyncDispatch);
640 }
641 
642 // Dispatch parallel compute functions by submitting all async compute tasks
643 // from a simple for loop in the caller thread.
644 static void
646  ParallelComputeFunction &parallelComputeFunction,
647  scf::ParallelOp op, Value blockSize, Value blockCount,
648  const SmallVector<Value> &tripCounts) {
649  MLIRContext *ctx = op->getContext();
650 
651  FuncOp compute = parallelComputeFunction.func;
652 
655 
656  // Create an async.group to wait on all async tokens from the concurrent
657  // execution of multiple parallel compute function. First block will be
658  // executed synchronously in the caller thread.
659  Value groupSize = b.create<arith::SubIOp>(blockCount, c1);
660  Value group = b.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
661 
662  // Call parallel compute function for all blocks.
663  using LoopBodyBuilder =
664  std::function<void(OpBuilder &, Location, Value, ValueRange)>;
665 
666  // Returns parallel compute function operands to process the given block.
667  auto computeFuncOperands = [&](Value blockIndex) -> SmallVector<Value> {
668  SmallVector<Value> computeFuncOperands = {blockIndex, blockSize};
669  computeFuncOperands.append(tripCounts);
670  computeFuncOperands.append(op.getLowerBound().begin(),
671  op.getLowerBound().end());
672  computeFuncOperands.append(op.getUpperBound().begin(),
673  op.getUpperBound().end());
674  computeFuncOperands.append(op.getStep().begin(), op.getStep().end());
675  computeFuncOperands.append(parallelComputeFunction.captures);
676  return computeFuncOperands;
677  };
678 
679  // Induction variable is the index of the block: [0, blockCount).
680  LoopBodyBuilder loopBuilder = [&](OpBuilder &loopBuilder, Location loc,
681  Value iv, ValueRange args) {
682  ImplicitLocOpBuilder nb(loc, loopBuilder);
683 
684  // Call parallel compute function inside the async.execute region.
685  auto executeBodyBuilder = [&](OpBuilder &executeBuilder,
686  Location executeLoc, ValueRange executeArgs) {
687  executeBuilder.create<CallOp>(executeLoc, compute.sym_name(),
688  compute.getCallableResults(),
689  computeFuncOperands(iv));
690  executeBuilder.create<async::YieldOp>(executeLoc, ValueRange());
691  };
692 
693  // Create async.execute operation to launch parallel computate function.
694  auto execute = nb.create<ExecuteOp>(TypeRange(), ValueRange(), ValueRange(),
695  executeBodyBuilder);
696  nb.create<AddToGroupOp>(rewriter.getIndexType(), execute.token(), group);
697  nb.create<scf::YieldOp>();
698  };
699 
700  // Iterate over all compute blocks and launch parallel compute operations.
701  b.create<scf::ForOp>(c1, blockCount, c1, ValueRange(), loopBuilder);
702 
703  // Call parallel compute function for the first block in the caller thread.
704  b.create<CallOp>(compute.sym_name(), compute.getCallableResults(),
705  computeFuncOperands(c0));
706 
707  // Wait for the completion of all async compute operations.
708  b.create<AwaitAllOp>(group);
709 }
710 
712 AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
713  PatternRewriter &rewriter) const {
714  // We do not currently support rewrite for parallel op with reductions.
715  if (op.getNumReductions() != 0)
716  return failure();
717 
718  ImplicitLocOpBuilder b(op.getLoc(), rewriter);
719 
720  // Computing minTaskSize emits IR and can be implemented as executing a cost
721  // model on the body of the scf.parallel. Thus it needs to be computed before
722  // the body of the scf.parallel has been manipulated.
723  Value minTaskSize = computeMinTaskSize(b, op);
724 
725  // Make sure that all constants will be inside the parallel operation body to
726  // reduce the number of parallel compute function arguments.
727  cloneConstantsIntoTheRegion(op.getLoopBody(), rewriter);
728 
729  // Compute trip count for each loop induction variable:
730  // tripCount = ceil_div(upperBound - lowerBound, step);
731  SmallVector<Value> tripCounts(op.getNumLoops());
732  for (size_t i = 0; i < op.getNumLoops(); ++i) {
733  auto lb = op.getLowerBound()[i];
734  auto ub = op.getUpperBound()[i];
735  auto step = op.getStep()[i];
736  auto range = b.createOrFold<arith::SubIOp>(ub, lb);
737  tripCounts[i] = b.createOrFold<arith::CeilDivSIOp>(range, step);
738  }
739 
740  // Compute a product of trip counts to get the 1-dimensional iteration space
741  // for the scf.parallel operation.
742  Value tripCount = tripCounts[0];
743  for (size_t i = 1; i < tripCounts.size(); ++i)
744  tripCount = b.create<arith::MulIOp>(tripCount, tripCounts[i]);
745 
746  // Short circuit no-op parallel loops (zero iterations) that can arise from
747  // the memrefs with dynamic dimension(s) equal to zero.
748  Value c0 = b.create<arith::ConstantIndexOp>(0);
749  Value isZeroIterations =
750  b.create<arith::CmpIOp>(arith::CmpIPredicate::eq, tripCount, c0);
751 
752  // Do absolutely nothing if the trip count is zero.
753  auto noOp = [&](OpBuilder &nestedBuilder, Location loc) {
754  nestedBuilder.create<scf::YieldOp>(loc);
755  };
756 
757  // Compute the parallel block size and dispatch concurrent tasks computing
758  // results for each block.
759  auto dispatch = [&](OpBuilder &nestedBuilder, Location loc) {
760  ImplicitLocOpBuilder nb(loc, nestedBuilder);
761 
762  // Collect statically known constants defining the loop nest in the parallel
763  // compute function. LLVM can't always push constants across the non-trivial
764  // async dispatch call graph, by providing these values explicitly we can
765  // choose to build more efficient loop nest, and rely on a better constant
766  // folding, loop unrolling and vectorization.
767  ParallelComputeFunctionBounds staticBounds = {
768  integerConstants(tripCounts),
769  integerConstants(op.getLowerBound()),
770  integerConstants(op.getUpperBound()),
771  integerConstants(op.getStep()),
772  };
773 
774  // Find how many inner iteration dimensions are statically known, and their
775  // product is smaller than the `512`. We align the parallel compute block
776  // size by the product of statically known dimensions, so that we can
777  // guarantee that the inner loops executes from 0 to the loop trip counts
778  // and we can elide dynamic loop boundaries, and give LLVM an opportunity to
779  // unroll the loops. The constant `512` is arbitrary, it should depend on
780  // how many iterations LLVM will typically decide to unroll.
781  static constexpr int64_t maxIterations = 512;
782 
783  // The number of inner loops with statically known number of iterations less
784  // than the `maxIterations` value.
785  int numUnrollableLoops = 0;
786 
787  auto getInt = [](IntegerAttr attr) { return attr ? attr.getInt() : 0; };
788 
789  SmallVector<int64_t> numIterations(op.getNumLoops());
790  numIterations.back() = getInt(staticBounds.tripCounts.back());
791 
792  for (int i = op.getNumLoops() - 2; i >= 0; --i) {
793  int64_t tripCount = getInt(staticBounds.tripCounts[i]);
794  int64_t innerIterations = numIterations[i + 1];
795  numIterations[i] = tripCount * innerIterations;
796 
797  // Update the number of inner loops that we can potentially unroll.
798  if (innerIterations > 0 && innerIterations <= maxIterations)
799  numUnrollableLoops++;
800  }
801 
802  // With large number of threads the value of creating many compute blocks
803  // is reduced because the problem typically becomes memory bound. For small
804  // number of threads it helps with stragglers.
805  float overshardingFactor = numWorkerThreads <= 4 ? 8.0
806  : numWorkerThreads <= 8 ? 4.0
807  : numWorkerThreads <= 16 ? 2.0
808  : numWorkerThreads <= 32 ? 1.0
809  : numWorkerThreads <= 64 ? 0.8
810  : 0.6;
811 
812  // Do not overload worker threads with too many compute blocks.
813  Value maxComputeBlocks = b.create<arith::ConstantIndexOp>(
814  std::max(1, static_cast<int>(numWorkerThreads * overshardingFactor)));
815 
816  // Compute parallel block size from the parallel problem size:
817  // blockSize = min(tripCount,
818  // max(ceil_div(tripCount, maxComputeBlocks),
819  // minTaskSize))
820  Value bs0 = b.create<arith::CeilDivSIOp>(tripCount, maxComputeBlocks);
821  Value bs1 = b.create<arith::MaxSIOp>(bs0, minTaskSize);
822  Value blockSize = b.create<arith::MinSIOp>(tripCount, bs1);
823 
824  ParallelComputeFunction notUnrollableParallelComputeFunction =
825  createParallelComputeFunction(op, staticBounds, 0, rewriter);
826 
827  // Dispatch parallel compute function using async recursive work splitting,
828  // or by submitting compute task sequentially from a caller thread.
829  auto doDispatch = asyncDispatch ? doAsyncDispatch : doSequentialDispatch;
830 
831  // Create a parallel compute function that takes a block id and computes
832  // the parallel operation body for a subset of iteration space.
833 
834  // Compute the number of parallel compute blocks.
835  Value blockCount = b.create<arith::CeilDivSIOp>(tripCount, blockSize);
836 
837  // Unroll when numUnrollableLoops > 0 && blockSize >= maxIterations.
838  bool staticShouldUnroll = numUnrollableLoops > 0;
839  auto dispatchNotUnrollable = [&](OpBuilder &nestedBuilder, Location loc) {
840  ImplicitLocOpBuilder nb(loc, nestedBuilder);
841  doDispatch(b, rewriter, notUnrollableParallelComputeFunction, op,
842  blockSize, blockCount, tripCounts);
843  nb.create<scf::YieldOp>();
844  };
845 
846  if (staticShouldUnroll) {
847  Value dynamicShouldUnroll = b.create<arith::CmpIOp>(
848  arith::CmpIPredicate::sge, blockSize,
849  b.create<arith::ConstantIndexOp>(maxIterations));
850 
851  ParallelComputeFunction unrollableParallelComputeFunction =
852  createParallelComputeFunction(op, staticBounds, numUnrollableLoops,
853  rewriter);
854 
855  auto dispatchUnrollable = [&](OpBuilder &nestedBuilder, Location loc) {
856  ImplicitLocOpBuilder nb(loc, nestedBuilder);
857  // Align the block size to be a multiple of the statically known
858  // number of iterations in the inner loops.
859  Value numIters = nb.create<arith::ConstantIndexOp>(
860  numIterations[op.getNumLoops() - numUnrollableLoops]);
861  Value alignedBlockSize = nb.create<arith::MulIOp>(
862  nb.create<arith::CeilDivSIOp>(blockSize, numIters), numIters);
863  doDispatch(b, rewriter, unrollableParallelComputeFunction, op,
864  alignedBlockSize, blockCount, tripCounts);
865  nb.create<scf::YieldOp>();
866  };
867 
868  b.create<scf::IfOp>(TypeRange(), dynamicShouldUnroll, dispatchUnrollable,
869  dispatchNotUnrollable);
870  nb.create<scf::YieldOp>();
871  } else {
872  dispatchNotUnrollable(nb, loc);
873  }
874  };
875 
876  // Replace the `scf.parallel` operation with the parallel compute function.
877  b.create<scf::IfOp>(TypeRange(), isZeroIterations, noOp, dispatch);
878 
879  // Parallel operation was replaced with a block iteration loop.
880  rewriter.eraseOp(op);
881 
882  return success();
883 }
884 
885 void AsyncParallelForPass::runOnOperation() {
886  MLIRContext *ctx = &getContext();
887 
888  RewritePatternSet patterns(ctx);
890  patterns, asyncDispatch, numWorkerThreads,
891  [&](ImplicitLocOpBuilder builder, scf::ParallelOp op) {
892  return builder.create<arith::ConstantIndexOp>(minTaskSize);
893  });
894  if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
895  signalPassFailure();
896 }
897 
898 std::unique_ptr<Pass> mlir::createAsyncParallelForPass() {
899  return std::make_unique<AsyncParallelForPass>();
900 }
901 
902 std::unique_ptr<Pass> mlir::createAsyncParallelForPass(bool asyncDispatch,
903  int32_t numWorkerThreads,
904  int32_t minTaskSize) {
905  return std::make_unique<AsyncParallelForPass>(asyncDispatch, numWorkerThreads,
906  minTaskSize);
907 }
908 
910  RewritePatternSet &patterns, bool asyncDispatch, int32_t numWorkerThreads,
911  const AsyncMinTaskSizeComputationFunction &computeMinTaskSize) {
912  MLIRContext *ctx = patterns.getContext();
913  patterns.add<AsyncParallelForRewrite>(ctx, asyncDispatch, numWorkerThreads,
914  computeMinTaskSize);
915 }
Include the generated interface declarations.
OpTy create(Location location, Args &&...args)
Create an operation of specific op type at the current insertion point.
Definition: Builders.h:430
static Operation * create(Location location, OperationName name, TypeRange resultTypes, ValueRange operands, ArrayRef< NamedAttribute > attributes, BlockRange successors, unsigned numRegions)
Create a new Operation with the specific fields.
Definition: Operation.cpp:27
MLIRContext * getContext() const
Definition: Builders.h:54
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:881
std::function< Value(ImplicitLocOpBuilder, scf::ParallelOp)> AsyncMinTaskSizeComputationFunction
Emit the IR to compute the minimum number of iterations of scf.parallel body that would be viable for...
Definition: Transforms.h:29
static SmallVector< IntegerAttr > integerConstants(ValueRange values)
Block represents an ordered list of Operations.
Definition: Block.h:29
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
Operation * clone(Operation &op, BlockAndValueMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
Definition: Builders.cpp:457
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
Definition: LogicalResult.h:72
virtual void notifyOperationInserted(Operation *op)
Notification handler for when an operation is inserted into the builder.
Definition: Builders.h:239
BlockArgument getArgument(unsigned i)
Definition: Block.h:120
StringAttr insert(Operation *symbol, Block::iterator insertPt={})
Insert a new symbol into the table, and rename it as necessary to avoid collisions.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:48
void map(Block *from, Block *to)
Inserts a new mapping for &#39;from&#39; to &#39;to&#39;.
void setInsertionPointAfter(Operation *op)
Sets the insertion point to the node after the specified operation, which will cause subsequent inser...
Definition: Builders.h:343
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
void getUsedValuesDefinedAbove(Region &region, Region &limit, SetVector< Value > &values)
Fill values with a list of values defined at the ancestors of the limit region and used within region...
Definition: RegionUtils.cpp:59
static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc, PatternRewriter &rewriter)
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:38
void populateAsyncParallelForPatterns(RewritePatternSet &patterns, bool asyncDispatch, int32_t numWorkerThreads, const AsyncMinTaskSizeComputationFunction &computeMinTaskSize)
Add a pattern to the given pattern list to lower scf.parallel to async operations.
Listener * getListener() const
Returns the current listener of this builder, or nullptr if this builder doesn&#39;t have a listener...
Definition: Builders.h:251
detail::constant_op_matcher m_Constant()
Matches a constant foldable operation.
Definition: Matchers.h:231
BlockArgListType getArguments()
Definition: Block.h:76
This class represents an argument of a Block.
Definition: Value.h:298
void cloneConstantsIntoTheRegion(Region &region)
Clone ConstantLike operations that are defined above the given region and have users in the region in...
Definition: PassDetail.cpp:15
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:72
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:84
SmallVector< int64_t, 4 > delinearize(ArrayRef< int64_t > strides, int64_t linearIndex)
Given the strides together with a linear index in the dimension space, returns the vector-space offse...
Definition: VectorUtils.cpp:75
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:355
OpTy create(Args &&...args)
Create an operation of specific op type at the current insertion point and location.
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:279
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&... args)
Add an instance of each of the pattern types &#39;Ts&#39; to the pattern list with the given arguments...
Definition: PatternMatch.h:930
IndexType getIndexType()
Definition: Builders.cpp:48
ImplicitLocOpBuilder maintains a &#39;current location&#39;, allowing use of the create<> method without spec...
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
Definition: Matchers.h:266
static void doSequentialDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter, ParallelComputeFunction &parallelComputeFunction, scf::ParallelOp op, Value blockSize, Value blockCount, const SmallVector< Value > &tripCounts)
Specialization of arith.constant op that returns an integer of index type.
Definition: Arithmetic.h:78
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:55
This class allows for representing and managing the symbol table used by operations with the &#39;SymbolT...
Definition: SymbolTable.h:23
void setInsertionPointToEnd(Block *block)
Sets the insertion point to the end of the specified block.
Definition: Builders.h:367
FunctionType getFunctionType(TypeRange inputs, TypeRange results)
Definition: Builders.cpp:67
Block * createBlock(Region *parent, Region::iterator insertPt={}, TypeRange argTypes=llvm::None, ArrayRef< Location > locs=llvm::None)
Add new block with &#39;argTypes&#39; arguments and set the insertion point to the end of it...
Definition: Builders.cpp:353
static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter, ParallelComputeFunction &parallelComputeFunction, scf::ParallelOp op, Value blockSize, Value blockCount, const SmallVector< Value > &tripCounts)
LogicalResult applyPatternsAndFoldGreedily(MutableArrayRef< Region > regions, const FrozenRewritePatternSet &patterns, GreedyRewriteConfig config=GreedyRewriteConfig())
Rewrite the regions of the specified operation, which must be isolated from above, by repeatedly applying the highest benefit patterns in a greedy work-list driven manner.
static ParallelComputeFunctionType getParallelComputeFunctionType(scf::ParallelOp op, PatternRewriter &rewriter)
static ParallelComputeFunction createParallelComputeFunction(scf::ParallelOp op, const ParallelComputeFunctionBounds &bounds, unsigned numBlockAlignedInnerLoops, PatternRewriter &rewriter)
This class helps build Operations.
Definition: Builders.h:177
This class provides an abstraction over the different types of ranges over Values.
MLIRContext * getContext() const
Definition: PatternMatch.h:906
std::unique_ptr< Pass > createAsyncParallelForPass()
static Value max(ImplicitLocOpBuilder &builder, Value a, Value b)