[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

diff --git a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
index 08f61c4..e202a40 100644 (file)
--- a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
+++ b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
@@ -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

diff --git a/mlir/lib/Dialect/LoopOps/LoopOps.cpp b/mlir/lib/Dialect/LoopOps/LoopOps.cpp
index 3d7ee38..c62507f 100644 (file)
--- a/mlir/lib/Dialect/LoopOps/LoopOps.cpp
+++ b/mlir/lib/Dialect/LoopOps/LoopOps.cpp
@@ -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);
 }
 
 //===----------------------------------------------------------------------===//

diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index 5873532..494bfa4 100644 (file)
--- 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 @@ 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
+}
+