[mlir][Affine] Fix AffineLoopInvariantCodeMotion

Make sure that memory ops that are defined inside the loop are registered
as such in 'defineOp'. In the test provided, the 'mulf' op was hoisted
outside the loop nest even when its 'affine.load' operand was not.

Reviewed By: bondhugula

Differential Revision: https://reviews.llvm.org/D86982

2 files changed, 68 insertions, 5 deletions

diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
index 364168c..df8d875 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
@@ -84,6 +84,11 @@ bool isOpLoopInvariant(Operation &op, Value indVar,
     // TODO: Support DMA ops.
     return false;
   } else if (!isa<ConstantOp>(op)) {
+    // Register op in the set of ops defined inside the loop. This set is used
+    // to prevent hoisting ops that depend on other ops defined inside the loop
+    // which are themselves not being hoisted.
+    definedOps.insert(&op);
+
     if (isMemRefDereferencingOp(op)) {
       Value memref = isa<AffineLoadOp>(op)
                          ? cast<AffineLoadOp>(op).getMemRef()
@@ -111,9 +116,6 @@ bool isOpLoopInvariant(Operation &op, Value indVar,
       }
     }
 
-    // Insert this op in the defined ops list.
-    definedOps.insert(&op);
-
     if (op.getNumOperands() == 0 && !isa<AffineYieldOp>(op)) {
       LLVM_DEBUG(llvm::dbgs() << "\nNon-constant op with 0 operands\n");
       return false;
diff --git a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
index 4c9c0db..8a2ed32 100644
--- a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
@@ -22,6 +22,8 @@ func @nested_loops_both_having_invariant_code() {
   return
 }
 
+// -----
+
 // The store-load forwarding can see through affine apply's since it relies on
 // dependence information.
 // CHECK-LABEL: func @store_affine_apply
@@ -36,12 +38,14 @@ func @store_affine_apply() -> memref<10xf32> {
 // CHECK:       %cst = constant 7.000000e+00 : f32
 // CHECK-NEXT:  %0 = alloc() : memref<10xf32>
 // CHECK-NEXT:  affine.for %arg0 = 0 to 10 {
-// CHECK-NEXT:      %1 = affine.apply #map3(%arg0)
+// CHECK-NEXT:      %1 = affine.apply #map{{[0-9]+}}(%arg0)
 // CHECK-NEXT:      affine.store %cst, %0[%1] : memref<10xf32>
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return %0 : memref<10xf32>
 }
 
+// -----
+
 func @nested_loops_code_invariant_to_both() {
   %m = alloc() : memref<10xf32>
   %cf7 = constant 7.0 : f32
@@ -61,6 +65,8 @@ func @nested_loops_code_invariant_to_both() {
   return
 }
 
+// -----
+
 func @single_loop_nothing_invariant() {
   %m1 = alloc() : memref<10xf32>
   %m2 = alloc() : memref<10xf32>
@@ -82,6 +88,8 @@ func @single_loop_nothing_invariant() {
   return
 }
 
+// -----
+
 func @invariant_code_inside_affine_if() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
@@ -98,7 +106,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 #map{{[0-9]+}}(%arg0)
   // CHECK-NEXT: affine.if #set0(%arg0, %1) {
   // CHECK-NEXT: %2 = addf %cst, %cst : f32
   // CHECK-NEXT: affine.store %2, %0[%arg0] : memref<10xf32>
@@ -108,6 +116,7 @@ func @invariant_code_inside_affine_if() {
   return
 }
 
+// -----
 
 func @dependent_stores() {
   %m = alloc() : memref<10xf32>
@@ -137,6 +146,8 @@ func @dependent_stores() {
   return
 }
 
+// -----
+
 func @independent_stores() {
   %m = alloc() : memref<10xf32>
   %cf7 = constant 7.0 : f32
@@ -165,6 +176,8 @@ func @independent_stores() {
   return
 }
 
+// -----
+
 func @load_dependent_store() {
   %m = alloc() : memref<10xf32>
   %cf7 = constant 7.0 : f32
@@ -192,6 +205,8 @@ func @load_dependent_store() {
   return
 }
 
+// -----
+
 func @load_after_load() {
   %m = alloc() : memref<10xf32>
   %cf7 = constant 7.0 : f32
@@ -219,6 +234,8 @@ func @load_after_load() {
   return
 }
 
+// -----
+
 func @invariant_affine_if() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
@@ -244,6 +261,8 @@ func @invariant_affine_if() {
   return
 }
 
+// -----
+
 func @invariant_affine_if2() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
@@ -271,6 +290,8 @@ func @invariant_affine_if2() {
   return
 }
 
+// -----
+
 func @invariant_affine_nested_if() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
@@ -303,6 +324,8 @@ func @invariant_affine_nested_if() {
   return
 }
 
+// -----
+
 func @invariant_affine_nested_if_else() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
@@ -339,6 +362,8 @@ func @invariant_affine_nested_if_else() {
   return
 }
 
+// -----
+
 func @invariant_affine_nested_if_else2() {
   %m = alloc() : memref<10xf32>
   %m2 = alloc() : memref<10xf32>
@@ -375,6 +400,7 @@ func @invariant_affine_nested_if_else2() {
   return
 }
 
+// -----
 
 func @invariant_affine_nested_if2() {
   %m = alloc() : memref<10xf32>
@@ -406,6 +432,8 @@ func @invariant_affine_nested_if2() {
   return
 }
 
+// -----
+
 func @invariant_affine_for_inside_affine_if() {
   %m = alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
@@ -438,6 +466,7 @@ func @invariant_affine_for_inside_affine_if() {
   return
 }
 
+// -----
 
 func @invariant_constant_and_load() {
   %m = alloc() : memref<100xf32>
@@ -459,6 +488,7 @@ func @invariant_constant_and_load() {
   return
 }
 
+// -----
 
 func @nested_load_store_same_memref() {
   %m = alloc() : memref<10xf32>
@@ -483,6 +513,7 @@ func @nested_load_store_same_memref() {
   return
 }
 
+// -----
 
 func @nested_load_store_same_memref2() {
   %m = alloc() : memref<10xf32>
@@ -505,3 +536,33 @@ func @nested_load_store_same_memref2() {
 
   return
 }
+
+// -----
+
+// CHECK-LABEL:   func @do_not_hoist_dependent_side_effect_free_op
+func @do_not_hoist_dependent_side_effect_free_op(%arg0: memref<10x512xf32>) {
+  %0 = alloca() : memref<1xf32>
+  %cst = constant 8.0 : f32
+  affine.for %i = 0 to 512 {
+    affine.for %j = 0 to 10 {
+      %5 = affine.load %arg0[%i, %j] : memref<10x512xf32>
+      %6 = affine.load %0[0] : memref<1xf32>
+      %add = addf %5, %6 : f32
+      affine.store %add, %0[0] : memref<1xf32>
+    }
+    %3 = affine.load %0[0] : memref<1xf32>
+    %4 = mulf %3, %cst : f32 // It shouldn't be hoisted.
+  }
+  return
+}
+
+// CHECK:       affine.for
+// CHECK-NEXT:    affine.for
+// CHECK-NEXT:      affine.load
+// CHECK-NEXT:      affine.load
+// CHECK-NEXT:      addf
+// CHECK-NEXT:      affine.store
+// CHECK-NEXT:    }
+// CHECK-NEXT:    affine.load
+// CHECK-NEXT:    mulf
+// CHECK-NEXT:  }