Enable AIRFuseChannel to fuse channels temporally, if there's no existing for loop (Xilinx#597)

* Fixup bug where channel fusion by for loop only checks puts but not gets

* Fixup ci tests which were wrong in the first place

* Enable temporal channel fusion without for loop in place

* Clang format

* Test

Author: erwei-xilinx

mlir/lib/Transform/AIRDependencyScheduleOpt.cpp

@@ -3009,23 +3009,49 @@ class AIRFuseChannels
     init_options();
     if (channelOps.empty())
       return;
+    // Rename symbols
+    // TODO: make this greedy
+    auto renameSymbols =
+        [](std::vector<air::ChannelOp> &channelOps,
+           std::map<air::ChannelOp, air::ChannelOp> chan_merge_map) {
+          for (unsigned i = 0; i < channelOps.size(); i++) {
+            for (auto chanKey : channelOps) {
+              if (!chan_merge_map.count(chanKey))
+                continue;
+              auto error = mlir::SymbolTable::replaceAllSymbolUses(
+                  chanKey.getOperation(),
+                  mlir::SymbolTable::getSymbolName(chan_merge_map[chanKey]),
+                  chanKey->getParentOfType<ModuleOp>());
+              // FIXME: what if this fails?
+              (void)error;
+            }
+          }
+        };
     std::map<air::ChannelOp, air::ChannelOp> chan_merge_map;
     for (unsigned i = 0; i < channelOps.size() - 1; i++) {
       for (unsigned j = i + 1; j < channelOps.size(); j++) {
         std::tuple<bool, std::string> checkScfForMergeableRes =
-            checkIfScfForMergeable(channelOps[i], channelOps[j]);
+            checkIfTemporalMergeable(channelOps[i], channelOps[j]);
         if (!std::get<0>(checkScfForMergeableRes))
           continue;
-        // Fuse air.channels by scf.for loop unpeeling, i.e. recovering any
-        // missing zeroth iterations in scf.for loops.
         air::ChannelOp chanA = channelOps[i];
         air::ChannelOp chanB = channelOps[j];
-        sortChannelsByLoopNests(chanA, chanB);
-        mergeChannelOpsByScfFor(chanA, chanB,
-                                std::get<1>(checkScfForMergeableRes));
-        chan_merge_map[chanB] = chanA;
+        if (std::get<1>(checkScfForMergeableRes) == "LB" ||
+            std::get<1>(checkScfForMergeableRes) == "UB") {
+          // Fuse air.channels by scf.for loop unpeeling, i.e. recovering any
+          // missing zeroth/last iterations in scf.for loops.
+          sortChannelsByLoopNests(chanA, chanB);
+          mergeChannelOpsTemporally(chanA, chanB,
+                                    std::get<1>(checkScfForMergeableRes));
+          chan_merge_map[chanB] = chanA;
+        } else if (std::get<1>(checkScfForMergeableRes) == "NFL") {
+          // Fuse air.channels temporally, if there isn't any for loop to fuse
+          // into.
+          chan_merge_map[chanB] = chanA;
+        }
       }
     }
+    renameSymbols(channelOps, chan_merge_map);
     if (!targetMemorySpaces.empty()) {
       for (unsigned i = 0; i < channelOps.size() - 1; i++) {
         for (unsigned j = i + 1; j < channelOps.size(); j++) {
@@ -3037,20 +3063,7 @@ class AIRFuseChannels
         }
       }
     }
-    // Rename symbols
-    // TODO: make this greedy
-    for (unsigned i = 0; i < channelOps.size(); i++) {
-      for (auto chanKey : channelOps) {
-        if (!chan_merge_map.count(chanKey))
-          continue;
-        auto error = mlir::SymbolTable::replaceAllSymbolUses(
-            chanKey.getOperation(),
-            mlir::SymbolTable::getSymbolName(chan_merge_map[chanKey]),
-            chanKey->getParentOfType<ModuleOp>());
-        // FIXME: what if this fails?
-        (void)error;
-      }
-    }
+    renameSymbols(channelOps, chan_merge_map);
   }
 
   void runOnOperation() override {
@@ -3183,8 +3196,8 @@ class AIRFuseChannels
   }
 
   std::tuple<bool, std::string>
