Skip to content

Commit

Permalink
resolve comments
Browse files Browse the repository at this point in the history
  • Loading branch information
@Yu-Zhewen
Yu-Zhewen committed Jan 8, 2025
1 parent ec53786 commit eba6384
Show file tree
Hide file tree
Showing 3 changed files with 132 additions and 47 deletions.
10 changes: 8 additions & 2 deletions compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignChannels.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,11 +70,14 @@ LogicalResult assignChannels(AMDAIE::WorkgroupOp workgroupOp) {
assert(tileToGeneratorMap.contains(tile) &&
"no channel generator found for tile");
std::optional<uint8_t> maybeChannel =
tileToGeneratorMap[tile].getAndAssignProducerDMAChannel(isPacketFlow);
tileToGeneratorMap[tile].getProducerDMAChannel();
if (!maybeChannel) {
return connectionOp.emitOpError()
<< "no producer DMA channel available";
}
// Only assign the channel if it is for circuit flow.
if (!isPacketFlow)
tileToGeneratorMap[tile].assignProducerDMAChannel(maybeChannel.value());
auto channelOp = rewriter.create<AMDAIE::ChannelOp>(
rewriter.getUnknownLoc(), tile, maybeChannel.value(),
StrmSwPortType::DMA, AMDAIE::DMAChannelDir::MM2S);
Expand All @@ -85,11 +88,14 @@ LogicalResult assignChannels(AMDAIE::WorkgroupOp workgroupOp) {
assert(tileToGeneratorMap.contains(tile) &&
"no channel generator found for tile");
std::optional<uint8_t> maybeChannel =
tileToGeneratorMap[tile].getAndAssignConsumerDMAChannel(isPacketFlow);
tileToGeneratorMap[tile].getConsumerDMAChannel();
if (!maybeChannel) {
return connectionOp.emitOpError()
<< "no consumer DMA channel available";
}
// Only assign the channel if it is for circuit flow.
if (!isPacketFlow)
tileToGeneratorMap[tile].assignConsumerDMAChannel(maybeChannel.value());
auto channelOp = rewriter.create<AMDAIE::ChannelOp>(
rewriter.getUnknownLoc(), tile, maybeChannel.value(),
StrmSwPortType::DMA, AMDAIE::DMAChannelDir::S2MM);
Expand Down
36 changes: 26 additions & 10 deletions runtime/src/iree-amd-aie/aie_runtime/Utils/ChannelGenerator.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ using namespace llvm;

namespace mlir::iree_compiler::AMDAIE {

enum class ChannelAssignmentMode { FirstAvailable, RoundRobin };

/// Utility to generate valid channels.
class ChannelGenerator {
public:
Expand All @@ -30,27 +32,41 @@ class ChannelGenerator {
lastUsedConsumerChannel = numConsumerChannels - 1;
}

/// Returns its next usable producer channel.
std::optional<uint8_t> getAndAssignProducerDMAChannel(bool isPacketFlow) {
/// Returns its next usable producer channel. By default, it uses round-robin
/// for load balancing.
std::optional<uint8_t> getProducerDMAChannel(
ChannelAssignmentMode mode = ChannelAssignmentMode::RoundRobin) {
for (uint8_t offset = 1; offset <= numProducerChannels; ++offset) {
uint8_t i = (lastUsedProducerChannel + offset) % numProducerChannels;
uint8_t i;
if (mode == ChannelAssignmentMode::FirstAvailable) {
i = offset - 1;
} else if (mode == ChannelAssignmentMode::RoundRobin) {
i = (lastUsedProducerChannel + offset) % numProducerChannels;
} else {
assert(false && "Unsupported ChannelAssignmentMode");
}
if (!assignedProducerChannels.count(i)) {
// Only assign the channel if it is for the circuit flow.
if (!isPacketFlow) assignedProducerChannels.insert(i);
lastUsedProducerChannel = i;
return i;
}
}
return std::nullopt;
}

/// Returns its next usable consumer channel.
std::optional<uint8_t> getAndAssignConsumerDMAChannel(bool isPacketFlow) {
/// Returns its next usable consumer channel. By default, it uses round-robin
/// for load balancing.
std::optional<uint8_t> getConsumerDMAChannel(
ChannelAssignmentMode mode = ChannelAssignmentMode::RoundRobin) {
for (uint8_t offset = 1; offset <= numConsumerChannels; ++offset) {
uint8_t i = (lastUsedConsumerChannel + offset) % numConsumerChannels;
uint8_t i;
if (mode == ChannelAssignmentMode::FirstAvailable) {
i = offset - 1;
} else if (mode == ChannelAssignmentMode::RoundRobin) {
i = (lastUsedConsumerChannel + offset) % numConsumerChannels;
} else {
assert(false && "Unsupported ChannelAssignmentMode");
}
if (!assignedConsumerChannels.count(i)) {
// Only assign the channel if it is for the circuit flow.
if (!isPacketFlow) assignedConsumerChannels.insert(i);
lastUsedConsumerChannel = i;
return i;
}
Expand Down
133 changes: 98 additions & 35 deletions runtime/src/iree-amd-aie/aie_runtime/Utils/test/ChannelGeneratorTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,53 +13,116 @@ namespace {

using namespace mlir::iree_compiler::AMDAIE;

TEST(ChannelGeneratorTest, GetAssign) {
TEST(ChannelGeneratorTest, GetFirstAvailable) {
ChannelGenerator generator(2, 2);
bool isPacketFlow = false;
EXPECT_EQ(generator.getAndAssignProducerDMAChannel(isPacketFlow).value(), 0);
EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(isPacketFlow).value(), 0);
EXPECT_EQ(generator.getAndAssignProducerDMAChannel(isPacketFlow).value(), 1);
EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(isPacketFlow).value(), 1);
EXPECT_EQ(generator.getAndAssignProducerDMAChannel(isPacketFlow),
std::nullopt);
EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(isPacketFlow),
std::nullopt);
EXPECT_EQ(
generator.getProducerDMAChannel(ChannelAssignmentMode::FirstAvailable)
.value(),
0);
EXPECT_EQ(
generator.getConsumerDMAChannel(ChannelAssignmentMode::FirstAvailable)
.value(),
0);
EXPECT_EQ(
generator.getProducerDMAChannel(ChannelAssignmentMode::FirstAvailable)
.value(),
0);
EXPECT_EQ(
generator.getConsumerDMAChannel(ChannelAssignmentMode::FirstAvailable)
.value(),
0);
EXPECT_EQ(
generator.getProducerDMAChannel(ChannelAssignmentMode::FirstAvailable)
.value(),
0);
EXPECT_EQ(
generator.getConsumerDMAChannel(ChannelAssignmentMode::FirstAvailable)
.value(),
0);
}

TEST(ChannelGeneratorTest, Occupied) {
ChannelGenerator generator(4, 4);
bool isPacketFlow = false;
TEST(ChannelGeneratorTest, GetRoundRobin) {
ChannelGenerator generator(2, 2);
EXPECT_EQ(generator.getProducerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
0);
EXPECT_EQ(generator.getConsumerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
0);
EXPECT_EQ(generator.getProducerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
1);
EXPECT_EQ(generator.getConsumerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
1);
EXPECT_EQ(generator.getProducerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
0);
EXPECT_EQ(generator.getConsumerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
0);
EXPECT_EQ(generator.getProducerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
1);
EXPECT_EQ(generator.getConsumerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
1);
}

TEST(ChannelGeneratorTest, GetAssign) {
ChannelGenerator generator(2, 2);
EXPECT_EQ(generator.getProducerDMAChannel().value(), 0);
generator.assignProducerDMAChannel(0);
EXPECT_EQ(generator.getConsumerDMAChannel().value(), 0);
generator.assignConsumerDMAChannel(0);
generator.assignProducerDMAChannel(2);
generator.assignConsumerDMAChannel(2);
EXPECT_EQ(generator.getAndAssignProducerDMAChannel(isPacketFlow).value(), 1);
EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(isPacketFlow).value(), 1);
EXPECT_EQ(generator.getAndAssignProducerDMAChannel(isPacketFlow).value(), 3);
EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(isPacketFlow).value(), 3);
EXPECT_EQ(generator.getAndAssignProducerDMAChannel(isPacketFlow),
std::nullopt);
EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(isPacketFlow),
std::nullopt);
EXPECT_EQ(generator.getProducerDMAChannel().value(), 1);
generator.assignProducerDMAChannel(1);
EXPECT_EQ(generator.getConsumerDMAChannel().value(), 1);
generator.assignConsumerDMAChannel(1);
EXPECT_EQ(generator.getProducerDMAChannel(), std::nullopt);
EXPECT_EQ(generator.getConsumerDMAChannel(), std::nullopt);
}

TEST(ChannelGeneratorTest, PacketFlow) {
TEST(ChannelGeneratorTest, Occupied) {
ChannelGenerator generator(4, 4);
generator.assignProducerDMAChannel(0);
generator.assignConsumerDMAChannel(0);
generator.assignProducerDMAChannel(2);
generator.assignConsumerDMAChannel(2);
bool isPacketFlow = true;
// Packet flow should not occupy the channel exclusively, and the available
// channel appears in a round-robin fashion.
EXPECT_EQ(generator.getAndAssignProducerDMAChannel(isPacketFlow).value(), 1);
EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(isPacketFlow).value(), 1);
EXPECT_EQ(generator.getAndAssignProducerDMAChannel(isPacketFlow).value(), 3);
EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(isPacketFlow).value(), 3);
EXPECT_EQ(generator.getAndAssignProducerDMAChannel(isPacketFlow).value(), 1);
EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(isPacketFlow).value(), 1);
EXPECT_EQ(generator.getAndAssignProducerDMAChannel(isPacketFlow).value(), 3);
EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(isPacketFlow).value(), 3);
EXPECT_EQ(
generator.getProducerDMAChannel(ChannelAssignmentMode::FirstAvailable)
.value(),
1);
EXPECT_EQ(
generator.getConsumerDMAChannel(ChannelAssignmentMode::FirstAvailable)
.value(),
1);
EXPECT_EQ(
generator.getProducerDMAChannel(ChannelAssignmentMode::FirstAvailable)
.value(),
1);
EXPECT_EQ(
generator.getConsumerDMAChannel(ChannelAssignmentMode::FirstAvailable)
.value(),
1);
EXPECT_EQ(generator.getProducerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
3);
EXPECT_EQ(generator.getConsumerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
3);
EXPECT_EQ(generator.getProducerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
1);
EXPECT_EQ(generator.getConsumerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
1);
EXPECT_EQ(generator.getProducerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
3);
EXPECT_EQ(generator.getConsumerDMAChannel(ChannelAssignmentMode::RoundRobin)
.value(),
3);
}

} // namespace
Expand Down

0 comments on commit eba6384

Please sign in to comment.