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 14 #include "mlir/Dialect/GPU/Transforms/Passes.h" 15 16 #include "mlir/Dialect/Async/IR/Async.h" 17 #include "mlir/Dialect/Func/IR/FuncOps.h" 18 #include "mlir/Dialect/GPU/IR/GPUDialect.h" 19 #include "mlir/Dialect/GPU/Transforms/Utils.h" 20 #include "mlir/IR/BlockAndValueMapping.h" 21 #include "mlir/IR/Builders.h" 22 #include "mlir/IR/PatternMatch.h" 23 #include "mlir/IR/SymbolTable.h" 24 #include "mlir/Support/LLVM.h" 25 #include "mlir/Transforms/RegionUtils.h" 26 #include "llvm/ADT/TypeSwitch.h" 27 28 namespace mlir { 29 #define GEN_PASS_DEF_GPUASYNCREGIONPASS 30 #include "mlir/Dialect/GPU/Transforms/Passes.h.inc" 31 } // namespace mlir 32 33 using namespace mlir; 34 35 namespace { 36 class GpuAsyncRegionPass 37 : public impl::GpuAsyncRegionPassBase<GpuAsyncRegionPass> { 38 struct ThreadTokenCallback; 39 struct DeferWaitCallback; 40 struct SingleTokenUseCallback; 41 void runOnOperation() override; 42 }; 43 } // namespace 44 45 static bool isTerminator(Operation *op) { 46 return op->mightHaveTrait<OpTrait::IsTerminator>(); 47 } 48 static bool hasSideEffects(Operation *op) { 49 return !MemoryEffectOpInterface::hasNoEffect(op); 50 } 51 52 // Region walk callback which makes GPU ops implementing the AsyncOpInterface 53 // execute asynchronously. 54 struct GpuAsyncRegionPass::ThreadTokenCallback { 55 ThreadTokenCallback(MLIRContext &context) : builder(&context) {} 56 57 WalkResult operator()(Block *block) { 58 for (Operation &op : make_early_inc_range(*block)) { 59 if (failed(visit(&op))) 60 return WalkResult::interrupt(); 61 } 62 return WalkResult::advance(); 63 } 64 65 private: 66 // If `op` implements the AsyncOpInterface, insert a `gpu.wait async` to 67 // create a current token (unless it already exists), and 'thread' that token 68 // through the `op` so that it executes asynchronously. 69 // 70 // If `op` is a terminator or an op with side-effects, insert a `gpu.wait` to 71 // host-synchronize execution. A `!gpu.async.token` will therefore only be 72 // used inside of its block and GPU execution will always synchronize with 73 // the host at block boundaries. 74 LogicalResult visit(Operation *op) { 75 if (isa<gpu::LaunchOp>(op)) 76 return op->emitOpError("replace with gpu.launch_func first"); 77 if (auto waitOp = llvm::dyn_cast<gpu::WaitOp>(op)) { 78 if (currentToken) 79 waitOp.addAsyncDependency(currentToken); 80 currentToken = waitOp.asyncToken(); 81 return success(); 82 } 83 builder.setInsertionPoint(op); 84 if (auto asyncOp = dyn_cast<gpu::AsyncOpInterface>(op)) 85 return rewriteAsyncOp(asyncOp); // Replace GPU op with async version. 86 if (!currentToken) 87 return success(); 88 // Insert host synchronization before terminator or op with side effects. 89 if (isTerminator(op) || hasSideEffects(op)) 90 currentToken = createWaitOp(op->getLoc(), Type(), {currentToken}); 91 return success(); 92 } 93 94 // Replaces asyncOp with a clone that returns a token. 95 LogicalResult rewriteAsyncOp(gpu::AsyncOpInterface asyncOp) { 96 auto *op = asyncOp.getOperation(); 97 auto tokenType = builder.getType<gpu::AsyncTokenType>(); 98 99 // If there is no current token, insert a `gpu.wait async` without 100 // dependencies to create one. 101 if (!currentToken) 102 currentToken = createWaitOp(op->getLoc(), tokenType, {}); 103 asyncOp.addAsyncDependency(currentToken); 104 105 // Return early if op returns a token already. 106 currentToken = asyncOp.getAsyncToken(); 107 if (currentToken) 108 return success(); 109 110 // Clone the op to return a token in addition to the other results. 111 SmallVector<Type, 1> resultTypes; 112 resultTypes.reserve(1 + op->getNumResults()); 113 copy(op->getResultTypes(), std::back_inserter(resultTypes)); 114 resultTypes.push_back(tokenType); 115 auto *newOp = Operation::create(op->getLoc(), op->getName(), resultTypes, 116 op->getOperands(), op->getAttrDictionary(), 117 op->getSuccessors(), op->getNumRegions()); 118 119 // Clone regions into new op. 120 BlockAndValueMapping 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).asyncToken(); 136 } 137 138 OpBuilder builder; 139 140 // The token that represents the current asynchronous dependency. It's valid 141 // range starts with a `gpu.wait async` op, and ends with a `gpu.wait` op. 142 // In between, each gpu::AsyncOpInterface depends on the current token and 143 // produces the new one. 144 Value currentToken = {}; 145 }; 146 147 /// Erases `executeOp` and returns a clone with additional `results`. 148 async::ExecuteOp addExecuteResults(async::ExecuteOp executeOp, 149 ValueRange results) { 150 // Add values to async.yield op. 151 Operation *yieldOp = executeOp.getBody()->getTerminator(); 152 yieldOp->insertOperands(yieldOp->getNumOperands(), results); 153 154 // Construct new result type list with additional types. 155 SmallVector<Type, 2> resultTypes; 156 resultTypes.reserve(executeOp.getNumResults() + results.size()); 157 transform(executeOp.getResultTypes(), std::back_inserter(resultTypes), 158 [](Type type) { 159 // Extract value type from !async.value. 160 if (auto valueType = type.dyn_cast<async::ValueType>()) 161 return valueType.getValueType(); 162 assert(type.isa<async::TokenType>() && "expected token type"); 163 return type; 164 }); 165 transform(results, std::back_inserter(resultTypes), 166 [](Value value) { return value.getType(); }); 167 168 // Clone executeOp with the extra results. 169 OpBuilder builder(executeOp); 170 auto newOp = builder.create<async::ExecuteOp>( 171 executeOp.getLoc(), TypeRange{resultTypes}.drop_front() /*drop token*/, 172 executeOp.getDependencies(), executeOp.getBodyOperands()); 173 BlockAndValueMapping mapper; 174 newOp.getRegion().getBlocks().clear(); 175 executeOp.getRegion().cloneInto(&newOp.getRegion(), mapper); 176 177 // Replace executeOp with cloned one. 178 executeOp.getOperation()->replaceAllUsesWith( 179 newOp.getResults().drop_back(results.size())); 180 executeOp.erase(); 181 182 return newOp; 183 } 184 185 // Callback for `async.execute` ops which tries to push the contained 186 // synchronous `gpu.wait` op to the dependencies of the `async.execute`. 187 struct GpuAsyncRegionPass::DeferWaitCallback { 188 // If the `executeOp`s token is used only in `async.execute` or `async.await` 189 // ops, add the region's last `gpu.wait` op to the worklist if it is 190 // synchronous and is the last op with side effects. 191 void operator()(async::ExecuteOp executeOp) { 192 if (!areAllUsersExecuteOrAwait(executeOp.getToken())) 193 return; 194 // async.execute's region is currently restricted to one block. 195 for (auto &op : llvm::reverse(executeOp.getBody()->without_terminator())) { 196 if (auto waitOp = dyn_cast<gpu::WaitOp>(op)) { 197 if (!waitOp.asyncToken()) 198 worklist.push_back(waitOp); 199 return; 200 } 201 if (hasSideEffects(&op)) 202 return; 203 } 204 } 205 206 // The destructor performs the actual rewrite work. 207 ~DeferWaitCallback() { 208 for (size_t i = 0; i < worklist.size(); ++i) { 209 auto waitOp = worklist[i]; 210 auto executeOp = waitOp->getParentOfType<async::ExecuteOp>(); 211 212 // Erase `gpu.wait` and return async dependencies from execute op instead. 213 SmallVector<Value, 4> dependencies = waitOp.getAsyncDependencies(); 214 waitOp.erase(); 215 executeOp = addExecuteResults(executeOp, dependencies); 216 217 // Add the async dependency to each user of the `async.execute` token. 218 auto asyncTokens = executeOp.getResults().take_back(dependencies.size()); 219 SmallVector<Operation *, 4> users(executeOp.getToken().user_begin(), 220 executeOp.getToken().user_end()); 221 for (Operation *user : users) 222 addAsyncDependencyAfter(asyncTokens, user); 223 } 224 } 225 226 private: 227 // Returns whether all token users are either 'async.execute' or 'async.await' 228 // ops. This is used as a requirement for pushing 'gpu.wait' ops from a 229 // 'async.execute' body to it's users. Specifically, we do not allow 230 // terminator users, because it could mean that the `async.execute` is inside 231 // control flow code. 232 static bool areAllUsersExecuteOrAwait(Value token) { 233 return !token.use_empty() && 234 llvm::all_of(token.getUsers(), [](Operation *user) { 235 return isa<async::ExecuteOp, async::AwaitOp>(user); 236 }); 237 } 238 239 // Add the `asyncToken` as dependency as needed after `op`. 240 void addAsyncDependencyAfter(ValueRange asyncTokens, Operation *op) { 241 OpBuilder builder(op->getContext()); 242 auto loc = op->getLoc(); 243 244 Block::iterator it; 245 SmallVector<Value, 1> tokens; 246 tokens.reserve(asyncTokens.size()); 247 TypeSwitch<Operation *>(op) 248 .Case<async::AwaitOp>([&](auto awaitOp) { 249 // Add async.await ops to wait for the !gpu.async.tokens. 250 builder.setInsertionPointAfter(op); 251 for (auto asyncToken : asyncTokens) 252 tokens.push_back( 253 builder.create<async::AwaitOp>(loc, asyncToken).getResult()); 254 // Set `it` after the inserted async.await ops. 255 it = builder.getInsertionPoint(); 256 }) 257 .Case<async::ExecuteOp>([&](auto executeOp) { 258 // Set `it` to the beginning of the region and add asyncTokens to the 259 // async.execute operands. 260 it = executeOp.getBody()->begin(); 261 executeOp.getBodyOperandsMutable().append(asyncTokens); 262 SmallVector<Type, 1> tokenTypes( 263 asyncTokens.size(), builder.getType<gpu::AsyncTokenType>()); 264 SmallVector<Location, 1> tokenLocs(asyncTokens.size(), 265 executeOp.getLoc()); 266 copy(executeOp.getBody()->addArguments(tokenTypes, tokenLocs), 267 std::back_inserter(tokens)); 268 }); 269 270 // Advance `it` to terminator or op with side-effects. 271 it = std::find_if(it, Block::iterator(), [](Operation &op) { 272 return isTerminator(&op) || hasSideEffects(&op); 273 }); 274 275 // If `op` implements the AsyncOpInterface, add `token` to the list of async 276 // dependencies. 277 if (auto asyncOp = dyn_cast<gpu::AsyncOpInterface>(*it)) { 278 for (auto token : tokens) 279 asyncOp.addAsyncDependency(token); 280 return; 281 } 282 283 // Otherwise, insert a gpu.wait before 'it'. 284 builder.setInsertionPoint(it->getBlock(), it); 285 auto waitOp = builder.create<gpu::WaitOp>(loc, Type{}, tokens); 286 287 // If the new waitOp is at the end of an async.execute region, add it to the 288 // worklist. 'operator()(executeOp)' would do the same, but this is faster. 289 auto executeOp = dyn_cast<async::ExecuteOp>(it->getParentOp()); 290 if (executeOp && areAllUsersExecuteOrAwait(executeOp.getToken()) && 291 !it->getNextNode()) 292 worklist.push_back(waitOp); 293 } 294 295 SmallVector<gpu::WaitOp, 8> worklist; 296 }; 297 298 // Callback for `async.execute` ops which repeats !gpu.async.token results 299 // so that each of them is only used once. 300 struct GpuAsyncRegionPass::SingleTokenUseCallback { 301 void operator()(async::ExecuteOp executeOp) { 302 // Extract !gpu.async.token results which have multiple uses. 303 auto multiUseResults = llvm::make_filter_range( 304 executeOp.getBodyResults(), [](OpResult result) { 305 if (result.use_empty() || result.hasOneUse()) 306 return false; 307 auto valueType = result.getType().dyn_cast<async::ValueType>(); 308 return valueType && 309 valueType.getValueType().isa<gpu::AsyncTokenType>(); 310 }); 311 if (multiUseResults.empty()) 312 return; 313 314 // Indices within !async.execute results (i.e. without the async.token). 315 SmallVector<int, 4> indices; 316 transform(multiUseResults, std::back_inserter(indices), 317 [](OpResult result) { 318 return result.getResultNumber() - 1; // Index without token. 319 }); 320 321 for (auto index : indices) { 322 assert(!executeOp.getBodyResults()[index].getUses().empty()); 323 // Repeat async.yield token result, one for each use after the first one. 324 auto uses = llvm::drop_begin(executeOp.getBodyResults()[index].getUses()); 325 auto count = std::distance(uses.begin(), uses.end()); 326 auto yieldOp = cast<async::YieldOp>(executeOp.getBody()->getTerminator()); 327 SmallVector<Value, 4> operands(count, yieldOp.getOperand(index)); 328 executeOp = addExecuteResults(executeOp, operands); 329 // Update 'uses' to refer to the new executeOp. 330 uses = llvm::drop_begin(executeOp.getBodyResults()[index].getUses()); 331 auto results = executeOp.getBodyResults().take_back(count); 332 for (auto pair : llvm::zip(uses, results)) 333 std::get<0>(pair).set(std::get<1>(pair)); 334 } 335 } 336 }; 337 338 // Replaces synchronous GPU ops in the op's region with asynchronous ones and 339 // inserts the necessary synchronization (as gpu.wait ops). Assumes sequential 340 // execution semantics and that no GPU ops are asynchronous yet. 341 void GpuAsyncRegionPass::runOnOperation() { 342 if (getOperation()->walk(ThreadTokenCallback(getContext())).wasInterrupted()) 343 return signalPassFailure(); 344 345 // Collect gpu.wait ops that we can move out of async.execute regions. 346 getOperation().getRegion().walk(DeferWaitCallback()); 347 // Makes each !gpu.async.token returned from async.execute op have single use. 348 getOperation().getRegion().walk(SingleTokenUseCallback()); 349 } 350 351 std::unique_ptr<OperationPass<func::FuncOp>> mlir::createGpuAsyncRegionPass() { 352 return std::make_unique<GpuAsyncRegionPass>(); 353 } 354