diff --git a/lib/Dialect/AIE/Transforms/AIEAssignBuffers.cpp b/lib/Dialect/AIE/Transforms/AIEAssignBuffers.cpp index d76d738e42..5c60ff74fd 100644 --- a/lib/Dialect/AIE/Transforms/AIEAssignBuffers.cpp +++ b/lib/Dialect/AIE/Transforms/AIEAssignBuffers.cpp @@ -26,7 +26,7 @@ using namespace xilinx::AIE; //===----------------------------------------------------------------------===// LogicalResult checkAndPrintOverflow(TileOp tile, int address, int maxDataMemorySize, int stacksize, - SmallVector buffers) { + SmallVector &buffers) { if (address > maxDataMemorySize) { InFlightDiagnostic error = tile.emitOpError("allocated buffers exceeded available memory\n"); @@ -65,17 +65,30 @@ LogicalResult basicAllocation(TileOp &tile) { else maxDataMemorySize = targetModel.getLocalMemorySize(); - SmallVector buffers; - // Collect all the buffers for this tile. + SmallVector buffers; + SmallVector allocated_buffers; + // Collect all the buffers for this tile. If the buffer has an address, add + // it to allocated_buffers. Otherwise, add it to buffers. device.walk([&](BufferOp buffer) { - if (buffer.getTileOp() == tile) - buffers.push_back(buffer); + if (buffer.getTileOp() == tile) { + if (buffer.getAddress()) + allocated_buffers.push_back(buffer); + else + buffers.push_back(buffer); + } }); - // Sort by allocation size. + + // Sort buffers by allocation size. std::sort(buffers.begin(), buffers.end(), [](BufferOp a, BufferOp b) { return a.getAllocationSize() > b.getAllocationSize(); }); + // Sort allocated_buffers by address + std::sort(allocated_buffers.begin(), allocated_buffers.end(), + [](BufferOp a, BufferOp b) { + return a.getAddress().value() < b.getAddress().value(); + }); + // Address range owned by the MemTile is 0x80000. // Address range owned by the tile is 0x8000 in // AIE1 and 0x10000 in AIE2, but we need room at @@ -87,9 +100,18 @@ LogicalResult basicAllocation(TileOp &tile) { address += stacksize; } + // As the next address to allocate is assigned, skip over any buffers + // from the allocated_buffers list. + auto current_alloc = allocated_buffers.begin(); for (auto buffer : buffers) { - if (buffer.getAddress()) - buffer->emitWarning("Overriding existing address"); + assert(!buffer.getAddress()); + while (current_alloc != allocated_buffers.end() && + address + buffer.getAllocationSize() > + current_alloc->getAddress().value()) { + address = current_alloc->getAddress().value() + + current_alloc->getAllocationSize(); + current_alloc++; + } buffer.setAddress(address); address += buffer.getAllocationSize(); } @@ -150,25 +172,29 @@ void setAndUpdateAddressInBank(BufferOp buffer, int64_t start_addr, // returns true and if not, the function emits a warning that the address // will be overwritten and returns false (which will cause the buffer to be // added to the list of buffers without addresses, to be completed later on). -bool checkAndAddBufferWithAddress(BufferOp buffer, int numBanks, - std::vector &nextAddrInBanks, - std::vector &bankLimits) { - if (auto addrAttr = buffer->getAttrOfType("address")) { - int addr = addrAttr.getInt(); - for (int i = 0; i < numBanks; i++) { - if (bankLimits[i].startAddr <= addr && addr < bankLimits[i].endAddr) { - if (addr >= nextAddrInBanks[i]) { - nextAddrInBanks[i] = addr + buffer.getAllocationSize(); - buffer.setMemBank(i); - } else { - buffer->emitWarning("Overriding existing address"); - return false; - } - } - } - return true; +FailureOr +checkAndAddBufferWithAddress(BufferOp buffer, int numBanks, + std::vector &nextAddrInBanks, + std::vector &bankLimits) { + auto addrAttr = buffer->getAttrOfType("address"); + if (!addrAttr) + return false; + + int addr = addrAttr.getInt(); + for (int i = 0; i < numBanks; i++) { + // if the address is not within the bank, continue + if (addr < bankLimits[i].startAddr || addr >= bankLimits[i].endAddr) + continue; + + // if the allocator already overwrote this address, fail + if (addr < nextAddrInBanks[i]) + return buffer->emitOpError("would override allocated address"); + + // the allocator can accomadate this existing allocation + nextAddrInBanks[i] = addr + buffer.getAllocationSize(); + buffer.setMemBank(i); } - return false; + return true; } // Function that checks whether the given buffer already has a set mem_bank @@ -177,22 +203,21 @@ bool checkAndAddBufferWithAddress(BufferOp buffer, int numBanks, // function emits a warning that the mem_bank will be overwritten and returns // false (which will cause the buffer to be added to the list of buffers // without addresses, to be completed later on). -bool checkAndAddBufferWithMemBank(BufferOp buffer, int numBanks, - std::vector &nextAddrInBanks, - std::vector &bankLimits) { - if (auto memBankAttr = buffer->getAttrOfType("mem_bank")) { - int mem_bank = memBankAttr.getInt(); - int64_t startAddr = nextAddrInBanks[mem_bank]; - int64_t endAddr = startAddr + buffer.getAllocationSize(); - if (endAddr <= bankLimits[mem_bank].endAddr) { - setAndUpdateAddressInBank(buffer, startAddr, endAddr, nextAddrInBanks); - } else { - buffer->emitWarning("Overriding existing mem_bank"); - return false; - } - return true; - } - return false; +FailureOr +checkAndAddBufferWithMemBank(BufferOp buffer, int numBanks, + std::vector &nextAddrInBanks, + std::vector &bankLimits) { + auto memBankAttr = buffer->getAttrOfType("mem_bank"); + if (!memBankAttr) + return false; + + int mem_bank = memBankAttr.getInt(); + int64_t startAddr = nextAddrInBanks[mem_bank]; + int64_t endAddr = startAddr + buffer.getAllocationSize(); + if (endAddr > bankLimits[mem_bank].endAddr) + return buffer->emitOpError("would override allocated address"); + setAndUpdateAddressInBank(buffer, startAddr, endAddr, nextAddrInBanks); + return true; } // Function that given a buffer will iterate over all the memory banks @@ -226,7 +251,7 @@ int setBufferAddress(BufferOp buffer, int numBanks, int startBankIndex, } LogicalResult checkAndPrintOverflow(TileOp tile, int numBanks, int stacksize, - SmallVector allBuffers, + SmallVector &allBuffers, std::vector &nextAddrInBanks, std::vector &bankLimits) { bool foundOverflow = false; @@ -311,8 +336,8 @@ LogicalResult simpleBankAwareAllocation(TileOp tile) { } fillBankLimits(numBanks, bankSize, bankLimits); - SmallVector buffersToAlloc; - SmallVector allBuffers; + SmallVector buffersToAlloc; + SmallVector allBuffers; // Collect all the buffers for this tile. device.walk([&](BufferOp buffer) { if (buffer.getTileOp() == tile) @@ -325,11 +350,13 @@ LogicalResult simpleBankAwareAllocation(TileOp tile) { // the above. for (auto buffer : allBuffers) { if (buffer.getTileOp() == tile) { - bool has_addr = checkAndAddBufferWithAddress(buffer, numBanks, + auto has_addr = checkAndAddBufferWithAddress(buffer, numBanks, nextAddrInBanks, bankLimits); - bool has_bank = checkAndAddBufferWithMemBank(buffer, numBanks, + auto has_bank = checkAndAddBufferWithMemBank(buffer, numBanks, nextAddrInBanks, bankLimits); - if (!has_addr && !has_bank) + if (failed(has_addr) || failed(has_bank)) + return failure(); + if (!has_addr.value() && !has_bank.value()) buffersToAlloc.push_back(buffer); } } diff --git a/python/compiler/aiecc/main.py b/python/compiler/aiecc/main.py index 6845481b4f..11f04dcc65 100644 --- a/python/compiler/aiecc/main.py +++ b/python/compiler/aiecc/main.py @@ -334,7 +334,12 @@ def run_passes(pass_pipeline, mlir_module_str, outputfile=None, verbose=False): print("Running:", pass_pipeline) with Context() as ctx, Location.unknown(): module = Module.parse(mlir_module_str) - PassManager.parse(pass_pipeline).run(module.operation) + pm = PassManager.parse(pass_pipeline) + try: + pm.run(module.operation) + except Exception as e: + print("Error running pass pipeline: ", pass_pipeline, e) + raise e mlir_module_str = str(module) if outputfile: with open(outputfile, "w") as g: