diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.h b/mlir/include/mlir/Interfaces/LoopLikeInterface.h index 9925fc6ce6ca9..d6c23131cf242 100644 --- a/mlir/include/mlir/Interfaces/LoopLikeInterface.h +++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.h @@ -14,6 +14,7 @@ #define MLIR_INTERFACES_LOOPLIKEINTERFACE_H_ #include "mlir/IR/OpDefinition.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" namespace mlir { class RewriterBase; diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td index e09b8672f2d08..644f096b69381 100644 --- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td +++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td @@ -329,6 +329,35 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { $_op->operand_begin() + firstOperandIndex + initsMutable.size()); } + /// Return whether the loop is known to have zero iterations. + /// Returns std::nullopt if not enough constant information is available. + ::std::optional isZeroTrip() { + auto lbs = $_op.getLoopLowerBounds(); + auto ubs = $_op.getLoopUpperBounds(); + auto steps = $_op.getLoopSteps(); + + if (!lbs || !ubs || !steps) + return ::std::nullopt; + + if (lbs->size() != ubs->size() || ubs->size() != steps->size()) + return ::std::nullopt; + + for (size_t i = 0; i < steps->size(); ++i) { + auto lb = ::mlir::getConstantIntValue((*lbs)[i]); + auto ub = ::mlir::getConstantIntValue((*ubs)[i]); + auto st = ::mlir::getConstantIntValue((*steps)[i]); + + if (!lb || !ub || !st) + return ::std::nullopt; // non-constant -> unknown + + if (*st >= 0 && *lb >= *ub) + return true; + if (*st < 0 && *lb <= *ub) + return true; + } + return false; + } + /// Return the region iter_arg that corresponds to the given init operand. /// Return an "empty" block argument if the given operand is not an init /// operand of this loop op. diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h index 9de20f0c69f1a..5442d652cc345 100644 --- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h +++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h @@ -14,6 +14,7 @@ #ifndef MLIR_INTERFACES_SIDEEFFECTINTERFACES_H #define MLIR_INTERFACES_SIDEEFFECTINTERFACES_H +#include "mlir/IR/Dominance.h" #include "mlir/IR/OpDefinition.h" namespace mlir { @@ -346,6 +347,20 @@ struct AlwaysSpeculatableImplTrait //===----------------------------------------------------------------------===// namespace MemoryEffects { +/// Defines the priority of the different memory effects. +/// +/// Sorting/ordering memory effects of an operation is done based on +/// their defined stage and priority, in that order. If stage values for two +/// effect instances are equal, they are then sorted by priority. Lower priority +/// values indicate higher precedence. +enum Priority { + DefaultPriority = 0, + AllocPriority = 1, + FreePriority = 2, + ReadPriority = 3, + WritePriority = 4 +}; + /// This class represents the base class used for memory effects. struct Effect : public SideEffects::Effect { using SideEffects::Effect::Effect; @@ -355,28 +370,50 @@ struct Effect : public SideEffects::Effect { using Base = SideEffects::Effect::Base; static bool classof(const SideEffects::Effect *effect); + + /// Return the priority associated with this memory effect. + Priority getPriority() const { return priority; } + +protected: + /// Priority value for this effect. Lower numbers indicate higher precedence. + Priority priority = Priority::DefaultPriority; }; using EffectInstance = SideEffects::EffectInstance; +/// Returns vector of the op's memory effects sorted in increasing stage order +/// and in increasing priority order within each stage. +llvm::SmallVector +getMemoryEffectsSorted(Operation *op); + /// The following effect indicates that the operation allocates from some /// resource. An 'allocate' effect implies only allocation of the resource, and /// not any visible mutation or dereference. -struct Allocate : public Effect::Base {}; +struct Allocate : public Effect::Base { + Allocate() : Effect::Base() { + this->priority = Priority::AllocPriority; + } +}; /// The following effect indicates that the operation frees some resource that /// has been allocated. An 'allocate' effect implies only de-allocation of the /// resource, and not any visible allocation, mutation or dereference. -struct Free : public Effect::Base {}; +struct Free : public Effect::Base { + Free() : Effect::Base() { this->priority = Priority::FreePriority; } +}; /// The following effect indicates that the operation reads from some resource. /// A 'read' effect implies only dereferencing of the resource, and not any /// visible mutation. -struct Read : public Effect::Base {}; +struct Read : public Effect::Base { + Read() : Effect::Base() { this->priority = Priority::ReadPriority; } +}; /// The following effect indicates that the operation writes to some resource. A /// 'write' effect implies only mutating a resource, and not any visible /// dereference or read. -struct Write : public Effect::Base {}; +struct Write : public Effect::Base { + Write() : Effect::Base() { this->priority = Priority::WritePriority; } +}; } // namespace MemoryEffects //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h index 3ceef44d799e8..f299c38439534 100644 --- a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h +++ b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h @@ -9,9 +9,11 @@ #ifndef MLIR_TRANSFORMS_LOOPINVARIANTCODEMOTIONUTILS_H #define MLIR_TRANSFORMS_LOOPINVARIANTCODEMOTIONUTILS_H +#include "mlir/Interfaces/SideEffectInterfaces.h" #include "mlir/Support/LLVM.h" +#include "mlir/Support/TypeID.h" -#include "llvm/ADT/SmallVector.h" +#include namespace mlir { @@ -21,6 +23,45 @@ class Region; class RewriterBase; class Value; +/// Alias for map used in LICM pass to track which memory resources have +/// conflicts due to sequence of memory effects applied to them in the region of +/// interest. +using MemoryConflictMap = DenseMap>; + +/// Gathers potential conflicts on all memory resources used within loop. +/// +/// Given a target loop and an op within it (or the loop op itself), +/// gathers op's memory effects and flags potential resource conflicts +/// in a map and then recurses into the op's regions to gather nested +/// resource conflicts. +/// +/// Typical usage: +/// \code +/// LoopLikeOpInterface myLoop = ...; +/// DenseMap> +/// myConflicts; gatherResourceConflicts(myLoop, myLoop.getOperation(), +/// resourceConflicts); +/// \endcode +/// +/// \param loop The loop to gather resource conflicts for. +/// \param op The operation to gather resource conflicts for, +/// typically the loop op itself via loop.getOperation(). +/// \param resourceConflicts Map to store potential resource conflicts. +/// Key is the resource ID that effects are applied to. Value is a pair of +/// a boolean, indicating if the resource has a conflict, and the last effect +/// that was applied to the resource (if no conflicts exist) or the effect +/// that caused the conflict (if conflicts exist). This was done so we +/// check the effect that causes the conflict for debugging purposes. +/// First call should use loop = someLoop and op = someLoop.getOperation() +/// +/// resourceConflicts is modified by the function and will be non-empty +/// as long as there are memory effects within the loop, even if there are +/// no conflicts. +void gatherResourceConflicts( + LoopLikeOpInterface loop, Operation *op, + DenseMap> + &resourceConflicts); + /// Given a list of regions, perform loop-invariant code motion. An operation is /// loop-invariant if it depends only of values defined outside of the loop. /// LICM moves these operations out of the loop body so that they are not @@ -63,9 +104,10 @@ class Value; /// /// Returns the number of operations moved. size_t moveLoopInvariantCode( - ArrayRef regions, + LoopLikeOpInterface loopLike, function_ref isDefinedOutsideRegion, - function_ref shouldMoveOutOfRegion, + function_ref shouldMoveSpeculatable, + function_ref shouldMoveMemoryEffect, function_ref moveOutOfRegion); /// Move side-effect free loop invariant code out of a loop-like op using diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp index b5a6888e5e1a4..6603f91936bc3 100644 --- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp +++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp @@ -8,7 +8,9 @@ #include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/IR/Dominance.h" #include "mlir/IR/SymbolTable.h" +#include #include using namespace mlir; @@ -317,14 +319,45 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) { return wouldOpBeTriviallyDeadImpl(op); } +llvm::SmallVector +mlir::MemoryEffects::getMemoryEffectsSorted(Operation *op) { + llvm::SmallVector effectsSorted; + + auto memInterface = dyn_cast(op); + + if (!memInterface) + return effectsSorted; // return empty vec + + memInterface.getEffects(effectsSorted); + + auto sortEffects = + [](llvm::SmallVectorImpl &effects) { + llvm::stable_sort(effects, [](const MemoryEffects::EffectInstance &a, + const MemoryEffects::EffectInstance &b) { + if (a.getStage() < b.getStage()) + return true; + + if (a.getStage() == b.getStage()) + return a.getEffect()->getPriority() < b.getEffect()->getPriority(); + + return false; // b before a + }); + }; + sortEffects(effectsSorted); + + return effectsSorted; +} + bool mlir::isMemoryEffectFree(Operation *op) { if (auto memInterface = dyn_cast(op)) { if (!memInterface.hasNoEffect()) return false; + // If the op does not have recursive side effects, then it is memory effect // free. if (!op->hasTrait()) return true; + } else if (!op->hasTrait()) { // Otherwise, if the op does not implement the memory effect interface and // it does not have recursive side effects, then it cannot be known that the @@ -338,6 +371,7 @@ bool mlir::isMemoryEffectFree(Operation *op) { for (Operation &op : region.getOps()) if (!isMemoryEffectFree(&op)) return false; + return true; } diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp index 111f58ef92f51..bd4eebf553a7e 100644 --- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp +++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp @@ -15,12 +15,14 @@ #include "mlir/IR/Operation.h" #include "mlir/IR/OperationSupport.h" #include "mlir/IR/PatternMatch.h" +#include "mlir/IR/Value.h" #include "mlir/Interfaces/LoopLikeInterface.h" #include "mlir/Interfaces/SideEffectInterfaces.h" #include "mlir/Interfaces/SubsetOpInterface.h" #include "llvm/Support/Debug.h" #include "llvm/Support/DebugLog.h" #include +#include #define DEBUG_TYPE "licm" @@ -58,62 +60,256 @@ static bool canBeHoisted(Operation *op, op, [&](OpOperand &operand) { return definedOutside(operand.get()); }); } -size_t mlir::moveLoopInvariantCode( - ArrayRef regions, - function_ref isDefinedOutsideRegion, - function_ref shouldMoveOutOfRegion, - function_ref moveOutOfRegion) { - size_t numMoved = 0; +/// Merges srcEffect's Memory Effect on its resource into the +/// resourceConflicts map, flagging the resource if the srcEffect +/// results in a conflict. +/// +/// \param resourceConflicts The map to store resources' conflicts status. +/// \param srcEffect The effect to merge into the resourceConflicts map. +/// \param srcHasConflict Whether the srcEffect results in a conflict based +/// on higher level analysis. +/// +/// resourceConflicts is modified by the function and will be non-empty +static void +mergeResource(DenseMap> + &resourceConflicts, + const MemoryEffects::EffectInstance &srcEffect, + bool srcHasConflict) { + + TypeID srcResourceID = srcEffect.getResource()->getResourceID(); + + bool srcIsAllocOrFree = isa(srcEffect.getEffect()) || + isa(srcEffect.getEffect()); + + LDBG() << " Merging Effect \"" << srcEffect.getEffect() << "<" << srcEffect.getResource()->getName() << ">\""; + + bool conflict = srcHasConflict || srcIsAllocOrFree; + + auto [dstIt, inserted] = resourceConflicts.insert( + std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect))); + if (inserted) { + LDBG() << " Effect inserted to map"; + return; + } - for (Region *region : regions) { - LDBG() << "Original loop:\n" << *region->getParentOp(); + // resource already in use + bool dstHasConflict = dstIt->second.first; + auto dstEffect = dstIt->second.second; - std::queue worklist; - // Add top-level operations in the loop body to the worklist. - for (Operation &op : region->getOps()) - worklist.push(&op); + if (dstHasConflict) { + LDBG() << " Resource has existing conflict from Effect \"" << dstEffect.getEffect() << "<" << dstEffect.getResource()->getName() << ">\""; + return; + } - auto definedOutside = [&](Value value) { - return isDefinedOutsideRegion(value, region); - }; + bool srcWrite = isa(srcEffect.getEffect()); + bool dstRead = isa(dstEffect.getEffect()); + bool readBeforeWrite = dstRead && srcWrite; - while (!worklist.empty()) { - Operation *op = worklist.front(); - worklist.pop(); - // Skip ops that have already been moved. Check if the op can be hoisted. - if (op->getParentRegion() != region) - continue; + conflict = conflict || readBeforeWrite; - LDBG() << "Checking op: " - << OpWithFlags(op, OpPrintingFlags().skipRegions()); - if (!shouldMoveOutOfRegion(op, region) || - !canBeHoisted(op, definedOutside)) - continue; + LDBG() << " Resource " << dstEffect.getResource()->getName() << " updated to conflict status = " << conflict; + dstIt->second = std::make_pair(conflict, srcEffect); +} - LDBG() << "Moving loop-invariant op: " << *op; - moveOutOfRegion(op, region); - ++numMoved; +/// Returns true if any of op's operands is defined inside the loop. +static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) { + return llvm::any_of(op->getOperands(), [&] (Value v) { + return !loopLike.isDefinedOutsideOfLoop(v); + }); +} - // Since the op has been moved, we need to check its users within the - // top-level of the loop body. - for (Operation *user : op->getUsers()) - if (user->getParentRegion() == region) - worklist.push(user); +/// Returns true if: +/// (a) any of the resources used by op's Memory Effects have been +/// flagged as having a conflict within the resourceConflicts map OR +/// (b) op doesn't have a MemoryEffectOpInterface or has one but +/// without any specific effects. +static bool mayHaveMemoryEffectConflict( + Operation *op, + DenseMap> + &resourceConflicts) { + LDBG() << "Checking for memory effect conflict on op: " + << OpWithFlags(op, OpPrintingFlags().skipRegions()); + + auto memInterface = dyn_cast(op); + + // op does not implement the memory effect op interface + // shouldn't be flagged as movable to be conservative + if (!memInterface) + return true; + + // gather all effects on op + llvm::SmallVector effects; + memInterface.getEffects(effects); + + // op has interface but no effects, be conservative + if (effects.empty()) + return true; + + // RFC moving ops with HasRecursiveMemoryEffects that have nested ops + // with effects is not supported + + // op has no conflicts IFF all resources are flagged as having no conflicts + for (const MemoryEffects::EffectInstance &effect : effects) { + auto resourceID = effect.getResource()->getResourceID(); + + auto resConIt = resourceConflicts.find(resourceID); + assert(resConIt != resourceConflicts.end()); + + bool hasConflict = resConIt->second.first; + if (hasConflict) { + LDBG() << " has conflict on resource " << effect.getResource()->getName() + << " from Memory Effect Mem" << effect.getEffect(); + return true; } } - return numMoved; + return false; +} + +void mlir::gatherResourceConflicts( + LoopLikeOpInterface loopLike, Operation *op, + DenseMap> + &resourceConflicts) { + + LDBG() << "Gather conflicts: on loop " + << OpWithFlags(loopLike.getOperation(), OpPrintingFlags().skipRegions()) + << ", visiting op: " << OpWithFlags(op, OpPrintingFlags().skipRegions()); + + if (auto memInterface = dyn_cast(op)) { + // gather all effects on op + SmallVector effects = + MemoryEffects::getMemoryEffectsSorted(op); + + if (!effects.empty()) { + LDBG() << " # of effects = " << effects.size(); + + // Any input to the op could be the input data source + // for write effects in the same op. E.g., + // scf.for ... { + // %0 = foo.bar(...) : ... + // foo.baz(%0) // foo.baz has a MemWrite effect + // } + // The input %0 could be the data source for the write effect in + // foo.baz. Since %0 is loop-variant, this may cause a conflict on + // SomeResource as the MemWrite contents may change between loop iterations. + // A more complex analysis would be needed to determine + // if this is a true conflict or not. + bool writesConflict = hasLoopVariantInput(loopLike, op); + LDBG() << " has loop-variant input? " + << (writesConflict ? "true" : "false"); + + for (const MemoryEffects::EffectInstance &effect : effects) { + bool conflict = writesConflict && isa(effect.getEffect()); + + mergeResource(resourceConflicts, effect, conflict); + + // All writes to a resource that follow a read on any other resource + // have to be considered as causing a conflict, similar + // to how we handle loop-variant inputs above; guaranteeing that the read + // is invariant and won't affect the write requires more robust logic. + if (isa(effect.getEffect())) { + LDBG() << " has MemRead; subsequent Writes will be treated as conflicting"; + writesConflict = true; + } + } + } + } + + for (Region ®ion : op->getRegions()) + for (Operation &opInner : region.getOps()) + gatherResourceConflicts(loopLike, &opInner, resourceConflicts); +} + +size_t mlir::moveLoopInvariantCode( + LoopLikeOpInterface loopLike, + function_ref isDefinedOutsideRegion, + function_ref shouldMoveOutOfRegion, + function_ref moveOutOfRegion) { + size_t numMovedTotal = 0; + + // TODO: see if this can be spec'd in a meaningful way to add back later. + // + // check that the loop isn't dead + auto isDead = loopLike.isZeroTrip(); + + LDBG() << "Loop " << OpWithFlags(loopLike.getOperation(), OpPrintingFlags().skipRegions()) + << " has constant bounds and steps? isZeroTrip()? " << (isDead.has_value() ? (isDead.value() ? "YES, YES" : "YES, NO") : "NO, NULL"); + + if (isDead.has_value() && isDead.value()) + return numMovedTotal; + + int numMoved = 0; + + do { + // reset value for iteration + numMoved = 0; + + // Go through loop body and map out resource usages. + // op->regions are essentially merged sequentially. + // E.g., an if's "then" and "else" regions are treated like one + // continuous region --> need to add fork checking. + // + // loop "do" and "then" regions also merged. + DenseMap> + resourceConflicts; + gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts); + + auto regions = loopLike.getLoopRegions(); + for (Region *region : regions) { + LDBG() << "Original loop:\n" << *region->getParentOp(); + + std::queue worklist; + // Add top-level operations in the loop body to the worklist. + for (Operation &op : region->getOps()) + worklist.push(&op); + + auto definedOutside = [&](Value value) { + return isDefinedOutsideRegion(value, region); + }; + + while (!worklist.empty()) { + Operation *op = worklist.front(); + worklist.pop(); + // Skip ops that have already been moved. Check if the op can be hoisted. + if (op->getParentRegion() != region) + continue; + + LDBG() << "Checking op: " + << OpWithFlags(op, OpPrintingFlags().skipRegions()); + + bool noMemoryConflicts = + isMemoryEffectFree(op) || + !mayHaveMemoryEffectConflict(op, resourceConflicts); + + if (!noMemoryConflicts || !shouldMoveOutOfRegion(op, region) || + !canBeHoisted(op, definedOutside)) + continue; + + LDBG() << "Moving loop-invariant op: " << *op; + moveOutOfRegion(op, region); + ++numMoved; + + // Since the op has been moved, we need to check its users within the + // top-level of the loop body. + for (Operation *user : op->getUsers()) + if (user->getParentRegion() == region) + worklist.push(user); + } + } + + numMovedTotal += numMoved; + } while (numMoved > 0); + + return numMovedTotal; } size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) { return moveLoopInvariantCode( - loopLike.getLoopRegions(), + loopLike, [&](Value value, Region *) { return loopLike.isDefinedOutsideOfLoop(value); }, - [&](Operation *op, Region *) { - return isMemoryEffectFree(op) && isSpeculatable(op); - }, + [&](Operation *op, Region *) { return isSpeculatable(op); }, [&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); }); } diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir index c1604e226a334..02bd3faca3f58 100644 --- a/mlir/test/Transforms/loop-invariant-code-motion.mlir +++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir @@ -1437,3 +1437,341 @@ func.func @do_not_hoist_vector_transfer_ops_memref( } func.return %final : vector<4x4xf32> } + +// ----- + +// CHECK-LABEL func.func @move_single_resource_basic +func.func @move_single_resource_basic() attributes {} { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 10 : i32 + %c2_i32 = arith.constant 1 : i32 + %c0_i32_0 = arith.constant 0 : i32 + + // Single write effect on one resource in a triple-nested loop + // No loop-variant inputs to op and no read effects -> movable + // CHECK: "test.test_effects_write_A"() : () -> () + // CHECK: scf.for + // CHECK: scf.for + // CHECK: scf.for + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + "test.test_effects_write_A"() : () -> () + } + } + } + return +} + +// ----- + +// CHECK-LABEL func.func @move_single_resource_write_dominant +func.func @move_single_resource_write_dominant() attributes {} { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 10 : i32 + %c2_i32 = arith.constant 1 : i32 + %c0_i32_0 = arith.constant 0 : i32 + + // Write effect on one resource followed by a Read. + // No loop-variant inputs to Write op, no conflict on + // "A" --> both ops movable + // CHECK: "test.test_effects_write_A"() : () -> () + // CHECK: "test.test_effects_read_A"() : () -> () + // CHECK: scf.for + // CHECK: scf.for + // CHECK: scf.for + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + "test.test_effects_write_A"() : () -> () + "test.test_effects_read_A"() : () -> () + } + } + } + return +} + +// ----- + +// CHECK-LABEL func.func @move_single_resource_read_dominant +func.func @move_single_resource_read_dominant() attributes {} { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 10 : i32 + %c2_i32 = arith.constant 1 : i32 + %c0_i32_0 = arith.constant 0 : i32 + + // CHECK: scf.for %arg0 + // CHECK: scf.for %arg1 + // CHECK: scf.for %arg2 + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + + // Read effect on "A" dominates write. + // Causes conflict "A" --> not movable. + // CHECK: "test.test_effects_read_A"() : () -> () + // CHECK: "test.test_effects_write_A"() : () -> () + + "test.test_effects_read_A"() : () -> () + "test.test_effects_write_A"() : () -> () + } + } + } + return +} + +// ----- + +// CHECK-LABEL func.func @move_single_resource_basic_conflict +func.func @move_single_resource_basic_conflict() attributes {} { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 10 : i32 + %c2_i32 = arith.constant 1 : i32 + %c0_i32_0 = arith.constant 0 : i32 + %c0_i32_1 = arith.constant 0 : i32 + %c0_i32_2 = arith.constant 0 : i32 + + // CHECK: scf.for + // CHECK: scf.for + // CHECK: scf.for + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + + // CHECK: "test.test_effects_write_A"() : () -> () + // CHECK: "test.test_effects_read_A"() : () -> () + // Read of "A" dominates Write "A" --> "A" is in conflict. + // "C" is not in conflict but, since all resources used + // by op aren't conflict free, they're not movable. + // CHECK: "test.test_effects_write_AC"() : () -> () + // CHECK: "test.test_effects_read_AC"() : () -> () + + "test.test_effects_write_A"() : () -> () + "test.test_effects_read_A"() : () -> () + "test.test_effects_write_AC"() : () -> () + "test.test_effects_read_AC"() : () -> () + } + } + } + return +} + +// ----- + +// CHECK-LABEL func.func @move_single_resource_if_region +func.func @move_single_resource_if_region() attributes {} { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 10 : i32 + %c2_i32 = arith.constant 1 : i32 + %c3_i32 = arith.constant 5 : i32 + + // CHECK: scf.for + // CHECK: scf.for + // CHECK: scf.for + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + %1 = arith.cmpi slt, %arg0, %c3_i32 : i32 + + scf.if %1 { + // Checking that we're not moving ops out of + // non-loop regions. + // CHECK: "test.test_effects_write_A"() : () -> () + // CHECK: "test.test_effects_read_A"() : () -> () + + "test.test_effects_write_A"() : () -> () + "test.test_effects_read_A"() : () -> () + } + } + } + } + return +} + +// ----- + +// CHECK-LABEL func.func @move_single_resource_for_inside_if_region +func.func @move_single_resource_for_inside_if_region() attributes {} { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 10 : i32 + %c2_i32 = arith.constant 1 : i32 + %c3_i32 = arith.constant 5 : i32 + + // CHECK: scf.for + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + %1 = arith.cmpi slt, %arg0, %c3_i32 : i32 + + // CHECK: scf.if + scf.if %1 { + %c0_i32_0 = arith.constant 0 : i32 + + // Checking that we can move ops out of loops nested + // within other regions, without moving ops out of + // the parent, non-loop region. + // CHECK: "test.test_effects_write_A"() : () -> () + // CHECK: scf.for + // CHECK: scf.for + + scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + "test.test_effects_write_A"() : () -> () + } + } + } + } + return +} + +// ----- + +// CHECK-LABEL func.func @move_multi_resource_comprehensive +func.func @move_multi_resource_comprehensive() attributes {} { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 10 : i32 + %c2_i32 = arith.constant 1 : i32 + %c0_i32_2 = arith.constant 0 : i32 + %c3_i32 = arith.constant 5 : i32 + + // CHECK: scf.for + // CHECK: scf.for + // CHECK: scf.for + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + + // CHECK: "test.test_effects_write_CD"() : () -> () + // CHECK: "test.test_effects_read_CD"() : () -> () + // CHECK: "test.test_effects_write_EF"() : () -> () + // CHECK: "test.test_effects_read_EF"() : () -> () + + // Both of these should be empty and moved out of their parent + scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + "test.test_effects_write_CD"() : () -> () + "test.test_effects_read_CD"() : () -> () + "test.test_effects_write_EF"() : () -> () + "test.test_effects_read_EF"() : () -> () + } + } + + %1 = arith.cmpi slt, %arg0, %c3_i32 : i32 + scf.if %1 { + %c0_i32_0 = arith.constant 0 : i32 + %c0_i32_1 = arith.constant 0 : i32 + + // CHECK: scf.for + // CHECK: "test.test_effects_write_B"() : () -> () + // CHECK: "test.test_effects_read_B"() : () -> () + // CHECK: scf.for + // CHECK: scf.for + // CHECK: scf.for + + scf.for %arg3 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + + // CHECK: "test.test_effects_write_A"() : () -> () + // CHECK: "test.test_effects_read_A"() : () -> () + + // emptyoop should be empty and moved out of parent + scf.for %arg4 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + "test.test_effects_write_A"() : () -> () + "test.test_effects_read_A"() : () -> () + } + + // Loop should be empty and moved out of parent + scf.for %arg5 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + "test.test_effects_write_B"() : () -> () + "test.test_effects_read_B"() : () -> () + } + + // CHECK: "test.test_effects_write_AC"() : () -> () + // CHECK: "test.test_effects_read_AC"() : () -> () + + // Loop should be empty and moved out of parent + scf.for %arg6 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + "test.test_effects_write_AC"() : () -> () + "test.test_effects_read_AC"() : () -> () + } + } + } + else { + // CHECK: "test.test_effects_write_F"() : () -> () + // CHECK: "test.test_effects_read_F"() : () -> () + + "test.test_effects_write_F"() : () -> () + "test.test_effects_read_F"() : () -> () + } + } + return +} + +// ----- + +// CHECK-LABEL func.func @move_multi_resource_write_conflicts +func.func @move_multi_resource_write_conflicts() attributes {} { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 10 : i32 + %c2_i32 = arith.constant 1 : i32 + + // CHECK: %{{.*}} = arith.constant 7 : index + + // CHECK: "test.test_effects_write_B"() : () -> () + // CHECK: "test.test_effects_read_B"() : () -> () + + // CHECK: "test.test_effects_write_A_with_input"(%c7) : (index) -> () + // CHECK: "test.test_effects_read_A"() : () -> () + + // CHECK: scf.for + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + %input = arith.constant 7 : index + "test.test_effects_write_A_with_input"(%input) : (index) -> () + "test.test_effects_read_A"() : () -> () + + "test.test_effects_write_B"() : () -> () + "test.test_effects_read_B"() : () -> () + } + + // CHECK: "test.test_effects_read_A"() : () -> () + // CHECK: scf.for + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + // CHECK: "test.test_effects_read_A_write_B"() : () -> () + // CHECK: "test.test_effects_read_B"() : () -> () + + "test.test_effects_read_A_write_B"() : () -> () + "test.test_effects_read_B"() : () -> () + + "test.test_effects_read_A"() : () -> () + } + + // CHECK: "test.test_effects_read_A"() : () -> () + // CHECK: scf.for + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + // CHECK: "test.test_effects_read_A_then_write_B"() : () -> () + // CHECK: "test.test_effects_read_B"() : () -> () + + "test.test_effects_read_A_then_write_B"() : () -> () + "test.test_effects_read_B"() : () -> () + + "test.test_effects_read_A"() : () -> () + } + + // CHECK: "test.test_effects_write_A_then_read_B"() : () -> () + // CHECK: "test.test_effects_read_B"() : () -> () + // CHECK: "test.test_effects_read_A"() : () -> () + // CHECK: scf.for + + scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 { + "test.test_effects_write_A_then_read_B"() : () -> () + "test.test_effects_read_B"() : () -> () + "test.test_effects_read_A"() : () -> () + } + + return +} \ No newline at end of file diff --git a/mlir/test/lib/Dialect/Test/TestOps.h b/mlir/test/lib/Dialect/Test/TestOps.h index 4201ade9795e7..2dba8d79cb52f 100644 --- a/mlir/test/lib/Dialect/Test/TestOps.h +++ b/mlir/test/lib/Dialect/Test/TestOps.h @@ -55,6 +55,29 @@ struct TestResource : public mlir::SideEffects::Resource::Base { llvm::StringRef getName() final { return ""; } }; +struct TestResourceA : public mlir::SideEffects::Resource::Base { + llvm::StringRef getName() final { return ""; } +}; + +struct TestResourceB : public mlir::SideEffects::Resource::Base { + llvm::StringRef getName() final { return ""; } +}; + +struct TestResourceC : public mlir::SideEffects::Resource::Base { + llvm::StringRef getName() final { return ""; } +}; + +struct TestResourceD : public mlir::SideEffects::Resource::Base { + llvm::StringRef getName() final { return ""; } +}; + +struct TestResourceE : public mlir::SideEffects::Resource::Base { + llvm::StringRef getName() final { return ""; } +}; + +struct TestResourceF : public mlir::SideEffects::Resource::Base { + llvm::StringRef getName() final { return ""; } +}; //===----------------------------------------------------------------------===// // PropertiesWithCustomPrint //===----------------------------------------------------------------------===// diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td index 5564264ed8b0b..1b3e3c7e21b08 100644 --- a/mlir/test/lib/Dialect/Test/TestOps.td +++ b/mlir/test/lib/Dialect/Test/TestOps.td @@ -2976,6 +2976,72 @@ def TestEffectsResult : TEST_Op<"test_effects_result"> { let results = (outs Res); } +//===----------------------------------------------------------------------===// +// Test Ops with multiple effects for Loop Invariant Code Motion. +//===----------------------------------------------------------------------===// + +def TestResourceA : Resource<"TestResourceA">; +def TestResourceB : Resource<"TestResourceB">; +def TestResourceC : Resource<"TestResourceC">; +def TestResourceD : Resource<"TestResourceD">; +def TestResourceE : Resource<"TestResourceE">; +def TestResourceF : Resource<"TestResourceF">; + +def TestEffectsWriteA : TEST_Op<"test_effects_write_A", + [MemoryEffects<[MemWrite]>]>; + +def TestEffectsWriteAWithInput : TEST_Op<"test_effects_write_A_with_input", + [MemoryEffects<[MemWrite]>]> { + + let arguments = (ins AnyType:$arg); +} + +def TestEffectsReadA : TEST_Op<"test_effects_read_A", + [MemoryEffects<[MemRead]>]>; + +def TestEffectsReadAWriteB : TEST_Op<"test_effects_read_A_write_B", + [MemoryEffects<[MemRead, + MemWrite]>]>; + +def TestEffectsReadAThenWriteB : TEST_Op<"test_effects_read_A_then_write_B", + [MemoryEffects<[MemRead, + MemWrite]>]>; + +def TestEffectsWriteAThenReadB : TEST_Op<"test_effects_write_A_then_read_B", + [MemoryEffects<[MemWrite, + MemRead]>]>; + +def TestEffectsWriteB : TEST_Op<"test_effects_write_B", + [MemoryEffects<[MemWrite]>]>; + +def TestEffectsReadB : TEST_Op<"test_effects_read_B", + [MemoryEffects<[MemRead]>]>; + +def TestEffectsWriteF : TEST_Op<"test_effects_write_F", + [MemoryEffects<[MemWrite]>]>; + +def TestEffectsReadF : TEST_Op<"test_effects_read_F", + [MemoryEffects<[MemRead]>]>; + +def TestEffectsWriteAC : TEST_Op<"test_effects_write_AC", + [MemoryEffects<[MemWrite, MemWrite]>]>; + +def TestEffectsReadAC : TEST_Op<"test_effects_read_AC", + [MemoryEffects<[MemRead, MemRead]>]>; + +def TestEffectsWriteCD : TEST_Op<"test_effects_write_CD", + [MemoryEffects<[MemWrite, MemWrite]>]>; + +def TestEffectsReadCD : TEST_Op<"test_effects_read_CD", + [MemoryEffects<[MemRead, MemRead]>]>; + +def TestEffectsWriteEF : TEST_Op<"test_effects_write_EF", + [MemoryEffects<[MemWrite, MemWrite]>]>; + +def TestEffectsReadEF : TEST_Op<"test_effects_read_EF", + [MemoryEffects<[MemRead, MemRead]>]>; + + //===----------------------------------------------------------------------===// // Test Ops with verifiers //===----------------------------------------------------------------------===//