[IR] Add keep_pkt_header attribute to amdaie.flow (#1338)

It was previously existing in aie.flow only.

Author: Yu-Zhewen

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.cpp

@@ -563,6 +563,13 @@ LogicalResult FlowOp::verify() {
     return emitOpError()
            << "multiple source and multiple targets is unsupported";
   }
+
+  if (!getIsPacketFlow() && getPacketId().has_value())
+    return emitOpError() << "packet ID can only be set for packet flows";
+
+  if (!getIsPacketFlow() && getKeepPktHeader().value_or(false))
+    return emitOpError() << "keep_pkt_header can only be set for packet flows";
+
   return success();
 }
 

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.td

@@ -132,15 +132,16 @@ def AMDAIE_FlowOp: AMDAIE_Op<"flow", [AttrSizedOperandSegments]>,
     %channel_1 = amdaie.channel(%tile_0_1, 0, port_type = DMA)
     %channel_2 = amdaie.channel(%tile_1_1, 0, port_type = DMA)
     %0 = amdaie.flow({%channel} -> {%channel_1, %channel_2})
-      {is_packet_flow = true, packet_id = 0 : ui8}
+      {is_packet_flow = true, keep_pkt_header = true, packet_id = 0 : ui8}
     ```
   }];
 
   let arguments = (
     ins Variadic<Index>:$sources,
         Variadic<Index>:$targets,
         BoolAttr:$is_packet_flow,
-        OptionalAttr<UI8Attr>:$packet_id
+        OptionalAttr<UI8Attr>:$packet_id,
+        OptionalAttr<BoolAttr>:$keep_pkt_header
   );
 
   let assemblyFormat = [{ `(` `{` $sources `}` `->` `{` $targets `}` `)` attr-dict }];

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/test/invalid.mlir

@@ -320,6 +320,50 @@ func.func @dma_cpy_nd_negative_source_stride(%arg0: !amdaie.logicalobjectfifo<me
 
 // -----
 
+func.func @flow_multi_sources_and_targets() {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %tile_0_0 = amdaie.tile(%c0, %c0)
+  %tile_0_1 = amdaie.tile(%c0, %c1)
+  %channel = amdaie.channel(%tile_0_0, 0, port_type = DMA, direction = MM2S)
+  %channel_1 = amdaie.channel(%tile_0_0, 1, port_type = DMA, direction = MM2S)
+  %channel_2 = amdaie.channel(%tile_0_1, 0, port_type = DMA, direction = S2MM)
+  %channel_3 = amdaie.channel(%tile_0_1, 1, port_type = DMA, direction = S2MM)
+  // expected-error @+1 {{multiple source and multiple targets is unsupported}}
+  %0 = amdaie.flow({%channel, %channel_1} -> {%channel_2, %channel_3}) {is_packet_flow = true}
+  return
+}
+
+// -----
+
+func.func @flow_invalid_packet_id() {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %tile_0_0 = amdaie.tile(%c0, %c0)
+  %tile_0_1 = amdaie.tile(%c0, %c1)
+  %channel = amdaie.channel(%tile_0_0, 0, port_type = DMA, direction = MM2S)
+  %channel_1 = amdaie.channel(%tile_0_1, 0, port_type = DMA, direction = S2MM)
+  // expected-error @+1 {{packet ID can only be set for packet flows}}
+  %0 = amdaie.flow({%channel} -> {%channel_1}) {is_packet_flow = false, packet_id = 0 : ui8}
+  return
+}
+
+// -----
+
+func.func @flow_invalid_keep_pkt_header() {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %tile_0_0 = amdaie.tile(%c0, %c0)
+  %tile_0_1 = amdaie.tile(%c0, %c1)
+  %channel = amdaie.channel(%tile_0_0, 0, port_type = DMA, direction = MM2S)
+  %channel_1 = amdaie.channel(%tile_0_1, 0, port_type = DMA, direction = S2MM)
+  // expected-error @+1 {{keep_pkt_header can only be set for packet flows}}
+  %0 = amdaie.flow({%channel} -> {%channel_1}) {is_packet_flow = false, keep_pkt_header = true}
+  return
+}
+
+// -----
+
 func.func @npu_dma_cpy_nd_invalid_src_offsets() {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/test/roundtrip.mlir

@@ -158,6 +158,7 @@ func.func @dma_cpy_nd_mixed(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32
 // CHECK:       %[[CHANNEL_1:.*]] = amdaie.channel(%[[TILE_0_1]], 0, port_type = DMA, direction = S2MM)
 // CHECK:       amdaie.flow({%[[CHANNEL]]} -> {%[[CHANNEL_1]]}) {is_packet_flow = false}
 // CHECK:       amdaie.flow({%[[CHANNEL]]} -> {%[[CHANNEL_1]]}) {is_packet_flow = true, packet_id = 1 : ui8}
+// CHECK:       amdaie.flow({%[[CHANNEL]]} -> {%[[CHANNEL_1]]}) {is_packet_flow = true, keep_pkt_header = true, packet_id = 2 : ui8}
 func.func @flow() {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
@@ -167,6 +168,7 @@ func.func @flow() {
   %channel_1 = amdaie.channel(%tile_0_1, 0, port_type = DMA, direction = S2MM)
   %0 = amdaie.flow({%channel} -> {%channel_1}) {is_packet_flow = false}
   %1 = amdaie.flow({%channel} -> {%channel_1}) {is_packet_flow = true, packet_id = 1 : ui8}
+  %2 = amdaie.flow({%channel} -> {%channel_1}) {is_packet_flow = true, keep_pkt_header = true, packet_id = 2 : ui8}
   return
 }
 

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignPacketIds.cpp

@@ -83,7 +83,7 @@ void AMDAIEAssignPacketIdsPass::runOnOperation() {
     rewriter.setInsertionPoint(flowOp);
     rewriter.replaceOpWithNewOp<AMDAIE::FlowOp>(
         flowOp, flowOp.getSources(), flowOp.getTargets(),
-        flowOp.getIsPacketFlow(), pktIdAttr);
+        flowOp.getIsPacketFlow(), pktIdAttr, /*keepPktHeader=*/nullptr);
     channelToPktFlowIndex[sourceChannelOp]++;
   }
 }

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEConnectionToFlow.cpp

@@ -37,7 +37,8 @@ void AMDAIEConnectionToFlowPass::runOnOperation() {
                                               AMDAIE::ConnectionType::Packet;
     auto flowOp = rewriter.create<AMDAIE::FlowOp>(
         rewriter.getUnknownLoc(), connectionOp.getSourceChannels(),
-        connectionOp.getTargetChannels(), isPacketFlow, /*packetId*/ nullptr);
+        connectionOp.getTargetChannels(), isPacketFlow, /*packetId=*/nullptr,
+        /*keepPktHeader=*/nullptr);
     rewriter.replaceOpWithNewOp<AMDAIE::ConnectionOp>(
         connectionOp, connectionOp.getTarget(),
         connectionOp.getTargetChannels(), connectionOp.getSource(),

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIELowerToAIE.cpp

@@ -206,7 +206,8 @@ SmallVector<Operation *> AIEDeviceBuilder::createFlowOps(
     if (pktId) {
       OpBuilder::InsertionGuard gg(rewriter);
       AIE::PacketFlowOp pktFlow = rewriter.create<AIE::PacketFlowOp>(
-          rewriter.getUnknownLoc(), pktId.value(), nullptr, nullptr);
+          rewriter.getUnknownLoc(), pktId.value(),
+          flowOp.getKeepPktHeaderAttr(), /*priority_route=*/nullptr);
       Region &r_pktFlow = pktFlow.getPorts();
       Block *b_pktFlow = rewriter.createBlock(&r_pktFlow);
       rewriter.setInsertionPointToStart(b_pktFlow);