MLIR  18.0.0git
AsyncRegionRewriter.cpp
Go to the documentation of this file.
1 //===- AsyncRegionRewriter.cpp - Implementation of GPU async rewriters ----===//
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 the GPU dialect pattern rewriters that make GPU op
10 // within a region execute asynchronously.
11 //
12 //===----------------------------------------------------------------------===//
13 
15 
20 #include "mlir/IR/Builders.h"
21 #include "mlir/IR/IRMapping.h"
22 #include "mlir/IR/PatternMatch.h"
23 #include "mlir/IR/SymbolTable.h"
25 #include "mlir/Support/LLVM.h"
27 #include "llvm/ADT/TypeSwitch.h"
28 
29 namespace mlir {
30 #define GEN_PASS_DEF_GPUASYNCREGIONPASS
31 #include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
32 } // namespace mlir
33 
34 using namespace mlir;
35 
36 namespace {
37 class GpuAsyncRegionPass
38  : public impl::GpuAsyncRegionPassBase<GpuAsyncRegionPass> {
39  struct ThreadTokenCallback;
40  struct DeferWaitCallback;
41  struct SingleTokenUseCallback;
42  void runOnOperation() override;
43 };
44 } // namespace
45 
46 static bool isTerminator(Operation *op) {
48 }
49 static bool hasSideEffects(Operation *op) { return !isMemoryEffectFree(op); }
50 
51 // Region walk callback which makes GPU ops implementing the AsyncOpInterface
52 // execute asynchronously.
54  ThreadTokenCallback(MLIRContext &context) : builder(&context) {}
55 
57  for (Operation &op : make_early_inc_range(*block)) {
58  if (failed(visit(&op)))
59  return WalkResult::interrupt();
60  }
61  return WalkResult::advance();
62  }
63 
64 private:
65  // If `op` implements the AsyncOpInterface, insert a `gpu.wait async` to
66  // create a current token (unless it already exists), and 'thread' that token
67  // through the `op` so that it executes asynchronously.
68  //
69  // If `op` is a terminator or an op with side-effects, insert a `gpu.wait` to
70  // host-synchronize execution. A `!gpu.async.token` will therefore only be
71  // used inside of its block and GPU execution will always synchronize with
72  // the host at block boundaries.
74  if (isa<gpu::LaunchOp>(op))
75  return op->emitOpError("replace with gpu.launch_func first");
76  if (auto waitOp = llvm::dyn_cast<gpu::WaitOp>(op)) {
77  if (currentToken)
78  waitOp.addAsyncDependency(currentToken);
79  currentToken = waitOp.getAsyncToken();
80  return success();
81  }
82  builder.setInsertionPoint(op);
83  if (auto asyncOp = dyn_cast<gpu::AsyncOpInterface>(op))
84  return rewriteAsyncOp(asyncOp); // Replace GPU op with async version.
85  if (!currentToken)
86  return success();
87  // Insert host synchronization before terminator or op with side effects.
88  if (isTerminator(op) || hasSideEffects(op))
89  currentToken = createWaitOp(op->getLoc(), Type(), {currentToken});
90  return success();
91  }
92 
93  // Replaces asyncOp with a clone that returns a token.
94  LogicalResult rewriteAsyncOp(gpu::AsyncOpInterface asyncOp) {
95  auto *op = asyncOp.getOperation();
96  auto tokenType = builder.getType<gpu::AsyncTokenType>();
97 
98  // If there is no current token, insert a `gpu.wait async` without
99  // dependencies to create one.
100  if (!currentToken)
101  currentToken = createWaitOp(op->getLoc(), tokenType, {});
102  asyncOp.addAsyncDependency(currentToken);
103 
104  // Return early if op returns a token already.
105  currentToken = asyncOp.getAsyncToken();
106  if (currentToken)
107  return success();
108 
109  // Clone the op to return a token in addition to the other results.
110  SmallVector<Type, 1> resultTypes;
111  resultTypes.reserve(1 + op->getNumResults());
112  copy(op->getResultTypes(), std::back_inserter(resultTypes));
113  resultTypes.push_back(tokenType);
114  auto *newOp = Operation::create(
115  op->getLoc(), op->getName(), resultTypes, op->getOperands(),
117  op->getSuccessors(), op->getNumRegions());
118 
119  // Clone regions into new op.
120  IRMapping mapping;
121  for (auto pair : llvm::zip_first(op->getRegions(), newOp->getRegions()))
122  std::get<0>(pair).cloneInto(&std::get<1>(pair), mapping);
123 
124  // Replace the op with the async clone.
125  auto results = newOp->getResults();
126  currentToken = results.back();
127  builder.insert(newOp);
128  op->replaceAllUsesWith(results.drop_back());
129  op->erase();
130 
131  return success();
132  }
133 
134  Value createWaitOp(Location loc, Type resultType, ValueRange operands) {
135  return builder.create<gpu::WaitOp>(loc, resultType, operands)
136  .getAsyncToken();
137  }
138 
139  OpBuilder builder;
140 
141  // The token that represents the current asynchronous dependency. It's valid
142  // range starts with a `gpu.wait async` op, and ends with a `gpu.wait` op.
143  // In between, each gpu::AsyncOpInterface depends on the current token and
144  // produces the new one.
145  Value currentToken = {};
146 };
147 
148 /// Erases `executeOp` and returns a clone with additional `results`.
149 async::ExecuteOp addExecuteResults(async::ExecuteOp executeOp,
150  ValueRange results) {
151  // Add values to async.yield op.
152  Operation *yieldOp = executeOp.getBody()->getTerminator();
153  yieldOp->insertOperands(yieldOp->getNumOperands(), results);
154 
155  // Construct new result type list with additional types.
156  SmallVector<Type, 2> resultTypes;
157  resultTypes.reserve(executeOp.getNumResults() + results.size());
158  transform(executeOp.getResultTypes(), std::back_inserter(resultTypes),
159  [](Type type) {
160  // Extract value type from !async.value.
161  if (auto valueType = dyn_cast<async::ValueType>(type))
162  return valueType.getValueType();
163  assert(isa<async::TokenType>(type) && "expected token type");
164  return type;
165  });
166  transform(results, std::back_inserter(resultTypes),
167  [](Value value) { return value.getType(); });
168 
169  // Clone executeOp with the extra results.
170  OpBuilder builder(executeOp);
171  auto newOp = builder.create<async::ExecuteOp>(
172  executeOp.getLoc(), TypeRange{resultTypes}.drop_front() /*drop token*/,
173  executeOp.getDependencies(), executeOp.getBodyOperands());
174  IRMapping mapper;
175  newOp.getRegion().getBlocks().clear();
176  executeOp.getRegion().cloneInto(&newOp.getRegion(), mapper);
177 
178  // Replace executeOp with cloned one.
179  executeOp.getOperation()->replaceAllUsesWith(
180  newOp.getResults().drop_back(results.size()));
181  executeOp.erase();
182 
183  return newOp;
184 }
185 
186 // Callback for `async.execute` ops which tries to push the contained
187 // synchronous `gpu.wait` op to the dependencies of the `async.execute`.
189  // If the `executeOp`s token is used only in `async.execute` or `async.await`
190  // ops, add the region's last `gpu.wait` op to the worklist if it is
191  // synchronous and is the last op with side effects.
192  void operator()(async::ExecuteOp executeOp) {
193  if (!areAllUsersExecuteOrAwait(executeOp.getToken()))
194  return;
195  // async.execute's region is currently restricted to one block.
196  for (auto &op : llvm::reverse(executeOp.getBody()->without_terminator())) {
197  if (auto waitOp = dyn_cast<gpu::WaitOp>(op)) {
198  if (!waitOp.getAsyncToken())
199  worklist.push_back(waitOp);
200  return;
201  }
202  if (hasSideEffects(&op))
203  return;
204  }
205  }
206 
207  // The destructor performs the actual rewrite work.
209  for (size_t i = 0; i < worklist.size(); ++i) {
210  auto waitOp = worklist[i];
211  auto executeOp = waitOp->getParentOfType<async::ExecuteOp>();
212 
213  // Erase `gpu.wait` and return async dependencies from execute op instead.
214  SmallVector<Value, 4> dependencies = waitOp.getAsyncDependencies();
215  waitOp.erase();
216  executeOp = addExecuteResults(executeOp, dependencies);
217 
218  // Add the async dependency to each user of the `async.execute` token.
219  auto asyncTokens = executeOp.getResults().take_back(dependencies.size());
220  SmallVector<Operation *, 4> users(executeOp.getToken().user_begin(),
221  executeOp.getToken().user_end());
222  for (Operation *user : users)
223  addAsyncDependencyAfter(asyncTokens, user);
224  }
225  }
226 
227 private:
228  // Returns whether all token users are either 'async.execute' or 'async.await'
229  // ops. This is used as a requirement for pushing 'gpu.wait' ops from a
230  // 'async.execute' body to it's users. Specifically, we do not allow
231  // terminator users, because it could mean that the `async.execute` is inside
232  // control flow code.
233  static bool areAllUsersExecuteOrAwait(Value token) {
234  return !token.use_empty() &&
235  llvm::all_of(token.getUsers(), [](Operation *user) {
236  return isa<async::ExecuteOp, async::AwaitOp>(user);
237  });
238  }
239 
240  // Add the `asyncToken` as dependency as needed after `op`.
241  void addAsyncDependencyAfter(ValueRange asyncTokens, Operation *op) {
242  OpBuilder builder(op->getContext());
243  auto loc = op->getLoc();
244 
245  Block::iterator it;
246  SmallVector<Value, 1> tokens;
247  tokens.reserve(asyncTokens.size());
249  .Case<async::AwaitOp>([&](auto awaitOp) {
250  // Add async.await ops to wait for the !gpu.async.tokens.
251  builder.setInsertionPointAfter(op);
252  for (auto asyncToken : asyncTokens)
253  tokens.push_back(
254  builder.create<async::AwaitOp>(loc, asyncToken).getResult());
255  // Set `it` after the inserted async.await ops.
256  it = builder.getInsertionPoint();
257  })
258  .Case<async::ExecuteOp>([&](auto executeOp) {
259  // Set `it` to the beginning of the region and add asyncTokens to the
260  // async.execute operands.
261  it = executeOp.getBody()->begin();
262  executeOp.getBodyOperandsMutable().append(asyncTokens);
263  SmallVector<Type, 1> tokenTypes(
264  asyncTokens.size(), builder.getType<gpu::AsyncTokenType>());
265  SmallVector<Location, 1> tokenLocs(asyncTokens.size(),
266  executeOp.getLoc());
267  copy(executeOp.getBody()->addArguments(tokenTypes, tokenLocs),
268  std::back_inserter(tokens));
269  });
270 
271  // Advance `it` to terminator or op with side-effects.
272  it = std::find_if(it, Block::iterator(), [](Operation &op) {
273  return isTerminator(&op) || hasSideEffects(&op);
274  });
275 
276  // If `op` implements the AsyncOpInterface, add `token` to the list of async
277  // dependencies.
278  if (auto asyncOp = dyn_cast<gpu::AsyncOpInterface>(*it)) {
279  for (auto token : tokens)
280  asyncOp.addAsyncDependency(token);
281  return;
282  }
283 
284  // Otherwise, insert a gpu.wait before 'it'.
285  builder.setInsertionPoint(it->getBlock(), it);
286  auto waitOp = builder.create<gpu::WaitOp>(loc, Type{}, tokens);
287 
288  // If the new waitOp is at the end of an async.execute region, add it to the
289  // worklist. 'operator()(executeOp)' would do the same, but this is faster.
290  auto executeOp = dyn_cast<async::ExecuteOp>(it->getParentOp());
291  if (executeOp && areAllUsersExecuteOrAwait(executeOp.getToken()) &&
292  !it->getNextNode())
293  worklist.push_back(waitOp);
294  }
295 
297 };
298 
299 // Callback for `async.execute` ops which repeats !gpu.async.token results
300 // so that each of them is only used once.
302  void operator()(async::ExecuteOp executeOp) {
303  // Extract !gpu.async.token results which have multiple uses.
304  auto multiUseResults = llvm::make_filter_range(
305  executeOp.getBodyResults(), [](OpResult result) {
306  if (result.use_empty() || result.hasOneUse())
307  return false;
308  auto valueType = dyn_cast<async::ValueType>(result.getType());
309  return valueType &&
310  isa<gpu::AsyncTokenType>(valueType.getValueType());
311  });
312  if (multiUseResults.empty())
313  return;
314 
315  // Indices within !async.execute results (i.e. without the async.token).
316  SmallVector<int, 4> indices;
317  transform(multiUseResults, std::back_inserter(indices),
318  [](OpResult result) {
319  return result.getResultNumber() - 1; // Index without token.
320  });
321 
322  for (auto index : indices) {
323  assert(!executeOp.getBodyResults()[index].getUses().empty());
324  // Repeat async.yield token result, one for each use after the first one.
325  auto uses = llvm::drop_begin(executeOp.getBodyResults()[index].getUses());
326  auto count = std::distance(uses.begin(), uses.end());
327  auto yieldOp = cast<async::YieldOp>(executeOp.getBody()->getTerminator());
328  SmallVector<Value, 4> operands(count, yieldOp.getOperand(index));
329  executeOp = addExecuteResults(executeOp, operands);
330  // Update 'uses' to refer to the new executeOp.
331  uses = llvm::drop_begin(executeOp.getBodyResults()[index].getUses());
332  auto results = executeOp.getBodyResults().take_back(count);
333  for (auto pair : llvm::zip(uses, results))
334  std::get<0>(pair).set(std::get<1>(pair));
335  }
336  }
337 };
338 
339 // Replaces synchronous GPU ops in the op's region with asynchronous ones and
340 // inserts the necessary synchronization (as gpu.wait ops). Assumes sequential
341 // execution semantics and that no GPU ops are asynchronous yet.
342 void GpuAsyncRegionPass::runOnOperation() {
343  if (getOperation()->walk(ThreadTokenCallback(getContext())).wasInterrupted())
344  return signalPassFailure();
345 
346  // Collect gpu.wait ops that we can move out of async.execute regions.
347  getOperation().getRegion().walk(DeferWaitCallback());
348  // Makes each !gpu.async.token returned from async.execute op have single use.
349  getOperation().getRegion().walk(SingleTokenUseCallback());
350 }
351 
352 std::unique_ptr<OperationPass<func::FuncOp>> mlir::createGpuAsyncRegionPass() {
353  return std::make_unique<GpuAsyncRegionPass>();
354 }
static bool isTerminator(Operation *op)
async::ExecuteOp addExecuteResults(async::ExecuteOp executeOp, ValueRange results)
Erases executeOp and returns a clone with additional results.
static bool hasSideEffects(Operation *op)
static void copy(Location loc, Value dst, Value src, Value size, OpBuilder &builder)
Copies the given number of bytes from src to dst pointers.
static void visit(Operation *op, DenseSet< Operation * > &visited)
Visits all the pdl.operand(s), pdl.result(s), and pdl.operation(s) connected to the given operation.
Definition: PDL.cpp:63
static MLIRContext * getContext(OpFoldResult val)
Block represents an ordered list of Operations.
Definition: Block.h:30
OpListType::iterator iterator
Definition: Block.h:133
This is a utility class for mapping one set of IR entities to another.
Definition: IRMapping.h:26
void clear()
Clears all mappings held by the mapper.
Definition: IRMapping.h:79
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
This class helps build Operations.
Definition: Builders.h:206
This is a value defined by a result of an operation.
Definition: Value.h:453
unsigned getResultNumber() const
Returns the number of this result.
Definition: Value.h:465
This class provides the API for ops that are known to be terminators.
Definition: OpDefinition.h:762
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
void insertOperands(unsigned index, ValueRange operands)
Insert the given operands into the operand list at the given 'index'.
Definition: Operation.cpp:255
bool mightHaveTrait()
Returns true if the operation might have the provided trait.
Definition: Operation.h:736
MLIRContext * getContext()
Return the context this operation is associated with.
Definition: Operation.h:216
unsigned getNumRegions()
Returns the number of regions held by this operation.
Definition: Operation.h:652
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
unsigned getNumOperands()
Definition: Operation.h:341
static Operation * create(Location location, OperationName name, TypeRange resultTypes, ValueRange operands, NamedAttrList &&attributes, OpaqueProperties properties, BlockRange successors, unsigned numRegions)
Create a new Operation with the specific fields.
Definition: Operation.cpp:66
MutableArrayRef< Region > getRegions()
Returns the regions held by this operation.
Definition: Operation.h:655
OperationName getName()
The name of an operation is the key identifier for it.
Definition: Operation.h:119
DictionaryAttr getDiscardableAttrDictionary()
Return all of the discardable attributes on this operation as a DictionaryAttr.
Definition: Operation.h:483
result_type_range getResultTypes()
Definition: Operation.h:423
operand_range getOperands()
Returns an iterator on the underlying Value's.
Definition: Operation.h:373
void replaceAllUsesWith(ValuesT &&values)
Replace all uses of results of this operation with the provided 'values'.
Definition: Operation.h:272
SuccessorRange getSuccessors()
Definition: Operation.h:682
InFlightDiagnostic emitOpError(const Twine &message={})
Emit an error with the op name prefixed, like "'dim' op " which is convenient for verifiers.
Definition: Operation.cpp:640
OpaqueProperties getPropertiesStorage()
Returns the properties storage.
Definition: Operation.h:879
void erase()
Remove this operation from its parent block and delete it.
Definition: Operation.cpp:538
unsigned getNumResults()
Return the number of results held by this operation.
Definition: Operation.h:399
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:36
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:378
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
bool use_empty() const
Returns true if this value has no uses.
Definition: Value.h:214
Type getType() const
Return the type of this value.
Definition: Value.h:125
user_range getUsers() const
Definition: Value.h:224
A utility result that is used to signal how to proceed with an ongoing walk:
Definition: Visitors.h:34
static WalkResult advance()
Definition: Visitors.h:52
static WalkResult interrupt()
Definition: Visitors.h:51
void walk(Operation *op, function_ref< void(Region *)> callback, WalkOrder order)
Walk all of the regions, blocks, or operations nested under (and including) the given operation.
Definition: Visitors.h:137
Include the generated interface declarations.
bool isMemoryEffectFree(Operation *op)
Returns true if the given operation is free of memory effects.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
std::unique_ptr< OperationPass< func::FuncOp > > createGpuAsyncRegionPass()
Rewrites a function region so that GPU ops execute asynchronously.
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:72
void operator()(async::ExecuteOp executeOp)
void operator()(async::ExecuteOp executeOp)
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26