Skip to content

Commit 2a09422

Browse files
nirvedhmeshramnirvedhmeshram
authored
Merge branch 'main' into nm_xclbin_merge
2 parents 34c47e3 + 5d91be6 commit 2a09422

22 files changed

+610
-433
lines changed

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

+3-2
Original file line numberDiff line numberDiff line change
@@ -349,8 +349,9 @@ LogicalObjectFifoAccessOp::getLogicalObjectFifo() {
349349
//===----------------------------------------------------------------------===//
350350

351351
void LogicalObjectFifoAcquire::build(OpBuilder &b, mlir::OperationState &result,
352-
Value dma, LogicalObjectFifoPort port) {
353-
build(b, result, dma, port, b.getI32IntegerAttr(1));
352+
mlir::TypeRange resultTypes, Value dma,
353+
LogicalObjectFifoPort port) {
354+
build(b, result, resultTypes, dma, port, b.getI32IntegerAttr(1));
354355
}
355356

356357
//===----------------------------------------------------------------------===//

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

+4-2
Original file line numberDiff line numberDiff line change
@@ -402,12 +402,14 @@ def AMDAIE_LogicalObjectFifoAcquire:
402402
OptionalAttr<I32Attr>:$size
403403
);
404404

405+
let results = (outs AnyAMDAIELogicalObjectFifoType:$output);
406+
405407
let assemblyFormat = [{
406-
`(` $dma `,` $port `)` attr-dict
408+
`(` $dma `,` $port `)` attr-dict `->` type($output)
407409
}];
408410

