11#include "mlir/Dialect/Arith/IR/Arith.h"
12#include "mlir/Dialect/MemRef/IR/MemRef.h"
13#include "mlir/Dialect/SCF/IR/SCF.h"
14#include "mlir/IR/Dominance.h"
15#include "mlir/IR/PatternMatch.h"
16#include "mlir/Interfaces/DataLayoutInterfaces.h"
17#include "mlir/Interfaces/MemorySlotInterfaces.h"
18#include "mlir/Pass/Pass.h"
19#include "mlir/Transforms/Mem2Reg.h"
20#include "mlir/Transforms/WalkPatternRewriteDriver.h"
29#define GEN_PASS_DEF_AIEOBJECTFIFOLOWERCORES
30#include "aie/Dialect/AIE/Transforms/AIEPasses.h.inc"
41constexpr llvm::StringLiteral kBookkeepingSlotAttrName =
42 "aie.objectfifo.bookkeeping_slot";
48Value buildRotatingSwitch(OpBuilder &builder, Location loc, Value idx,
50 llvm::function_ref<Value(
int)> elem) {
51 SmallVector<int64_t, 4> caseValues;
52 for (
int c = 0; c < n; ++c) {
53 caseValues.push_back(c);
55 auto cases = DenseI64ArrayAttr::get(builder.getContext(), caseValues);
56 auto switchOp = scf::IndexSwitchOp::create(builder, loc, TypeRange{resultTy},
58 builder.createBlock(&switchOp.getDefaultRegion());
59 builder.setInsertionPointToStart(&switchOp.getDefaultBlock());
60 scf::YieldOp::create(builder, loc, elem(0));
61 for (
int c = 0; c < n; ++c) {
62 builder.createBlock(&switchOp.getCaseRegions()[c]);
63 builder.setInsertionPointToStart(&switchOp.getCaseBlock(c));
64 scf::YieldOp::create(builder, loc, elem(c));
66 builder.setInsertionPointAfter(switchOp);
67 return switchOp.getResult(0);
72void emitRotateObjects(OpBuilder &builder, Location loc,
73 ArrayRef<memref::AllocaOp> slots,
int released) {
74 if (slots.empty() || released %
static_cast<int>(slots.size()) == 0) {
77 SmallVector<Value> objects;
78 for (memref::AllocaOp slot : slots) {
80 memref::LoadOp::create(builder, loc, slot.getResult(), ValueRange{}));
82 for (
size_t k = 0; k < slots.size(); ++k) {
83 memref::AllocaOp slot = slots[k];
84 memref::StoreOp::create(builder, loc,
85 objects[(k + released) % objects.size()], slot,
92void emitAdvanceObjectIndex(OpBuilder &builder, Location loc, Value slot,
93 int depth,
int released) {
95 arith::ConstantOp::create(builder, loc, builder.getI32IntegerAttr(depth));
96 Value old = memref::LoadOp::create(builder, loc, slot, ValueRange{});
97 Value step = arith::ConstantOp::create(builder, loc,
98 builder.getI32IntegerAttr(released));
99 Value sum = arith::AddIOp::create(builder, loc, old, step);
101 arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::sge, sum, size);
104 Value wrapped = released == 1
105 ? Value(arith::ConstantOp::create(
106 builder, loc, builder.getI32IntegerAttr(0)))
107 : Value(
arith::SubIOp::create(builder, loc, sum, size));
108 Value next = arith::SelectOp::create(builder, loc, wrapping, wrapped, sum);
109 memref::StoreOp::create(builder, loc, next, slot, ValueRange{});
115SmallVector<FlatSymbolRefAttr>
116segmentLocksFor(ObjectFifoCoreEndpointOp endpoint, LockAction action) {
117 bool acquiring = action == LockAction::AcquireGreaterEqual;
118 bool wantsProduce = acquiring != endpoint.drains();
120 SmallVector<FlatSymbolRefAttr> locks;
121 for (ObjectFifoSegmentOp segment : endpoint.getSelectedSegments()) {
122 if (
auto lock = wantsProduce ? segment.getProduceLockAttr()
123 : segment.getConsumeLockAttr()) {
124 locks.push_back(lock);
132struct LoweringContext {
134 bool usesSemaphoreLocks;
135 DenseMap<Operation *, SmallVector<Value>> poolBuffers{};
141 DenseMap<std::pair<Operation *, Operation *>, memref::AllocaOp> objectIndex{};
142 DenseMap<std::pair<Operation *, Operation *>, SmallVector<memref::AllocaOp>>
144 DenseMap<std::pair<Operation *, Operation *>, Value> soleObject{};
145 DenseMap<std::pair<Operation *, Operation *>, memref::AllocaOp> heldCount{};
146 bool sawError =
false;
148 ObjectFifoCoreEndpointOp endpointOf(Operation *op, StringRef name) {
149 return SymbolTable::lookupNearestSymbolFrom<ObjectFifoCoreEndpointOp>(
150 device, StringAttr::get(op->getContext(), name));
153 LockOp lockOf(FlatSymbolRefAttr name) {
162void emitBinaryUseLocks(OpBuilder &builder, Location loc,
163 MutableArrayRef<LockOp> locks,
int depth,
bool drains,
164 Value counterSlot,
int count, LockAction action) {
165 if (count == 0 || locks.empty()) {
170 bool full = (!drains && action == LockAction::Release) ||
171 (drains && action == LockAction::Acquire);
173 memref::LoadOp::create(builder, loc, counterSlot, ValueRange{});
175 arith::IndexCastOp::create(builder, loc, builder.getIndexType(), counter);
176 Type lockTy = locks[0].getType();
177 for (
int i = 0; i < count; i++) {
179 buildRotatingSwitch(builder, loc, index, lockTy, depth, [&](
int c) {
180 return locks[(i + c) % depth].getResult();
182 UseLockOp::create(builder, loc, lock, action, full ? 1 : 0);
187 LoweringContext &ctx;
188 LowerRelease(MLIRContext *context, LoweringContext &ctx)
189 : OpRewritePattern(context), ctx(ctx) {}
191 LogicalResult matchAndRewrite(ObjectFifoReleaseOp releaseOp,
192 PatternRewriter &rewriter)
const override {
193 auto endpoint = ctx.endpointOf(releaseOp, releaseOp.getObjFifoName());
197 ObjectFifoPoolOp pool = endpoint.getPoolOp();
198 CoreOp core = releaseOp->getParentOfType<CoreOp>();
199 auto key = std::make_pair(core.getOperation(), endpoint.getOperation());
200 memref::AllocaOp index = ctx.objectIndex.lookup(key);
202 int count = releaseOp.relNumber() * pool.getRepeatCount().value_or(1);
203 rewriter.setInsertionPointAfter(releaseOp);
204 Location loc = releaseOp.getLoc();
206 if (!ctx.usesSemaphoreLocks) {
210 releaseOp->emitOpError(
211 "objectFifo release has no corresponding acquire on this tile");
215 SmallVector<LockOp> locks = pool.getLockOps();
216 emitBinaryUseLocks(rewriter, loc, locks, pool.getDepth(),
217 endpoint.drains(), index.getResult(), count,
218 LockAction::Release);
220 Value amount = arith::ConstantOp::create(
221 rewriter, loc, rewriter.getI32IntegerAttr(count));
222 for (FlatSymbolRefAttr name :
223 segmentLocksFor(endpoint, LockAction::Release)) {
224 UseLockOp::create(rewriter, loc, ctx.lockOf(name), LockAction::Release,
228 memref::AllocaOp held = ctx.heldCount.lookup(key);
229 assert(held &&
"a held counter is created for every released endpoint");
231 memref::LoadOp::create(rewriter, loc, held.getResult(), ValueRange{});
232 Value next = arith::SubIOp::create(rewriter, loc, current, amount);
233 memref::StoreOp::create(rewriter, loc, next, held.getResult(),
238 emitAdvanceObjectIndex(rewriter, loc, index.getResult(), pool.getDepth(),
239 releaseOp.getSize());
241 auto slots = ctx.objectSlots.find(key);
242 if (slots != ctx.objectSlots.end()) {
243 emitRotateObjects(rewriter, loc, slots->second, releaseOp.getSize());
250 LoweringContext &ctx;
251 LowerAcquire(MLIRContext *context, LoweringContext &ctx)
252 : OpRewritePattern(context), ctx(ctx) {}
254 LogicalResult matchAndRewrite(ObjectFifoAcquireOp acquireOp,
255 PatternRewriter &rewriter)
const override {
256 auto endpoint = ctx.endpointOf(acquireOp, acquireOp.getObjFifoName());
260 ObjectFifoPoolOp pool = endpoint.getPoolOp();
261 CoreOp core = acquireOp->getParentOfType<CoreOp>();
262 auto key = std::make_pair(core.getOperation(), endpoint.getOperation());
264 int wanted = acquireOp.acqNumber();
265 if (wanted > pool.getDepth()) {
266 acquireOp->emitOpError(
"acquires ")
267 << wanted <<
" objects from a pool of " << pool.getDepth();
271 int repeat = pool.getRepeatCount().value_or(1);
272 rewriter.setInsertionPointAfter(acquireOp);
273 Location loc = acquireOp.getLoc();
275 if (!ctx.usesSemaphoreLocks) {
280 memref::AllocaOp index = ctx.objectIndex.lookup(key);
281 assert(index &&
"a rotation counter is created for every acquire");
282 SmallVector<LockOp> locks = pool.getLockOps();
283 emitBinaryUseLocks(rewriter, loc, locks, pool.getDepth(),
284 endpoint.drains(), index.getResult(), wanted * repeat,
285 LockAction::Acquire);
286 return replaceWithObjects(acquireOp, key, rewriter);
291 memref::AllocaOp held = ctx.heldCount.lookup(key);
292 assert(held &&
"a held counter is created for every acquire");
294 memref::LoadOp::create(rewriter, loc, held.getResult(), ValueRange{});
295 Value target = arith::ConstantOp::create(
296 rewriter, loc, rewriter.getI32IntegerAttr(wanted));
298 arith::ConstantOp::create(rewriter, loc, rewriter.getI32IntegerAttr(0));
299 Value shortfall = arith::SubIOp::create(rewriter, loc, target, current);
300 Value delta = arith::MaxSIOp::create(rewriter, loc, shortfall, zero);
302 Value repeatVal = arith::ConstantOp::create(
303 rewriter, loc, rewriter.getI32IntegerAttr(repeat));
304 delta = arith::MulIOp::create(rewriter, loc, delta, repeatVal);
306 for (FlatSymbolRefAttr name :
307 segmentLocksFor(endpoint, LockAction::AcquireGreaterEqual)) {
308 UseLockOp::create(rewriter, loc, ctx.lockOf(name),
309 LockAction::AcquireGreaterEqual, delta);
312 Value next = arith::AddIOp::create(rewriter, loc, current, delta);
313 memref::StoreOp::create(rewriter, loc, next, held.getResult(),
315 return replaceWithObjects(acquireOp, key, rewriter);
319 LogicalResult replaceWithObjects(ObjectFifoAcquireOp acquireOp,
320 std::pair<Operation *, Operation *> key,
321 PatternRewriter &rewriter)
const {
322 if (Value
object = ctx.soleObject.lookup(key)) {
323 for (Value result : acquireOp.getObjects()) {
324 rewriter.replaceAllUsesWith(result,
object);
328 auto slots = ctx.objectSlots.find(key);
329 if (slots == ctx.objectSlots.end() || slots->second.empty()) {
332 for (
auto [k, result] : llvm::enumerate(acquireOp.getObjects())) {
334 memref::LoadOp::create(rewriter, acquireOp.getLoc(),
335 slots->second[k].getResult(), ValueRange{});
336 rewriter.replaceAllUsesWith(result,
object);
342struct AIEObjectFifoLowerCoresPass
343 :
public xilinx::AIE::impl::AIEObjectFifoLowerCoresBase<
344 AIEObjectFifoLowerCoresPass> {
349 memref::AllocaOp makeSlot(OpBuilder &builder, Location loc, Value initial,
350 StringRef marker = kBookkeepingSlotAttrName) {
351 auto slotTy = MemRefType::get(SmallVector<int64_t>{}, initial.getType());
352 auto slot = memref::AllocaOp::create(builder, loc, slotTy);
353 slot->setAttr(marker, builder.getUnitAttr());
354 memref::StoreOp::create(builder, loc, initial, slot.getResult(),
362 void trackObjects(OpBuilder &builder, Location loc,
363 ObjectFifoCoreEndpointOp endpoint, int64_t start,
364 std::pair<Operation *, Operation *> key,
365 LoweringContext &ctx) {
366 ObjectFifoPoolOp pool = endpoint.getPoolOp();
367 SmallVector<Value> buffers = ctx.poolBuffers.lookup(pool.getOperation());
370 MemRefType accessType = endpoint.getAccessType();
371 if (accessType != pool.getElemType()) {
372 auto [offset, size] = endpoint.getExtent();
373 for (Value &buffer : buffers) {
374 buffer = memref::SubViewOp::create(
375 builder, loc, accessType, buffer,
376 ArrayRef<OpFoldResult>{builder.getIndexAttr(offset)},
377 {builder.getIndexAttr(size)}, {builder.getIndexAttr(1)});
380 if (buffers.size() == 1) {
381 ctx.soleObject[key] = buffers.front();
384 SmallVector<memref::AllocaOp> &slots = ctx.objectSlots[key];
385 for (
size_t k = 0; k < buffers.size(); ++k) {
386 slots.push_back(makeSlot(builder, loc,
387 buffers[(start + k) % buffers.size()],
392 void emitCounters(CoreOp coreOp, LoweringContext &ctx, OpBuilder &builder) {
393 builder.setInsertionPointToStart(&coreOp.getBody().front());
394 Operation *core = coreOp.getOperation();
399 auto firstObject = [](ObjectFifoCoreEndpointOp endpoint) -> int64_t {
400 ObjectFifoPoolOp pool = endpoint.getPoolOp();
401 auto initValues = pool.getInitValues();
402 if (endpoint.drains() || !initValues || pool.getDepth() <= 0) {
405 return static_cast<int64_t
>(initValues->size()) % pool.getDepth();
408 auto record = [&](DenseMap<std::pair<Operation *, Operation *>,
409 memref::AllocaOp> &slots,
410 StringRef name,
bool rotating) {
411 auto endpoint = ctx.endpointOf(coreOp, name);
415 auto key = std::make_pair(core, endpoint.getOperation());
416 if (slots.count(key)) {
419 int64_t start = rotating ? firstObject(endpoint) : 0;
422 initial = arith::ConstantOp::create(builder, coreOp.getLoc(),
423 builder.getI32IntegerAttr(start));
426 zero = arith::ConstantOp::create(builder, coreOp.getLoc(),
427 builder.getI32IntegerAttr(0));
431 slots[key] = makeSlot(builder, coreOp.getLoc(), initial);
435 coreOp.walk([&](ObjectFifoAcquireOp a) {
436 auto endpoint = ctx.endpointOf(coreOp, a.getObjFifoName());
440 ObjectFifoPoolOp pool = endpoint.getPoolOp();
441 auto [buffers, inserted] =
442 ctx.poolBuffers.try_emplace(pool.getOperation());
444 for (BufferLike buffer : pool.getBufferOps()) {
445 buffers->second.push_back(buffer.getBuffer());
447 if (
static_cast<int64_t
>(buffers->second.size()) != pool.getDepth()) {
449 "expects every one of its 'depth' buffers to resolve");
453 if (ctx.sawError || a->use_empty()) {
456 auto key = std::make_pair(core, endpoint.getOperation());
457 if (!ctx.objectSlots.count(key) && !ctx.soleObject.count(key)) {
458 trackObjects(builder, coreOp.getLoc(), endpoint, firstObject(endpoint),
462 if (!ctx.usesSemaphoreLocks) {
464 coreOp.walk([&](ObjectFifoAcquireOp a) {
465 record(ctx.objectIndex, a.getObjFifoName(),
true);
469 coreOp.walk([&](ObjectFifoAcquireOp a) {
470 record(ctx.heldCount, a.getObjFifoName(),
false);
472 coreOp.walk([&](ObjectFifoReleaseOp r) {
473 record(ctx.heldCount, r.getObjFifoName(),
false);
481 void annotateUnrollHints(DeviceOp device, LoweringContext &ctx,
482 OpBuilder &builder) {
483 for (
auto coreOp : device.getOps<CoreOp>()) {
484 coreOp.walk([&](Operation *loop) {
489 bool touched =
false;
490 auto account = [&](Operation *op, StringRef name) {
494 auto endpoint = ctx.endpointOf(op, name);
499 period = std::lcm(period, (int64_t)endpoint.getPoolOp().getDepth());
502 [&](ObjectFifoAcquireOp a) { account(a, a.getObjFifoName()); });
504 [&](ObjectFifoReleaseOp r) { account(r, r.getObjFifoName()); });
507 builder.getI64IntegerAttr(period));
513 LogicalResult promoteBookkeepingSlots(DeviceOp device) {
514 SmallVector<PromotableAllocationOpInterface> allocators;
515 device.walk([&](memref::AllocaOp allocaOp) {
516 if (allocaOp->hasAttr(kBookkeepingSlotAttrName)) {
517 allocators.push_back(
518 cast<PromotableAllocationOpInterface>(allocaOp.getOperation()));
521 if (allocators.empty()) {
525 DataLayout dataLayout = DataLayout::closest(device);
526 DominanceInfo dominance(device);
527 OpBuilder builder(device.getContext());
528 (void)tryToPromoteMemorySlots(allocators, builder, dataLayout, dominance);
530 WalkResult leftover = device.walk([&](memref::AllocaOp allocaOp) {
531 if (!allocaOp->hasAttr(kBookkeepingSlotAttrName)) {
532 return WalkResult::advance();
534 allocaOp.emitOpError()
535 <<
"objectFifo bookkeeping slot could not be promoted to SSA "
536 "(mem2reg left it in place); the objectFifo lowering requires "
537 "all bookkeeping counters to become loop-carried SSA values";
538 return WalkResult::interrupt();
540 return failure(leftover.wasInterrupted());
543 void runOnOperation()
override {
544 DeviceOp device = getOperation();
545 OpBuilder builder(device.getContext());
547 LoweringContext ctx{device, device.getTargetModel().hasProperty(
550 annotateUnrollHints(device, ctx, builder);
551 for (
auto coreOp : device.getOps<CoreOp>()) {
552 emitCounters(coreOp, ctx, builder);
555 return signalPassFailure();
558 RewritePatternSet patterns(device.getContext());
559 patterns.insert<LowerAcquire, LowerRelease>(device.getContext(), ctx);
560 walkAndApplyPatterns(device, std::move(patterns));
562 return signalPassFailure();
565 SmallVector<Operation *> toErase;
566 device.walk([&](Operation *op) {
567 if (isa<ObjectFifoAcquireOp, ObjectFifoReleaseOp>(op)) {
568 toErase.push_back(op);
571 for (Operation *op : toErase) {
576 for (
auto endpoint : llvm::make_early_inc_range(
577 device.getOps<ObjectFifoCoreEndpointOp>())) {
581 if (failed(promoteBookkeepingSlots(device))) {
582 return signalPassFailure();
589std::unique_ptr<OperationPass<DeviceOp>>
591 return std::make_unique<AIEObjectFifoLowerCoresPass>();
std::unique_ptr< mlir::OperationPass< DeviceOp > > createAIEObjectFifoLowerCoresPass()
constexpr llvm::StringLiteral kObjectFifoObjectSlotAttrName
Discardable attribute set by --aie-objectfifo-lower-cores on the rank-0 memref.allocas that hold a co...
mlir::Operation * lookupNamedOp(mlir::Operation *from, mlir::StringAttr name)
constexpr llvm::StringLiteral kObjectFifoUnrollHintAttrName
Discardable attribute set by --aie-objectfifo-lower-cores on scf.for and forever scf....
mlir::Operation * getObjectFifoLoopParent(mlir::Operation *op)
The innermost loop around op that isObjectFifoLoop accepts, or null.
bool isObjectFifoLoop(mlir::Operation *op)
Whether op is a loop that objectFifo lowering unrolls by the rotation period of the objectFifos acces...