Skip to content

Commit

Permalink
Merge branch 'main' into noalias
Browse files Browse the repository at this point in the history
  • Loading branch information
newling authored Jan 10, 2025
2 parents b41597b + f367180 commit 922fca7
Show file tree
Hide file tree
Showing 4 changed files with 375 additions and 309 deletions.
322 changes: 322 additions & 0 deletions compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/AMDAIERT.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,322 @@
// 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 "AMDAIERT.h"

#include "iree-amd-aie/aie_runtime/iree_aie_configure.h"

using namespace mlir;

using xilinx::AIE::AMSelOp;
using xilinx::AIE::BDDimLayoutAttr;
using xilinx::AIE::BDPadLayoutAttr;
using xilinx::AIE::BufferOp;
using xilinx::AIE::ConnectOp;
using xilinx::AIE::CoreOp;
using xilinx::AIE::DeviceOp;
using xilinx::AIE::DMABDOp;
using xilinx::AIE::DMABDPACKETOp;
using xilinx::AIE::DMAChannelDir;
using xilinx::AIE::DMAStartOp;
using xilinx::AIE::LockAction;
using xilinx::AIE::LockOp;
using xilinx::AIE::MasterSetOp;
using xilinx::AIE::MemOp;
using xilinx::AIE::MemTileDMAOp;
using xilinx::AIE::PacketRuleOp;
using xilinx::AIE::PacketRulesOp;
using xilinx::AIE::ShimMuxOp;
using xilinx::AIE::SwitchboxOp;
using xilinx::AIE::TileOp;
using xilinx::AIE::UseLockOp;

using Path = std::filesystem::path;

