[Router] Fix packet group id assignment (#1335)

Background: The group ID is used inside the router, where only packet
flows that share any source or destination port are assigned the same
group ID. The router then restricts channel sharing to occur only within
the same group, which helps prevent arbiter deadlocks.

Consider three packet flows: a, b, and c, where a and c have overlapping
source or destination ports, and b is independent. The expected group
IDs should be: 0, 1, 0.

However, the previous assignment logic was incorrect: `packetGroupId =
std::max(packetGroupId, existingId);`, which produced the wrong results:
0, 1, 1.

This issue was observed when using out-of-order mode DMA for
MatmulThinBias.

Author: Yu-Zhewen

compiler/plugins/target/AMD-AIE/aie/test/unit_out_of_order_routing.mlir

@@ -0,0 +1,54 @@
+// RUN: iree-opt --amdaie-create-pathfinder-flows %s | FileCheck %s
+
+// Routing test for three packet flows from memtile (0,1) to coretile (0,2).
+// Two of the flows target the same destination port, configured as out-of-order mode.
+
+// CHECK-LABEL:   aie.device(npu1_4col) {
+// CHECK:           %[[TILE_0_1:.*]] = aie.tile(0, 1)
+// CHECK:           %[[SWITCHBOX_0_1:.*]] = aie.switchbox(%[[TILE_0_1]]) {
+// CHECK:             %[[AMSEL_0:.*]] = aie.amsel<0> (0)
+// CHECK:             %[[AMSEL_1:.*]] = aie.amsel<1> (0)
+// CHECK:             %[[MASTERSET_NORTH_0:.*]] = aie.masterset(NORTH : 0, %[[AMSEL_0]])
+// CHECK:             %[[MASTERSET_NORTH_5:.*]] = aie.masterset(NORTH : 5, %[[AMSEL_1]])
+// CHECK:             aie.packet_rules(DMA : 1) {
+// CHECK:               aie.rule(31, 0, %[[AMSEL_0]]) {packet_ids = array<i32: 0>}
+// CHECK:             }
+// CHECK:             aie.packet_rules(DMA : 2) {
+// CHECK:               aie.rule(31, 0, %[[AMSEL_1]]) {packet_ids = array<i32: 0>}
+// CHECK:             }
+// CHECK:             aie.packet_rules(DMA : 3) {
+// CHECK:               aie.rule(31, 0, %[[AMSEL_0]]) {packet_ids = array<i32: 0>}
+// CHECK:             }
+// CHECK:           }
+// CHECK:           %[[TILE_0_2:.*]] = aie.tile(0, 2)
+// CHECK:           %[[SWITCHBOX_0_2:.*]] = aie.switchbox(%[[TILE_0_2]]) {
+// CHECK:             %[[AMSEL_0:.*]] = aie.amsel<0> (0)
+// CHECK:             %[[AMSEL_1:.*]] = aie.amsel<1> (0)
+// CHECK:             %[[MASTERSET_DMA_0:.*]] = aie.masterset(DMA : 0, %[[AMSEL_0]]) {keep_pkt_header = "true"}
+// CHECK:             %[[MASTERSET_DMA_1:.*]] = aie.masterset(DMA : 1, %[[AMSEL_1]])
+// CHECK:             aie.packet_rules(SOUTH : 0) {
+// CHECK:               aie.rule(31, 0, %[[AMSEL_0]]) {packet_ids = array<i32: 0>}
+// CHECK:             }
+// CHECK:             aie.packet_rules(SOUTH : 5) {
+// CHECK:               aie.rule(31, 0, %[[AMSEL_1]]) {packet_ids = array<i32: 0>}
+// CHECK:             }
+// CHECK:           }
+// CHECK:         }
+module {
+  aie.device(npu1_4col) {
+    %tile_0_1 = aie.tile(0, 1)
+    %tile_0_2 = aie.tile(0, 2)
+    aie.packet_flow(0) {
+      aie.packet_source<%tile_0_1, DMA : 1>
+      aie.packet_dest<%tile_0_2, DMA : 0>
+    } {keep_pkt_header = true}
+    aie.packet_flow(0) {
+      aie.packet_source<%tile_0_1, DMA : 2>
+      aie.packet_dest<%tile_0_2, DMA : 1>
+    }
+    aie.packet_flow(0) {
+      aie.packet_source<%tile_0_1, DMA : 3>
+      aie.packet_dest<%tile_0_2, DMA : 0>
+    } {keep_pkt_header = true}
+  }
+}

runtime/src/iree-amd-aie/aie_runtime/iree_aie_router.cc

@@ -273,32 +273,23 @@ void Router::addFlow(TileLoc srcCoords, Port srcPort, TileLoc dstCoords,
     }
   }
 
-  // Assign a group ID for packet flows
-  // any overlapping in source/destination will lead to the same group ID
-  // channel sharing will happen within the same group ID
-  // for circuit flows, group ID is always -1, and no channel sharing
-  int packetGroupId = -1;
-  if (isPacketFlow) {
-    bool found = false;
+  // Assign a group ID to packet flows:
+  // - If the source or any destination overlaps with an existing flow, reuse
+  // its group ID.
+  // - Flows sharing the same group ID can share channels.
+  // - For circuit flows, always assign group ID -1 (no channel sharing).
+  int32_t packetGroupId = [&]() -> int32_t {
+    if (!isPacketFlow) return -1;
     for (auto &[existingId, src, dsts] : impl->flows) {
-      if (src.tileLoc == srcCoords && src.port == srcPort) {
-        packetGroupId = existingId;
-        found = true;
-        break;
-      }
-      for (auto &dst : dsts) {
-        if (dst.tileLoc == dstCoords && dst.port == dstPort) {
-          packetGroupId = existingId;
-          found = true;
-          break;
-        }
+      if (src.tileLoc == srcCoords && src.port == srcPort) return existingId;
+      if (llvm::any_of(dsts, [&](const auto &dst) {
+            return dst.tileLoc == dstCoords && dst.port == dstPort;
+          })) {
+        return existingId;
       }
-      packetGroupId = std::max(packetGroupId, existingId);
-    }
-    if (!found) {
-      packetGroupId++;
     }
-  }
+    return nextAvailablePacketGroupId++;
+  }();
 
   // If no existing flow was found with this source, create a new flow.
   PhysPort srcPhysPort{srcCoords, srcPort, PhysPort::Direction::SRC};

runtime/src/iree-amd-aie/aie_runtime/iree_aie_router.h

@@ -140,6 +140,7 @@ struct Router {
   std::map<PhysPort, PhysPort> dijkstraShortestPaths(PhysPort src);
   std::optional<std::map<PhysPort, SwitchSettings>> findPaths(
       int maxIterations = 1000);
+  int32_t nextAvailablePacketGroupId = 0;
 };
 
 // A map from a switchbox output (physical) port to the number of that port.