Fix for bank-aware allocation (#1958)

Author: pvasireddy-amd

lib/Dialect/AIE/Transforms/AIEAssignBuffers.cpp

@@ -262,7 +262,7 @@ void printMemMap(TileOp tile, SmallVector<BufferOp> &allocatedBuffers,
 // Returns true if the buffer was successfully allocated, false otherwise.
 // If no bank has enough space to accommodate the buffer, an error is emitted.
 
-int setBufferAddress(BufferOp buffer, int numBanks, int startBankIndex,
+int setBufferAddress(BufferOp buffer, int numBanks, int &startBankIndex,
                      std::vector<int64_t> &nextAddrInBanks,
                      std::vector<BankLimits> &bankLimits) {
   assert(startBankIndex < numBanks &&
@@ -276,13 +276,12 @@ int setBufferAddress(BufferOp buffer, int numBanks, int startBankIndex,
       buffer.setMemBank(bankIndex);
       setAndUpdateAddressInBank(buffer, startAddr, endAddr, nextAddrInBanks);
       allocated = true;
-      bankIndex++;
-      bankIndex %= numBanks;
+      bankIndex = (bankIndex + 1) % numBanks;
+      startBankIndex = bankIndex;
       break;
     }
     // Move to the next bank
-    bankIndex++;
-    bankIndex %= numBanks;
+    bankIndex = (bankIndex + 1) % numBanks;
   }
   // If no bank has enough space, throws error
   if (!allocated) {

test/assign-buffer-addresses/bank_aware_alloc_simple.mlir

@@ -9,12 +9,10 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: aie-opt --aie-assign-buffer-addresses="alloc-scheme=bank-aware" %s | FileCheck %s
-// CHECK: %a = aie.buffer(%tile_3_3) {address = 3104 : i32, mem_bank = 0 : i32, sym_name = "a"} : memref<16xi8> 
+// CHECK: %a = aie.buffer(%tile_3_3) {address = 16384 : i32, mem_bank = 2 : i32, sym_name = "a"} : memref<16xi8> 
 // CHECK: %b = aie.buffer(%tile_3_3) {address = 1024 : i32, mem_bank = 0 : i32, sym_name = "b"} : memref<512xi32> 
-// CHECK: %c = aie.buffer(%tile_3_3) {address = 3072 : i32, mem_bank = 0 : i32, sym_name = "c"} : memref<16xi16> 
+// CHECK: %c = aie.buffer(%tile_3_3) {address = 8192 : i32, mem_bank = 1 : i32, sym_name = "c"} : memref<16xi16> 
 // CHECK: %_anonymous0 = aie.buffer(%tile_4_4) {address = 1024 : i32, mem_bank = 0 : i32, sym_name = "_anonymous0"} : memref<500xi32> 
-    
-
 
 module @test {
   aie.device(xcvc1902) {