From 01aba906394b2a89b232e47ed2494dbf89654d9e Mon Sep 17 00:00:00 2001 From: James Newling Date: Tue, 5 Dec 2023 14:06:43 -0800 Subject: [PATCH] Fixes for --assign-lock-ids pass (#772) --- .../AIE/Transforms/AIEAssignLockIDs.cpp | 109 ++++++++++-------- test/assign-lock-ids/assign-lockIDs.mlir | 37 +++++- test/assign-lock-ids/invalid.mlir | 65 +++++++++++ 3 files changed, 164 insertions(+), 47 deletions(-) create mode 100644 test/assign-lock-ids/invalid.mlir diff --git a/lib/Dialect/AIE/Transforms/AIEAssignLockIDs.cpp b/lib/Dialect/AIE/Transforms/AIEAssignLockIDs.cpp index 303a972c96..b0350ce667 100644 --- a/lib/Dialect/AIE/Transforms/AIEAssignLockIDs.cpp +++ b/lib/Dialect/AIE/Transforms/AIEAssignLockIDs.cpp @@ -10,18 +10,17 @@ // This pass aims to assign lockIDs to AIE.lock operations. The lockID is // numbered from the most recent AIE.lock within the same tile. If the lockID -// exceeds 15 then the pass generates an error and terminates. AIE.lock -// operations for different tiles are numbered independently. If there are -// existing lock IDs, this pass is idempotent and only assign lock ids to locks -// without an ID. +// exceeds the number of locks on the tile, the pass generates an error and +// terminates. AIE.lock operations for different tiles are numbered +// independently. If there are existing lock IDs, this pass is idempotent +// and only assigns lock IDs to locks without an ID. #include "aie/Dialect/AIE/IR/AIEDialect.h" #include "aie/Dialect/AIE/Transforms/AIEPasses.h" #include "mlir/IR/Attributes.h" #include "mlir/Pass/Pass.h" - -#include +#include "llvm/ADT/DenseMap.h" #define DEBUG_TYPE "aie-assign-lock-ids" @@ -40,50 +39,68 @@ struct AIEAssignLockIDsPass : AIEAssignLockIDsBase { DeviceOp device = getOperation(); OpBuilder rewriter = OpBuilder::atBlockEnd(device.getBody()); - std::map>> tileToLastID; - - // The first pass scans for and stores the existing lockIDs. This data is - // stored in a map with the lock’s tile operation as the key, while the - // value to the key is a pair with the current potential lockID and a set - // that stores the currently assigned lockIDs. - for (auto lock : device.getOps()) { - if (lock.getLockID().has_value()) { - Operation *lock_tile = lock.getTile().getDefiningOp(); - tileToLastID[lock_tile].first = 0; - tileToLastID[lock_tile].second.insert(lock.getLockIDValue()); - } - } + // All of the lock ops on a tile, separated into ops which have been + // assigned to a lock, and ops which have not. + struct TileLockOps { + DenseSet assigned; + SmallVector unassigned; + }; + + DenseMap tileToLocks; + + // Construct data structure storing locks by tile. + for (LockOp lockOp : device.getOps()) { + + TileOp tileOp = lockOp.getTileOp(); + bool isAssigned = lockOp.getLockID().has_value(); + + if (isAssigned) { + auto lockID = lockOp.getLockID().value(); + auto iter = tileToLocks.find(tileOp); - // The second pass scans for locks with no lockIDs and assigns locks. - for (auto lock : device.getOps()) { - Operation *lock_tile = lock.getTile().getDefiningOp(); - if (!lock.getLockID().has_value()) { - if (tileToLastID.find(lock_tile) == tileToLastID.end()) { - // If the tile operation corresponding to the lock does not exist in - // the data structure, initialize the lockID with 0 with an empty set. - tileToLastID[lock_tile].first = 0; - } else if (tileToLastID[lock_tile].first < 15) { - // If the tile operation of the lock exists, the potential lockID is - // checked with the set containing occupied lockIDs until a lockID - // that is free is found. - int potential_ID = tileToLastID[lock_tile].first; - while (true) { - if (tileToLastID[lock_tile].second.find(potential_ID) != - tileToLastID[lock_tile].second.end()) - potential_ID++; - else - break; + if (iter == tileToLocks.end()) + tileToLocks.insert({tileOp, {{lockID}, /* unassigned = */ {}}}); + else { + if (iter->second.assigned.find(lockID) != + iter->second.assigned.end()) { + auto diag = lockOp->emitOpError("is assigned to the same lock (") + << lockID << ") as another op."; + diag.attachNote(tileOp.getLoc()) + << "tile has lock ops assigned to same lock."; + return signalPassFailure(); } - tileToLastID[lock_tile].first = potential_ID; - } else { - lock->emitError() << "Exceeded the number of unique LockIDs"; - return; + iter->second.assigned.insert(lockID); } + } else { + auto iter = tileToLocks.find(tileOp); + if (iter == tileToLocks.end()) + tileToLocks.insert({tileOp, {/* assigned = */ {}, {lockOp}}}); + else + iter->second.unassigned.push_back(lockOp); + } + } + + // IR mutation: assign locks to all unassigned lock ops. + for (auto [tileOp, locks] : tileToLocks) { - // The lockID is assigned and is stored in the set. - lock->setAttr("lockID", rewriter.getI32IntegerAttr( - tileToLastID[lock_tile].first)); - tileToLastID[lock_tile].second.insert(tileToLastID[lock_tile].first); + const auto locksPerTile = + getTargetModel(tileOp).getNumLocks(tileOp.getCol(), tileOp.getRow()); + + uint32_t nextID = 0; + for (auto lockOp : locks.unassigned) { + while (nextID < locksPerTile && + (locks.assigned.find(nextID) != locks.assigned.end())) { + ++nextID; + } + if (nextID == locksPerTile) { + mlir::InFlightDiagnostic diag = + lockOp->emitOpError("not allocated a lock."); + diag.attachNote(tileOp.getLoc()) << "because only " << locksPerTile + << " locks available in this tile."; + return signalPassFailure(); + } + lockOp->setAttr("lockID", rewriter.getI32IntegerAttr(nextID)); + ++nextID; } } } diff --git a/test/assign-lock-ids/assign-lockIDs.mlir b/test/assign-lock-ids/assign-lockIDs.mlir index 39be48682d..d58415e509 100644 --- a/test/assign-lock-ids/assign-lockIDs.mlir +++ b/test/assign-lock-ids/assign-lockIDs.mlir @@ -8,7 +8,7 @@ // //===----------------------------------------------------------------------===// -// RUN: aie-opt --aie-assign-lock-ids %s | FileCheck %s +// RUN: aie-opt --aie-assign-lock-ids --split-input-file %s | FileCheck %s // CHECK: %[[VAL_0:.*]] = AIE.tile(2, 2) // CHECK: %[[VAL_1:.*]] = AIE.tile(2, 3) // CHECK: %[[VAL_2:.*]] = AIE.tile(3, 3) @@ -85,3 +85,38 @@ module @test_assign_lockIDs { %l60 = AIE.lock(%t60) } } + +// ----- + +module @memTileTest { + AIE.device(xcve2802) { + + // Memory tiles on xcve have 64 locks. + %tmemtile = AIE.tile(1,1) + %l0 = AIE.lock(%tmemtile, 1) + %l1 = AIE.lock(%tmemtile, 0) + %l2 = AIE.lock(%tmemtile) + %l3 = AIE.lock(%tmemtile) + %l4 = AIE.lock(%tmemtile) + %l5 = AIE.lock(%tmemtile) + %l6 = AIE.lock(%tmemtile) + %l7 = AIE.lock(%tmemtile) + %l8 = AIE.lock(%tmemtile) + %l9 = AIE.lock(%tmemtile) + %l10 = AIE.lock(%tmemtile) + %l11 = AIE.lock(%tmemtile) + %l12 = AIE.lock(%tmemtile) + %l13 = AIE.lock(%tmemtile) + %l14 = AIE.lock(%tmemtile,33) + %l15 = AIE.lock(%tmemtile) + %l16 = AIE.lock(%tmemtile) + %l17 = AIE.lock(%tmemtile) + %l18 = AIE.lock(%tmemtile) + %l19 = AIE.lock(%tmemtile,2) + } +} + + +// CHECK-LABEL: memTileTest +// CHECK-COUNT-20: AIE.lock +// CHECK-NOT: AIE.lock diff --git a/test/assign-lock-ids/invalid.mlir b/test/assign-lock-ids/invalid.mlir new file mode 100644 index 0000000000..65cf456460 --- /dev/null +++ b/test/assign-lock-ids/invalid.mlir @@ -0,0 +1,65 @@ +// RUN: aie-opt --aie-assign-lock-ids %s -split-input-file -verify-diagnostics + +AIE.device(xcve2802) { + // expected-note @below {{because only 16 locks available in this tile}} + %tMemTile = AIE.tile(4,4) + %l0 = AIE.lock(%tMemTile) + %l1 = AIE.lock(%tMemTile) + %l2 = AIE.lock(%tMemTile) + %l3 = AIE.lock(%tMemTile) + %l4 = AIE.lock(%tMemTile) + %l5 = AIE.lock(%tMemTile) + %l6 = AIE.lock(%tMemTile) + %l7 = AIE.lock(%tMemTile) + %l8 = AIE.lock(%tMemTile) + %l9 = AIE.lock(%tMemTile) + %l10 = AIE.lock(%tMemTile) + %l11 = AIE.lock(%tMemTile) + %l12 = AIE.lock(%tMemTile) + %l13 = AIE.lock(%tMemTile) + %l14 = AIE.lock(%tMemTile) + %l15 = AIE.lock(%tMemTile) + // expected-error @below {{not allocated a lock}} + %l16 = AIE.lock(%tMemTile) + %l17 = AIE.lock(%tMemTile) + %l18 = AIE.lock(%tMemTile) + %l19 = AIE.lock(%tMemTile) +} + +// ----- + +AIE.device(xcve2802) { + // expected-note @below {{because only 16 locks available in this tile}} + %tMemTile = AIE.tile(4,4) + %l0 = AIE.lock(%tMemTile) + %l1 = AIE.lock(%tMemTile, 1) + %l2 = AIE.lock(%tMemTile) + %l3 = AIE.lock(%tMemTile) + %l4 = AIE.lock(%tMemTile) + %l5 = AIE.lock(%tMemTile) + %l6 = AIE.lock(%tMemTile) + %l7 = AIE.lock(%tMemTile, 2) + %l8 = AIE.lock(%tMemTile) + %l9 = AIE.lock(%tMemTile) + %l10 = AIE.lock(%tMemTile, 15) + %l11 = AIE.lock(%tMemTile) + %l12 = AIE.lock(%tMemTile) + %l13 = AIE.lock(%tMemTile) + %l14 = AIE.lock(%tMemTile) + // expected-error @below {{not allocated a lock}} + %l15 = AIE.lock(%tMemTile) + %l16 = AIE.lock(%tMemTile) + %l17 = AIE.lock(%tMemTile, 3) + %l18 = AIE.lock(%tMemTile) + %l19 = AIE.lock(%tMemTile) +} + +// ----- + +AIE.device(xcve2802) { + //expected-note @below {{tile has lock ops assigned to same lock}} + %t22 = AIE.tile(2,2) + %l0 = AIE.lock(%t22, 7) + // expected-error @below {{is assigned to the same lock (7) as another op}} + %l1 = AIE.lock(%t22, 7) +}