[MLIR] Vector store to load forwarding

[anandkodnani]

Summary

The MemRefDataFlow pass does store to load forwarding
only for affine store/loads. This patch updates the pass
to use affine read/write interface which enables vector
forwarding.

[aartbik]

just a few fly by comment on test style

mlir/test/Transforms/memref-vector-dataflow.mlir

4

single white space line

20

why are these comment here?

54

if you want to match closing, then perhaps also add opening, and all closing patterns?

[bondhugula]

This looks good - thanks for contributing. A bunch of minor doc related comments.

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

728–729

Nit: stored value -> value to be stored

mlir/test/Transforms/memref-vector-dataflow.mlir

3

Please add a line on what these test cases are testing.

33–36

Please indent suitably.

51–52

You may want to capture the store value and check if the add operand is that.

54

CHECK-NEXT here?

[bondhugula]

The second line in the commit summary is broken. Please rephrase. Also,

only for scalars -> only for affine store/loads

[bondhugula]

Sorry, I've taken the liberty to add the MLIR tag to the commit title - to avoid any ambiguity w.r.t LLVM proper.

[anandkodnani]

Addressed comments on test and commit message.

[anandkodnani]

mlir/test/Transforms/memref-vector-dataflow.mlir

20

Updated the test case.

[ftynse]

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

85

I don't understand what "underlying value" means in this context. Do you mean the value read by this operation?

91

Nit: return cast<ConcreteOp>(this->getOperaiton()); looks more concise. return this->getOperation()->getResult(0); looks equally concise and arguably clearer about the intent as opposed to an explicit cast to ConcreteOp followed by an implicit cast to Value.

mlir/test/Transforms/memref-vector-dataflow.mlir

2

Why can't this live in mlir/test/Transforms/memref-dataflow-opt.mlir?

[anandkodnani]

• Updated getValue method in interface.
• Moved the test to memref-dataflow-opt.mlir

[anandkodnani]

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

85

I meant the value represented by the interface.

91

Using return cast<ConcreteOp>(this->getOperaiton()); just to keep consistent style with other methods in the interface.

[ftynse]

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

85

These are _operation_ interfaces, they don't represent values. The fact that single-result operations are implicitly convertible to values should not affect this. So "the value read by this operation" sounds like the right formulation here.

[dcaballe]

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

87

getValue is something we could get from the OneResult trait. However, I was playing with this, adding OneResult trait to this interfaces and I couldn't make it work. @ftynse or anybody else, do you know if this is possible?

[bondhugula]

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

87

The OneResult should automatically be added by the generator. You could verify this. Doing .getResult() should already give you the loaded value, right? Why do we need this?

[bondhugula]

mlir/test/Transforms/memref-dataflow-opt.mlir

284 ↗

(On Diff #280081)

Nit: ... value forwarding from vector stores to vector loads.

[ftynse]

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

87

@ftynse or anybody else, do you know if this is possible?

I don't think there is an easy way. Some time ago, I had a prototype that inherited the interface trait class from another trait class (for FunctionLike), but the code to support that was quite convoluted. That being said, having cast<ConcreteOp>(this->getOperation().getResult() should dispatch properly, without having to know if ConcreteOp has the trait or if it got the method from elsewhere. If you want just a check, something like static_assert(std::is_base_of<InterfaceName::Trait, ConcreteOp>::value, "") should do the trick.

[dcaballe]

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

87

The OneResult should automatically be added by the generator. You could verify this. Doing .getResult() should already give you the loaded value, right? Why do we need this?

The OneResult trait is automatically added to the concrete ops but not to the interfaces. My hope was that we could add OneResult to the interface and then get rid of the getValue() method altogether and replace it with .getResult() instead. However, that doesn't seem to be possible right now so... fair enough! Thanks for the info!

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

InterfaceMethod<
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         assert(memref == getMemRef());
         return {Identifier::get(op.getMapAttrName(), op.getContext()),
                 op.getAffineMapAttr()};
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/"Returns the underlying value.",
+      /*retTy=*/"Value",
+      /*methodName=*/"getValue",
+      /*args=*/(ins),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op;
+      }]
+    >,
   ];
 }
 
 def AffineWriteOpInterface : OpInterface<"AffineWriteOpInterface"> {
   let description = [{
       Interface to query characteristics of write-like ops with affine
       restrictions.
   }];
InterfaceMethod<
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         assert(memref == getMemRef());
         return {Identifier::get(op.getMapAttrName(), op.getContext()),
                 op.getAffineMapAttr()};
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/"Returns the value to store.",
+      /*retTy=*/"Value",
+      /*methodName=*/"getValueToStore",
+      /*args=*/(ins),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getOperand(op.getStoredValOperandIndex());
+      }]
+    >,
   ];
 }
 
 #endif // AFFINEMEMORYOPINTERFACES

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

