[MLIR] Implement LoopLikeInterface for loop.parallel

Summary:
This is to allow optimizations like loop invariant code motion to work
on the ParallelOp.

Additional small cleanup on the ForOp implementation of
LoopLikeInterface and the test file of loop-invariant-code-motion.

Differential Revision: https://reviews.llvm.org/D77128
This commit is contained in:
Tres Popp 2020-03-31 11:41:51 +02:00
parent 51d594d33c
commit a67cd71acd
3 changed files with 70 additions and 8 deletions

View File

@ -246,7 +246,9 @@ def IfOp : Loop_Op<"if",
}
def ParallelOp : Loop_Op<"parallel",
[AttrSizedOperandSegments, SingleBlockImplicitTerminator<"YieldOp">]> {
[AttrSizedOperandSegments,
DeclareOpInterfaceMethods<LoopLikeOpInterface>,
SingleBlockImplicitTerminator<"YieldOp">]> {
let summary = "parallel for operation";
let description = [{
The "loop.parallel" operation represents a loop nest taking 4 groups of SSA

View File

@ -182,7 +182,7 @@ bool ForOp::isDefinedOutsideOfLoop(Value value) {
LogicalResult ForOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
for (auto op : ops)
op->moveBefore(this->getOperation());
op->moveBefore(*this);
return success();
}
@ -191,8 +191,8 @@ ForOp mlir::loop::getForInductionVarOwner(Value val) {
if (!ivArg)
return ForOp();
assert(ivArg.getOwner() && "unlinked block argument");
auto *containingInst = ivArg.getOwner()->getParentOp();
return dyn_cast_or_null<ForOp>(containingInst);
auto *containingOp = ivArg.getOwner()->getParentOp();
return dyn_cast_or_null<ForOp>(containingOp);
}
//===----------------------------------------------------------------------===//
@ -459,13 +459,25 @@ static void print(OpAsmPrinter &p, ParallelOp op) {
op.getAttrs(), /*elidedAttrs=*/ParallelOp::getOperandSegmentSizeAttr());
}
Region &ParallelOp::getLoopBody() { return region(); }
bool ParallelOp::isDefinedOutsideOfLoop(Value value) {
return !region().isAncestor(value.getParentRegion());
}
LogicalResult ParallelOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
for (auto op : ops)
op->moveBefore(*this);
return success();
}
ParallelOp mlir::loop::getParallelForInductionVarOwner(Value val) {
auto ivArg = val.dyn_cast<BlockArgument>();
if (!ivArg)
return ParallelOp();
assert(ivArg.getOwner() && "unlinked block argument");
auto *containingInst = ivArg.getOwner()->getParentOp();
return dyn_cast<ParallelOp>(containingInst);
auto *containingOp = ivArg.getOwner()->getParentOp();
return dyn_cast<ParallelOp>(containingOp);
}
//===----------------------------------------------------------------------===//

View File

@ -1,4 +1,4 @@
// RUN: mlir-opt %s -loop-invariant-code-motion -split-input-file | FileCheck %s
// RUN: mlir-opt %s -split-input-file -loop-invariant-code-motion | FileCheck %s
func @nested_loops_both_having_invariant_code() {
%m = alloc() : memref<10xf32>
@ -25,6 +25,8 @@ func @nested_loops_both_having_invariant_code() {
return
}
// -----
func @nested_loops_code_invariant_to_both() {
%m = alloc() : memref<10xf32>
%cf7 = constant 7.0 : f32
@ -44,6 +46,8 @@ func @nested_loops_code_invariant_to_both() {
return
}
// -----
func @single_loop_nothing_invariant() {
%m1 = alloc() : memref<10xf32>
%m2 = alloc() : memref<10xf32>
@ -65,6 +69,8 @@ func @single_loop_nothing_invariant() {
return
}
// -----
func @invariant_code_inside_affine_if() {
%m = alloc() : memref<10xf32>
%cf8 = constant 8.0 : f32
@ -81,7 +87,7 @@ func @invariant_code_inside_affine_if() {
// CHECK: %0 = alloc() : memref<10xf32>
// CHECK-NEXT: %cst = constant 8.000000e+00 : f32
// CHECK-NEXT: affine.for %arg0 = 0 to 10 {
// CHECK-NEXT: %1 = affine.apply #map3(%arg0)
// CHECK-NEXT: %1 = affine.apply #map0(%arg0)
// CHECK-NEXT: affine.if #set0(%arg0, %1) {
// CHECK-NEXT: %2 = addf %cst, %cst : f32
// CHECK-NEXT: affine.store %2, %0[%arg0] : memref<10xf32>
@ -91,6 +97,8 @@ func @invariant_code_inside_affine_if() {
return
}
// -----
func @invariant_affine_if() {
%m = alloc() : memref<10xf32>
%cf8 = constant 8.0 : f32
@ -114,6 +122,8 @@ func @invariant_affine_if() {
return
}
// -----
func @invariant_affine_if2() {
%m = alloc() : memref<10xf32>
%cf8 = constant 8.0 : f32
@ -139,6 +149,8 @@ func @invariant_affine_if2() {
return
}
// -----
func @invariant_affine_nested_if() {
%m = alloc() : memref<10xf32>
%cf8 = constant 8.0 : f32
@ -169,6 +181,8 @@ func @invariant_affine_nested_if() {
return
}
// -----
func @invariant_affine_nested_if_else() {
%m = alloc() : memref<10xf32>
%cf8 = constant 8.0 : f32
@ -205,6 +219,8 @@ func @invariant_affine_nested_if_else() {
return
}
// -----
func @invariant_loop_dialect() {
%ci0 = constant 0 : index
%ci10 = constant 10 : index
@ -226,6 +242,8 @@ func @invariant_loop_dialect() {
return
}
// -----
func @variant_loop_dialect() {
%ci0 = constant 0 : index
%ci10 = constant 10 : index
@ -244,3 +262,33 @@ func @variant_loop_dialect() {
return
}
// -----
func @parallel_loop_with_invariant() {
%c0 = constant 0 : index
%c10 = constant 10 : index
%c1 = constant 1 : index
%c7 = constant 7 : i32
%c8 = constant 8 : i32
loop.parallel (%arg0, %arg1) = (%c0, %c0) to (%c10, %c10) step (%c1, %c1) {
%v0 = addi %c7, %c8 : i32
%v3 = addi %arg0, %arg1 : index
}
// CHECK-LABEL: func @parallel_loop_with_invariant
// CHECK: %c0 = constant 0 : index
// CHECK-NEXT: %c10 = constant 10 : index
// CHECK-NEXT: %c1 = constant 1 : index
// CHECK-NEXT: %c7_i32 = constant 7 : i32
// CHECK-NEXT: %c8_i32 = constant 8 : i32
// CHECK-NEXT: addi %c7_i32, %c8_i32 : i32
// CHECK-NEXT: loop.parallel (%arg0, %arg1) = (%c0, %c0) to (%c10, %c10) step (%c1, %c1)
// CHECK-NEXT: addi %arg0, %arg1 : index
// CHECK-NEXT: yield
// CHECK-NEXT: }
// CHECK-NEXT: return
return
}