Make pre-allocated buffers work again (#1901)

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<BufferOp, 4> buffers) {
+                                    SmallVector<BufferOp> &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<BufferOp, 4> buffers;
-  // Collect all the buffers for this tile.
+  SmallVector<BufferOp> buffers;
+  SmallVector<BufferOp> 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<WalkOrder::PreOrder>([&](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();
   }
@@ -142,25 +164,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<int64_t> &nextAddrInBanks,
-                                  std::vector<BankLimits> &bankLimits) {
-  if (auto addrAttr = buffer->getAttrOfType<IntegerAttr>("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<bool>
+checkAndAddBufferWithAddress(BufferOp buffer, int numBanks,
+                             std::vector<int64_t> &nextAddrInBanks,
+                             std::vector<BankLimits> &bankLimits) {
+  auto addrAttr = buffer->getAttrOfType<IntegerAttr>("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 allocated 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
@@ -169,27 +195,26 @@ 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<int64_t> &nextAddrInBanks,
-                                  std::vector<BankLimits> &bankLimits) {
-  if (auto memBankAttr = buffer->getAttrOfType<IntegerAttr>("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<bool>
+checkAndAddBufferWithMemBank(BufferOp buffer, int numBanks,
+                             std::vector<int64_t> &nextAddrInBanks,
+                             std::vector<BankLimits> &bankLimits) {
+  auto memBankAttr = buffer->getAttrOfType<IntegerAttr>("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 existing mem_bank");
+  setAndUpdateAddressInBank(buffer, startAddr, endAddr, nextAddrInBanks);
+  return true;
 }
 
 // Prints the memory map across banks
-void printMemMap(TileOp tile, SmallVector<BufferOp, 4> allocatedBuffers,
-                 SmallVector<BufferOp, 4> preAllocatedBuffers, int numBanks,
+void printMemMap(TileOp tile, SmallVector<BufferOp> &allocatedBuffers,
+                 SmallVector<BufferOp> &preAllocatedBuffers, int numBanks,
                  std::vector<BankLimits> &bankLimits, int stacksize) {
   InFlightDiagnostic error = tile.emitOpError(
       "Not all requested buffers fit in the available memory.\n");
@@ -270,7 +295,7 @@ int setBufferAddress(BufferOp buffer, int numBanks, int startBankIndex,
 }
 
 LogicalResult checkAndPrintOverflow(TileOp tile, int numBanks, int stacksize,
-                                    SmallVector<BufferOp, 4> allBuffers,
+                                    SmallVector<BufferOp> &allBuffers,
                                     std::vector<int64_t> &nextAddrInBanks,
                                     std::vector<BankLimits> &bankLimits) {
   bool foundOverflow = false;
@@ -320,7 +345,7 @@ LogicalResult checkAndPrintOverflow(TileOp tile, int numBanks, int stacksize,
 }
 
 // Function to deallocate attributes of buffers in case of a failure
-void deAllocationBuffers(SmallVector<BufferOp, 4> &buffers) {
+void deAllocationBuffers(SmallVector<BufferOp> &buffers) {
   for (auto buffer : buffers) {
     buffer->removeAttr("address");
     buffer->removeAttr("mem_bank");
@@ -363,9 +388,9 @@ LogicalResult simpleBankAwareAllocation(TileOp tile) {
   }
   fillBankLimits(numBanks, bankSize, bankLimits);
 
-  SmallVector<BufferOp, 4> buffersToAlloc;
-  SmallVector<BufferOp, 4> preAllocatedBuffers;
-  SmallVector<BufferOp, 4> allBuffers;
+  SmallVector<BufferOp> preAllocatedBuffers;
+  SmallVector<BufferOp> buffersToAlloc;
+  SmallVector<BufferOp> allBuffers;
   // Collect all the buffers for this tile.
   device.walk<WalkOrder::PreOrder>([&](BufferOp buffer) {
     if (buffer.getTileOp() == tile)
@@ -378,11 +403,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);
       else
         preAllocatedBuffers.push_back(buffer);
@@ -396,7 +423,7 @@ LogicalResult simpleBankAwareAllocation(TileOp tile) {
             });
 
   // Set addresses for remaining buffers.
-  SmallVector<BufferOp, 4> allocatedBuffers;
+  SmallVector<BufferOp> allocatedBuffers;
   int bankIndex = 0;
   for (auto buffer : buffersToAlloc) {
     // If the buffer doesn't fit in any of the bank space then

python/compiler/aiecc/cl_arguments.py

@@ -147,7 +147,7 @@ def parse_args(args=None):
     parser.add_argument(
         "--alloc-scheme",
         dest="alloc_scheme",
-        default=None,
+        default="bank-aware",
         help="Allocation scheme for AIE buffers: basic-sequential, bank-aware (default).",
     )
     parser.add_argument(

python/compiler/aiecc/main.py

@@ -34,32 +34,30 @@
 from aie.ir import Context, Location, Module
 from aie.passmanager import PassManager
 
-INPUT_WITH_ADDRESSES_PIPELINE = (
-    lambda scheme="", dynamic_objFifos=False, ctrl_pkt_overlay=False: (
+INPUT_WITH_ADDRESSES_PIPELINE = lambda scheme, dynamic_objFifos, ctrl_pkt_overlay: (
+    Pipeline()
+    .lower_affine()
+    .add_pass("aie-canonicalize-device")
+    .Nested(
+        "aie.device",
         Pipeline()
-        .lower_affine()
-        .add_pass("aie-canonicalize-device")
-        .Nested(
-            "aie.device",
-            Pipeline()
-            .add_pass("aie-assign-lock-ids")
-            .add_pass("aie-register-objectFifos")
-            .add_pass(
-                "aie-objectFifo-stateful-transform", dynamic_objFifos=dynamic_objFifos
-            )
-            .add_pass("aie-assign-bd-ids")
-            .add_pass("aie-lower-cascade-flows")
-            .add_pass("aie-lower-broadcast-packet")
-            .add_pass("aie-lower-multicast")
-            .add_pass("aie-assign-tile-controller-ids")
-            .add_pass(
-                "aie-generate-column-control-overlay",
-                route_shim_to_tile_ctrl=ctrl_pkt_overlay,
-            )
-            .add_pass("aie-assign-buffer-addresses", alloc_scheme=scheme),
+        .add_pass("aie-assign-lock-ids")
+        .add_pass("aie-register-objectFifos")
+        .add_pass(
+            "aie-objectFifo-stateful-transform", dynamic_objFifos=dynamic_objFifos
         )
-        .convert_scf_to_cf()
+        .add_pass("aie-assign-bd-ids")
+        .add_pass("aie-lower-cascade-flows")
+        .add_pass("aie-lower-broadcast-packet")
+        .add_pass("aie-lower-multicast")
+        .add_pass("aie-assign-tile-controller-ids")
+        .add_pass(
+            "aie-generate-column-control-overlay",
+            route_shim_to_tile_ctrl=ctrl_pkt_overlay,
+        )
+        .add_pass("aie-assign-buffer-addresses", alloc_scheme=scheme),
     )
+    .convert_scf_to_cf()
 )
 
 LOWER_TO_LLVM_PIPELINE = (
@@ -334,7 +332,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:
@@ -1062,12 +1065,9 @@ async def run_flow(self):
 
             file_with_addresses = self.prepend_tmp("input_with_addresses.mlir")
 
-            if opts.alloc_scheme:
-                pass_pipeline = INPUT_WITH_ADDRESSES_PIPELINE(
-                    opts.alloc_scheme, opts.dynamic_objFifos, opts.ctrl_pkt_overlay
-                ).materialize(module=True)
-            else:
-                pass_pipeline = INPUT_WITH_ADDRESSES_PIPELINE().materialize(module=True)
+            pass_pipeline = INPUT_WITH_ADDRESSES_PIPELINE(
+                opts.alloc_scheme, opts.dynamic_objFifos, opts.ctrl_pkt_overlay
+            ).materialize(module=True)
 
             run_passes(
                 pass_pipeline,

test/assign-buffer-addresses/bank_aware_alloc_error.mlir

@@ -1,4 +1,4 @@
-//===- bank_aware_alloc_error.mlir ---------------------------------------------*- MLIR -*-===//
+//===- bank_aware_alloc_error.mlir -----------------------------*- MLIR -*-===//
 //
 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -26,18 +26,18 @@
 // CHECK:         bank : 3        0x6000-0x7FFF
 
 module @test {
- aie.device(xcvc1902) {
-  %0 = aie.tile(3, 3)
-  %b1 = aie.buffer(%0) { sym_name = "a" } : memref<16xi8>
-  %1 = aie.buffer(%0) { sym_name = "b" } : memref<8192xi32>
-  %b2 = aie.buffer(%0) { sym_name = "c" } : memref<16xi16>
-  %3 = aie.tile(4, 4)
-  %4 = aie.buffer(%3) : memref<500xi32>
-  aie.core(%0) {
-    aie.end
+  aie.device(xcvc1902) {
+    %0 = aie.tile(3, 3)
+    %b1 = aie.buffer(%0) { sym_name = "a" } : memref<16xi8>
+    %1 = aie.buffer(%0) { sym_name = "b" } : memref<8192xi32>
+    %b2 = aie.buffer(%0) { sym_name = "c" } : memref<16xi16>
+    %3 = aie.tile(4, 4)
+    %4 = aie.buffer(%3) : memref<500xi32>
+    aie.core(%0) {
+      aie.end
+    }
+    aie.core(%3) {
+      aie.end
+    }
   }
-  aie.core(%3) {
-    aie.end
-  }
- }
-}
+}

test/assign-buffer-addresses/bank_aware_alloc_memory_exhausted.mlir

@@ -1,4 +1,4 @@
-//===- bank_aware_alloc_memory_exhausted.mlir ---------------------------------------------*- MLIR -*-===//
+//===- bank_aware_alloc_memory_exhausted.mlir ------------------*- MLIR -*-===//
 //
 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

test/assign-buffer-addresses/bank_aware_alloc_memtile_error.mlir

@@ -1,4 +1,4 @@
-//===- bank_aware_alloc_memtile_error.mlir ---------------------------------------*- MLIR -*-===//
+//===- bank_aware_alloc_memtile_error.mlir ---------------------*- MLIR -*-===//
 //
 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -29,13 +29,12 @@
 // CHECK:         bank : 6        0x60000-0x6FFFF
 // CHECK:         bank : 7        0x70000-0x7FFFF
 
-
 module @test {
- aie.device(xcve2302) {
-  %0 = aie.tile(3, 1)
-  %b1 = aie.buffer(%0) { sym_name = "a" } : memref<132000xi32>
-  aie.memtile_dma(%0) {
-    aie.end
+  aie.device(xcve2302) {
+    %0 = aie.tile(3, 1)
+    %b1 = aie.buffer(%0) { sym_name = "a" } : memref<132000xi32>
+    aie.memtile_dma(%0) {
+      aie.end
+    }
   }
- }
-}
+}