[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

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;

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:  }