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