[Mlir-commits] [mlir] 3430ae1 - [mlir] Update LICM to support Graph Regions
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Apr 15 10:30:31 PDT 2022
Author: Mogball
Date: 2022-04-15T17:30:27Z
New Revision: 3430ae1e7baa11992084e3624eb28f9dfd941ec7
URL: https://github.com/llvm/llvm-project/commit/3430ae1e7baa11992084e3624eb28f9dfd941ec7
DIFF: https://github.com/llvm/llvm-project/commit/3430ae1e7baa11992084e3624eb28f9dfd941ec7.diff
LOG: [mlir] Update LICM to support Graph Regions
Changes the algorithm of LICM to support graph regions (no guarantee of topologically sorted order). Also fixes an issue where ops with recursive side effects and regions would not be hoisted if any nested ops used operands that were defined within the nested region.
Reviewed By: rriddle
Differential Revision: https://reviews.llvm.org/D122465
Added:
Modified:
mlir/lib/Interfaces/LoopLikeInterface.cpp
mlir/test/Transforms/loop-invariant-code-motion.mlir
mlir/test/lib/Dialect/Test/TestOps.td
Removed:
################################################################################
diff --git a/mlir/lib/Interfaces/LoopLikeInterface.cpp b/mlir/lib/Interfaces/LoopLikeInterface.cpp
index 554bea1e61677..a0ea47eaaf975 100644
--- a/mlir/lib/Interfaces/LoopLikeInterface.cpp
+++ b/mlir/lib/Interfaces/LoopLikeInterface.cpp
@@ -10,6 +10,7 @@
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/Debug.h"
+#include <queue>
using namespace mlir;
@@ -26,75 +27,88 @@ using namespace mlir;
// LoopLike Utilities
//===----------------------------------------------------------------------===//
-// Checks whether the given op can be hoisted by checking that
-// - the op and any of its contained operations do not depend on SSA values
-// defined inside of the loop (by means of calling definedOutside).
-// - the op has no side-effects. If sideEffecting is Never, sideeffects of this
-// op and its nested ops are ignored.
-static bool canBeHoisted(Operation *op,
- function_ref<bool(Value)> definedOutside) {
- // Check that dependencies are defined outside of loop.
- if (!llvm::all_of(op->getOperands(), definedOutside))
- return false;
- // Check whether this op is side-effect free. If we already know that there
- // can be no side-effects because the surrounding op has claimed so, we can
- // (and have to) skip this step.
+/// Returns true if the given operation is side-effect free as are all of its
+/// nested operations.
+///
+/// TODO: There is a duplicate function in ControlFlowSink. Move
+/// `moveLoopInvariantCode` to TransformUtils and then factor out this function.
+static bool isSideEffectFree(Operation *op) {
if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+ // If the op has side-effects, it cannot be moved.
if (!memInterface.hasNoEffect())
return false;
- // If the operation doesn't have side effects and it doesn't recursively
- // have side effects, it can always be hoisted.
+ // If the op does not have recursive side effects, then it can be moved.
if (!op->hasTrait<OpTrait::HasRecursiveSideEffects>())
return true;
-
- // Otherwise, if the operation doesn't provide the memory effect interface
- // and it doesn't have recursive side effects we treat it conservatively as
- // side-effecting.
} else if (!op->hasTrait<OpTrait::HasRecursiveSideEffects>()) {
+ // 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
+ // op is moveable.
return false;
}
- // Recurse into the regions for this op and check whether the contained ops
- // can be hoisted.
- for (auto ®ion : op->getRegions()) {
- for (auto &block : region) {
- for (auto &innerOp : block)
- if (!canBeHoisted(&innerOp, definedOutside))
- return false;
- }
- }
+ // Recurse into the regions and ensure that all nested ops can also be moved.
+ for (Region ®ion : op->getRegions())
+ for (Operation &op : region.getOps())
+ if (!isSideEffectFree(&op))
+ return false;
return true;
}
+/// Checks whether the given op can be hoisted by checking that
+/// - the op and none of its contained operations depend on values inside of the
+/// loop (by means of calling definedOutside).
+/// - the op has no side-effects.
+static bool canBeHoisted(Operation *op,
+ function_ref<bool(Value)> definedOutside) {
+ if (!isSideEffectFree(op))
+ return false;
+
+ // Do not move terminators.
+ if (op->hasTrait<OpTrait::IsTerminator>())
+ return false;
+
+ // Walk the nested operations and check that all used values are either
+ // defined outside of the loop or in a nested region, but not at the level of
+ // the loop body.
+ auto walkFn = [&](Operation *child) {
+ for (Value operand : child->getOperands()) {
+ // Ignore values defined in a nested region.
+ if (op->isAncestor(operand.getParentRegion()->getParentOp()))
+ continue;
+ if (!definedOutside(operand))
+ return WalkResult::interrupt();
+ }
+ return WalkResult::advance();
+ };
+ return !op->walk(walkFn).wasInterrupted();
+}
+
void mlir::moveLoopInvariantCode(LoopLikeOpInterface looplike) {
- auto &loopBody = looplike.getLoopBody();
-
- // We use two collections here as we need to preserve the order for insertion
- // and this is easiest.
- SmallPtrSet<Operation *, 8> willBeMovedSet;
- SmallVector<Operation *, 8> opsToMove;
-
- // Helper to check whether an operation is loop invariant wrt. SSA properties.
- auto isDefinedOutsideOfBody = [&](Value value) {
- auto *definingOp = value.getDefiningOp();
- return (definingOp && !!willBeMovedSet.count(definingOp)) ||
- looplike.isDefinedOutsideOfLoop(value);
+ Region *loopBody = &looplike.getLoopBody();
+
+ std::queue<Operation *> worklist;
+ // Add top-level operations in the loop body to the worklist.
+ for (Operation &op : loopBody->getOps())
+ worklist.push(&op);
+
+ auto definedOutside = [&](Value value) {
+ return looplike.isDefinedOutsideOfLoop(value);
};
- // Do not use walk here, as we do not want to go into nested regions and hoist
- // operations from there. These regions might have semantics unknown to this
- // rewriting. If the nested regions are loops, they will have been processed.
- for (auto &block : loopBody) {
- for (auto &op : block.without_terminator()) {
- if (canBeHoisted(&op, isDefinedOutsideOfBody)) {
- opsToMove.push_back(&op);
- willBeMovedSet.insert(&op);
- }
- }
- }
+ 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() != loopBody || !canBeHoisted(op, definedOutside))
+ continue;
- // For all instructions that we found to be invariant, move outside of the
- // loop.
- for (Operation *op : opsToMove)
looplike.moveOutOfLoop(op);
+
+ // 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() == loopBody)
+ worklist.push(user);
+ }
}
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index 488726321da40..1127711d758f1 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -157,10 +157,10 @@ func @invariant_affine_nested_if() {
affine.for %arg0 = 0 to 10 {
affine.for %arg1 = 0 to 10 {
affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
- %cf9 = arith.addf %cf8, %cf8 : f32
- affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
- %cf10 = arith.addf %cf9, %cf9 : f32
- }
+ %cf9 = arith.addf %cf8, %cf8 : f32
+ affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
+ %cf10 = arith.addf %cf9, %cf9 : f32
+ }
}
}
}
@@ -168,6 +168,7 @@ func @invariant_affine_nested_if() {
// CHECK: memref.alloc
// CHECK-NEXT: arith.constant
// CHECK-NEXT: affine.for
+ // CHECK-NEXT: }
// CHECK-NEXT: affine.for
// CHECK-NEXT: affine.if
// CHECK-NEXT: arith.addf
@@ -175,7 +176,6 @@ func @invariant_affine_nested_if() {
// CHECK-NEXT: arith.addf
// CHECK-NEXT: }
// CHECK-NEXT: }
- // CHECK-NEXT: }
return
@@ -319,6 +319,110 @@ func @nested_uses_inside(%lb: index, %ub: index, %step: index) {
scf.yield %val2: index
}
}
- return
+ return
}
+// -----
+
+// Test that two ops that feed into each other are moved without violating
+// dominance in non-graph regions.
+// CHECK-LABEL: func @invariant_subgraph
+// CHECK-SAME: %{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %[[ARG:.*]]: i32
+func @invariant_subgraph(%lb: index, %ub: index, %step: index, %arg: i32) {
+ // CHECK: %[[V0:.*]] = arith.addi %[[ARG]], %[[ARG]]
+ // CHECK-NEXT: %[[V1:.*]] = arith.addi %[[ARG]], %[[V0]]
+ // CHECK-NEXT: scf.for
+ scf.for %i = %lb to %ub step %step {
+ // CHECK-NEXT: "test.sink"(%[[V1]])
+ %v0 = arith.addi %arg, %arg : i32
+ %v1 = arith.addi %arg, %v0 : i32
+ "test.sink"(%v1) : (i32) -> ()
+ }
+ return
+}
+
+// -----
+
+// Test invariant nested loop is hoisted.
+// CHECK-LABEL: func @test_invariant_nested_loop
+func @test_invariant_nested_loop() {
+ // CHECK: %[[C:.*]] = arith.constant
+ %0 = arith.constant 5 : i32
+ // CHECK: %[[V0:.*]] = arith.addi %[[C]], %[[C]]
+ // CHECK-NEXT: %[[V1:.*]] = arith.addi %[[V0]], %[[C]]
+ // CHECK-NEXT: test.graph_loop
+ // CHECK-NEXT: ^bb0(%[[ARG0:.*]]: i32)
+ // CHECK-NEXT: %[[V2:.*]] = arith.subi %[[ARG0]], %[[ARG0]]
+ // CHECK-NEXT: test.region_yield %[[V2]]
+ // CHECK: test.graph_loop
+ // CHECK-NEXT: test.region_yield %[[V1]]
+ test.graph_loop {
+ %1 = arith.addi %0, %0 : i32
+ %2 = arith.addi %1, %0 : i32
+ test.graph_loop {
+ ^bb0(%arg0: i32):
+ %3 = arith.subi %arg0, %arg0 : i32
+ test.region_yield %3 : i32
+ } : () -> ()
+ test.region_yield %2 : i32
+ } : () -> ()
+ return
+}
+
+
+// -----
+
+// Test ops in a graph region are hoisted.
+// CHECK-LABEL: func @test_invariants_in_graph_region
+func @test_invariants_in_graph_region() {
+ // CHECK: test.single_no_terminator_op
+ test.single_no_terminator_op : {
+ // CHECK-NEXT: %[[C:.*]] = arith.constant
+ // CHECK-NEXT: %[[V1:.*]] = arith.addi %[[C]], %[[C]]
+ // CHECK-NEXT: %[[V0:.*]] = arith.addi %[[C]], %[[V1]]
+ test.graph_loop {
+ %v0 = arith.addi %c0, %v1 : i32
+ %v1 = arith.addi %c0, %c0 : i32
+ %c0 = arith.constant 5 : i32
+ test.region_yield %v0 : i32
+ } : () -> ()
+ }
+ return
+}
+
+// -----
+
+// Test ops in a graph region are hoisted in topological order into non-graph
+// regions and that dominance is preserved.
+// CHECK-LABEL: func @test_invariant_backedge
+func @test_invariant_backedge() {
+ // CHECK-NEXT: %[[C:.*]] = arith.constant
+ // CHECK-NEXT: %[[V1:.*]] = arith.addi %[[C]], %[[C]]
+ // CHECK-NEXT: %[[V0:.*]] = arith.addi %[[C]], %[[V1]]
+ // CHECK-NEXT: test.graph_loop
+ test.graph_loop {
+ // CHECK-NEXT: test.region_yield %[[V0]]
+ %v0 = arith.addi %c0, %v1 : i32
+ %v1 = arith.addi %c0, %c0 : i32
+ %c0 = arith.constant 5 : i32
+ test.region_yield %v0 : i32
+ } : () -> ()
+ return
+}
+
+// -----
+
+// Test that cycles aren't hoisted from graph regions to non-graph regions.
+// CHECK-LABEL: func @test_invariant_cycle_not_hoisted
+func @test_invariant_cycle_not_hoisted() {
+ // CHECK: test.graph_loop
+ test.graph_loop {
+ // CHECK-NEXT: %[[A:.*]] = "test.a"(%[[B:.*]]) :
+ // CHECK-NEXT: %[[B]] = "test.b"(%[[A]]) :
+ // CHECK-NEXT: test.region_yield %[[A]]
+ %a = "test.a"(%b) : (i32) -> i32
+ %b = "test.b"(%a) : (i32) -> i32
+ test.region_yield %a : i32
+ } : () -> ()
+ return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 6405b1a78c258..644c6fd9e58d2 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -10,7 +10,9 @@
#define TEST_OPS
include "TestDialect.td"
+include "TestInterfaces.td"
include "mlir/Dialect/DLTI/DLTIBase.td"
+include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
include "mlir/IR/EnumAttr.td"
include "mlir/IR/OpBase.td"
include "mlir/IR/OpAsmInterface.td"
@@ -22,9 +24,8 @@ include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/CopyOpInterface.td"
include "mlir/Interfaces/DataLayoutInterfaces.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/LoopLikeInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
-include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
-include "TestInterfaces.td"
// Include the attribute definitions.
@@ -2759,14 +2760,14 @@ def : Pat<(TestDefaultStrAttrNoValueOp $value),
def TestResource : Resource<"TestResource">;
def TestEffectsOpA : TEST_Op<"op_with_effects_a"> {
- let arguments = (ins
- Arg<Variadic<AnyMemRef>, "", [MemRead]>,
- Arg<FlatSymbolRefAttr, "", [MemRead]>:$first,
- Arg<SymbolRefAttr, "", [MemWrite]>:$second,
- Arg<OptionalAttr<SymbolRefAttr>, "", [MemRead]>:$optional_symbol
- );
+ let arguments = (ins
+ Arg<Variadic<AnyMemRef>, "", [MemRead]>,
+ Arg<FlatSymbolRefAttr, "", [MemRead]>:$first,
+ Arg<SymbolRefAttr, "", [MemWrite]>:$second,
+ Arg<OptionalAttr<SymbolRefAttr>, "", [MemRead]>:$optional_symbol
+ );
- let results = (outs Res<AnyMemRef, "", [MemAlloc<TestResource>]>);
+ let results = (outs Res<AnyMemRef, "", [MemAlloc<TestResource>]>);
}
def TestEffectsOpB : TEST_Op<"op_with_effects_b",
@@ -2792,4 +2793,26 @@ def TestVerifiersOp : TEST_Op<"verifiers",
let hasRegionVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// Test Loop Op with a graph region
+//===----------------------------------------------------------------------===//
+
+// Test loop op with a graph region.
+def TestGraphLoopOp : TEST_Op<"graph_loop",
+ [LoopLikeOpInterface, NoSideEffect,
+ RecursiveSideEffects, SingleBlock,
+ RegionKindInterface, HasOnlyGraphRegion]> {
+ let arguments = (ins Variadic<AnyType>:$args);
+ let results = (outs Variadic<AnyType>:$rets);
+ let regions = (region SizedRegion<1>:$body);
+
+ let assemblyFormat = [{
+ $args $body attr-dict `:` functional-type(operands, results)
+ }];
+
+ let extraClassDeclaration = [{
+ mlir::Region &getLoopBody() { return getBody(); }
+ }];
+}
+
#endif // TEST_OPS
More information about the Mlir-commits
mailing list