620 lines · cpp
1//===- EliminateBarriers.cpp - Eliminate extra barriers --===//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-exception6//7//===----------------------------------------------------------------------===//8//9// Barrier elimination pattern and pass. If a barrier does not enforce any10// conflicting pair of memory effects, including a pair that is enforced by11// another barrier, it is unnecessary and can be removed. Adapted from12// "High-Performance GPU-to-CPU Transpilation and Optimization via High-Level13// Parallel Constructs" by Moses, Ivanov, Domke, Endo, Doerfert, and Zinenko in14// PPoPP 2023 and implementation in Polygeist.15//16//===----------------------------------------------------------------------===//17 18#include "mlir/Dialect/Func/IR/FuncOps.h"19#include "mlir/Dialect/GPU/IR/GPUDialect.h"20#include "mlir/Dialect/GPU/Transforms/Passes.h"21#include "mlir/Dialect/MemRef/IR/MemRef.h"22#include "mlir/Dialect/SCF/IR/SCF.h"23#include "mlir/Dialect/Vector/IR/VectorOps.h"24#include "mlir/IR/Operation.h"25#include "mlir/Transforms/GreedyPatternRewriteDriver.h"26#include "llvm/ADT/TypeSwitch.h"27#include "llvm/Support/Debug.h"28#include "llvm/Support/DebugLog.h"29 30namespace mlir {31#define GEN_PASS_DEF_GPUELIMINATEBARRIERS32#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"33} // namespace mlir34 35using namespace mlir;36using namespace mlir::gpu;37 38#define DEBUG_TYPE "gpu-erase-barriers"39#define DEBUG_TYPE_ALIAS "gpu-erase-barries-alias"40 41// The functions below provide interface-like verification, but are too specific42// to barrier elimination to become interfaces.43 44/// Returns `true` if the op is defines the parallel region that is subject to45/// barrier synchronization.46static bool isParallelRegionBoundary(Operation *op) {47 if (op->hasAttr("__parallel_region_boundary_for_test"))48 return true;49 50 return isa<GPUFuncOp, LaunchOp>(op);51}52 53/// Returns `true` if the op behaves like a sequential loop, e.g., the control54/// flow "wraps around" from the end of the body region back to its start.55static bool isSequentialLoopLike(Operation *op) { return isa<scf::ForOp>(op); }56 57/// Returns `true` if the regions of the op are guaranteed to be executed at58/// most once. Thus, if an operation in one of the nested regions of `op` is59/// executed than so are all the other operations in this region.60static bool hasSingleExecutionBody(Operation *op) {61 return isa<FunctionOpInterface, scf::IfOp, memref::AllocaScopeOp>(op);62}63 64/// Returns `true` if the operation is known to produce a pointer-like object65/// distinct from any other object produced by a similar operation. For example,66/// an allocation produces such an object.67static bool producesDistinctBase(Operation *op) {68 return isa_and_nonnull<memref::AllocOp, memref::AllocaOp>(op);69}70 71/// Populates `effects` with all memory effects without associating them to a72/// specific value.73static void addAllValuelessEffects(74 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {75 effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Read>());76 effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Write>());77 effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Allocate>());78 effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Free>());79}80 81/// Collect the memory effects of the given op in 'effects'. Returns 'true' if82/// it could extract the effect information from the op, otherwise returns83/// 'false' and conservatively populates the list with all possible effects84/// associated with no particular value or symbol.85static bool86collectEffects(Operation *op,87 SmallVectorImpl<MemoryEffects::EffectInstance> &effects,88 bool ignoreBarriers = true) {89 // Skip over barriers to avoid infinite recursion (those barriers would ask90 // this barrier again).91 if (ignoreBarriers && isa<BarrierOp>(op))92 return true;93 94 // Collect effect instances the operation. Note that the implementation of95 // getEffects erases all effect instances that have the type other than the96 // template parameter so we collect them first in a local buffer and then97 // copy.98 if (auto iface = dyn_cast<MemoryEffectOpInterface>(op)) {99 SmallVector<MemoryEffects::EffectInstance> localEffects;100 iface.getEffects(localEffects);101 llvm::append_range(effects, localEffects);102 return true;103 }104 if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>()) {105 for (auto ®ion : op->getRegions()) {106 for (auto &block : region) {107 for (auto &innerOp : block)108 if (!collectEffects(&innerOp, effects, ignoreBarriers))109 return false;110 }111 }112 return true;113 }114 115 // We need to be conservative here in case the op doesn't have the interface116 // and assume it can have any possible effect.117 addAllValuelessEffects(effects);118 return false;119}120 121/// Get all effects before the given operation caused by other operations in the122/// same block. That is, this will not consider operations beyond the block.123static bool124getEffectsBeforeInBlock(Operation *op,125 SmallVectorImpl<MemoryEffects::EffectInstance> &effects,126 bool stopAtBarrier) {127 if (op == &op->getBlock()->front())128 return true;129 130 for (Operation *it = op->getPrevNode(); it != nullptr;131 it = it->getPrevNode()) {132 if (isa<BarrierOp>(it)) {133 if (stopAtBarrier)134 return true;135 continue;136 }137 138 if (!collectEffects(it, effects))139 return false;140 }141 return true;142}143 144/// Collects memory effects from operations that may be executed before `op` in145/// a trivial structured control flow, e.g., without branches. Stops at the146/// parallel region boundary or at the barrier operation if `stopAtBarrier` is147/// set. Returns `true` if the memory effects added to `effects` are exact,148/// `false` if they are a conservative over-approximation. The latter means that149/// `effects` contain instances not associated with a specific value.150static bool151getEffectsBefore(Operation *op,152 SmallVectorImpl<MemoryEffects::EffectInstance> &effects,153 bool stopAtBarrier) {154 if (!op->getBlock())155 return true;156 157 // If there is a non-structured control flow, bail.158 Region *region = op->getBlock()->getParent();159 if (region && !region->hasOneBlock()) {160 addAllValuelessEffects(effects);161 return false;162 }163 164 // Collect all effects before the op.165 getEffectsBeforeInBlock(op, effects, stopAtBarrier);166 167 // Stop if reached the parallel region boundary.168 if (isParallelRegionBoundary(op->getParentOp()))169 return true;170 171 Operation *parent = op->getParentOp();172 // Otherwise, keep collecting above the parent operation.173 if (!parent->hasTrait<OpTrait::IsIsolatedFromAbove>() &&174 !getEffectsBefore(parent, effects, stopAtBarrier))175 return false;176 177 // If the op is loop-like, collect effects from the trailing operations until178 // we hit a barrier because they can executed before the current operation by179 // the previous iteration of this loop. For example, in the following loop180 //181 // for i = ... {182 // op1183 // ...184 // barrier185 // op2186 // }187 //188 // the operation `op2` at iteration `i` is known to be executed before the189 // operation `op1` at iteration `i+1` and the side effects must be ordered190 // appropriately.191 if (isSequentialLoopLike(parent)) {192 // Assuming loop terminators have no side effects.193 return getEffectsBeforeInBlock(op->getBlock()->getTerminator(), effects,194 /*stopAtBarrier=*/true);195 }196 197 // If the parent operation is not guaranteed to execute its (single-block)198 // region once, walk the block.199 bool conservative = false;200 if (!hasSingleExecutionBody(op->getParentOp()))201 op->getParentOp()->walk([&](Operation *in) {202 if (conservative)203 return WalkResult::interrupt();204 if (!collectEffects(in, effects)) {205 conservative = true;206 return WalkResult::interrupt();207 }208 return WalkResult::advance();209 });210 211 return !conservative;212}213 214/// Get all effects after the given operation caused by other operations in the215/// same block. That is, this will not consider operations beyond the block.216static bool217getEffectsAfterInBlock(Operation *op,218 SmallVectorImpl<MemoryEffects::EffectInstance> &effects,219 bool stopAtBarrier) {220 if (op == &op->getBlock()->back())221 return true;222 223 for (Operation *it = op->getNextNode(); it != nullptr;224 it = it->getNextNode()) {225 if (isa<BarrierOp>(it)) {226 if (stopAtBarrier)227 return true;228 continue;229 }230 if (!collectEffects(it, effects))231 return false;232 }233 return true;234}235 236/// Collects memory effects from operations that may be executed after `op` in237/// a trivial structured control flow, e.g., without branches. Stops at the238/// parallel region boundary or at the barrier operation if `stopAtBarrier` is239/// set. Returns `true` if the memory effects added to `effects` are exact,240/// `false` if they are a conservative over-approximation. The latter means that241/// `effects` contain instances not associated with a specific value.242static bool243getEffectsAfter(Operation *op,244 SmallVectorImpl<MemoryEffects::EffectInstance> &effects,245 bool stopAtBarrier) {246 if (!op->getBlock())247 return true;248 249 // If there is a non-structured control flow, bail.250 Region *region = op->getBlock()->getParent();251 if (region && !region->hasOneBlock()) {252 addAllValuelessEffects(effects);253 return false;254 }255 256 // Collect all effects after the op.257 getEffectsAfterInBlock(op, effects, stopAtBarrier);258 259 Operation *parent = op->getParentOp();260 // Stop if reached the parallel region boundary.261 if (isParallelRegionBoundary(parent))262 return true;263 264 // Otherwise, keep collecting below the parent operation.265 // Don't look into, for example, neighboring functions266 if (!parent->hasTrait<OpTrait::IsIsolatedFromAbove>() &&267 !getEffectsAfter(parent, effects, stopAtBarrier))268 return false;269 270 // If the op is loop-like, collect effects from the leading operations until271 // we hit a barrier because they can executed after the current operation by272 // the next iteration of this loop. For example, in the following loop273 //274 // for i = ... {275 // op1276 // ...277 // barrier278 // op2279 // }280 //281 // the operation `op1` at iteration `i` is known to be executed after the282 // operation `op2` at iteration `i-1` and the side effects must be ordered283 // appropriately.284 if (isSequentialLoopLike(parent)) {285 if (isa<BarrierOp>(op->getBlock()->front()))286 return true;287 288 bool exact = collectEffects(&op->getBlock()->front(), effects);289 return getEffectsAfterInBlock(&op->getBlock()->front(), effects,290 /*stopAtBarrier=*/true) &&291 exact;292 }293 294 // If the parent operation is not guaranteed to execute its (single-block)295 // region once, walk the block.296 bool conservative = false;297 if (!hasSingleExecutionBody(op->getParentOp()))298 op->getParentOp()->walk([&](Operation *in) {299 if (conservative)300 return WalkResult::interrupt();301 if (!collectEffects(in, effects)) {302 conservative = true;303 return WalkResult::interrupt();304 }305 return WalkResult::advance();306 });307 308 return !conservative;309}310 311/// Looks through known "view-like" ops to find the base memref.312static Value getBase(Value v) {313 while (true) {314 Operation *definingOp = v.getDefiningOp();315 if (!definingOp)316 break;317 318 bool shouldContinue =319 TypeSwitch<Operation *, bool>(v.getDefiningOp())320 .Case<memref::CastOp, memref::SubViewOp, memref::ViewOp>(321 [&](auto op) {322 v = op.getSource();323 return true;324 })325 .Case<memref::TransposeOp>([&](auto op) {326 v = op.getIn();327 return true;328 })329 .Case<memref::CollapseShapeOp, memref::ExpandShapeOp>([&](auto op) {330 v = op.getSrc();331 return true;332 })333 .Default(false);334 if (!shouldContinue)335 break;336 }337 return v;338}339 340/// Returns `true` if the value is defined as a function argument.341static bool isFunctionArgument(Value v) {342 auto arg = dyn_cast<BlockArgument>(v);343 return arg && isa<FunctionOpInterface>(arg.getOwner()->getParentOp());344}345 346/// Returns the operand that the operation "propagates" through it for capture347/// purposes. That is, if the value produced by this operation is captured, then348/// so is the returned value.349static Value propagatesCapture(Operation *op) {350 return llvm::TypeSwitch<Operation *, Value>(op)351 .Case(352 [](ViewLikeOpInterface viewLike) { return viewLike.getViewSource(); })353 .Case([](CastOpInterface castLike) { return castLike->getOperand(0); })354 .Case([](memref::TransposeOp transpose) { return transpose.getIn(); })355 .Case<memref::ExpandShapeOp, memref::CollapseShapeOp>(356 [](auto op) { return op.getSrc(); })357 .Default(nullptr);358}359 360/// Returns `true` if the given operation is known to capture the given value,361/// `false` if it is known not to capture the given value, `nullopt` if neither362/// is known.363static std::optional<bool> getKnownCapturingStatus(Operation *op, Value v) {364 return llvm::TypeSwitch<Operation *, std::optional<bool>>(op)365 // Store-like operations don't capture the destination, but do capture366 // the value.367 .Case<memref::StoreOp, vector::TransferWriteOp>(368 [&](auto op) { return op.getValue() == v; })369 .Case<vector::StoreOp, vector::MaskedStoreOp>(370 [&](auto op) { return op.getValueToStore() == v; })371 // These operations are known not to capture.372 .Case([](memref::DeallocOp) { return false; })373 // By default, we don't know anything.374 .Default(std::nullopt);375}376 377/// Returns `true` if the value may be captured by any of its users, i.e., if378/// the user may be storing this value into memory. This makes aliasing analysis379/// more conservative as it cannot assume the pointer-like value is only passed380/// around through SSA use-def.381static bool maybeCaptured(Value v) {382 SmallVector<Value> todo = {v};383 while (!todo.empty()) {384 Value v = todo.pop_back_val();385 for (Operation *user : v.getUsers()) {386 // A user that is known to only read cannot capture.387 auto iface = dyn_cast<MemoryEffectOpInterface>(user);388 if (iface) {389 SmallVector<MemoryEffects::EffectInstance> effects;390 iface.getEffects(effects);391 if (llvm::all_of(effects,392 [](const MemoryEffects::EffectInstance &effect) {393 return isa<MemoryEffects::Read>(effect.getEffect());394 })) {395 continue;396 }397 }398 399 // When an operation is known to create an alias, consider if the400 // source is captured as well.401 if (Value v = propagatesCapture(user)) {402 todo.push_back(v);403 continue;404 }405 406 std::optional<bool> knownCaptureStatus = getKnownCapturingStatus(user, v);407 if (!knownCaptureStatus || *knownCaptureStatus)408 return true;409 }410 }411 412 return false;413}414 415/// Returns true if two values may be referencing aliasing memory. This is a416/// rather naive and conservative analysis. Values defined by different417/// allocation-like operations as well as values derived from those by casts and418/// views cannot alias each other. Similarly, values defined by allocations419/// inside a function cannot alias function arguments. Global values cannot420/// alias each other or local allocations. Values that are captured, i.e.421/// themselves potentially stored in memory, are considered as aliasing with422/// everything. This seems sufficient to achieve barrier removal in structured423/// control flow, more complex cases would require a proper dataflow analysis.424static bool mayAlias(Value first, Value second) {425 LDBG(DEBUG_TYPE_ALIAS, 1)426 << "checking aliasing between " << first << " and " << second;427 428 first = getBase(first);429 second = getBase(second);430 431 LDBG(DEBUG_TYPE_ALIAS, 1) << "base " << first << " and " << second;432 433 // Values derived from the same base memref do alias (unless we do a more434 // advanced analysis to prove non-overlapping accesses).435 if (first == second) {436 LDBG(DEBUG_TYPE_ALIAS, 1) << "-> do alias!";437 return true;438 }439 440 // Different globals cannot alias.441 if (auto globFirst = first.getDefiningOp<memref::GetGlobalOp>()) {442 if (auto globSecond = second.getDefiningOp<memref::GetGlobalOp>()) {443 return globFirst.getNameAttr() == globSecond.getNameAttr();444 }445 }446 447 // Two function arguments marked as noalias do not alias.448 auto isNoaliasFuncArgument = [](Value value) {449 auto bbArg = dyn_cast<BlockArgument>(value);450 if (!bbArg)451 return false;452 auto iface = dyn_cast<FunctionOpInterface>(bbArg.getOwner()->getParentOp());453 if (!iface)454 return false;455 // TODO: we need a way to not depend on the LLVM dialect here.456 return iface.getArgAttr(bbArg.getArgNumber(), "llvm.noalias") != nullptr;457 };458 if (isNoaliasFuncArgument(first) && isNoaliasFuncArgument(second))459 return false;460 461 bool isDistinct[] = {producesDistinctBase(first.getDefiningOp()),462 producesDistinctBase(second.getDefiningOp())};463 bool isGlobal[] = {first.getDefiningOp<memref::GetGlobalOp>() != nullptr,464 second.getDefiningOp<memref::GetGlobalOp>() != nullptr};465 466 // Non-equivalent distinct bases and globals cannot alias. At this point, we467 // have already filtered out based on values being equal and global name being468 // equal.469 if ((isDistinct[0] || isGlobal[0]) && (isDistinct[1] || isGlobal[1]))470 return false;471 472 bool isArg[] = {isFunctionArgument(first), isFunctionArgument(second)};473 474 // Distinct bases (allocations) cannot have been passed as an argument.475 if ((isDistinct[0] && isArg[1]) || (isDistinct[1] && isArg[0]))476 return false;477 478 // Non-captured base distinct values cannot conflict with another base value.479 if (isDistinct[0] && !maybeCaptured(first))480 return false;481 if (isDistinct[1] && !maybeCaptured(second))482 return false;483 484 // Otherwise, conservatively assume aliasing.485 LDBG(DEBUG_TYPE_ALIAS, 1) << "-> may alias!";486 return true;487}488 489/// Returns `true` if the effect may be affecting memory aliasing the value. If490/// the effect is not associated with any value, it is assumed to affect all491/// memory and therefore aliases with everything.492static bool mayAlias(MemoryEffects::EffectInstance a, Value v2) {493 if (Value v = a.getValue()) {494 return mayAlias(v, v2);495 }496 return true;497}498 499/// Returns `true` if the two effects may be affecting aliasing memory. If500/// an effect is not associated with any value, it is assumed to affect all501/// memory and therefore aliases with everything. Effects on different resources502/// cannot alias.503static bool mayAlias(MemoryEffects::EffectInstance a,504 MemoryEffects::EffectInstance b) {505 if (a.getResource()->getResourceID() != b.getResource()->getResourceID())506 return false;507 if (Value v2 = b.getValue()) {508 return mayAlias(a, v2);509 } else if (Value v = a.getValue()) {510 return mayAlias(b, v);511 }512 return true;513}514 515/// Returns `true` if any of the "before" effect instances has a conflict with516/// any "after" instance for the purpose of barrier elimination. The effects are517/// supposed to be limited to a barrier synchronization scope. A conflict exists518/// if effects instances affect aliasing memory locations and at least on of519/// then as a write. As an exception, if the non-write effect is an allocation520/// effect, there is no conflict since we are only expected to see the521/// allocation happening in the same thread and it cannot be accessed from522/// another thread without capture (which we do handle in alias analysis).523static bool524haveConflictingEffects(ArrayRef<MemoryEffects::EffectInstance> beforeEffects,525 ArrayRef<MemoryEffects::EffectInstance> afterEffects) {526 for (const MemoryEffects::EffectInstance &before : beforeEffects) {527 for (const MemoryEffects::EffectInstance &after : afterEffects) {528 // If cannot alias, definitely no conflict.529 if (!mayAlias(before, after))530 continue;531 532 // Read/read is not a conflict.533 if (isa<MemoryEffects::Read>(before.getEffect()) &&534 isa<MemoryEffects::Read>(after.getEffect())) {535 continue;536 }537 538 // Allocate/* is not a conflict since the allocation happens within the539 // thread context.540 // TODO: This is not the case for */Free unless the allocation happened in541 // the thread context, which we could also check for.542 if (isa<MemoryEffects::Allocate>(before.getEffect()) ||543 isa<MemoryEffects::Allocate>(after.getEffect())) {544 continue;545 }546 547 // In the particular case that the before effect is a free, we only have 2548 // possibilities:549 // 1. either the program is well-formed and there must be an interleaved550 // alloc that must limit the scope of effect lookback and we can551 // safely ignore the free -> read / free -> write and free -> free552 // conflicts.553 // 2. either the program is ill-formed and we are in undefined behavior554 // territory.555 if (isa<MemoryEffects::Free>(before.getEffect()))556 continue;557 558 // Other kinds of effects create a conflict, e.g. read-after-write.559 LDBG() << "found a conflict between (before): " << before.getValue()560 << " read:" << isa<MemoryEffects::Read>(before.getEffect())561 << " write:" << isa<MemoryEffects::Write>(before.getEffect())562 << " alloc:" << isa<MemoryEffects::Allocate>(before.getEffect())563 << " free:" << isa<MemoryEffects::Free>(before.getEffect());564 LDBG() << "and (after): " << after.getValue()565 << " read:" << isa<MemoryEffects::Read>(after.getEffect())566 << " write:" << isa<MemoryEffects::Write>(after.getEffect())567 << " alloc:" << isa<MemoryEffects::Allocate>(after.getEffect())568 << " free:" << isa<MemoryEffects::Free>(after.getEffect());569 return true;570 }571 }572 573 return false;574}575 576namespace {577class BarrierElimination final : public OpRewritePattern<BarrierOp> {578public:579 using OpRewritePattern<BarrierOp>::OpRewritePattern;580 581 LogicalResult matchAndRewrite(BarrierOp barrier,582 PatternRewriter &rewriter) const override {583 LDBG() << "checking the necessity of: " << barrier << " "584 << barrier.getLoc();585 586 SmallVector<MemoryEffects::EffectInstance> beforeEffects;587 getEffectsBefore(barrier, beforeEffects, /*stopAtBarrier=*/true);588 589 SmallVector<MemoryEffects::EffectInstance> afterEffects;590 getEffectsAfter(barrier, afterEffects, /*stopAtBarrier=*/true);591 592 if (!haveConflictingEffects(beforeEffects, afterEffects)) {593 LDBG() << "the surrounding barriers are sufficient, removing " << barrier;594 rewriter.eraseOp(barrier);595 return success();596 }597 598 LDBG() << "barrier is necessary: " << barrier << " " << barrier.getLoc();599 return failure();600 }601};602 603class GpuEliminateBarriersPass604 : public impl::GpuEliminateBarriersBase<GpuEliminateBarriersPass> {605 void runOnOperation() override {606 auto funcOp = getOperation();607 RewritePatternSet patterns(&getContext());608 mlir::populateGpuEliminateBarriersPatterns(patterns);609 if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) {610 return signalPassFailure();611 }612 }613};614 615} // namespace616 617void mlir::populateGpuEliminateBarriersPatterns(RewritePatternSet &patterns) {618 patterns.insert<BarrierElimination>(patterns.getContext());619}620