[Mlir-commits] [mlir] [MLIR][Affine] Fix affine loop fusion with vector ops #115849, #120227 (PR #122799)

[Uday Bondhugula]

================
@@ -177,13 +184,68 @@ gatherProducerConsumerMemrefs(unsigned srcId, unsigned dstId,
                                 producerConsumerMemrefs);
 }
 
+/// Performs two checks:
+/// Firstly, checks if both src/dst ops are vector operations.
+/// Secondly, the shapes of the loads and stores of each memref in
+/// the producer/consumer chains. If the load shapes are larger
+/// than the stores then we cannot fuse the loops. The loads
+/// would have a dependency on the values stored.
+static bool checkVectorLoadStoreOps(unsigned srcId, unsigned dstId,
+                                    DenseSet<Value> &producerConsumerMemrefs,
+                                    MemRefDependenceGraph *mdg) {
+  SmallVector<Operation *> storeOps;
+  SmallVector<Operation *> loadOps;
+
+  auto *srcNode = mdg->getNode(srcId);
+  auto *dstNode = mdg->getNode(dstId);
+
+  for (Value memref : producerConsumerMemrefs) {
+    srcNode->getStoreOpsForMemref(memref, &storeOps);
+    dstNode->getLoadOpsForMemref(memref, &loadOps);
+
+    for (Operation *storeOp : storeOps) {
+      auto vectorStoreOp = dyn_cast<AffineVectorStoreOp>(storeOp);
+
+      if (!vectorStoreOp)
+        continue;
+
+      auto storeVecType = vectorStoreOp.getVectorType();
+
+      for (Operation *loadOp : loadOps) {
+        auto vectorLoadOp = dyn_cast<AffineVectorLoadOp>(loadOp);
+
+        if (!vectorLoadOp)
+          return false;
+
+        auto loadVecType = vectorLoadOp.getVectorType();
+
+        if (loadVecType.getRank() != storeVecType.getRank())
+          return false;
+
+        for (int i = 0; i < loadVecType.getRank(); ++i) {
----------------
bondhugula wrote:

Use `unsigned i = 0, e = ...` form to prevent repeated evaluation.

https://github.com/llvm/llvm-project/pull/122799