[mlir] Update LICM to support Graph Regions

[Mogball]

Summary

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.

[Mogball]

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

171

@nicolasvasilache Can you verify whether this is acceptable behaviour?

[Mogball]

change to worklist algorithm

[mehdi_amini]

Adds default implementations of isDefinedOutsideOfLoop and moveOutOfLoop since 99% of all implementations of these functions were identical

Can you split this as a NFC commit?

[Mogball]

rebase (fix bazel build maybe)

[Mogball]

Yeah let me split the patch

[Mogball]

split out simplifications to LoopLikeOpInterface

[bondhugula]

mlir/lib/Interfaces/LoopLikeInterface.cpp

35

This would be a useful function and can directly be added in MLIRAnalysis library?

mlir/test/lib/Dialect/Test/TestOps.td

25

Sorted order here.

2782

Nit: ... with a graph region.

2801

Nit: extra blank line.

[Mogball]

review comments

mlir/lib/Interfaces/LoopLikeInterface.cpp

35

I will be submitting a patch after this one that moves moveLoopInvariantCode to TransformUtils. At the same time, I'll move this function to MLIRAnalysis (the same function exists elsewhere in the codebase and I will refactor those as well)

[rriddle]

mlir/lib/Interfaces/LoopLikeInterface.cpp

58

-> /

58

Sigh, I hate phab. Use /// comments here.

mlir/test/lib/Dialect/Test/TestOps.td

15

Do you know why this dialect is depending on Linalg?

2802–2804

[Mogball]

review comments

[Mogball]

mlir/test/lib/Dialect/Test/TestOps.td

15

There's a test op for LinalgConvolutionOpInterface

[rriddle]

mlir/lib/Interfaces/LoopLikeInterface.cpp

68

Probably more idiomatic to check the trait for IsTerminator.

107

nit: Add a newline before the comment.

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

362

Can you add a test that checks that we don't hoist from graph regions into non-graph regions?

366

Do we have a test in a non-graph region where operations that can be hoisted feed into each other? Something like:

%v0 = arith.addi %c0, %c0 : i32
%v1 = arith.addi %c0, %v0 : i32

to make sure that dominance still holds after hoisting?

[Mogball]

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

362

Is the behaviour that we want that ops don't cross between non-graph and graph regions (in either direction)?

Right now, LICM can move ops from graph regions to non-graph regions but only in topological order (dominance will be preserved). Ops that form cycles won't be moved because it doesn't detect them.

366

Oddly enough, I don't see one. I'll add one.

[rriddle]

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

362

It's more about being conservatively correct. We can hoist from a graph into a non-graph(and vice-versa), but we, for example, have to ensure that the dominance of the hoisted operations is valid. We have to be careful about the contract changes between the two regions.

[Mogball]

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

362

We will be okay to move ops from graph regions to non-graph regions; since LICM only considers one operation at a time, ops are moved only when all their operands are defined outside the graph region body. If the outer region is a non-graph region, the defining ops must precede the loop op.

I'll add some test cases.

[Mogball]

review comments -- adding some tests

mlir/lib/Interfaces/LoopLikeInterface.cpp

 //===- LoopLikeInterface.cpp - Loop-like operations in MLIR ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
+#include <queue>
 
 using namespace mlir;
 
 #define DEBUG_TYPE "loop-like"
 
 //===----------------------------------------------------------------------===//
 // LoopLike Interfaces
 //===----------------------------------------------------------------------===//
 
 /// Include the definitions of the loop-like interfaces.
 #include "mlir/Interfaces/LoopLikeInterface.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // 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 &region : 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 &region : op->getRegions())
+    for (Operation &op : region.getOps())
+      if (!isSideEffectFree(&op))
+  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 true;
+    return WalkResult::advance();
+  };
+  return !op->walk(walkFn).wasInterrupted();
 }
 
 void mlir::moveLoopInvariantCode(LoopLikeOpInterface looplike) {
-  auto &loopBody = looplike.getLoopBody();
+  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);
 
-  // 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);
+  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);
+  }
 }

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

 // -----
 
 func @invariant_affine_nested_if() {
   %m = memref.alloc() : memref<10xf32>
   %cf8 = arith.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 = arith.addf %cf8, %cf8 : f32
-          affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
-            %cf10 = arith.addf %cf9, %cf9 : f32
-          }
       }
     }
   }
 
   // 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
   // CHECK-NEXT: affine.if
   // CHECK-NEXT: arith.addf
   // CHECK-NEXT: }
   // CHECK-NEXT: }
-  // CHECK-NEXT: }
 
 
   return
 }
 
 // -----
 
 func @invariant_affine_nested_if_else() {
scf.for %i = %lb to %ub step %step {
     %val = call @make_val() : () -> (index)
     %val2 = call @make_val() : () -> (index)
     %r = scf.if %true -> (index) {
       scf.yield %val: index
     } else {
       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
+}

mlir/test/lib/Dialect/Test/TestOps.td

 //===-- TestOps.td - Test dialect operation definitions ----*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TEST_OPS
 #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"
 include "mlir/IR/PatternBase.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 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.
 include "TestAttrDefs.td"
 // Include the type definitions.
 include "TestTypeDefs.td"
 
 

 
 //===----------------------------------------------------------------------===//
 // Test Ops with effects
 //===----------------------------------------------------------------------===//
 
 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 results = (outs Res<AnyMemRef, "", [MemAlloc<TestResource>]>);
 }
 
 def TestEffectsOpB : TEST_Op<"op_with_effects_b",
     [MemoryEffects<[MemWrite<TestResource>]>]>;
 
 def TestEffectsRead : TEST_Op<"op_with_memread",
     [MemoryEffects<[MemRead]>]> {
   let results = (outs AnyInteger);
 }
 
 def TestEffectsWrite : TEST_Op<"op_with_memwrite",
     [MemoryEffects<[MemWrite]>]>;
 
+//===----------------------------------------------------------------------===//
+// 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