Select Shim Tile Used for Trace Packet Extraction

include/aie/Dialect/AIE/IR/AIETraceOps.td

@@ -487,12 +487,14 @@ def AIE_TraceHostConfigOp : AIE_Op<"trace.host_config", [
     - arg_idx (default=4): XRT argument index for trace buffer. Set to -1 to
       append trace data after the last tensor argument in the runtime_sequence.
     - routing (default=single): Shim routing strategy. Currently only 'single'
-      is supported, which routes all traces to column 0's shim.
+      is supported, which routes all traces to a single shim tile.
+    - egress_shim_col (default=0): Column index of the shim tile that receives
+      the trace packet flows.
 
     Example:
     ```mlir
     aie.runtime_sequence(%arg0: memref<16xi32>) {
-      aie.trace.host_config buffer_size=65536
+      aie.trace.host_config buffer_size=65536 egress_shim_col=4
       aie.trace.start_config @trace1
     }
     ```
@@ -501,7 +503,8 @@ def AIE_TraceHostConfigOp : AIE_Op<"trace.host_config", [
   let arguments = (ins
     AIEI32Attr:$buffer_size,
     DefaultValuedAttr<AIEI32Attr, "4">:$arg_idx,
-    DefaultValuedAttr<TraceShimRoutingAttr, "TraceShimRouting::Single">:$routing
+    DefaultValuedAttr<TraceShimRoutingAttr, "TraceShimRouting::Single">:$routing,
+    DefaultValuedAttr<AIEI32Attr, "0">:$egress_shim_col
   );
 
   let hasCustomAssemblyFormat = 1;

lib/Dialect/AIE/IR/AIETraceOps.cpp

@@ -753,9 +753,13 @@ void TraceHostConfigOp::print(OpAsmPrinter &p) {
   if (getRouting() != TraceShimRouting::Single)
     p << " routing = " << stringifyTraceShimRouting(getRouting());
 
+  if (getEgressShimCol() != 0)
+    p << " egress_shim_col = " << getEgressShimCol();
+
   p.printOptionalAttrDict(
       (*this)->getAttrs(),
-      /*elidedAttrs=*/{"buffer_size", "arg_idx", "routing"});
+      /*elidedAttrs=*/{"buffer_size", "arg_idx", "routing",
+                       "egress_shim_col"});
 }
 
 ParseResult TraceHostConfigOp::parse(OpAsmParser &parser,
@@ -798,6 +802,21 @@ ParseResult TraceHostConfigOp::parse(OpAsmParser &parser,
   result.attributes.set(
       "routing", TraceShimRoutingAttr::get(parser.getContext(), routingVal));
 
+  // Parse egress_shim_col (default: 0 = column 0)
+  int32_t egressShimColVal = 0;
+  if (succeeded(parser.parseOptionalKeyword("egress_shim_col"))) {
+    IntegerAttr egressShimCol;
+    if (parser.parseEqual() ||
+        parser.parseAttribute(egressShimCol,
+                              parser.getBuilder().getI32Type(),
+                              "egress_shim_col", result.attributes))
+      return failure();
+  } else {
+    result.attributes.set(
+        "egress_shim_col",
+        parser.getBuilder().getI32IntegerAttr(egressShimColVal));
+  }
+
   if (parser.parseOptionalAttrDict(result.attributes))
     return failure();
 
@@ -819,6 +838,11 @@ LogicalResult TraceHostConfigOp::verify() {
     return emitOpError("buffer_size must be positive");
   }
 
+  // Validate Shim col id
+  if (getEgressShimCol() < 0) {
+    return emitOpError("egress_shim_col must be >= 0");
+  }
+
   return success();
 }
 

lib/Dialect/AIE/Transforms/AIEInsertTraceFlows.cpp

@@ -123,6 +123,7 @@ struct AIEInsertTraceFlowsPass
     int bufferSizeBytes = hostConfig.getBufferSize();
     int traceArgIdx = hostConfig.getArgIdx();
     auto routing = hostConfig.getRouting();
+    int egressShimColFromIR = hostConfig.getEgressShimCol();
 
     // arg_idx=-1 means "append trace after last tensor"
     int traceBufferOffset = 0; // in bytes
@@ -323,8 +324,15 @@ struct AIEInsertTraceFlowsPass
     std::map<int, ShimInfo> shimInfos; // col -> ShimInfo
 
     if (routing == TraceShimRouting::Single) {
-      // All traces route to column 0 shim
-      int targetCol = 0;
+      // All traces route to a single shim, controlled by the egress_shim_col parameter (default is 0).
+      int targetCol = egressShimColFromIR;
+      if (targetCol >= targetModel.columns() ||
+          !targetModel.isShimNOCTile(targetCol, 0)) {
+        device.emitError() << "egress_shim_col " << targetCol
+                           << " is not a valid shim NOC tile (device has "
+                           << targetModel.columns() << " columns)";
+        return signalPassFailure();
+      }
       TileOp shimTile = nullptr;
       for (auto tile : device.getOps<TileOp>()) {
         if (tile.getCol() == targetCol && tile.getRow() == 0) {

programming_examples/basic/event_trace/aie_trace.mlir

@@ -183,8 +183,9 @@ module {
       // TRACE INITIALIZATION
       // ========================================================================
 
-      // Configure trace output buffer (8192 bytes, default arg_idx=4)
-      aie.trace.host_config buffer_size = 8192
+      // Configure trace output buffer (8192 bytes, default arg_idx=4).
+      // Select which shim should be used by the egress trace transfer (default is the shim tile from column 0). 
+      aie.trace.host_config buffer_size = 8192 egress_shim_col = 4
 
       // Start trace configuration
       aie.trace.start_config @core_trace

python/dialects/aie.py

@@ -637,6 +637,7 @@ def trace_host_config(
     *,
     arg_idx=4,
     routing=TraceShimRouting.Single,
+    egress_shim_col=0,
     loc=None,
     ip=None,
 ):
@@ -649,6 +650,7 @@ def trace_host_config(
         buffer_size=buffer_size,
         arg_idx=arg_idx,
         routing=routing,
+        egress_shim_col=egress_shim_col,
         loc=loc,
         ip=ip,
     )