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

[dcaballe]

Summary

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.

[ftynse]

This should be a canonical example of why analyses/passes should use interfaces

[mehdi_amini]

In D80971#2068940, @ftynse wrote:

This should be a canonical example of why analyses/passes should use interfaces

Indeed! Nice :)

[bondhugula]

This is really great - the same could be done for memref store to load forwarding, and it should work out of the box with affine.vector_load/store. But there's an issue with this (here and in general) if you mix vector load/stores with regular load/stores on the same memref. The slices computed by fusion aren't aware of the larger underlying loaded/stored data with the vectors. So the regions will be inaccurate at the boundaries. You are fine as long as you don't have memrefs that have both vector load/stores and scalar load/stores in different loop nests. With this patch, at the moment, if you mix the two, you'll get incorrect output from fusion say when the producer nest has scalar stores and the consumer vector loads. The slice won't have all the data needed.

This can be fixed by fixing the dependence information and memref region computation to account for vector load/stores. They would be inaccurate after this revision (and of course anyway inaccurate even if completely ignoring vector load/store).

[dcaballe]

Thanks Uday! Good point. I haven't looked too much into the slice computation but if you can send me some pointers and some examples on how the regions should be computed I could have a look. Thanks!

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

InterfaceMethod<
       /*methodName=*/"getAffineMap",
       /*args=*/(ins),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         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()};
+      }]
+    >,
   ];
 }
 
 def AffineWriteOpInterface : OpInterface<"AffineWriteOpInterface"> {
   let description = [{
       Interface to query characteristics of write-like ops with affine
       restrictions.
   }];
InterfaceMethod<
       /*methodName=*/"getAffineMap",
       /*args=*/(ins),
       /*methodName=*/[{}],
       /*defaultImplementation=*/[{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         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()};
+      }]
+    >,
   ];
 }
 
 #endif // AFFINEMEMORYOPINTERFACES

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

code extraClassDeclarationBase = [{
 
     void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
 
     /// Returns the affine map used to index the memref for this operation.
     AffineMapAttr getAffineMapAttr() {
       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"; }
   }];
 }
 
 def AffineLoadOp : AffineLoadOpBase<"load"> {
   let summary = "affine load operation";
   let description = [{
     The "affine.load" op reads an element from a memref, where the index

mlir/lib/Transforms/Utils/LoopFusionUtils.cpp

 
 using namespace mlir;
 
 // Gathers all load and store memref accesses in 'opA' into 'values', where
 // 'values[memref] == true' for each store operation.
 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;
     }
   });
 }
 
 // Returns true if 'op' is a load or store operation which access an memref
 // accessed 'values' and at least one of the access is a store operation.
 // 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;
 }
 
 // Returns the first operation in range ('opA', 'opB') which has a data
 // dependence on 'opA'. Returns 'nullptr' of no dependence exists.
 static Operation *getFirstDependentOpInRange(Operation *opA, Operation *opB) {
static Operation *getLastDependentOpInRange(Operation *opA, Operation *opB) {
   // *) 'opX' and 'opB' access the same memref and at least one of the accesses
   //    is a store.
   // *) 'opX' produces an SSA Value which is used by 'opB'.
   Operation *lastDepOp = nullptr;
   for (Block::reverse_iterator it = std::next(Block::reverse_iterator(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();
         }
         return WalkResult::advance();
       }
       for (auto value : op->getResults()) {
         for (auto user : value.getUsers()) {

 
 // Gathers all load and store ops in loop nest rooted at 'forOp' into
 // 'loadAndStoreOps'.
 static bool
 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;
   });
   return !hasIfOp;
 }
 
 // TODO(andydavis) Prevent fusion of loop nests with side-effecting operations.
bool mlir::getFusionComputeCost(AffineForOp srcForOp, LoopNestStats &srcStats,
   // The store and loads to this memref will disappear.
   // TODO(andydavis) Add load coalescing to memref data flow opt pass.
   if (storeLoadFwdGuaranteed) {
     // Subtract from operation count the loads/store we expect load/store
     // forwarding to remove.
     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;
       }
     });
     // Subtract out any store ops in single-iteration src slice loop nest.
     if (storeCount > 0)
       computeCostMap[insertPointParent] = -storeCount;
     // Subtract out any load users of 'storeMemrefs' nested below
     // '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'.
           getLoopIVs(*user, &loops);
           if (llvm::is_contained(loops, cast<AffineForOp>(insertPointParent))) {
             if (auto forOp =
                     dyn_cast_or_null<AffineForOp>(user->getParentOp())) {
               if (computeCostMap.count(forOp) == 0)

mlir/lib/Transforms/Utils/Utils.cpp

 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/TypeSwitch.h"
 using namespace mlir;
 
 /// Return true if this operation dereferences one or more memref's.
 // 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;
 }
 
 /// 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); });
 }
 
 // Perform the replacement in `op`.
 LogicalResult mlir::replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
                                              Operation *op,
                                              ArrayRef<Value> extraIndices,
                                              AffineMap indexRemap,

mlir/test/Transforms/loop-fusion.mlir

func @reshape_into_matmul(%lhs : memref<1024x1024xf32>,
   return
 }
 // MAXIMAL-NEXT: alloc
 // MAXIMAL-NEXT: affine.for
 // MAXIMAL-NEXT:   affine.for
 // 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