[mlir][Affine] Fix AffineLoopInvariantCodeMotion

[dcaballe]

Summary

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.

[bondhugula]

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

87

The comment is a bit redundant. Either drop or replace it with something that indicates what it means to add it to the definedOps.

mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir

542

This has in general nothing to do with mul specifically, but any side-effect free op consuming the load result?

[dcaballe]

Addressed feedback. Thanks!

[bondhugula]

Thanks!

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

89

are are -> are

89

which are not -> which are themselves not
?

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

bool isOpLoopInvariant(Operation &op, Value indVar,
   } else if (isa<AffineForOp>(op)) {
     // If the body of a predicated region has a for loop, we don't hoist the
     // 'affine.if'.
     return false;
   } else if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) {
     // 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()
                          : cast<AffineStoreOp>(op).getMemRef();
       for (auto *user : memref.getUsers()) {
         // If this memref has a user that is a DMA, give up because these
         // operations write to this memref.
         if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) {
if (isMemRefDereferencingOp(op)) {
             if (llvm::is_contained(userIVs, getForInductionVarOwner(indVar))) {
               return false;
             }
           }
         }
       }
     }
 
-    // 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;
     }
     for (unsigned int i = 0; i < op.getNumOperands(); ++i) {
       auto *operandSrc = op.getOperand(i).getDefiningOp();
 
       LLVM_DEBUG(

mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir

func @nested_loops_both_having_invariant_code() {
   // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
   // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
   // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32>
 
   return
 }
 
+// -----
+
 // The store-load forwarding can see through affine apply's since it relies on
 // dependence information.
 // CHECK-LABEL: func @store_affine_apply
 func @store_affine_apply() -> memref<10xf32> {
   %cf7 = constant 7.0 : f32
   %m = alloc() : memref<10xf32>
   affine.for %arg0 = 0 to 10 {
       %t0 = affine.apply affine_map<(d1) -> (d1 + 1)>(%arg0)
       affine.store %cf7, %m[%t0] : memref<10xf32>
   }
   return %m : 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
   %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 #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>
   // CHECK-NEXT: }
 
 
   return
 }
 
+// -----
 
 func @dependent_stores() {
   %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 @dependent_stores() {
 
   // CHECK-NEXT: affine.for %arg1 = 0 to 10 {
   // CHECK-NEXT:   affine.store %2, %0[%arg1] : memref<10xf32>
   // CHECK-NEXT:   affine.store %1, %0[%arg0] : memref<10xf32>
 
   return
 }
 
+// -----
+
 func @independent_stores() {
   %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
     affine.for %arg1 = 0 to 10 {
func @independent_stores() {
   // CHECK-NEXT:   affine.for %arg1 = 0 to 10 {
   // CHECK-NEXT:     affine.store %1, %0[%arg0] : memref<10xf32>
   // CHECK-NEXT:     affine.store %2, %0[%arg1] : memref<10xf32>
   // CHECK-NEXT:    }
 
   return
 }
 
+// -----
+
 func @load_dependent_store() {
   %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
     affine.for %arg1 = 0 to 10 {
func @load_dependent_store() {
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
   // CHECK-NEXT: affine.for %arg1 = 0 to 10 {
   // CHECK-NEXT:   affine.store %1, %0[%arg1] : memref<10xf32>
   // CHECK-NEXT:   %3 = affine.load %0[%arg0] : memref<10xf32>
 
   return
 }
 
+// -----
+
 func @load_after_load() {
   %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
     affine.for %arg1 = 0 to 10 {
func @load_after_load() {
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
   // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32>
   // CHECK-NEXT: affine.for %arg1 = 0 to 10 {
   // CHECK-NEXT: %4 = affine.load %0[%arg1] : memref<10xf32>
 
   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
           affine.store %cf9, %m[%arg0] : memref<10xf32>
func @invariant_affine_if() {
   // CHECK-NEXT: %1 = addf %cst, %cst : f32
   // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32>
   // 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: affine.store %1, %0[%arg1] : memref<10xf32>
   // 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.store %cf9, %m[%arg0] : memref<10xf32>
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_affine_nested_if_else2() {
   %m = alloc() : memref<10xf32>
   %m2 = 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
func @invariant_affine_nested_if_else2() {
   // CHECK-NEXT: %4 = affine.load %0[%arg0] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: }
 
 
   return
 }
 
+// -----
 
 func @invariant_affine_nested_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
func @invariant_affine_nested_if2() {
   // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: }
 
 
   return
 }
 
+// -----
+
 func @invariant_affine_for_inside_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
           affine.store %cf9, %m[%arg0] : memref<10xf32>
func @invariant_affine_for_inside_affine_if() {
   // CHECK-NEXT: }
   // CHECK-NEXT: }
   // CHECK-NEXT: }
 
 
   return
 }
 
+// -----
 
 func @invariant_constant_and_load() {
   %m = alloc() : memref<100xf32>
   %m2 = alloc() : memref<100xf32>
   affine.for %arg0 = 0 to 5 {
     %c0 = constant 0 : index
     %v = affine.load %m2[%c0] : memref<100xf32>
     affine.store %v, %m[%arg0] : memref<100xf32>
   }
 
   // CHECK: %0 = alloc() : memref<100xf32>
   // CHECK-NEXT: %1 = alloc() : memref<100xf32>
   // CHECK-NEXT: %c0 = constant 0 : index
   // CHECK-NEXT: %2 = affine.load %1[%c0] : memref<100xf32>
   // CHECK-NEXT: affine.for %arg0 = 0 to 5 {
   // CHECK-NEXT:  affine.store %2, %0[%arg0] : memref<100xf32>
 
 
   return
 }
 
+// -----
 
 func @nested_load_store_same_memref() {
   %m = alloc() : memref<10xf32>
   %cst = constant 8.0 : f32
   %c0 = constant 0 : index
    affine.for %arg0 = 0 to 10 {
     %v0 = affine.load %m[%c0] : memref<10xf32>
     affine.for %arg1 = 0 to 10 {
       affine.store %cst, %m[%arg1] : memref<10xf32>
     }
   }
 
   // CHECK: %0 = alloc() : memref<10xf32>
   // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
   // CHECK-NEXT: %c0 = constant 0 : index
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
   // CHECK-NEXT:  %1 = affine.load %0[%c0] : memref<10xf32>
   // CHECK-NEXT:   affine.for %arg1 = 0 to 10 {
   // CHECK-NEXT:    affine.store %cst, %0[%arg1] : memref<10xf32>
 
 
   return
 }
 
+// -----
 
 func @nested_load_store_same_memref2() {
   %m = alloc() : memref<10xf32>
   %cst = constant 8.0 : f32
   %c0 = constant 0 : index
    affine.for %arg0 = 0 to 10 {
      affine.store %cst, %m[%c0] : memref<10xf32>
       affine.for %arg1 = 0 to 10 {
         %v0 = affine.load %m[%arg0] : memref<10xf32>
     }
   }
 
   // CHECK: %0 = alloc() : memref<10xf32>
   // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
   // CHECK-NEXT: %c0 = constant 0 : index
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
   // CHECK-NEXT:   affine.store %cst, %0[%c0] : memref<10xf32>
   // CHECK-NEXT:   %1 = affine.load %0[%arg0] : memref<10xf32>
 
 
   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:  }