[mlir][Affine] Enable fusion of loops with vector loads/stores

This patch enables affine loop fusion for loops with affine vector loads
and stores. For that, we only had to use affine memory op interfaces in
LoopFusionUtils.cpp and Utils.cpp so that vector loads and stores are
also taken into account.

Reviewed By: andydavis1, ftynse

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

mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td

@@ -68,6 +68,19 @@ def AffineReadOpInterface : OpInterface<"AffineReadOpInterface"> {
         return op.getAffineMapAttr().getValue();
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/"Returns the AffineMapAttr associated with 'memref'.",
+      /*retTy=*/"NamedAttribute",
+      /*methodName=*/"getAffineMapAttrForMemRef",
+      /*args=*/(ins "Value":$memref),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        assert(memref == getMemRef());
+        return {Identifier::get(op.getMapAttrName(), op.getContext()),
+                op.getAffineMapAttr()};
+      }]
+    >,
   ];
 }
 
@@ -124,6 +137,19 @@ def AffineWriteOpInterface : OpInterface<"AffineWriteOpInterface"> {
         return op.getAffineMapAttr().getValue();
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/"Returns the AffineMapAttr associated with 'memref'.",
+      /*retTy=*/"NamedAttribute",
+      /*methodName=*/"getAffineMapAttrForMemRef",
+      /*args=*/(ins "Value":$memref),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        assert(memref == getMemRef());
+        return {Identifier::get(op.getMapAttrName(), op.getContext()),
+                op.getAffineMapAttr()};
+      }]
+    >,
   ];
 }
 

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

@@ -389,13 +389,6 @@ class AffineLoadOpBase<string mnemonic, list<OpTrait> traits = []> :
       return getAttr(getMapAttrName()).cast<AffineMapAttr>();
     }
 
-    /// Returns the AffineMapAttr associated with 'memref'.
-    NamedAttribute getAffineMapAttrForMemRef(Value memref) {
-      assert(memref == getMemRef());
-      return {Identifier::get(getMapAttrName(), getContext()),
-              getAffineMapAttr()};
-    }
-
     static StringRef getMapAttrName() { return "map"; }
   }];
 }

mlir/lib/Transforms/Utils/LoopFusionUtils.cpp

@@ -38,10 +38,10 @@ using namespace mlir;
 static void getLoadAndStoreMemRefAccesses(Operation *opA,
                                           DenseMap<Value, bool> &values) {
   opA->walk([&](Operation *op) {
-    if (auto loadOp = dyn_cast<AffineLoadOp>(op)) {
+    if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
       if (values.count(loadOp.getMemRef()) == 0)
         values[loadOp.getMemRef()] = false;
-    } else if (auto storeOp = dyn_cast<AffineStoreOp>(op)) {
+    } else if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
       values[storeOp.getMemRef()] = true;
     }
   });