-  checkIfScfForMergeableImpl(std::vector<Operation *> a_loop_nest,
-                             std::vector<Operation *> b_loop_nest) {
+  checkIfTemporalMergeableByScfForImpl(std::vector<Block *> a_loop_nest,
+                                       std::vector<Block *> b_loop_nest) {
     std::tuple<bool, std::string> notMergeable = {false, ""};
     std::tuple<bool, std::string> mergeableToLB = {true, "LB"};
     std::tuple<bool, std::string> mergeableToUB = {true, "UB"};
@@ -3194,20 +3207,22 @@ class AIRFuseChannels
         std::max(a_loop_nest.size(), b_loop_nest.size());
     // Skip over the unequal scf.for loop, and check equality for the rest of
     // the loops first.
-    unsigned outermostScfFor = -1;
+    int outermostScfFor = -1;
     for (unsigned i = 0; i < max_loop_nest_count; i++) {
       unsigned scfForCount = 0;
-      if ((i < a_loop_nest.size()) && isa<scf::ForOp>(a_loop_nest[i]))
+      if ((i < a_loop_nest.size()) &&
+          isa<scf::ForOp>(a_loop_nest[i]->getParentOp()))
         scfForCount++;
-      if ((i < b_loop_nest.size()) && isa<scf::ForOp>(b_loop_nest[i]))
+      if ((i < b_loop_nest.size()) &&
+          isa<scf::ForOp>(b_loop_nest[i]->getParentOp()))
         scfForCount++;
       if (scfForCount == 1)
-        outermostScfFor = i;
+        outermostScfFor = (int)i;
     }
     if (outermostScfFor < 0)
       return notMergeable;
     SmallVector<unsigned> controlLoopIndices;
-    for (unsigned i = 0; i < max_loop_nest_count; i++)
+    for (int i = 0; i < (int)max_loop_nest_count; i++)
       if (i != outermostScfFor)
         controlLoopIndices.push_back(i);
     // TODO: Assuming b_loop_nest is before a_loop_nest. Always true? TODO:
@@ -3220,7 +3235,8 @@ class AIRFuseChannels
       }
       // Check if the skipped scf.for loop has LB >= 1. This is a sign of
       // peeling, indicating opportunity of merge by unpeeling into LB.
-      auto outerMostScfFor = dyn_cast<scf::ForOp>(a_loop_nest[outermostScfFor]);
+      auto outerMostScfFor =
+          dyn_cast<scf::ForOp>(a_loop_nest[outermostScfFor]->getParentOp());
       assert(outerMostScfFor);
       if (auto constLB = getConstantIntValue(outerMostScfFor.getLowerBound()))
         if (*constLB < 1)
@@ -3232,12 +3248,31 @@ class AIRFuseChannels
           return notMergeable;
       }
       // Merge by unpeeling into UB.
-      auto outerMostScfFor = dyn_cast<scf::ForOp>(b_loop_nest[outermostScfFor]);
+      auto outerMostScfFor =
+          dyn_cast<scf::ForOp>(b_loop_nest[outermostScfFor]->getParentOp());
       assert(outerMostScfFor);
       return mergeableToUB;
     }
     return mergeableToLB;
   }
+  std::tuple<bool, std::string>
+  checkIfTemporalMergeableImpl(std::vector<Block *> a_loop_nest,
+                               std::vector<Block *> b_loop_nest) {
+    std::tuple<bool, std::string> notMergeable = {false, ""};
+    std::tuple<bool, std::string> mergeableToLB = {true, "LB"};
+    std::tuple<bool, std::string> mergeableToUB = {true, "UB"};
+    std::tuple<bool, std::string> mergeableNoForLoop = {true, "NFL"};
+    if (std::abs((int)a_loop_nest.size() - (int)b_loop_nest.size()) == 1)
+      return checkIfTemporalMergeableByScfForImpl(a_loop_nest, b_loop_nest);
+    else if (a_loop_nest.size() != b_loop_nest.size())
+      return notMergeable;
+
+    for (unsigned i = 0; i < a_loop_nest.size(); i++) {
+      if (!areEquivalentControlLoops(a_loop_nest[i], b_loop_nest[i]))
+        return notMergeable;
+    }
+    return mergeableNoForLoop;
+  }
   Value getHierOperandFromHierBlockArgument(BlockArgument arg) {
     if (!arg)
       return nullptr;
@@ -3284,9 +3319,10 @@ class AIRFuseChannels
   }
   // Check of two air.channels are mergeable in time, by fusing into a shared
   // scf.for loop. Returns a tuple of bool of whether mergeable, and string of
