diff --git a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td --- a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td +++ b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td @@ -246,7 +246,9 @@ } def ParallelOp : Loop_Op<"parallel", - [AttrSizedOperandSegments, SingleBlockImplicitTerminator<"YieldOp">]> { + [AttrSizedOperandSegments, + DeclareOpInterfaceMethods, + SingleBlockImplicitTerminator<"YieldOp">]> { let summary = "parallel for operation"; let description = [{ The "loop.parallel" operation represents a loop nest taking 4 groups of SSA diff --git a/mlir/lib/Dialect/LoopOps/LoopOps.cpp b/mlir/lib/Dialect/LoopOps/LoopOps.cpp --- a/mlir/lib/Dialect/LoopOps/LoopOps.cpp +++ b/mlir/lib/Dialect/LoopOps/LoopOps.cpp @@ -182,7 +182,7 @@ LogicalResult ForOp::moveOutOfLoop(ArrayRef ops) { for (auto op : ops) - op->moveBefore(this->getOperation()); + op->moveBefore(*this); return success(); } @@ -191,8 +191,8 @@ if (!ivArg) return ForOp(); assert(ivArg.getOwner() && "unlinked block argument"); - auto *containingInst = ivArg.getOwner()->getParentOp(); - return dyn_cast_or_null(containingInst); + auto *containingOp = ivArg.getOwner()->getParentOp(); + return dyn_cast_or_null(containingOp); } //===----------------------------------------------------------------------===// @@ -459,13 +459,25 @@ op.getAttrs(), /*elidedAttrs=*/ParallelOp::getOperandSegmentSizeAttr()); } +Region &ParallelOp::getLoopBody() { return region(); } + +bool ParallelOp::isDefinedOutsideOfLoop(Value value) { + return !region().isAncestor(value.getParentRegion()); +} + +LogicalResult ParallelOp::moveOutOfLoop(ArrayRef ops) { + for (auto op : ops) + op->moveBefore(*this); + return success(); +} + ParallelOp mlir::loop::getParallelForInductionVarOwner(Value val) { auto ivArg = val.dyn_cast(); if (!ivArg) return ParallelOp(); assert(ivArg.getOwner() && "unlinked block argument"); - auto *containingInst = ivArg.getOwner()->getParentOp(); - return dyn_cast(containingInst); + auto *containingOp = ivArg.getOwner()->getParentOp(); + return dyn_cast(containingOp); } //===----------------------------------------------------------------------===// diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir --- a/mlir/test/Transforms/loop-invariant-code-motion.mlir +++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir @@ -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 @@ return } +// ----- + func @nested_loops_code_invariant_to_both() { %m = alloc() : memref<10xf32> %cf7 = constant 7.0 : f32 @@ -44,6 +46,8 @@ return } +// ----- + func @single_loop_nothing_invariant() { %m1 = alloc() : memref<10xf32> %m2 = alloc() : memref<10xf32> @@ -65,6 +69,8 @@ return } +// ----- + func @invariant_code_inside_affine_if() { %m = alloc() : memref<10xf32> %cf8 = constant 8.0 : f32 @@ -81,7 +87,7 @@ // 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 @@ return } +// ----- + func @invariant_affine_if() { %m = alloc() : memref<10xf32> %cf8 = constant 8.0 : f32 @@ -114,6 +122,8 @@ return } +// ----- + func @invariant_affine_if2() { %m = alloc() : memref<10xf32> %cf8 = constant 8.0 : f32 @@ -139,6 +149,8 @@ return } +// ----- + func @invariant_affine_nested_if() { %m = alloc() : memref<10xf32> %cf8 = constant 8.0 : f32 @@ -169,6 +181,8 @@ return } +// ----- + func @invariant_affine_nested_if_else() { %m = alloc() : memref<10xf32> %cf8 = constant 8.0 : f32 @@ -205,6 +219,8 @@ return } +// ----- + func @invariant_loop_dialect() { %ci0 = constant 0 : index %ci10 = constant 10 : index @@ -226,6 +242,8 @@ return } +// ----- + func @variant_loop_dialect() { %ci0 = constant 0 : index %ci10 = constant 10 : index @@ -244,3 +262,33 @@ 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 +} +