Insert trace flows pass (#2962)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Joseph Melber <jgmelber@gmail.com>

Author: 3 people

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

@@ -159,4 +159,22 @@ def EdgeTriggerAttr : I32EnumAttr<"EdgeTrigger",
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Trace Shim Routing Strategy Enumeration
+//===----------------------------------------------------------------------===//
+
+def TraceShimRoutingSingle    : I32EnumAttrCase<"Single", 0, "single">;
+
+def TraceShimRoutingAttr : I32EnumAttr<"TraceShimRouting",
+    "Shim tile routing strategy for trace collection",
+    [
+      TraceShimRoutingSingle
+    ]> {
+  let cppNamespace = "::xilinx::AIE";
+  let description = [{
+    Determines how traces are routed to shim tiles:
+    - single: All traces route to a single shim tile (column 0)
+  }];
+}
+
 #endif // AIE_TRACE_ATTRS

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

@@ -54,15 +54,13 @@ def AIE_TraceOp : AIE_Op<"trace", [
 
   let arguments = (ins
     Index:$tile,
-    SymbolNameAttr:$sym_name,
-    OptionalAttr<AIEI32Attr>:$buffer_size
+    SymbolNameAttr:$sym_name
   );
 
   let regions = (region SizedRegion<1>:$body);
 
   let assemblyFormat = [{
-    $sym_name `(` $tile `)` (`buffer_size` `=` $buffer_size^)?
-    $body attr-dict
+    $sym_name `(` $tile `)` $body attr-dict
   }];
 
   let hasVerifier = 1;
@@ -464,6 +462,47 @@ def AIE_TraceEdgeEventOp : AIE_Op<"trace.edge_event", [
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// AIE_TraceHostConfigOp - Host Buffer Configuration
+//===----------------------------------------------------------------------===//
+
+def AIE_TraceHostConfigOp : AIE_Op<"trace.host_config", [
+    HasParent<"RuntimeSequenceOp">
+  ]> {
+  let summary = "Configure host buffer for trace data collection";
+
+  let description = [{
+    Configures host-side trace buffer.
+
+    Attributes:
+    - buffer_size: Trace buffer size in bytes
+    - arg_idx (default=4): XRT argument index for trace buffer
+    - routing (default=single): Shim routing strategy. Currently only 'single'
+      is supported, which routes all traces to column 0's shim.
+    - trace_after_last_tensor (default=false): Append trace data to the last
+      tensor argument in the runtime_sequence. When true, buffer_offset is computed
+      from the tensor size. Only valid with routing=single.
+
+    Example:
+    ```mlir
+    aie.runtime_sequence(%arg0: memref<16xi32>) {
+      aie.trace.host_config buffer_size=65536
+      aie.trace.start_config @trace1
+    }
+    ```
+  }];
+
+  let arguments = (ins
+    AIEI32Attr:$buffer_size,
+    DefaultValuedAttr<AIEI32Attr, "4">:$arg_idx,
+    DefaultValuedAttr<TraceShimRoutingAttr, "TraceShimRouting::Single">:$routing,
+    DefaultValuedAttr<BoolAttr, "false">:$trace_after_last_tensor
+  );
+
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // AIE_TraceStartConfigOp - Runtime Configuration Invocation
 //===----------------------------------------------------------------------===//

include/aie/Dialect/AIE/Transforms/AIEPasses.h

@@ -68,6 +68,7 @@ std::unique_ptr<mlir::OperationPass<DeviceOp>> createAIEAssignTileCtrlIDsPass();
 std::unique_ptr<mlir::OperationPass<DeviceOp>> createAIETraceToConfigPass();
 std::unique_ptr<mlir::OperationPass<DeviceOp>>
 createAIETraceRegPackWritesPass();
+std::unique_ptr<mlir::OperationPass<DeviceOp>> createAIEInsertTraceFlowsPass();
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION

include/aie/Dialect/AIE/Transforms/AIEPasses.td

@@ -424,4 +424,41 @@ def AIETraceRegPackWrites : Pass<"aie-trace-pack-reg-writes", "DeviceOp"> {
   ];
 }
 
+def AIEInsertTraceFlows : Pass<"aie-insert-trace-flows", "DeviceOp"> {
+  let summary = "Insert packet flows and runtime sequence trace setup";
+  let description = [{
+    For each aie.trace operation, this pass:
+    - Creates ONE packet flow from trace port to shim DMA
+    - Groups traces by target shim (minimizes shim usage, ideally 1)
+    - Inserts ONE shim buffer descriptor per shim tile for all traces
+    - Inserts per-tile timer control register writes
+    - Inserts per-shim broadcast and DMA control setup
+    
+    All trace configuration is injected at the beginning of the runtime
+    sequence, before user data transfer operations.
+    
+    Multiple trace streams (from different tiles or different trace units
+    on the same tile) are routed to the same shim DMA channel and drained
+    by a single buffer descriptor.
+  }];
+  
+  let constructor = "xilinx::AIE::createAIEInsertTraceFlowsPass()";
+
+  let dependentDialects = [
+    "xilinx::AIE::AIEDialect",
+    "xilinx::AIEX::AIEXDialect",
+  ];
+
+  let options = [
+    Option<"clShimChannel", "shim-channel", "int", "1",
+           "S2MM DMA channel to use for trace (default: 1)">,
+    Option<"clDefaultBdId", "default-bd-id", "int", "15",
+           "Buffer descriptor ID for trace (default: 15)">,
+    Option<"clPacketIdStart", "packet-id-start", "int", "1",
+           "Starting packet ID for trace flows (default: 1)">,
+    Option<"clTraceBurstLength", "burst-length", "int", "64",
+           "DMA burst length for trace transfers (default: 64 bytes)">
+  ];
+}
+
 #endif

lib/Dialect/AIE/IR/AIETraceOps.cpp

@@ -727,6 +727,111 @@ ParseResult TraceEdgeEventOp::parse(OpAsmParser &parser,
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// TraceHostConfigOp
+//===----------------------------------------------------------------------===//
+
+void TraceHostConfigOp::print(OpAsmPrinter &p) {
+  p << " buffer_size = " << getBufferSize();
+
+  // Only print non-default values
+  if (getArgIdx() != 4)
+    p << " arg_idx = " << getArgIdx();
+
+  if (getRouting() != TraceShimRouting::Single)
+    p << " routing = " << stringifyTraceShimRouting(getRouting());
+
+  if (getTraceAfterLastTensor())
+    p << " trace_after_last_tensor = true";
+
+  p.printOptionalAttrDict((*this)->getAttrs(),
+                          /*elidedAttrs=*/{"buffer_size", "arg_idx", "routing",
+                                           "trace_after_last_tensor"});
+}
+
+ParseResult TraceHostConfigOp::parse(OpAsmParser &parser,
+                                     OperationState &result) {
+  // Parse required buffer_size
+  IntegerAttr bufferSize;
+  if (parser.parseKeyword("buffer_size") || parser.parseEqual() ||
+      parser.parseAttribute(bufferSize, parser.getBuilder().getI32Type(),
+                            "buffer_size", result.attributes))
+    return failure();
+
+  // Parse arg_idx (default: 4)
+  int32_t argIdxVal = 4;
+  if (succeeded(parser.parseOptionalKeyword("arg_idx"))) {
+    IntegerAttr argIdx;
+    if (parser.parseEqual() ||
+        parser.parseAttribute(argIdx, parser.getBuilder().getI32Type(),
+                              "arg_idx", result.attributes))
+      return failure();
+  } else {
+    result.attributes.set("arg_idx",
+                          parser.getBuilder().getI32IntegerAttr(argIdxVal));
+  }
+
+  // Parse routing (default: single)
+  TraceShimRouting routingVal = TraceShimRouting::Single;
+  if (succeeded(parser.parseOptionalKeyword("routing"))) {
+    if (parser.parseEqual())
+      return failure();
+    StringRef routingStr;
+    if (failed(parser.parseKeyword(&routingStr)))
+      return failure();
+    auto routing = symbolizeTraceShimRouting(routingStr);
+    if (!routing)
+      return parser.emitError(parser.getCurrentLocation(),
+                              "unknown routing strategy: ")
+             << routingStr;
+    routingVal = *routing;
+  }
+  result.attributes.set(
+      "routing", TraceShimRoutingAttr::get(parser.getContext(), routingVal));
+
+  // Parse trace_after_last_tensor (default: false)
+  bool traceAfterVal = false;
+  if (succeeded(parser.parseOptionalKeyword("trace_after_last_tensor"))) {
+    if (parser.parseEqual())
+      return failure();
+    StringRef boolStr;
+    if (failed(parser.parseKeyword(&boolStr)))
+      return failure();
+    if (boolStr == "true")
+      traceAfterVal = true;
+    else if (boolStr != "false")
+      return parser.emitError(parser.getCurrentLocation(),
+                              "expected 'true' or 'false'");
+  }
+  result.attributes.set("trace_after_last_tensor",
+                        parser.getBuilder().getBoolAttr(traceAfterVal));
+
+  if (parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+
+  return success();
+}
+
+LogicalResult TraceHostConfigOp::verify() {
+  bool useTraceAfterLastTensor = getTraceAfterLastTensor();
+
+  // Appending trace data after the last tensor only works with single shim
+  if (useTraceAfterLastTensor) {
+    if (getRouting() != TraceShimRouting::Single) {
+      return emitOpError("appending trace data to the last tensor argument "
+                         "only works with single shim destination strategy "
+                         "(routing=single)");
+    }
+  }
+
+  // Validate buffer_size is positive
+  if (getBufferSize() <= 0) {
+    return emitOpError("buffer_size must be positive");
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TraceStartConfigOp
 //===----------------------------------------------------------------------===//