-  // fusing into for loop lower bound (LB) or upper bound (UB).
-  std::tuple<bool, std::string> checkIfScfForMergeable(air::ChannelOp chan_a,
-                                                       air::ChannelOp chan_b) {
+  // fusing into for loop lower bound (LB) or upper bound (UB), or fuse with no
+  // for loop (NFL).
+  std::tuple<bool, std::string>
+  checkIfTemporalMergeable(air::ChannelOp chan_a, air::ChannelOp chan_b) {
     std::vector<air::ChannelPutOp> a_puts =
         getChannelPutOpThroughSymbol(chan_a);
     std::vector<air::ChannelPutOp> b_puts =
@@ -3297,6 +3333,8 @@ class AIRFuseChannels
         getChannelGetOpThroughSymbol(chan_b);
     std::tuple<bool, std::string> notMergeable = {false, ""};
     std::tuple<bool, std::string> mergeableToLB = {true, "LB"};
+    std::tuple<bool, std::string> mergeableToUB = {true, "UB"};
+    std::tuple<bool, std::string> mergeableNoForLoop = {true, "NFL"};
     if (a_puts.size() != b_puts.size())
       return notMergeable;
     if (a_puts.size() != 1)
@@ -3356,35 +3394,73 @@ class AIRFuseChannels
         (!areTheSameSSAValueLists(aSizes, bSizes)) ||
         (!areTheSameSSAValueLists(aStrides, bStrides)))
       return notMergeable;
+    std::vector<std::tuple<bool, std::string>> putResults;
     for (unsigned i = 0; i < a_puts.size(); i++) {
       auto a_put_loop_nest = getParentLoopNest(a_puts[i].getOperation());
       auto b_put_loop_nest = getParentLoopNest(b_puts[i].getOperation());
-      return checkIfScfForMergeableImpl(a_put_loop_nest, b_put_loop_nest);
+      putResults.push_back(
+          checkIfTemporalMergeableImpl(a_put_loop_nest, b_put_loop_nest));
     }
+    std::vector<std::tuple<bool, std::string>> getResults;
     for (unsigned i = 0; i < a_gets.size(); i++) {
       auto a_get_loop_nest = getParentLoopNest(a_gets[i].getOperation());
       auto b_get_loop_nest = getParentLoopNest(b_gets[i].getOperation());
-      return checkIfScfForMergeableImpl(a_get_loop_nest, b_get_loop_nest);
-    }
-    return mergeableToLB;
+      getResults.push_back(
+          checkIfTemporalMergeableImpl(a_get_loop_nest, b_get_loop_nest));
+    }
+    bool overallUBMergeable = true;
+    bool overallLBMergeable = true;
+    bool overallNFLMergeable = true;
+    for (auto putRes : putResults) {
+      if (!std::get<0>(putRes))
+        return notMergeable;
+      overallUBMergeable &= (std::get<1>(putRes) == "UB");
+      overallLBMergeable &= (std::get<1>(putRes) == "LB");
+      overallNFLMergeable &= (std::get<1>(putRes) == "NFL");
+    }
+    for (auto getRes : getResults) {
+      if (!std::get<0>(getRes))
+        return notMergeable;
+      overallUBMergeable &= (std::get<1>(getRes) == "UB");
+      overallLBMergeable &= (std::get<1>(getRes) == "LB");
+      overallNFLMergeable &= (std::get<1>(getRes) == "NFL");
+    }
+    if (overallNFLMergeable)
+      return mergeableNoForLoop;
+    else if (overallLBMergeable)
+      return mergeableToLB;
+    else if (overallUBMergeable)
+      return mergeableToUB;
+    return notMergeable;
   }
