[MLIR] Implement LoopLikeInterface for loop.parallel

[tpopp]

Summary

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

[bondhugula]

mlir/lib/Dialect/LoopOps/LoopOps.cpp

470

You don't need ->getOperation()

mlir/test/Transforms/loop-invariant-code-motion.mlir

266

Please always include CHECK-LABEL's when you aren't using -split-input-file. There is otherwise the chance of CHECK's matching with previous function's output while developing / changing / adding tests and this could be quite painful. Looks like the other tests here aren't following the right convention as well and should be fixed.

[bondhugula]

mlir/lib/Dialect/LoopOps/LoopOps.cpp

478–479

Nit: containingOp

mlir/test/Transforms/loop-invariant-code-motion.mlir

282

You don't need %0 =

284

You don't need %1 =

[tpopp]

Replace ->getOperation() with dereference and run split-input-file on test file.

[tpopp]

mlir/test/Transforms/loop-invariant-code-motion.mlir

266

I also split most test cases here. One test fails then, so I will have to follow up on that.

[tpopp]

Update map# on test case to allow splitting of test cases.

[herhut]

Thanks!

[tpopp]

Uday, I landed this thinking you were doing more of a drive-by review, but I'll do a follow up commit if you have any more concerns.

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

OpBuilder getElseBodyBuilder() {
       Block &body = elseRegion().front();
       return OpBuilder(&body,
                        results().empty() ? std::prev(body.end()) : body.end());
     }
   }];
 }
 
 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
     values as operands that represent the lower bounds, upper bounds, steps and
     initial values, respectively. The operation defines a variadic number of
     SSA values for its induction variables. It has one region capturing the
     loop body. The induction variables are represented as an argument of this
     region. These SSA values always have type index, which is the size of the

mlir/lib/Dialect/LoopOps/LoopOps.cpp

 Region &ForOp::getLoopBody() { return region(); }
 
 bool ForOp::isDefinedOutsideOfLoop(Value value) {
   return !region().isAncestor(value.getParentRegion());
 }
 
 LogicalResult ForOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
   for (auto op : ops)
-    op->moveBefore(this->getOperation());
+    op->moveBefore(*this);
   return success();
 }
 
 ForOp mlir::loop::getForInductionVarOwner(Value val) {
   auto ivArg = val.dyn_cast<BlockArgument>();
   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);
 }
 
 //===----------------------------------------------------------------------===//
 // IfOp
 //===----------------------------------------------------------------------===//
 
 void IfOp::build(Builder *builder, OperationState &result, Value cond,
                  bool withElseRegion) {
static void print(OpAsmPrinter &p, ParallelOp op) {
   if (!op.initVals().empty())
     p << " init (" << op.initVals() << ")";
   p.printOptionalArrowTypeList(op.getResultTypes());
   p.printRegion(op.region(), /*printEntryBlockArgs=*/false);
   p.printOptionalAttrDict(
       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);
 }
 
 //===----------------------------------------------------------------------===//
 // ReduceOp
 //===----------------------------------------------------------------------===//
 
 void ReduceOp::build(Builder *builder, OperationState &result, Value operand) {
   auto type = operand.getType();

mlir/test/Transforms/loop-invariant-code-motion.mlir

-// 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>
   %cf7 = constant 7.0 : f32
   %cf8 = constant 8.0 : f32
 
   affine.for %arg0 = 0 to 10 {
     %v0 = addf %cf7, %cf8 : f32
func @nested_loops_both_having_invariant_code() {
   // CHECK-NEXT: addf %[[ADD0]], %[[CST1]] : f32
   // CHECK-NEXT: affine.for
   // CHECK-NEXT: affine.for
   // CHECK-NEXT: affine.store
 
   return
 }
 
+// -----
+
 func @nested_loops_code_invariant_to_both() {
   %m = alloc() : memref<10xf32>
   %cf7 = constant 7.0 : f32
   %cf8 = constant 8.0 : f32
 
   affine.for %arg0 = 0 to 10 {
     affine.for %arg1 = 0 to 10 {
       %v0 = addf %cf7, %cf8 : f32
     }
   }
 
   // CHECK: %0 = alloc() : memref<10xf32>
   // CHECK-NEXT: %cst = constant 7.000000e+00 : f32
   // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
   // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
 
   return
 }
 
+// -----
+
 func @single_loop_nothing_invariant() {
   %m1 = alloc() : memref<10xf32>
   %m2 = alloc() : memref<10xf32>
   affine.for %arg0 = 0 to 10 {
     %v0 = affine.load %m1[%arg0] : memref<10xf32>
     %v1 = affine.load %m2[%arg0] : memref<10xf32>
     %v2 = addf %v0, %v1 : f32
     affine.store %v2, %m1[%arg0] : memref<10xf32>
   }
 
   // CHECK: %0 = alloc() : memref<10xf32>
   // CHECK-NEXT: %1 = alloc() : memref<10xf32>
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
   // CHECK-NEXT: %2 = affine.load %0[%arg0] : memref<10xf32>
   // CHECK-NEXT: %3 = affine.load %1[%arg0] : memref<10xf32>
   // CHECK-NEXT: %4 = addf %2, %3 : f32
   // CHECK-NEXT: affine.store %4, %0[%arg0] : memref<10xf32>
 
   return
 }
 
+// -----
+
 func @invariant_code_inside_affine_if() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
 
   affine.for %arg0 = 0 to 10 {
     %t0 = affine.apply affine_map<(d1) -> (d1 + 1)>(%arg0)
     affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %t0) {
         %cf9 = addf %cf8, %cf8 : f32
         affine.store %cf9, %m[%arg0] : memref<10xf32>
 
     }
   }
 
   // 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>
   // CHECK-NEXT: }
 
 
   return
 }
 
+// -----
+
 func @invariant_affine_if() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
   affine.for %arg0 = 0 to 10 {
     affine.for %arg1 = 0 to 10 {
       affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
           %cf9 = addf %cf8, %cf8 : f32
       }
     }
   }
 
   // CHECK: %0 = alloc() : memref<10xf32>
   // CHECK-NEXT: %[[CST:.*]] = constant 8.000000e+00 : f32
   // CHECK-NEXT: affine.for %[[ARG:.*]] = 0 to 10 {
   // CHECK-NEXT: }
   // CHECK-NEXT: affine.for %[[ARG:.*]] = 0 to 10 {
   // CHECK-NEXT: affine.if #set0(%[[ARG]], %[[ARG]]) {
   // CHECK-NEXT: addf %[[CST]], %[[CST]] : f32
   // CHECK-NEXT: }
 
   return
 }
 