namespace mlir::iree_compiler::AMDAIE {

LogicalResult addAllAieElfs(const AMDAIEDeviceModel &deviceModel,
DeviceOp &device, const Path &workDirPath,
bool aieSim) {
for (auto tileOp : device.getOps<TileOp>()) {
TileLoc tileLoc{tileOp.getCol(), tileOp.getRow()};
if (deviceModel.isShimNOCorPLTile(tileLoc.col, tileLoc.row)) continue;
if (CoreOp coreOp = getCoreOp(tileOp)) {
std::string fileName;
std::optional<StringRef> elfFile = coreOp.getElfFile();
if (elfFile.has_value()) {
fileName = *elfFile;
} else {
fileName = "core_" + std::to_string(tileLoc.col) + "_" +
std::to_string(tileLoc.row) + ".elf";
}
if (failed(addElfToTile(deviceModel, tileLoc, workDirPath / fileName,
aieSim))) {
return failure();
}
}
}
return success();
}

LogicalResult addAllCoreEnable(const AMDAIEDeviceModel &deviceModel,
DeviceOp &device) {
// Start execution of all the cores.
for (auto tileOp : device.getOps<TileOp>()) {
TileLoc tileLoc = {tileOp.getCol(), tileOp.getRow()};
if (auto coreOp = getCoreOp(tileOp);
coreOp && failed(coreEnable(deviceModel, tileLoc)))
return failure();
}
return success();
}

Lock::Action toLock(LockAction l) {
// Convert `xilinx::AIE::LockAction` to
// `mlir::iree_compiler::AMDAIE::LockAction`.
switch (l) {
case LockAction::Acquire:
return Lock::Action::Acquire;
case LockAction::AcquireGreaterEqual:
return Lock::Action::AcquireGreaterEqual;
case LockAction::Release:
return Lock::Action::Release;
}
llvm::report_fatal_error("unhandled lock action");
}

LogicalResult configureLocksAndBd(Block &block, const TileLoc &tileLoc,
const AMDAIEDeviceModel &deviceModel) {
FailureOr<XAie_DmaDesc> dmaTileBd = initDMADesc(deviceModel, tileLoc);
if (failed(dmaTileBd)) return failure();
std::optional<int> acqValue, relValue, acqLockId, relLockId;
bool acqEn;
for (UseLockOp op : block.getOps<UseLockOp>()) {
// Only dyn_cast if you are going to check if it was of the type
// expected; if you aren't checking use cast instead as it will at
// least assert in debug mode with an easier to understand error than
// dereferencing.
LockOp lock = cast<LockOp>(op.getLock().getDefiningOp());
switch (toLock(op.getAction())) {
case Lock::Action::Acquire:
case Lock::Action::AcquireGreaterEqual:
acqEn = op.getAcqEn();
acqLockId = lock.getLockID();
acqValue = op.getValue().value_or(1);
if (op.getAction() == LockAction::AcquireGreaterEqual)
acqValue.value() = -acqValue.value();
break;
case Lock::Action::Release:
relLockId = lock.getLockID();
relValue = op.getValue().value_or(1);
break;
}
}
// Disable acquire and release locks if not set.
if (!acqLockId) {
acqLockId = 0;
acqValue = 0;
acqEn = false;
}
if (!relLockId) {
relLockId = 0;
relValue = 0;
}
assert(acqValue && relValue && acqLockId && relLockId &&
"expected both use_lock(acquire) and use_lock(release) with bd");
if (failed(configureDMALocks(deviceModel, dmaTileBd.value(), tileLoc,
*acqValue, *relValue, *acqLockId, *relLockId,
acqEn))) {
return failure();
}

// Pull metadata related to packet routing, bdId, buffer length, size, stride
// to pass to aie-rt.
DMABDOp bdOp = *block.getOps<DMABDOp>().begin();
assert(bdOp.getBdId().has_value() &&
"DMABDOp must have assigned bd_id; did you forget to run "
"aie-assign-bd-ids?");
bool validBd = true;
std::optional<uint8_t> packetType;
std::optional<uint8_t> packetID;
bool enablePacket = false;
auto maybePacketOps = block.getOps<DMABDPACKETOp>();
if (!maybePacketOps.empty()) {
assert(llvm::range_size(maybePacketOps) == 1 &&
"expected only one dma_bd_packet");
DMABDPACKETOp packetOp = *maybePacketOps.begin();
packetType = packetOp.getPacketType();
packetID = packetOp.getPacketId();
enablePacket = true;
}

BufferOp bufferOp = cast<BufferOp>(bdOp.getBuffer().getDefiningOp());
if (!bufferOp.getAddress())
return bufferOp.emitError("buffer must have address assigned");
// Convert `xilinx::AIE::BDDimLayoutAttr` to
// `mlir::iree_compiler::AMDAIE::BDDimLayout`.
std::optional<std::vector<BDDimLayout>> maybeDims;
if (std::optional<std::vector<BDDimLayoutAttr>> dims = bdOp.getDimensions()) {
maybeDims = std::vector<BDDimLayout>{};
for (const BDDimLayoutAttr &dim : (*dims)) {
maybeDims->emplace_back(BDDimLayout{dim.getSize(), dim.getStride()});
}
}

// Convert `xilinx::AIE::BDPadLayoutAttr` to
// `mlir::iree_compiler::AMDAIE::BDPadLayout`.
std::optional<std::vector<BDPadLayout>> maybePadDims;
if (std::optional<std::vector<BDPadLayoutAttr>> dims =
bdOp.getPadDimensions()) {
maybePadDims = std::vector<BDPadLayout>{};
for (const BDPadLayoutAttr &dim : (*dims)) {
maybePadDims->emplace_back(
BDPadLayout{dim.getConstPadBefore(), dim.getConstPadAfter()});
}
}

bool enableNextBd = bdOp.getNextBdId().has_value();
std::optional<uint8_t> nextBdId =
enableNextBd
? std::optional<uint8_t>{static_cast<uint8_t>(*bdOp.getNextBdId())}
: std::nullopt;
std::optional<BDIterLayout> maybeIter = std::nullopt;
if (failed(configureDMABD(deviceModel, dmaTileBd.value(), tileLoc, validBd,
static_cast<uint8_t>(*bdOp.getBdId()), enableNextBd,
nextBdId, enablePacket, packetType, packetID,
*bufferOp.getAddress(), getLenInBytes(bdOp),
getOffsetInBytes(bdOp),
getBufferElementTypeWidthInBytes(bdOp), maybeDims,
maybePadDims, maybeIter))) {
return failure();
}
return success();
}

LogicalResult addInitConfig(const AMDAIEDeviceModel &deviceModel,
DeviceOp &device) {
for (auto tileOp : device.getOps<TileOp>()) {
TileLoc tileLoc = {tileOp.getCol(), tileOp.getRow()};
if (deviceModel.isShimTile(tileOp.getCol(), tileOp.getRow())) {
continue;
}
if (auto coreOp = getCoreOp(tileOp);
coreOp && failed(resetUnResetCore(deviceModel, tileLoc))) {
return failure();
}
}

// Set locks with explicit initializers.
WalkResult r;
r = device.walk<WalkOrder::PreOrder>([&](LockOp lockOp) {
if (lockOp.getLockID() && lockOp.getInit()) {
TileOp t = xilinx::AIE::getTileOp(*lockOp.getOperation());
TileLoc tileLoc = {t.getCol(), t.getRow()};
Lock lock{tileLoc, static_cast<uint8_t>(*lockOp.getLockID()),
static_cast<int8_t>(*lockOp.getInit())};
if (failed(initializeLock(deviceModel, lock)))
return WalkResult::interrupt();
}
return WalkResult::advance();
});
if (r.wasInterrupted()) return failure();

auto memOps = llvm::to_vector_of<Operation *>(device.getOps<MemOp>());
llvm::append_range(memOps, device.getOps<MemTileDMAOp>());
for (Operation *memOp : memOps) {
TileOp t = xilinx::AIE::getTileOp(*memOp);
TileLoc tileLoc = {t.getCol(), t.getRow()};
if (deviceModel.isShimNOCorPLTile(tileLoc.col, tileLoc.row)) {
continue;
}

// Handle DMA ops separately.
for (Block &block : memOp->getRegion(0)) {
if (block.getOps<DMABDOp>().empty()) continue;
if (failed(configureLocksAndBd(block, tileLoc, deviceModel)))
return failure();
}

for (Block &block : memOp->getRegion(0)) {
for (auto op : block.getOps<DMAStartOp>()) {
DMABDOp bd = *op.getDest()->getOps<DMABDOp>().begin();
int chNum = op.getChannelIndex();
auto channelDir = static_cast<DMAChannelDir>(op.getChannelDir());
bool issueToken = tileLoc.row == 0 && channelDir == DMAChannelDir::MM2S;
bool setChannelEnable = true;
if (failed(configurePushToBdQueue(
deviceModel, tileLoc, chNum, channelDir, bd.getBdId().value(),
op.getRepeatCount(), issueToken, setChannelEnable)))
return failure();
}
}
}

// StreamSwitch (switchbox) configuration.
for (auto switchboxOp : device.getOps<SwitchboxOp>()) {
TileOp t = xilinx::AIE::getTileOp(*switchboxOp.getOperation());
TileLoc tileLoc = {t.getCol(), t.getRow()};
std::vector<Connect> connects;
for (auto connectOp : switchboxOp.getOps<ConnectOp>()) {
connects.emplace_back(
Port{connectOp.getSourceBundle(), connectOp.getSourceChannel()},
Port{connectOp.getDestBundle(), connectOp.getDestChannel()},
Connect::Interconnect::SWB, tileLoc.col, tileLoc.row);
}
if (failed(configureStreamSwitch(deviceModel, tileLoc, connects))) {
return failure();
}

Block &b = switchboxOp.getConnections().front();
for (auto masterSetOp : b.getOps<MasterSetOp>()) {
std::vector<AMSel> amSels;
for (Value val : masterSetOp.getAmsels()) {
AMSelOp amsel = cast<AMSelOp>(val.getDefiningOp());
amSels.push_back({static_cast<uint8_t>(amsel.getArbiterID()),
static_cast<uint8_t>(amsel.getMsel())});
}
if (failed(configureSwitchPacketMasters(
deviceModel, tileLoc, masterSetOp.getDestBundle(),
masterSetOp.getDestChannel(), amSels,
masterSetOp->hasAttr("keep_pkt_header"))))
return failure();
}

for (auto packetRulesOp : b.getOps<PacketRulesOp>()) {
int slot = 0;
Block &block = packetRulesOp.getRules().front();
for (auto packetRuleOp : block.getOps<PacketRuleOp>()) {
AMSelOp amselOp =
cast<AMSelOp>(packetRuleOp.getAmsel().getDefiningOp());
if (failed(configureSwitchPacketSlaves(
deviceModel, tileLoc, packetRulesOp.getSourceBundle(),
packetRulesOp.getSourceChannel(),
AMSel{amselOp.getArbiterID(), amselOp.getMsel()},
packetRuleOp.getValue(), packetRuleOp.getMask(), slot)))
return failure();
slot++;
}
}
}

for (auto muxOp : device.getOps<ShimMuxOp>()) {
TileOp t = xilinx::AIE::getTileOp(*muxOp.getOperation());
TileLoc tileLoc = {t.getCol(), t.getRow()};
std::vector<Connect> connects;
for (auto connectOp : muxOp.getOps<ConnectOp>()) {
connects.emplace_back(
Port{connectOp.getSourceBundle(), connectOp.getSourceChannel()},
Port{connectOp.getDestBundle(), connectOp.getDestChannel()},
Connect::Interconnect::SHIMMUX, tileLoc.col, tileLoc.row);
}
if (failed(configureStreamSwitch(deviceModel, tileLoc, connects))) {
return failure();
}
}

return success();
}

} // namespace mlir::iree_compiler::AMDAIE
34 changes: 34 additions & 0 deletions compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/AMDAIERT.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
// 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

#ifndef IREE_AMD_AIE_TARGET_AMDAIERT_H_
#define IREE_AMD_AIE_TARGET_AMDAIERT_H_

#include <filesystem>

#include "aie/AIEDialect.h"

namespace mlir::iree_compiler::AMDAIE {

/// Load ELF files for all cores within the device operation.
LogicalResult addAllAieElfs(const AMDAIEDeviceModel &deviceModel,
xilinx::AIE::DeviceOp &device,
const std::filesystem::path &workDirPath,
bool aieSim);

/// Update core control registers to enable all cores within the device
/// operation.
LogicalResult addAllCoreEnable(const AMDAIEDeviceModel &deviceModel,
xilinx::AIE::DeviceOp &device);

/// Utility function to reset all cores, initialize hardware locks,
/// and configure all switchboxes.
LogicalResult addInitConfig(const AMDAIEDeviceModel &deviceModel,
xilinx::AIE::DeviceOp &device);

} // namespace mlir::iree_compiler::AMDAIE

#endif // IREE_AMD_AIE_TARGET_AMDAIERT_H_
Loading

0 comments on commit 922fca7

Please sign in to comment.