[mlir][Affine] Support affine vector loads/stores in LICM

[dcaballe]

Summary

Introduce affine memory op interfaces in AffineLoopInvariantCodeMotion so
that it can also work on affine.vector_load and affine.vector_store ops.

[bondhugula]

Looks good.

[bondhugula]

I'd reword the commit summary a bit to replace "Introduce" - instead, something like "Make use of ..."

[dcaballe]

In D86986#2251245, @bondhugula wrote:

I'd reword the commit summary a bit to replace "Introduce" - instead, something like "Make use of ..."

Will do. Thanks!

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

 
 static bool
 areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar,
                                  SmallPtrSetImpl<Operation *> &definedOps,
                                  SmallPtrSetImpl<Operation *> &opsToHoist);
 
 static bool isMemRefDereferencingOp(Operation &op) {
   // TODO: Support DMA Ops.
-  return isa<AffineLoadOp, AffineStoreOp>(op);
+  return isa<AffineReadOpInterface, AffineWriteOpInterface>(op);
 }
 
 // Returns true if the individual op is loop invariant.
 bool isOpLoopInvariant(Operation &op, Value indVar,
                        SmallPtrSetImpl<Operation *> &definedOps,
                        SmallPtrSetImpl<Operation *> &opsToHoist) {
   LLVM_DEBUG(llvm::dbgs() << "iterating on op: " << op;);
 
if (isa<AffineIfOp>(op)) {
     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();
+      Value memref = isa<AffineReadOpInterface>(op)
+                         ? cast<AffineReadOpInterface>(op).getMemRef()
+                         : cast<AffineWriteOpInterface>(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)) {
           return false;
         }
         // If the memref used by the load/store is used in a store elsewhere in
         // the loop nest, we do not hoist. Similarly, if the memref used in a
         // load is also being stored too, we do not hoist the load.
-        if (isa<AffineStoreOp>(user) ||
-            (isa<AffineLoadOp>(user) && isa<AffineStoreOp>(op))) {
+        if (isa<AffineWriteOpInterface>(user) ||
+            (isa<AffineReadOpInterface>(user) &&
+             isa<AffineWriteOpInterface>(op))) {
           if (&op != user) {
             SmallVector<AffineForOp, 8> userIVs;
             getLoopIVs(*user, &userIVs);
             // Check that userIVs don't contain the for loop around the op.
             if (llvm::is_contained(userIVs, getForInductionVarOwner(indVar))) {
               return false;
             }
           }

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

 // 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:  }
+
+// -----
+
+// CHECK-LABEL: func @vector_loop_nothing_invariant
+func @vector_loop_nothing_invariant() {
+  %m1 = alloc() : memref<40xf32>
+  %m2 = alloc() : memref<40xf32>
+  affine.for %arg0 = 0 to 10 {
+    %v0 = affine.vector_load %m1[%arg0*4] : memref<40xf32>, vector<4xf32>
+    %v1 = affine.vector_load %m2[%arg0*4] : memref<40xf32>, vector<4xf32>
+    %v2 = addf %v0, %v1 : vector<4xf32>
+    affine.vector_store %v2, %m1[%arg0*4] : memref<40xf32>, vector<4xf32>
+  }
+  return
+}
+
+// CHECK:       affine.for
+// CHECK-NEXT:    affine.vector_load
+// CHECK-NEXT:    affine.vector_load
+// CHECK-NEXT:    addf
+// CHECK-NEXT:    affine.vector_store
+// CHECK-NEXT:  }
+
+// -----
+
+// CHECK-LABEL: func @vector_loop_all_invariant
+func @vector_loop_all_invariant() {
+  %m1 = alloc() : memref<4xf32>
+  %m2 = alloc() : memref<4xf32>
+  %m3 = alloc() : memref<4xf32>
+  affine.for %arg0 = 0 to 10 {
+    %v0 = affine.vector_load %m1[0] : memref<4xf32>, vector<4xf32>
+    %v1 = affine.vector_load %m2[0] : memref<4xf32>, vector<4xf32>
+    %v2 = addf %v0, %v1 : vector<4xf32>
+    affine.vector_store %v2, %m3[0] : memref<4xf32>, vector<4xf32>
+  }
+  return
+}
+
+// CHECK:       alloc()
+// CHECK-NEXT:  alloc()
+// CHECK-NEXT:  alloc()
+// CHECK-NEXT:  affine.vector_load
+// CHECK-NEXT:  affine.vector_load
+// CHECK-NEXT:  addf
+// CHECK-NEXT:  affine.vector_store
+// CHECK-NEXT:  affine.for