+// -----
+
 func @invariant_affine_if2() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
   affine.for %arg0 = 0 to 10 {
     affine.for %arg1 = 0 to 10 {
       affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
           %cf9 = addf %cf8, %cf8 : f32
           affine.store %cf9, %m[%arg1] : memref<10xf32>
func @invariant_affine_if2() {
   // CHECK-NEXT: addf
   // CHECK-NEXT: affine.store
   // CHECK-NEXT: }
   // CHECK-NEXT: }
 
   return
 }
 
+// -----
+
 func @invariant_affine_nested_if() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
   affine.for %arg0 = 0 to 10 {
     affine.for %arg1 = 0 to 10 {
       affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
           %cf9 = addf %cf8, %cf8 : f32
           affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
func @invariant_affine_nested_if() {
   // CHECK-NEXT: }
   // CHECK-NEXT: }
   // CHECK-NEXT: }
 
 
   return
 }
 
+// -----
+
 func @invariant_affine_nested_if_else() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
   affine.for %arg0 = 0 to 10 {
     affine.for %arg1 = 0 to 10 {
       affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
           %cf9 = addf %cf8, %cf8 : f32
           affine.store %cf9, %m[%arg0] : memref<10xf32>
func @invariant_affine_nested_if_else() {
   // CHECK-NEXT: }
   // CHECK-NEXT: }
   // CHECK-NEXT: }
 
 
   return
 }
 
+// -----
+
 func @invariant_loop_dialect() {
   %ci0 = constant 0 : index
   %ci10 = constant 10 : index
   %ci1 = constant 1 : index
   %m = alloc() : memref<10xf32>
   %cf7 = constant 7.0 : f32
   %cf8 = constant 8.0 : f32
   loop.for %arg0 = %ci0 to %ci10 step %ci1 {
     loop.for %arg1 = %ci0 to %ci10 step %ci1 {
       %v0 = addf %cf7, %cf8 : f32
     }
   }
 
   // CHECK: %0 = alloc() : memref<10xf32>
   // CHECK-NEXT: %cst = constant 7.000000e+00 : f32
   // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
   // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
 
   return
 }
 
+// -----
+
 func @variant_loop_dialect() {
   %ci0 = constant 0 : index
   %ci10 = constant 10 : index
   %ci1 = constant 1 : index
   %m = alloc() : memref<10xf32>
   loop.for %arg0 = %ci0 to %ci10 step %ci1 {
     loop.for %arg1 = %ci0 to %ci10 step %ci1 {
       %v0 = addi %arg0, %arg1 : index
     }
   }
 
   // CHECK: %0 = alloc() : memref<10xf32>
   // CHECK-NEXT: loop.for
   // CHECK-NEXT: loop.for
   // CHECK-NEXT: addi
 
   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
+}
+