-  std::vector<Operation *> getParentLoopNest(Operation *op) {
-    std::vector<Operation *> parent_loop_nest;
+  std::vector<Block *> getParentLoopNest(Operation *op) {
+    std::vector<Block *> parent_loop_nest;
     Operation *parent = op;
     while (parent) {
-      if (isa<scf::ForOp>(parent))
-        parent_loop_nest.push_back(parent);
-      else if (isa<scf::ParallelOp>(parent))
-        parent_loop_nest.push_back(parent);
-      else if (isa<air::HierarchyInterface>(parent))
-        parent_loop_nest.push_back(parent);
-      else if (isa<affine::AffineIfOp>(parent))
-        parent_loop_nest.push_back(parent);
+      if (auto forOp = dyn_cast<scf::ForOp>(parent))
+        parent_loop_nest.push_back(forOp.getBody());
+      else if (auto parOp = dyn_cast<scf::ParallelOp>(parent))
+        parent_loop_nest.push_back(parOp.getBody());
+      else if (auto hierOp = dyn_cast<air::HierarchyInterface>(parent))
+        parent_loop_nest.push_back(&hierOp->getRegion(0).front());
+      else if (auto aifOp = dyn_cast<affine::AffineIfOp>(parent)) {
+        if (aifOp.getThenBlock()->findAncestorOpInBlock(*op))
+          parent_loop_nest.push_back(aifOp.getThenBlock());
+        else if (aifOp.hasElse() &&
+                 aifOp.getElseBlock()->findAncestorOpInBlock(*op))
+          parent_loop_nest.push_back(aifOp.getElseBlock());
+      }
       parent = parent->getParentOp();
     }
     return parent_loop_nest;
   }
-  bool areEquivalentControlLoops(Operation *a, Operation *b) {
+  bool areEquivalentControlLoops(Block *aBlock, Block *bBlock) {
+    Operation *a = aBlock->getParentOp();
+    Operation *b = bBlock->getParentOp();
+    if (!a)
+      return false;
+    if (!b)
+      return false;
     if (isa<scf::ForOp>(a) && isa<scf::ForOp>(b)) {
       auto a_for = dyn_cast<scf::ForOp>(a);
       auto b_for = dyn_cast<scf::ForOp>(b);
@@ -3437,11 +3513,27 @@ class AIRFuseChannels
       }
       return true;
     } else if (isa<air::HierarchyInterface>(a) &&
-               isa<air::HierarchyInterface>(a)) {
+               isa<air::HierarchyInterface>(b)) {
       if (a == b)
         return true;
-    } else if (isa<affine::AffineIfOp>(a) || isa<affine::AffineIfOp>(b))
-      return false;
+      auto aHierSym =
+          a->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
+      auto bHierSym =
+          b->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
+      if (aHierSym && bHierSym && aHierSym.str() == bHierSym.str())
+        return true;
+    } else if (isa<affine::AffineIfOp>(a) || isa<affine::AffineIfOp>(b)) {
+      if (a == b)
+        return false; // Sharing the same affine.if means spatially parallel
+                      // ops. Cannot merge by for loop (i.e. in time).
+      auto aIf = dyn_cast<affine::AffineIfOp>(a);
+      auto bIf = dyn_cast<affine::AffineIfOp>(b);
+      if (aBlock == aIf.getThenBlock() && bBlock == bIf.getThenBlock())
+        return true;
+      if (aIf.hasElse() && bIf.hasElse() && aBlock == aIf.getElseBlock() &&
+          bBlock == bIf.getElseBlock())
+        return true;
+    }
     return false;
   }
   void mergeChannelOps(air::ChannelInterface a, air::ChannelInterface b) {
@@ -3459,8 +3551,9 @@ class AIRFuseChannels
       async_a.addAsyncDependency(
           dyn_cast<air::AsyncOpInterface>(new_b).getAsyncToken());
   }
