Skip to content

Commit af67ee2

Browse files
jtuylsjtuyls
authored
[AssignPacketIds] Base packet ID assignment on source channels (#930)
1 parent 458945b commit af67ee2

File tree

2 files changed

+141
-19
lines changed

2 files changed

+141
-19
lines changed

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

+39-15
Original file line numberDiff line numberDiff line change
@@ -40,21 +40,45 @@ void AMDAIEAssignPacketIdsPass::runOnOperation() {
4040
AMDAIE::getDeviceModel(maybeDevice.value());
4141
auto ui8ty =
4242
IntegerType::get(rewriter.getContext(), 8, IntegerType::Unsigned);
43-
int pktFlowIndex{0};
44-
WalkResult res = parentOp->walk([&](AMDAIE::FlowOp flowOp) {
45-
if (pktFlowIndex > deviceModel.getPacketIdMaxIdx()) {
46-
flowOp.emitOpError() << "ran out of packet IDs to assign";
47-
return WalkResult::interrupt();
48-
}
49-
rewriter.setInsertionPoint(flowOp);
50-
IntegerAttr pktIdAttr = flowOp.getIsPacketFlow()
51-
? IntegerAttr::get(ui8ty, pktFlowIndex++)
52-
: nullptr;
53-
rewriter.replaceOpWithNewOp<AMDAIE::FlowOp>(
54-
flowOp, flowOp.getSources(), flowOp.getTargets(),
55-
flowOp.getIsPacketFlow(), pktIdAttr);
56-
return WalkResult::advance();
57-
});
43+
44+
// Perform assignment of packet IDs based on the source channels of the flow
45+
// ops. I.e. `amdaie.flow` ops with the same source channel will get a
46+
// different packet IDs assigned to accommodate multiple data packets being
47+
// routed through the same ports.
48+
DenseMap<AMDAIE::ChannelOp, size_t> channelToPktFlowIndex;
49+
WalkResult res =
50+
parentOp->walk([&](AMDAIE::FlowOp flowOp) {
51+
if (!flowOp.getIsPacketFlow()) return WalkResult::advance();
52+
SmallVector<Value> sourceChannels = flowOp.getSources();
53+
if (sourceChannels.size() == 0) {
54+
flowOp.emitOpError() << "with no source channel is unsupported";
55+
return WalkResult::interrupt();
56+
}
57+
if (sourceChannels.size() > 1) {
58+
flowOp.emitOpError()
59+
<< "with multiple source channels is unsupported";
60+
return WalkResult::interrupt();
61+
}
62+
auto sourceChannelOp = dyn_cast_if_present<AMDAIE::ChannelOp>(
63+
sourceChannels[0].getDefiningOp());
64+
if (!sourceChannelOp) {
65+
flowOp.emitOpError() << "source should be an `amdaie.channel` op";
66+
return WalkResult::interrupt();
67+
}
68+
size_t pktFlowIndex = channelToPktFlowIndex[sourceChannelOp];
69+
if (pktFlowIndex > deviceModel.getPacketIdMaxIdx()) {
70+
flowOp.emitOpError()
71+
<< "ran out of packet IDs to assign for source channel";
72+
return WalkResult::interrupt();
73+
}
74+
IntegerAttr pktIdAttr = IntegerAttr::get(ui8ty, pktFlowIndex);
75+
rewriter.setInsertionPoint(flowOp);
76+
rewriter.replaceOpWithNewOp<AMDAIE::FlowOp>(
77+
flowOp, flowOp.getSources(), flowOp.getTargets(),
78+
flowOp.getIsPacketFlow(), pktIdAttr);
79+
channelToPktFlowIndex[sourceChannelOp]++;
80+
return WalkResult::advance();
81+
});
5882
if (res.wasInterrupted()) return signalPassFailure();
5983
}
6084

‎compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/assign_packet_ids.mlir

+102-4
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
// expected-error @+1 {{has no AMDAIEDevice in the target attribute configuration}}
44
module {
5-
func.func @assign_packet_ids(%arg0: memref<8x16xi32>, %arg1: memref<1x1x8x16xi32, 1>, %arg2: memref<1x1x8x16xi32, 2>) {
5+
func.func @no_device() {
66
%c0 = arith.constant 0 : index
77
%c1 = arith.constant 1 : index
88
amdaie.workgroup {
@@ -21,6 +21,69 @@ module {
2121

2222
// -----
2323

24+
#executable_target_amdaie_xclbin_fb = #hal.executable.target<"amd-aie", "amdaie-xclbin-fb", {target_device = "npu1_4col", ukernels = "none"}>
25+
module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} {
26+
func.func @no_source_channel() {
27+
%c0 = arith.constant 0 : index
28+
%c1 = arith.constant 1 : index
29+
amdaie.workgroup {
30+
%tile_0_1 = amdaie.tile(%c0, %c1)
31+
%channel = amdaie.channel(%tile_0_1, 0, port_type = DMA, direction = S2MM)
32+
// expected-error @+1 {{with no source channel is unsupported}}
33+
%0 = amdaie.flow({} -> {%channel}) {is_packet_flow = true}
34+
amdaie.controlcode {
35+
amdaie.end
36+
}
37+
}
38+
return
39+
}
40+
}
41+
42+
// -----
43+
44+
#executable_target_amdaie_xclbin_fb = #hal.executable.target<"amd-aie", "amdaie-xclbin-fb", {target_device = "npu1_4col", ukernels = "none"}>
45+
module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} {
46+
func.func @no_channel() {
47+
%c0 = arith.constant 0 : index
48+
%c1 = arith.constant 1 : index
49+
amdaie.workgroup {
50+
%tile_0_1 = amdaie.tile(%c0, %c1)
51+
%channel = amdaie.channel(%tile_0_1, 0, port_type = DMA, direction = S2MM)
52+
// expected-error @+1 {{source should be an `amdaie.channel` op}}
53+
%0 = amdaie.flow({%c0} -> {%channel}) {is_packet_flow = true}
54+
amdaie.controlcode {
55+
amdaie.end
56+
}
57+
}
58+
return
59+
}
60+
}
61+
62+
// -----
63+
64+
#executable_target_amdaie_xclbin_fb = #hal.executable.target<"amd-aie", "amdaie-xclbin-fb", {target_device = "npu1_4col", ukernels = "none"}>
65+
module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} {
66+
func.func @multiple_source_channels() {
67+
%c0 = arith.constant 0 : index
68+
%c1 = arith.constant 1 : index
69+
amdaie.workgroup {
70+
%tile_0_0 = amdaie.tile(%c0, %c0)
71+
%tile_0_1 = amdaie.tile(%c0, %c1)
72+
%channel = amdaie.channel(%tile_0_0, 0, port_type = DMA, direction = MM2S)
73+
%channel_1 = amdaie.channel(%tile_0_0, 1, port_type = DMA, direction = MM2S)
74+
%channel_2 = amdaie.channel(%tile_0_1, 0, port_type = DMA, direction = S2MM)
75+
// expected-error @+1 {{with multiple source channels is unsupported}}
76+
%0 = amdaie.flow({%channel, %channel_1} -> {%channel_2}) {is_packet_flow = true}
77+
amdaie.controlcode {
78+
amdaie.end
79+
}
80+
}
81+
return
82+
}
83+
}
84+
85+
// -----
86+
2487
// CHECK-LABEL: @assign_packet_ids
2588
// CHECK: %[[C0:.*]] = arith.constant 0 : index
2689
// CHECK: %[[C1:.*]] = arith.constant 1 : index
@@ -35,10 +98,10 @@ module {
3598
// CHECK: %[[CHANNEL_3:.*]] = amdaie.channel(%[[TILE_0_1]], 1, port_type = DMA, direction = S2MM)
3699
// CHECK: amdaie.flow({%[[CHANNEL]]} -> {%[[CHANNEL_1]]}) {is_packet_flow = false}
37100
// CHECK: amdaie.flow({%[[CHANNEL]]} -> {%[[CHANNEL_1]]}) {is_packet_flow = true, packet_id = 0 : ui8}
38-
// CHECK: amdaie.flow({%[[CHANNEL_2]]} -> {%[[CHANNEL_3]]}) {is_packet_flow = true, packet_id = 1 : ui8}
101+
// CHECK: amdaie.flow({%[[CHANNEL_2]]} -> {%[[CHANNEL_3]]}) {is_packet_flow = true, packet_id = 0 : ui8}
39102
#executable_target_amdaie_xclbin_fb = #hal.executable.target<"amd-aie", "amdaie-xclbin-fb", {target_device = "npu1_4col", ukernels = "none"}>
40103
module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} {
41-
func.func @assign_packet_ids(%arg0: memref<8x16xi32>, %arg1: memref<1x1x8x16xi32, 1>, %arg2: memref<1x1x8x16xi32, 2>) {
104+
func.func @assign_packet_ids() {
42105
%c0 = arith.constant 0 : index
43106
%c1 = arith.constant 1 : index
44107
%c2 = arith.constant 2 : index
@@ -63,9 +126,44 @@ module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb}
63126

64127
// -----
65128

129+
// Test that different packet IDs are used for flows with the same source channel.
130+
// CHECK-LABEL: @assign_packet_ids_same_source_channel
131+
// CHECK: %[[C0:.*]] = arith.constant 0 : index
132+
// CHECK: %[[C1:.*]] = arith.constant 1 : index
133+
// CHECK: amdaie.workgroup
134+
// CHECK: %[[TILE_0_0:.*]] = amdaie.tile(%[[C0]], %[[C0]])
135+
// CHECK: %[[TILE_0_1:.*]] = amdaie.tile(%[[C0]], %[[C1]])
136+
// CHECK: %[[CHANNEL:.*]] = amdaie.channel(%[[TILE_0_0]], 0, port_type = DMA, direction = MM2S)
137+
// CHECK: %[[CHANNEL_1:.*]] = amdaie.channel(%[[TILE_0_1]], 0, port_type = DMA, direction = S2MM)
138+
// CHECK: %[[CHANNEL_2:.*]] = amdaie.channel(%[[TILE_0_1]], 1, port_type = DMA, direction = S2MM)
139+
// CHECK: amdaie.flow({%[[CHANNEL]]} -> {%[[CHANNEL_1]]}) {is_packet_flow = true, packet_id = 0 : ui8}
140+
// CHECK: amdaie.flow({%[[CHANNEL]]} -> {%[[CHANNEL_2]]}) {is_packet_flow = true, packet_id = 1 : ui8}
141+
#executable_target_amdaie_xclbin_fb = #hal.executable.target<"amd-aie", "amdaie-xclbin-fb", {target_device = "npu1_4col", ukernels = "none"}>
142+
module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} {
143+
func.func @assign_packet_ids_same_source_channel() {
144+
%c0 = arith.constant 0 : index
145+
%c1 = arith.constant 1 : index
146+
amdaie.workgroup {
147+
%tile_0_0 = amdaie.tile(%c0, %c0)
148+
%tile_0_1 = amdaie.tile(%c0, %c1)
149+
%channel = amdaie.channel(%tile_0_0, 0, port_type = DMA, direction = MM2S)
150+
%channel_1 = amdaie.channel(%tile_0_1, 0, port_type = DMA, direction = S2MM)
151+
%channel_2 = amdaie.channel(%tile_0_1, 1, port_type = DMA, direction = S2MM)
152+
%0 = amdaie.flow({%channel} -> {%channel_1}) {is_packet_flow = true}
153+
%1 = amdaie.flow({%channel} -> {%channel_2}) {is_packet_flow = true}
154+
amdaie.controlcode {
155+
amdaie.end
156+
}
157+
}
158+
return
159+
}
160+
}
161+
162+
// -----
163+
66164
#executable_target_amdaie_xclbin_fb = #hal.executable.target<"amd-aie", "amdaie-xclbin-fb", {target_device = "npu1_4col", ukernels = "none"}>
67165
module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} {
68-
func.func @assign_packet_ids(%arg0: memref<8x16xi32>, %arg1: memref<1x1x8x16xi32, 1>, %arg2: memref<1x1x8x16xi32, 2>) {
166+
func.func @assign_packet_ids() {
69167
%c0 = arith.constant 0 : index
70168
%c1 = arith.constant 1 : index
71169
amdaie.workgroup {

0 commit comments

Comments
 (0)