support trace packet routing with corner turns

lib/Dialect/AIE/Transforms/AIECreatePacketFlows.cpp

@@ -157,10 +157,17 @@ void updateCoordinates(int &xCur, int &yCur, WireBundle move) {
 // Build a packet-switched route from the sourse to the destination with the
 // given ID. The route is recorded in the given map of switchboxes.
 void buildPSRoute(
-    int xSrc, int ySrc, Port sourcePort, int xDest, int yDest, Port destPort,
-    int flowID,
+    TileOp srcTile, Port sourcePort, TileOp destTile, Port destPort, int flowID,
     DenseMap<TileID, SmallVector<std::pair<Connect, int>, 8>> &switchboxes,
     bool reverseOrder = false) {
+
+  int xSrc = srcTile.colIndex();
+  int ySrc = srcTile.rowIndex();
+  int xDest = destTile.colIndex();
+  int yDest = destTile.rowIndex();
+
+  const auto &targetModel = getTargetModel(srcTile);
+
   int xCur = xSrc;
   int yCur = ySrc;
   WireBundle curBundle = {};
@@ -213,6 +220,13 @@ void buildPSRoute(
       if (move == lastBundle)
         continue;
 
+      // If the source port is a trace port, we need to validate the destination
+      if (xCur == xSrc && yCur == ySrc &&
+          sourcePort.bundle == WireBundle::Trace &&
+          !targetModel.isValidTraceMaster(xSrc, ySrc, move, curChannel)) {
+        continue;
+      }
+
       updateCoordinates(xCur, yCur, move);
 
       if (std::find(congestion.begin(), congestion.end(), TileID{xCur, yCur}) !=
@@ -320,22 +334,18 @@ struct AIERoutePacketFlowsPass
       Region &r = pktflow.getPorts();
       Block &b = r.front();
       int flowID = pktflow.IDInt();
-      int xSrc = 0, ySrc = 0;
-      Port sourcePort;
+      Port sourcePort, destPort;
+      TileOp srcTile, destTile;
 
       for (Operation &Op : b.getOperations()) {
         if (auto pktSource = dyn_cast<PacketSourceOp>(Op)) {
-          auto srcTile = dyn_cast<TileOp>(pktSource.getTile().getDefiningOp());
-          xSrc = srcTile.colIndex();
-          ySrc = srcTile.rowIndex();
+          srcTile = dyn_cast<TileOp>(pktSource.getTile().getDefiningOp());
           sourcePort = pktSource.port();
         } else if (auto pktDest = dyn_cast<PacketDestOp>(Op)) {
-          auto destTile = dyn_cast<TileOp>(pktDest.getTile().getDefiningOp());
-          int xDest = destTile.colIndex();
-          int yDest = destTile.rowIndex();
-          Port destPort = pktDest.port();
+          destTile = dyn_cast<TileOp>(pktDest.getTile().getDefiningOp());
+          destPort = pktDest.port();
 
-          buildPSRoute(xSrc, ySrc, sourcePort, xDest, yDest, destPort, flowID,
+          buildPSRoute(srcTile, sourcePort, destTile, destPort, flowID,
                        switchboxes, true);
 
           // Assign "keep_pkt_header flag"

test/create-packet-flows/trace_packet_routing.mlir

@@ -0,0 +1,28 @@
+//===- trace_packet_routing.mlir ------------------------------------------------*- MLIR -*-===//
+//
+// Copyright (C) 2024, Advanced Micro Devices, Inc.
+// SPDX-License-Identifier: MIT
+//
+//===----------------------------------------------------------------------===//
+// REQUIRES: ryzen_ai, chess
+
+// RUN: aie-opt --aie-create-packet-flows %s | FileCheck %s
+// CHECK-LABEL: module @trace_packet_routing {
+
+module @trace_packet_routing {
+ aie.device(npu1_4col) {
+  %tile_0_0 = aie.tile(0, 0)
+  %tile_1_0 = aie.tile(1, 0)
+  %tile_0_2 = aie.tile(0, 2)
+  %tile_0_3 = aie.tile(0, 3)
+
+  aie.packet_flow(0) { 
+    aie.packet_source<%tile_0_2, Trace : 0> // core trace
+    aie.packet_dest<%tile_0_0, DMA : 1>
+  } {keep_pkt_header = true}
+  aie.packet_flow(1) { 
+    aie.packet_source<%tile_0_3, Trace : 0> // core trace
+    aie.packet_dest<%tile_1_0, DMA : 1>
+  } {keep_pkt_header = true}
+ }
+}