def AffinePrefetchOp : Affine_Op<"prefetch"> {
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 class AffineStoreOpBase<string mnemonic, list<OpTrait> traits = []> :
     Affine_Op<mnemonic, !listconcat(traits,
         [DeclareOpInterfaceMethods<AffineWriteOpInterface>])> {
   code extraClassDeclarationBase = [{
-    /// Get value to be stored by store operation.
-    Value getValueToStore() { return getOperand(0); }
+    /// Returns the operand index of the stored value.
+    unsigned getStoredValOperandIndex() { return 0; }
 
     /// Returns the operand index of the memref.
     unsigned getMemRefOperandIndex() { return 1; }
 
     void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
 
     /// Returns the affine map used to index the memref for this operation.
     AffineMapAttr getAffineMapAttr() {

mlir/lib/Transforms/MemRefDataFlowOpt.cpp

 // loop/conditional live-out SSA values is available.
 // TODO: do general dead store elimination for memref's. This pass
 // currently only eliminates the stores only if no other loads/uses (other
 // than dealloc) remain.
 //
 struct MemRefDataFlowOpt : public MemRefDataFlowOptBase<MemRefDataFlowOpt> {
   void runOnFunction() override;
 
-  void forwardStoreToLoad(AffineLoadOp loadOp);
+  void forwardStoreToLoad(AffineReadOpInterface loadOp);
 
   // A list of memref's that are potentially dead / could be eliminated.
   SmallPtrSet<Value, 4> memrefsToErase;
   // Load op's whose results were replaced by those forwarded from stores.
   SmallVector<Operation *, 8> loadOpsToErase;
 
   DominanceInfo *domInfo = nullptr;
   PostDominanceInfo *postDomInfo = nullptr;
 };
 
 } // end anonymous namespace
 
 /// Creates a pass to perform optimizations relying on memref dataflow such as
 /// store to load forwarding, elimination of dead stores, and dead allocs.
 std::unique_ptr<OperationPass<FuncOp>> mlir::createMemRefDataFlowOptPass() {
   return std::make_unique<MemRefDataFlowOpt>();
 }
 
 // This is a straightforward implementation not optimized for speed. Optimize
 // if needed.
-void MemRefDataFlowOpt::forwardStoreToLoad(AffineLoadOp loadOp) {
+void MemRefDataFlowOpt::forwardStoreToLoad(AffineReadOpInterface loadOp) {
   // First pass over the use list to get the minimum number of surrounding
   // loops common between the load op and the store op, with min taken across
   // all store ops.
   SmallVector<Operation *, 8> storeOps;
   unsigned minSurroundingLoops = getNestingDepth(loadOp);
   for (auto *user : loadOp.getMemRef().getUsers()) {
-    auto storeOp = dyn_cast<AffineStoreOp>(user);
+    auto storeOp = dyn_cast<AffineWriteOpInterface>(user);
     if (!storeOp)
       continue;
     unsigned nsLoops = getNumCommonSurroundingLoops(*loadOp, *storeOp);
     minSurroundingLoops = std::min(nsLoops, minSurroundingLoops);
     storeOps.push_back(storeOp);
   }
 
   // The list of store op candidates for forwarding that satisfy conditions
if (llvm::all_of(depSrcStores, [&](Operation *depStore) {
       lastWriteStoreOp = storeOp;
       break;
     }
   }
   if (!lastWriteStoreOp)
     return;
 
   // Perform the actual store to load forwarding.
-  Value storeVal = cast<AffineStoreOp>(lastWriteStoreOp).getValueToStore();
-  loadOp.replaceAllUsesWith(storeVal);
+  Value storeVal = cast<AffineWriteOpInterface>(lastWriteStoreOp).getValueToStore();
+  loadOp.getValue().replaceAllUsesWith(storeVal);
   // Record the memref for a later sweep to optimize away.
   memrefsToErase.insert(loadOp.getMemRef());
   // Record this to erase later.
   loadOpsToErase.push_back(loadOp);
 }
 
 void MemRefDataFlowOpt::runOnFunction() {
   // Only supports single block functions at the moment.
   FuncOp f = getFunction();
   if (!llvm::hasSingleElement(f)) {
     markAllAnalysesPreserved();
     return;
   }
 
   domInfo = &getAnalysis<DominanceInfo>();
   postDomInfo = &getAnalysis<PostDominanceInfo>();
 
   loadOpsToErase.clear();
   memrefsToErase.clear();
 
   // Walk all load's and perform store to load forwarding.
-  f.walk([&](AffineLoadOp loadOp) { forwardStoreToLoad(loadOp); });
+  f.walk([&](AffineReadOpInterface loadOp) { forwardStoreToLoad(loadOp); });
 
   // Erase all load op's whose results were replaced with store fwd'ed ones.
   for (auto *loadOp : loadOpsToErase)
     loadOp->erase();
 
   // Check if the store fwd'ed memrefs are now left with only stores and can
   // thus be completely deleted. Note: the canonicalize pass should be able
   // to do this as well, but we'll do it here since we collected these anyway.
   for (auto memref : memrefsToErase) {
     // If the memref hasn't been alloc'ed in this function, skip.
     Operation *defOp = memref.getDefiningOp();
     if (!defOp || !isa<AllocOp>(defOp))
       // TODO: if the memref was returned by a 'call' operation, we
       // could still erase it if the call had no side-effects.
       continue;
     if (llvm::any_of(memref.getUsers(), [&](Operation *ownerOp) {
-          return !isa<AffineStoreOp, DeallocOp>(ownerOp);
+          return !isa<AffineWriteOpInterface, DeallocOp>(ownerOp);
         }))
       continue;
 
     // Erase all stores, the dealloc, and the alloc on the memref.
     for (auto *user : llvm::make_early_inc_range(memref.getUsers()))
       user->erase();
     defOp->erase();
   }
 }

mlir/test/Transforms/memref-vector-dataflow.mlir

+// RUN: mlir-opt %s -memref-dataflow-opt -split-input-file | FileCheck %s
+
+
+func @eltwise_filter_2d(%ifm0 : memref<10x20x30x512xf32> {tpc.tensor}, // Reused for eltwise output.
+                        %ifm1 : memref<10x20x30x512xf32> {tpc.tensor},
+                        %filter : memref<5x5x512xf32> {tpc.tensor},
+                        %ofm : memref<10x20x30x512xf32> {tpc.tensor},
+                        %auxBias : memref<512xf32> {tpc.tensor},
+                        %padW : index,
+                        %padH : index
+                        //%kernelW : index, -> 5
+                        //%kernelH : index, -> 5
+                        //%strideW : index, -> 2
+                        //%strideH : index, -> 2
+                        //%dilationW : index, -> 3
+                        //%dilationH : index -> 3
+                        ) {
+  // Dims
+  //const int depth = 0;
+  //const int width = 1;
+  //const int height = 2;
+  //const int batch = 3;
+
+  affine.for %d = 0 to 16 {
+    affine.for %b = 0 to 10 {
+      affine.for %h = 0 to 20 {
+        affine.for %w = 0 to 30 {
+          %lhs = affine.vector_load %ifm0[%b, %h, %w, 32*%d] : memref<10x20x30x512xf32>, vector<32xf32>
+          %rhs = affine.vector_load %ifm1[%b, %h, %w, 32*%d] : memref<10x20x30x512xf32>, vector<32xf32>
+          %add = addf %lhs, %rhs : vector<32xf32>
+          affine.vector_store %add, %ifm0[%b, %h, %w, 32*%d] : memref<10x20x30x512xf32>, vector<32xf32>
+    %lhs1 = affine.vector_load %ifm0[%b, %h, %w, 32*%d] : memref<10x20x30x512xf32>, vector<32xf32>
+    %add1 = addf %lhs1, %rhs : vector<32xf32>
+    affine.vector_store %add1, %ofm[%b, %h, %w, 32*%d] : memref<10x20x30x512xf32>, vector<32xf32>
+        }
+      }
+    }
+  }
+
+  return
+}
+
+// CHECK:      affine.for %{{.*}} = 0 to 16
+// CHECK-NEXT:   affine.for %{{.*}} = 0 to 10
+// CHECK-NEXT:     affine.for %{{.*}} = 0 to 20
+// CHECK-NEXT:       affine.for %{{.*}} = 0 to 30
+// CHECK:                  affine.vector_load
+// CHECK:                  affine.vector_load
+// CHECK:                  addf
+// CHECK:                  affine.vector_store
+// CHECK:                  addf
+// CHECK:                  affine.vector_store
+// CHECK:            }
+