@@ -52,10 +52,10 @@ static void getLoadAndStoreMemRefAccesses(Operation *opA,
 // Returns false otherwise.
 static bool isDependentLoadOrStoreOp(Operation *op,
                                      DenseMap<Value, bool> &values) {
-  if (auto loadOp = dyn_cast<AffineLoadOp>(op)) {
+  if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
     return values.count(loadOp.getMemRef()) > 0 &&
            values[loadOp.getMemRef()] == true;
-  } else if (auto storeOp = dyn_cast<AffineStoreOp>(op)) {
+  } else if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
     return values.count(storeOp.getMemRef()) > 0;
   }
   return false;
@@ -105,7 +105,7 @@ static Operation *getLastDependentOpInRange(Operation *opA, Operation *opB) {
        it != Block::reverse_iterator(opA); ++it) {
     Operation *opX = &(*it);
     opX->walk([&](Operation *op) {
-      if (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op)) {
+      if (isa<AffineReadOpInterface>(op) || isa<AffineWriteOpInterface>(op)) {
         if (isDependentLoadOrStoreOp(op, values)) {
           lastDepOp = opX;
           return WalkResult::interrupt();
@@ -179,7 +179,7 @@ gatherLoadsAndStores(AffineForOp forOp,
                      SmallVectorImpl<Operation *> &loadAndStoreOps) {
   bool hasIfOp = false;
   forOp.walk([&](Operation *op) {
-    if (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op))
+    if (isa<AffineReadOpInterface>(op) || isa<AffineWriteOpInterface>(op))
       loadAndStoreOps.push_back(op);
     else if (isa<AffineIfOp>(op))
       hasIfOp = true;
@@ -464,7 +464,7 @@ bool mlir::getFusionComputeCost(AffineForOp srcForOp, LoopNestStats &srcStats,
     unsigned storeCount = 0;
     llvm::SmallDenseSet<Value, 4> storeMemrefs;
     srcForOp.walk([&](Operation *op) {
-      if (auto storeOp = dyn_cast<AffineStoreOp>(op)) {
+      if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
         storeMemrefs.insert(storeOp.getMemRef());
         ++storeCount;
       }
@@ -476,7 +476,7 @@ bool mlir::getFusionComputeCost(AffineForOp srcForOp, LoopNestStats &srcStats,
     // 'insertPointParent'.
     for (auto value : storeMemrefs) {
       for (auto *user : value.getUsers()) {
-        if (auto loadOp = dyn_cast<AffineLoadOp>(user)) {
+        if (auto loadOp = dyn_cast<AffineReadOpInterface>(user)) {
           SmallVector<AffineForOp, 4> loops;
           // Check if any loop in loop nest surrounding 'user' is
           // 'insertPointParent'.

mlir/lib/Transforms/Utils/Utils.cpp

@@ -30,7 +30,7 @@ using namespace mlir;
 // Temporary utility: will be replaced when this is modeled through
 // side-effects/op traits. TODO(b/117228571)
 static bool isMemRefDereferencingOp(Operation &op) {
-  if (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op) ||
+  if (isa<AffineReadOpInterface>(op) || isa<AffineWriteOpInterface>(op) ||
       isa<AffineDmaStartOp>(op) || isa<AffineDmaWaitOp>(op))
     return true;
   return false;
@@ -39,8 +39,8 @@ static bool isMemRefDereferencingOp(Operation &op) {
 /// Return the AffineMapAttr associated with memory 'op' on 'memref'.
 static NamedAttribute getAffineMapAttrForMemRef(Operation *op, Value memref) {
   return TypeSwitch<Operation *, NamedAttribute>(op)
-      .Case<AffineDmaStartOp, AffineLoadOp, AffinePrefetchOp, AffineStoreOp,
-            AffineDmaWaitOp>(
+      .Case<AffineDmaStartOp, AffineReadOpInterface, AffinePrefetchOp,
+            AffineWriteOpInterface, AffineDmaWaitOp>(
           [=](auto op) { return op.getAffineMapAttrForMemRef(memref); });
 }
 

mlir/test/Transforms/loop-fusion.mlir

@@ -2464,3 +2464,32 @@ func @reshape_into_matmul(%lhs : memref<1024x1024xf32>,
 // MAXIMAL-NEXT:     affine.for
 // MAXIMAL-NOT:      affine.for
 // MAXIMAL:      return
+
+// -----
+
+// CHECK-LABEL: func @vector_loop
+func @vector_loop(%a : memref<10x20xf32>, %b : memref<10x20xf32>,
+                  %c : memref<10x20xf32>) {
+  affine.for %j = 0 to 10 {
+    affine.for %i = 0 to 5 {
+      %ld0 = affine.vector_load %a[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
+      affine.vector_store %ld0, %b[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
+    }
+  }
+
+  affine.for %j = 0 to 10 {
+    affine.for %i = 0 to 5 {
+      %ld0 = affine.vector_load %b[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
+      affine.vector_store %ld0, %c[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
+    }
+  }
+
+  return
+}
+// CHECK:      affine.for
+// CHECK-NEXT:   affine.for
+// CHECK-NEXT:     affine.vector_load
+// CHECK-NEXT:     affine.vector_store
+// CHECK-NEXT:     affine.vector_load
+// CHECK-NEXT:     affine.vector_store
+// CHECK-NOT:  affine.for