409411
let builders = [
410-
OpBuilder<(ins "mlir::Value":$dma, "LogicalObjectFifoPort":$port)>,
412+
OpBuilder<(ins "mlir::TypeRange":$resultTypes, "mlir::Value":$dma, "LogicalObjectFifoPort":$port)>,
411413
];
412414
}
413415

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ func.func @logicalobjectfifo_access(%arg0: !amdaie.logicalobjectfifo<memref<1x1x
131131
// CHECK-SAME: %[[DMA]]
132132
func.func @logicalobjectfifo_acquire(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>) {
133133
%0 = amdaie.dma_cpy_nd(%arg0[0, 0, 0, 0] [1, 1, 8, 16] [128, 128, 16, 1], %arg1[0, 0, 0, 0] [1, 1, 8, 16] [128, 16, 16, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>)
134-
amdaie.logicalobjectfifo.acquire(%0, Consume) {size = 1 : i32}
134+
%1 = amdaie.logicalobjectfifo.acquire(%0, Consume) {size = 1 : i32} -> !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>
135135
return
136136
}
137137

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

+248
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,248 @@
1+
// Copyright 2024 The IREE Authors
2+
//
3+
// Licensed under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
7+
#include "iree-amd-aie/IR/AMDAIEOps.h"
8+
#include "iree-amd-aie/Transforms/Passes.h"
9+
#include "mlir/IR/IRMapping.h"
10+
#include "mlir/IR/Iterators.h"
11+
12+
#define DEBUG_TYPE "iree-amdaie-access-to-acquire-release"
13+
14+
namespace mlir::iree_compiler::AMDAIE {
15+
16+
namespace {
17+
18+
/// Walk all read access operations within the core operations and insert
19+
/// semaphore acquire and release stubs. Acquire operations will be inserted
20+
/// at the location of the access operation and release operations will be
21+
/// inserted before the next access or at the end of the block.
22+
LogicalResult readAccessToAcquireRelease(Operation *parentOp) {
23+
IRRewriter rewriter(parentOp->getContext());
24+
25+
SmallVector<AMDAIE::CoreOp> coreOps;
26+
parentOp->walk([&](AMDAIE::CoreOp coreOp) { coreOps.push_back(coreOp); });
27+
28+
// Map from DMA source/target logical objectFifos to those respective DMA
29+
// operations.
30+
DenseMap<Value, AMDAIE::CircularDmaCpyNdOp> logicalObjectFifoToDma;
31+
parentOp->walk([&](AMDAIE::CircularDmaCpyNdOp dmaOp) {
32+
logicalObjectFifoToDma[dmaOp.getSource()] = dmaOp;
33+
logicalObjectFifoToDma[dmaOp.getTarget()] = dmaOp;
34+
});
35+
36+
for (AMDAIE::CoreOp coreOp : coreOps) {
37+
DenseMap<Value, AMDAIE::LogicalObjectFifoAccessOp>
38+
logicalObjectFifoToLastAccess;
39+
WalkResult res =
40+
coreOp->walk([&](AMDAIE::LogicalObjectFifoAccessOp accessOp) {
41+
if (accessOp.getAccessType() != AMDAIE::MemoryAccess::Read)
42+
return WalkResult::advance();
43+
44+
if (logicalObjectFifoToLastAccess.contains(accessOp.getInput())) {
45+
rewriter.setInsertionPoint(accessOp);
46+
rewriter.create<AMDAIE::LogicalObjectFifoRelease>(
47+
rewriter.getUnknownLoc(),
48+
logicalObjectFifoToDma[accessOp.getInput()].getResult(),
49+
LogicalObjectFifoPort::Consume);
50+
}
51+
52+
if (!logicalObjectFifoToDma.contains(accessOp.getInput())) {
53+
accessOp.emitOpError()
54+
<< "read access not found as source of DMA operation";
55+
return WalkResult::interrupt();
56+
}
57+
rewriter.setInsertionPoint(accessOp);
58+
auto acquireOp = rewriter.create<AMDAIE::LogicalObjectFifoAcquire>(
59+
rewriter.getUnknownLoc(),
60+
llvm::cast<LogicalObjectFifoType>(accessOp.getInput().getType()),
61+
logicalObjectFifoToDma[accessOp.getInput()].getResult(),
62+
LogicalObjectFifoPort::Consume);
63+
auto newAccessOp = rewriter.create<AMDAIE::LogicalObjectFifoAccessOp>(
64+
rewriter.getUnknownLoc(), acquireOp.getResult(),
65+
AMDAIE::MemoryAccess::Read);
66+
rewriter.replaceAllUsesWith(accessOp.getResult(),
67+
newAccessOp.getResult());
68+
logicalObjectFifoToLastAccess[accessOp.getInput()] = accessOp;
69+
return WalkResult::advance();
70+
});
71+
if (res.wasInterrupted()) return failure();
72+
73+
// Insert release for remaining read access operations at end of block.
74+
for (auto &&[value, accessOp] : logicalObjectFifoToLastAccess) {
75+
Block *parentBlock = accessOp->getBlock();
76+
if (!parentBlock->back().hasTrait<OpTrait::IsTerminator>()) {
77+
rewriter.setInsertionPointToEnd(parentBlock);
78+
} else {
79+
rewriter.setInsertionPoint(parentBlock->getTerminator());
80+
}
81+
if (!logicalObjectFifoToDma.contains(accessOp.getInput())) {
82+
accessOp.emitOpError()
83+
<< "read access not found as source of DMA operation";
84+
return failure();
85+
}
86+
rewriter.create<AMDAIE::LogicalObjectFifoRelease>(
87+
rewriter.getUnknownLoc(), logicalObjectFifoToDma[accessOp.getInput()],
88+
LogicalObjectFifoPort::Consume);
89+
}
90+
}
91+
return success();
92+
}
93+
94+
/// Walk all write access operations within the core operations and insert
95+
/// semaphore operations. Release operations will be inserted
96+
/// at the location of the access operation and acquire operations will be
97+
/// inserted after the preceding access or at the beginning of the block.
98+
LogicalResult writeAccessToAcquireRelease(Operation *parentOp) {
99+
IRRewriter rewriter(parentOp->getContext());
100+
101+
SmallVector<AMDAIE::CoreOp> coreOps;
102+
parentOp->walk([&](AMDAIE::CoreOp coreOp) { coreOps.push_back(coreOp); });
103+
104+
// Map from DMA source/target logical objectFifos to those respective DMA
105+
// operations.
106+
DenseMap<Value, AMDAIE::CircularDmaCpyNdOp> logicalObjectFifoToDma;
107+
parentOp->walk([&](AMDAIE::CircularDmaCpyNdOp dmaOp) {
108+
logicalObjectFifoToDma[dmaOp.getSource()] = dmaOp;
109+
logicalObjectFifoToDma[dmaOp.getTarget()] = dmaOp;
110+
});
111+
112+
for (AMDAIE::CoreOp coreOp : coreOps) {
113+
DenseMap<Value, SmallVector<AMDAIE::LogicalObjectFifoAccessOp>>
114+
logicalObjectFifoToAccesses;
115+
DenseMap<Value, AMDAIE::LogicalObjectFifoAccessOp>
116+
logicalObjectFifoLastWriteAccesses;
117+
WalkResult res = coreOp->walk([&](AMDAIE::LogicalObjectFifoAccessOp
118+
accessOp) {
119+
// If there is another write op on the same logical objectFifo,
120+
// release it before the acquire.
121+
if (logicalObjectFifoLastWriteAccesses.contains(accessOp.getInput())) {
122+
AMDAIE::LogicalObjectFifoAccessOp prevAccess =
123+
logicalObjectFifoLastWriteAccesses[accessOp.getInput()];
124+
if (!logicalObjectFifoToDma.contains(prevAccess.getInput())) {
125+
prevAccess.emitOpError()
126+
<< "write access not found as source of DMA operation";
127+
return WalkResult::interrupt();
128+
}
129+
rewriter.setInsertionPoint(accessOp);
130+
rewriter.create<AMDAIE::LogicalObjectFifoRelease>(
131+
rewriter.getUnknownLoc(),
132+
logicalObjectFifoToDma[prevAccess.getInput()],
133+
LogicalObjectFifoPort::Produce);
134+
// Remove from last access as settled.
135+
logicalObjectFifoLastWriteAccesses.erase(prevAccess.getInput());
136+
}
137+
// Insert acquire for write access at first `Any` access or at start
138+
// of block.
139+
if (accessOp.getAccessType() == AMDAIE::MemoryAccess::Write) {
140+
if (!logicalObjectFifoToAccesses.contains(accessOp.getInput())) {
141+
rewriter.setInsertionPointToStart(accessOp->getBlock());
142+
} else {
143+
AMDAIE::LogicalObjectFifoAccessOp firstAccess =
144+
logicalObjectFifoToAccesses[accessOp.getInput()][0];
145+
rewriter.setInsertionPoint(firstAccess);
146+
}
147+
if (!logicalObjectFifoToDma.contains(accessOp.getInput())) {
148+
accessOp.emitOpError()
149+
<< "write access not found as source of DMA operation";
150+
return WalkResult::interrupt();
151+
}
152+
auto acquireOp = rewriter.create<AMDAIE::LogicalObjectFifoAcquire>(
153+
rewriter.getUnknownLoc(),
154+
llvm::cast<LogicalObjectFifoType>(accessOp.getInput().getType()),
155+
logicalObjectFifoToDma[accessOp.getInput()].getResult(),
156+
LogicalObjectFifoPort::Produce);
157+
auto newAccessOp = rewriter.create<AMDAIE::LogicalObjectFifoAccessOp>(
158+
rewriter.getUnknownLoc(), acquireOp.getResult(),
159+
AMDAIE::MemoryAccess::Write);
160+
161+
// Update uses of this access operation and the preceding ones.
162+
rewriter.replaceAllUsesWith(accessOp.getResult(),
163+
newAccessOp.getResult());
164+
if (logicalObjectFifoToAccesses.contains(accessOp.getInput())) {
165+
for (AMDAIE::LogicalObjectFifoAccessOp precedingAccessOp :
166+
logicalObjectFifoToAccesses[accessOp.getInput()]) {
167+
rewriter.replaceAllUsesWith(precedingAccessOp.getResult(),
168+
newAccessOp.getResult());
169+
}
170+
}
171+
172+
// Insert into last access map
173+
logicalObjectFifoLastWriteAccesses[accessOp.getInput()] = accessOp;
174+
}
175+
// Insert any access operation into first access map.
176+
if (!logicalObjectFifoToAccesses.contains(accessOp.getInput())) {
177+
logicalObjectFifoToAccesses[accessOp.getInput()] = {accessOp};
178+
} else {
179+
logicalObjectFifoToAccesses[accessOp.getInput()].push_back(accessOp);
180+
}
181+
return WalkResult::advance();
182+
});
183+
if (res.wasInterrupted()) return failure();
184+
185+
// Insert release for remaining access operations at end of block.
186+
for (auto &&[value, writeAccessOp] : logicalObjectFifoLastWriteAccesses) {
187+
Block *parentBlock = writeAccessOp->getBlock();
188+
if (!parentBlock->back().hasTrait<OpTrait::IsTerminator>()) {
189+
rewriter.setInsertionPointToEnd(parentBlock);
190+
} else {
191+
rewriter.setInsertionPoint(parentBlock->getTerminator());
192+
}
193+
if (!logicalObjectFifoToDma.contains(writeAccessOp.getInput())) {
194+
writeAccessOp.emitOpError()
195+
<< "write access not found as source of DMA operation";
196+
return failure();
197+
}
198+
rewriter.create<AMDAIE::LogicalObjectFifoRelease>(
199+
rewriter.getUnknownLoc(),
200+
logicalObjectFifoToDma[writeAccessOp.getInput()],
201+
LogicalObjectFifoPort::Produce);
202+
}
203+
}
204+
return success();
205+
}
206+
207+
class AMDAIEAccessToAcquireReleasePass
208+
: public impl::AMDAIEAccessToAcquireReleaseBase<
209+
AMDAIEAccessToAcquireReleasePass> {
210+
public:
211+
void getDependentDialects(DialectRegistry &registry) const override {
212+
registry.insert<AMDAIEDialect>();
213+
}
214+
215+
AMDAIEAccessToAcquireReleasePass() = default;
216+
AMDAIEAccessToAcquireReleasePass(
217+
const AMDAIEAccessToAcquireReleasePass &pass){};
218+
void runOnOperation() override;
219+
};
220+
221+
void AMDAIEAccessToAcquireReleasePass::runOnOperation() {
222+
Operation *parentOp = getOperation();
223+
if (failed(readAccessToAcquireRelease(parentOp))) {
224+
parentOp->emitOpError() << "failed to convert read access operations to "
225+
"acquire-release semaphore stubs";
226+
return signalPassFailure();
227+
}
228+
if (failed(writeAccessToAcquireRelease(parentOp))) {
229+
parentOp->emitOpError() << "failed to convert write access operations to "
230+
"acquire-release semaphore stubs";
231+
return signalPassFailure();
232+
}
233+
// Erase old access operations.
234+
IRRewriter rewriter(parentOp->getContext());
235+
parentOp->walk([&](AMDAIE::LogicalObjectFifoAccessOp accessOp) {
236+
if (accessOp->getUses().empty()) {
237+
rewriter.eraseOp(accessOp);
238+
}
239+
});
240+
}
241+
242+
} // namespace
243+
244+
std::unique_ptr<Pass> createAMDAIEAccessToAcquireReleasePass() {
245+
return std::make_unique<AMDAIEAccessToAcquireReleasePass>();
246+
}
247+
248+
} // namespace mlir::iree_compiler::AMDAIE

0 commit comments

Comments
 (0)