-  void mergeChannelOpsByScfFor(air::ChannelInterface a, air::ChannelInterface b,
-                               std::string mergeByLBOrUB) {
+  void mergeChannelOpsTemporally(air::ChannelInterface a,
+                                 air::ChannelInterface b,
+                                 std::string mergeByLBOrUB) {
     scf::ForOp parentForOp = a->getParentOfType<scf::ForOp>();
     while (parentForOp && parentForOp->getParentOfType<scf::ForOp>()) {
       parentForOp = parentForOp->getParentOfType<scf::ForOp>();
@@ -3498,8 +3591,8 @@ class AIRFuseChannels
       mergeChannelOps(a_gets[i], b_gets[i]);
     }
   }
-  void mergeChannelOpsByScfFor(air::ChannelOp chan_a, air::ChannelOp chan_b,
-                               std::string mergeByLBOrUB) {
+  void mergeChannelOpsTemporally(air::ChannelOp chan_a, air::ChannelOp chan_b,
+                                 std::string mergeByLBOrUB) {
     std::vector<air::ChannelPutOp> a_puts =
         getChannelPutOpThroughSymbol(chan_a);
     std::vector<air::ChannelPutOp> b_puts =
@@ -3509,10 +3602,10 @@ class AIRFuseChannels
     std::vector<air::ChannelGetOp> b_gets =
         getChannelGetOpThroughSymbol(chan_b);
     if (!b_puts[0]->getParentOfType<air::HerdOp>()) {
-      mergeChannelOpsByScfFor(a_puts[0], b_puts[0], mergeByLBOrUB);
+      mergeChannelOpsTemporally(a_puts[0], b_puts[0], mergeByLBOrUB);
     }
     if (!b_gets[0]->getParentOfType<air::HerdOp>()) {
-      mergeChannelOpsByScfFor(a_gets[0], b_gets[0], mergeByLBOrUB);
+      mergeChannelOpsTemporally(a_gets[0], b_gets[0], mergeByLBOrUB);
     }
   }
   Operation *cloneOpAndOperands(OpBuilder builder, IRMapping remap,
@@ -3550,16 +3643,18 @@ class AIRFuseChannels
     if (a_loop_nest.size() != b_loop_nest.size())
       return;
     for (unsigned i = 0; i < a_loop_nest.size(); i++) {
-      if (auto a_for = dyn_cast<scf::ForOp>(a_loop_nest[i])) {
-        if (auto b_for = dyn_cast<scf::ForOp>(b_loop_nest[i])) {
+      if (auto a_for = dyn_cast<scf::ForOp>(a_loop_nest[i]->getParentOp())) {
+        if (auto b_for = dyn_cast<scf::ForOp>(b_loop_nest[i]->getParentOp())) {
           for (unsigned j = 0; j < a_for.getBody()->getNumArguments(); j++) {
             remap.map(b_for.getBody()->getArgument(j),
                       a_for.getBody()->getArgument(j));
           }
         }
       }
-      if (auto a_par = dyn_cast<scf::ParallelOp>(a_loop_nest[i])) {
-        if (auto b_par = dyn_cast<scf::ParallelOp>(b_loop_nest[i])) {
+      if (auto a_par =
+              dyn_cast<scf::ParallelOp>(a_loop_nest[i]->getParentOp())) {
+        if (auto b_par =
+                dyn_cast<scf::ParallelOp>(b_loop_nest[i]->getParentOp())) {
           for (unsigned j = 0; j < a_par.getBody()->getNumArguments(); j++) {
             remap.map(b_par.getBody()->getArgument(j),
                       a_par.getBody()->getArgument(j));

mlir/test/Conversion/AIRToAIE/async_gemm_w_pingpong_to_locks_aie2.mlir

@@ -17,10 +17,10 @@
 // CHECK:   %[[VAL_7:.*]] = aie.tile(6, 4)
 // CHECK-COUNT-8:    aie.lock(%[[VAL_3]], {{.*}})
 // CHECK-COUNT-2:    aie.lock(%[[VAL_2]], {{.*}})
-// CHECK-COUNT-8:    aie.lock(%[[VAL_4]], {{.*}})
-// CHECK-COUNT-8:    aie.lock(%[[VAL_5]], {{.*}})
-// CHECK-COUNT-8:    aie.lock(%[[VAL_6]], {{.*}})
-// CHECK-COUNT-8:    aie.lock(%[[VAL_7]], {{.*}})
+// CHECK-COUNT-6:    aie.lock(%[[VAL_4]], {{.*}})
+// CHECK-COUNT-6:    aie.lock(%[[VAL_5]], {{.*}})
+// CHECK-COUNT-6:    aie.lock(%[[VAL_6]], {{.*}})
+// CHECK-COUNT-6:    aie.lock(%[[VAL_7]], {{.*}})
 // CHECK:    aie.buffer(%[[VAL_2]]) {{{.*}}} : memref<64x64xi32, 1>
 // CHECK:    aie.buffer(%[[VAL_3]]) {{{.*}}} : memref<64x128xi32, 1>
 // CHECK:    aie.buffer(%[[VAL_3]]) {{{.*}}} : memref<128x64xi32, 1>