From 59a0b32c290a65d619466329b99a0896859d2684 Mon Sep 17 00:00:00 2001 From: Yu-Zhewen Date: Fri, 3 Jan 2025 14:04:04 +0000 Subject: [PATCH 1/3] init commit --- .../Transforms/AMDAIEAssignChannels.cpp | 47 ++++++++++++++++--- .../Transforms/test/assign_channels.mlir | 44 +++++++++++++++-- .../aie_runtime/Utils/ChannelGenerator.h | 44 +++++++++++++---- .../aie_runtime/Utils/test/CMakeLists.txt | 10 ++++ .../Utils/test/ChannelGeneratorTest.cpp | 45 ++++++++++++++++++ 5 files changed, 169 insertions(+), 21 deletions(-) create mode 100644 runtime/src/iree-amd-aie/aie_runtime/Utils/test/ChannelGeneratorTest.cpp diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignChannels.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignChannels.cpp index 22337a6c1..c437b473b 100644 --- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignChannels.cpp +++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignChannels.cpp @@ -19,7 +19,26 @@ namespace { /// Assign channels to `amdaie.connection` ops. LogicalResult assignChannels(AMDAIE::WorkgroupOp workgroupOp) { IRRewriter rewriter(workgroupOp->getContext()); - ChannelGenerator generator; + + // Get the device model. + std::optional device = getConfigAMDAIEDevice(workgroupOp); + if (!device) + return workgroupOp->emitOpError() + << "could not find an AMDAIEDevice attribute"; + AMDAIEDeviceModel deviceModel = AMDAIE::getDeviceModel(device.value()); + + // Get the number of producer and consumer channels for each tile. + DenseMap tileToGeneratorMap; + workgroupOp.walk([&](AMDAIE::TileOp tileOp) { + uint32_t col = getConstantIndexOrAssert(tileOp.getCol()); + uint32_t row = getConstantIndexOrAssert(tileOp.getRow()); + AMDAIETileType tileType = deviceModel.getTileType(col, row); + uint8_t numDmaChannels = + deviceModel.getDmaProp(tileType, AMDAIEDmaProp::NumChannels); + tileToGeneratorMap[tileOp.getResult()] = + ChannelGenerator(numDmaChannels, numDmaChannels); + }); + SmallVector connectionOps; workgroupOp->walk([&](AMDAIE::ConnectionOp connectionOp) { connectionOps.push_back(connectionOp); @@ -43,18 +62,32 @@ LogicalResult assignChannels(AMDAIE::WorkgroupOp workgroupOp) { rewriter.setInsertionPoint(connectionOp); SmallVector sourceChannels; for (Value tile : sourceLogicalObjFifo.getTiles()) { - uint8_t channel = generator.getProducerDMAChannel(tile); + assert(tileToGeneratorMap.contains(tile) && + "no channel generator found for tile"); + std::optional maybeChannel = + tileToGeneratorMap[tile].getAndAssignProducerDMAChannel(); + if (!maybeChannel) { + return connectionOp.emitOpError() + << "no producer DMA channel available"; + } auto channelOp = rewriter.create( - rewriter.getUnknownLoc(), tile, channel, StrmSwPortType::DMA, - AMDAIE::DMAChannelDir::MM2S); + rewriter.getUnknownLoc(), tile, maybeChannel.value(), + StrmSwPortType::DMA, AMDAIE::DMAChannelDir::MM2S); sourceChannels.push_back(channelOp.getResult()); } SmallVector targetChannels; for (Value tile : targetLogicalObjFifo.getTiles()) { - uint8_t channel = generator.getConsumerDMAChannel(tile); + assert(tileToGeneratorMap.contains(tile) && + "no channel generator found for tile"); + std::optional maybeChannel = + tileToGeneratorMap[tile].getAndAssignConsumerDMAChannel(); + if (!maybeChannel) { + return connectionOp.emitOpError() + << "no consumer DMA channel available"; + } auto channelOp = rewriter.create( - rewriter.getUnknownLoc(), tile, channel, StrmSwPortType::DMA, - AMDAIE::DMAChannelDir::S2MM); + rewriter.getUnknownLoc(), tile, maybeChannel.value(), + StrmSwPortType::DMA, AMDAIE::DMAChannelDir::S2MM); targetChannels.push_back(channelOp.getResult()); } rewriter.replaceOpWithNewOp( diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/assign_channels.mlir b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/assign_channels.mlir index 1dbf73473..2e5e9f78e 100644 --- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/assign_channels.mlir +++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/assign_channels.mlir @@ -1,5 +1,16 @@ // RUN: iree-opt --pass-pipeline="builtin.module(iree-amdaie-assign-channels)" --split-input-file --verify-diagnostics %s | FileCheck %s +module { + // expected-error @+1 {{could not find an AMDAIEDevice attribute}} + amdaie.workgroup { + amdaie.controlcode { + amdaie.end + } + } +} + +// ----- + // CHECK-LABEL: @assign_channels // CHECK: %[[C0:.+]] = arith.constant 0 : index // CHECK: %[[C1:.+]] = arith.constant 1 : index @@ -12,10 +23,8 @@ // CHECK: %[[CHANNEL_2:.+]] = amdaie.channel(%[[tile_0_0]], 1, port_type = DMA, direction = MM2S) // CHECK: %[[CHANNEL_3:.+]] = amdaie.channel(%[[tile_0_1]], 1, port_type = DMA, direction = S2MM) // CHECK: amdaie.connection(%{{.+}} {%[[CHANNEL_3]]}, %{{.+}} {%[[CHANNEL_2]]}) -// CHECK: %[[CHANNEL_4:.+]] = amdaie.channel(%[[tile_0_0]], 2, port_type = DMA, direction = MM2S) -// CHECK: %[[CHANNEL_5:.+]] = amdaie.channel(%[[tile_0_1]], 2, port_type = DMA, direction = S2MM) -// CHECK: amdaie.connection(%{{.+}} {%[[CHANNEL_5]]}, %{{.+}} {%[[CHANNEL_4]]}) -module { +#executable_target_amdaie_xclbin_fb = #hal.executable.target<"amd-aie", "amdaie-xclbin-fb", {target_device = "npu1_4col", ukernels = "none"}> +module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} { func.func @assign_channels(%arg0: memref<1x1x8x16xi32, 1>, %arg1: memref<8x16xi32>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -26,6 +35,30 @@ module { %1 = amdaie.logicalobjectfifo.from_memref %arg1, {%tile_0_0} : memref<8x16xi32> -> !amdaie.logicalobjectfifo> %2 = amdaie.connection(%0, %1) : (!amdaie.logicalobjectfifo>, !amdaie.logicalobjectfifo>) %3 = amdaie.connection(%0, %1) : (!amdaie.logicalobjectfifo>, !amdaie.logicalobjectfifo>) + amdaie.controlcode { + amdaie.end + } + } + return + } +} + +// ----- + +// Each tile has two MM2S channels and two S2MM channels. +#executable_target_amdaie_xclbin_fb = #hal.executable.target<"amd-aie", "amdaie-xclbin-fb", {target_device = "npu1_4col", ukernels = "none"}> +module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} { + func.func @run_out_of_channel(%arg0: memref<1x1x8x16xi32, 1>, %arg1: memref<8x16xi32>) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + amdaie.workgroup { + %tile_0_0 = amdaie.tile(%c0, %c0) + %tile_0_1 = amdaie.tile(%c0, %c1) + %0 = amdaie.logicalobjectfifo.from_memref %arg0, {%tile_0_1} : memref<1x1x8x16xi32, 1> -> !amdaie.logicalobjectfifo> + %1 = amdaie.logicalobjectfifo.from_memref %arg1, {%tile_0_0} : memref<8x16xi32> -> !amdaie.logicalobjectfifo> + %2 = amdaie.connection(%0, %1) : (!amdaie.logicalobjectfifo>, !amdaie.logicalobjectfifo>) + %3 = amdaie.connection(%0, %1) : (!amdaie.logicalobjectfifo>, !amdaie.logicalobjectfifo>) + // expected-error @+1 {{no producer DMA channel available}} %4 = amdaie.connection(%0, %1) : (!amdaie.logicalobjectfifo>, !amdaie.logicalobjectfifo>) amdaie.controlcode { amdaie.end @@ -37,7 +70,8 @@ module { // ----- -module { +#executable_target_amdaie_xclbin_fb = #hal.executable.target<"amd-aie", "amdaie-xclbin-fb", {target_device = "npu1_4col", ukernels = "none"}> +module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} { func.func @no_source(%arg0: memref<1x1x8x16xi32, 1>, %arg1: !amdaie.logicalobjectfifo>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index diff --git a/runtime/src/iree-amd-aie/aie_runtime/Utils/ChannelGenerator.h b/runtime/src/iree-amd-aie/aie_runtime/Utils/ChannelGenerator.h index 3fca74e7c..52b0c3bb1 100644 --- a/runtime/src/iree-amd-aie/aie_runtime/Utils/ChannelGenerator.h +++ b/runtime/src/iree-amd-aie/aie_runtime/Utils/ChannelGenerator.h @@ -17,24 +17,50 @@ using namespace llvm; namespace mlir::iree_compiler::AMDAIE { /// Utility to generate valid channels. -/// TODO(jornt): add physical feasibility checks on channels. class ChannelGenerator { public: ChannelGenerator() {} + ChannelGenerator(uint8_t numProducerChannels, uint8_t numConsumerChannels) + : numProducerChannels(numProducerChannels), + numConsumerChannels(numConsumerChannels) {} - /// Given a tile, returns its next usable producer channel. - uint8_t getProducerDMAChannel(Value tile) { - return producerChannelsPerTile[tile]++; + /// Returns its next usable producer channel. + std::optional getAndAssignProducerDMAChannel() { + for (uint8_t i = 0; i < numProducerChannels; i++) { + if (!assignedProducerChannels.count(i)) { + assignedProducerChannels.insert(i); + return i; + } + } + return std::nullopt; } - /// Given a tile, returns its next usable consumer channel. - uint8_t getConsumerDMAChannel(Value tile) { - return consumerChannelsPerTile[tile]++; + /// Returns its next usable consumer channel. + std::optional getAndAssignConsumerDMAChannel() { + for (uint8_t i = 0; i < numConsumerChannels; i++) { + if (!assignedConsumerChannels.count(i)) { + assignedConsumerChannels.insert(i); + return i; + } + } + return std::nullopt; + } + + /// Assigns the provided producer channel. + void assignProducerDMAChannel(uint8_t channel) { + assignedProducerChannels.insert(channel); + } + + /// Assigns the provided consumer channel. + void assignConsumerDMAChannel(uint8_t channel) { + assignedConsumerChannels.insert(channel); } private: - DenseMap producerChannelsPerTile; - DenseMap consumerChannelsPerTile; + uint8_t numProducerChannels = 0; + uint8_t numConsumerChannels = 0; + DenseSet assignedProducerChannels; + DenseSet assignedConsumerChannels; }; } // namespace mlir::iree_compiler::AMDAIE diff --git a/runtime/src/iree-amd-aie/aie_runtime/Utils/test/CMakeLists.txt b/runtime/src/iree-amd-aie/aie_runtime/Utils/test/CMakeLists.txt index 5146a78ba..462f02461 100644 --- a/runtime/src/iree-amd-aie/aie_runtime/Utils/test/CMakeLists.txt +++ b/runtime/src/iree-amd-aie/aie_runtime/Utils/test/CMakeLists.txt @@ -16,6 +16,16 @@ iree_cc_test( iree-amd-aie::aie_runtime::Utils::Utils ) +iree_cc_test( + NAME + ChannelGeneratorTest + SRCS + "ChannelGeneratorTest.cpp" + DEPS + gtest + iree-amd-aie::aie_runtime::Utils::Utils +) + iree_cc_test( NAME LockIdGeneratorTest diff --git a/runtime/src/iree-amd-aie/aie_runtime/Utils/test/ChannelGeneratorTest.cpp b/runtime/src/iree-amd-aie/aie_runtime/Utils/test/ChannelGeneratorTest.cpp new file mode 100644 index 000000000..85c106f9b --- /dev/null +++ b/runtime/src/iree-amd-aie/aie_runtime/Utils/test/ChannelGeneratorTest.cpp @@ -0,0 +1,45 @@ +// Copyright 2025 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include + +#include "gtest/gtest.h" +#include "iree-amd-aie/aie_runtime/Utils/ChannelGenerator.h" + +namespace { + +using namespace mlir::iree_compiler::AMDAIE; + +TEST(ChannelGeneratorTest, GetAssign) { + ChannelGenerator generator(2, 2); + EXPECT_EQ(generator.getAndAssignProducerDMAChannel().value(), 0); + EXPECT_EQ(generator.getAndAssignConsumerDMAChannel().value(), 0); + EXPECT_EQ(generator.getAndAssignProducerDMAChannel().value(), 1); + EXPECT_EQ(generator.getAndAssignConsumerDMAChannel().value(), 1); + EXPECT_EQ(generator.getAndAssignProducerDMAChannel(), std::nullopt); + EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(), std::nullopt); +} + +TEST(ChannelGeneratorTest, Occupied) { + ChannelGenerator generator(4, 4); + generator.assignProducerDMAChannel(0); + generator.assignConsumerDMAChannel(0); + generator.assignProducerDMAChannel(2); + generator.assignConsumerDMAChannel(2); + EXPECT_EQ(generator.getAndAssignProducerDMAChannel().value(), 1); + EXPECT_EQ(generator.getAndAssignConsumerDMAChannel().value(), 1); + EXPECT_EQ(generator.getAndAssignProducerDMAChannel().value(), 3); + EXPECT_EQ(generator.getAndAssignConsumerDMAChannel().value(), 3); + EXPECT_EQ(generator.getAndAssignProducerDMAChannel(), std::nullopt); + EXPECT_EQ(generator.getAndAssignConsumerDMAChannel(), std::nullopt); +} + +} // namespace + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} From 9c0efbda8f5df2d85089c3cdc73b40d3042b4377 Mon Sep 17 00:00:00 2001 From: Yu-Zhewen Date: Fri, 3 Jan 2025 14:37:40 +0000 Subject: [PATCH 2/3] update comment --- .../AMD-AIE/iree-amd-aie/Transforms/test/assign_channels.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/assign_channels.mlir b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/assign_channels.mlir index 2e5e9f78e..095cc09f7 100644 --- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/assign_channels.mlir +++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/assign_channels.mlir @@ -45,7 +45,7 @@ module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} // ----- -// Each tile has two MM2S channels and two S2MM channels. +// Shim tile (0, 0) has only two producer (MM2S) channels. #executable_target_amdaie_xclbin_fb = #hal.executable.target<"amd-aie", "amdaie-xclbin-fb", {target_device = "npu1_4col", ukernels = "none"}> module attributes {hal.executable.target = #executable_target_amdaie_xclbin_fb} { func.func @run_out_of_channel(%arg0: memref<1x1x8x16xi32, 1>, %arg1: memref<8x16xi32>) { From 076bff6dcc0b74010b98461504ee31bd11014c99 Mon Sep 17 00:00:00 2001 From: Yu-Zhewen Date: Fri, 3 Jan 2025 19:54:28 +0000 Subject: [PATCH 3/3] add brackets --- .../AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignChannels.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignChannels.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignChannels.cpp index c437b473b..4bf66f282 100644 --- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignChannels.cpp +++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignChannels.cpp @@ -22,9 +22,10 @@ LogicalResult assignChannels(AMDAIE::WorkgroupOp workgroupOp) { // Get the device model. std::optional device = getConfigAMDAIEDevice(workgroupOp); - if (!device) + if (!device) { return workgroupOp->emitOpError() << "could not find an AMDAIEDevice attribute"; + } AMDAIEDeviceModel deviceModel = AMDAIE::getDeviceModel(device.value()); // Get the number of producer and consumer channels for each tile.