Fall back routine when bank-aware allocation fails

include/aie-c/TargetModel.h

@@ -137,6 +137,9 @@ aieTargetModelGetNumMemTileRows(AieTargetModel targetModel);
 MLIR_CAPI_EXPORTED uint32_t
 aieTargetModelGetMemTileSize(AieTargetModel targetModel);
 
+MLIR_CAPI_EXPORTED uint32_t
+aieTargetModelGetNumBanks(AieTargetModel targetModel, int col, int row);
+
 /// Returns true if this is an NPU target model.
 MLIR_CAPI_EXPORTED bool aieTargetModelIsNPU(AieTargetModel targetModel);
 

include/aie/Dialect/AIE/IR/AIETargetModel.h

@@ -176,6 +176,8 @@ class AIETargetModel {
   virtual uint32_t getNumMemTileRows() const = 0;
   /// Return the size (in bytes) of a MemTile.
   virtual uint32_t getMemTileSize() const = 0;
+  /// Return the number of memory banks of a given tile.
+  virtual uint32_t getNumBanks(int col, int row) const = 0;
   /// Return the number of destinations of connections inside a switchbox. These
   /// are the targets of connect operations in the switchbox.
   virtual uint32_t getNumDestSwitchboxConnections(int col, int row,
@@ -259,6 +261,7 @@ class AIE1TargetModel : public AIETargetModel {
   uint32_t getNumBDs(int col, int row) const override { return 16; }
   uint32_t getNumMemTileRows() const override { return 0; }
   uint32_t getMemTileSize() const override { return 0; }
+  uint32_t getNumBanks(int col, int row) const override { return 4; }
 
   uint32_t getNumDestSwitchboxConnections(int col, int row,
                                           WireBundle bundle) const override;
@@ -320,6 +323,10 @@ class AIE2TargetModel : public AIETargetModel {
 
   uint32_t getMemTileSize() const override { return 0x00080000; }
 
+  uint32_t getNumBanks(int col, int row) const override {
+    return isMemTile(col, row) ? 8 : 4;
+  }
+
   uint32_t getNumDestSwitchboxConnections(int col, int row,
                                           WireBundle bundle) const override;
   uint32_t getNumSourceSwitchboxConnections(int col, int row,

include/aie/Dialect/AIE/Transforms/AIEPasses.td

@@ -27,8 +27,8 @@ def AIEAssignBufferAddresses : Pass<"aie-assign-buffer-addresses", "DeviceOp"> {
   let constructor = "xilinx::AIE::createAIEAssignBufferAddressesPass()";
 
   let options = [
-    Option<"clBasicAlloc", "basic-alloc", "bool", /*default=*/"false",
-            "Flag to enable the basic sequential allocation scheme (not bank-aware).">
+    Option<"clAllocScheme", "alloc-scheme", "std::string", /*default=*/"",
+           "Choose allocation scheme; possibilities: basic-sequential, bank-aware. By default, bank-aware is selected and if it fails, will fall back to basic-sequential scheme.">,
   ];
 }
 

lib/CAPI/TargetModel.cpp

@@ -149,6 +149,11 @@ uint32_t aieTargetModelGetMemTileSize(AieTargetModel targetModel) {
   return unwrap(targetModel).getMemTileSize();
 }
 
+uint32_t aieTargetModelGetNumBanks(AieTargetModel targetModel, int col,
+                                   int row) {
+  return unwrap(targetModel).getNumBanks(col, row);
+}
+
 bool aieTargetModelIsNPU(AieTargetModel targetModel) {
   return unwrap(targetModel).isNPU();
 }

lib/Dialect/AIE/Transforms/AIEAssignBuffers.cpp

@@ -28,8 +28,8 @@ LogicalResult checkAndPrintOverflow(TileOp tile, int address,
                                     int maxDataMemorySize, int stacksize,
                                     SmallVector<BufferOp, 4> buffers) {
   if (address > maxDataMemorySize) {
-    InFlightDiagnostic error =
-        tile.emitOpError("allocated buffers exceeded available memory\n");
+    InFlightDiagnostic error = tile.emitOpError(
+        "allocated buffers exceeded available memory: Sequential\n");
     auto &note = error.attachNote() << "MemoryMap:\n";
     auto printbuffer = [&](StringRef name, int address, int size) {
       note << "\t" << name << " \t"
@@ -53,7 +53,7 @@ LogicalResult checkAndPrintOverflow(TileOp tile, int address,
   return success();
 }
 
-LogicalResult basicAllocation(TileOp &tile) {
+LogicalResult basicAllocation(TileOp tile) {
   auto device = tile->getParentOfType<AIE::DeviceOp>();
   if (!device)
     return failure();
@@ -113,14 +113,6 @@ typedef struct BankLimits {
   int64_t endAddr;
 } BankLimits;
 
-// TODO: add to target model
-int getNumBanks(TileOp tile) {
-  if (tile.isMemTile())
-    return 1;
-  else
-    return 4;
-}
-
 // Function that given a number of banks and their size, computes
 // the start and end addresses for each bank and fills in the entry
 // in the bankLimits vector.
@@ -195,33 +187,84 @@ bool checkAndAddBufferWithMemBank(BufferOp buffer, int numBanks,
   return false;
 }
 
+// Prints the memory map across banks
+void printMemMap(TileOp tile, SmallVector<BufferOp, 4> allocatedBuffers,
+                 SmallVector<BufferOp, 4> preAllocatedBuffers, int numBanks,
+                 std::vector<BankLimits> &bankLimits, int stacksize) {
+  InFlightDiagnostic error =
+      tile.emitOpError("All requested buffers doesn't fit in the available "
+                       "memory: Bank aware\n");
+  auto &note = error.attachNote()
+               << "Current configuration of buffers in bank(s) : ";
+  note << "MemoryMap:\n";
+  auto printbuffer = [&](StringRef name, int address, int size) {
+    note << "\t"
+         << "\t" << name << " \t"
+         << ": 0x" << llvm::utohexstr(address) << "-0x"
+         << llvm::utohexstr(address + size - 1) << " \t(" << size
+         << " bytes)\n";
+  };
+  for (int i = 0; i < numBanks; i++) {
+    if (i == 0) {
+      if (stacksize > 0)
+        printbuffer("(stack)", 0, stacksize);
+      else
+        note << "(no stack allocated)\n";
+    }
+    note << "\t"
+         << "bank : " << i << "\t"
+         << "0x" << llvm::utohexstr(bankLimits[i].startAddr) << "-0x"
+         << llvm::utohexstr(bankLimits[i].endAddr - 1) << "\n";
+    for (auto buffer : preAllocatedBuffers) {
+      auto addr = buffer.getAddress().value();
+      auto mem_bank = buffer.getMemBank().value();
+      if (mem_bank == i)
+        printbuffer(buffer.name(), addr, buffer.getAllocationSize());
+    }
+    for (auto buffer : allocatedBuffers) {
+      auto addr = buffer.getAddress().value();
+      auto mem_bank = buffer.getMemBank().value();
+      if (mem_bank == i)
+        printbuffer(buffer.name(), addr, buffer.getAllocationSize());
+    }
+  }
+}
+
 // Function that given a buffer will iterate over all the memory banks
 // starting from the given index to try and find a bank with enough
 // space. If it does, it will set the buffer's address and mem_bank
 // attributes and update the nextAddrInBanks vector.
-// If it does not find one with enough space, it will allocate the
-// buffer in the last checked bank (this will be picked up during
-// overflow error checking). Finally, the function returns the index
-// of the next bank to search (which should be given to subsequent
+// If it does not find one with enough space, it will throw an error.
+// Finally, the function returns a pass or a fail.
+// The index of the next bank to search (which should be given to subsequent
 // calls of this function to ensure a round-robin allocation scheme
 // over the available banks).
-int setBufferAddress(BufferOp buffer, int numBanks, int startBankIndex,
-                     std::vector<int64_t> &nextAddrInBanks,
-                     std::vector<BankLimits> &bankLimits) {
-  int bankIndex = startBankIndex;
+bool setBufferAddress(BufferOp buffer, int numBanks, int &bankIndex,
+                      std::vector<int64_t> &nextAddrInBanks,
+                      std::vector<BankLimits> &bankLimits) {
+  bool allocated = false;
   for (int i = 0; i < numBanks; i++) {
     int64_t startAddr = nextAddrInBanks[bankIndex];
     int64_t endAddr = startAddr + buffer.getAllocationSize();
-    if (endAddr <= bankLimits[bankIndex].endAddr || i == numBanks - 1) {
+    if (endAddr <= bankLimits[bankIndex].endAddr) {
       buffer.setMemBank(bankIndex);
       setAndUpdateAddressInBank(buffer, startAddr, endAddr, nextAddrInBanks);
+      allocated = true;
       bankIndex++;
       break;
     }
+    // Move to the next bank
     bankIndex++;
+    bankIndex %= numBanks;
+  }
+  // If no bank has enough space, throws error
+  if (!allocated) {
+    buffer.emitError("Failed to allocate buffer: ")
+        << buffer.name() << " with size: " << buffer.getAllocationSize()
+        << " bytes.";
+    return false;
   }
-  bankIndex %= numBanks;
-  return bankIndex;
+  return true;
 }
 
 LogicalResult checkAndPrintOverflow(TileOp tile, int numBanks, int stacksize,
@@ -237,8 +280,8 @@ LogicalResult checkAndPrintOverflow(TileOp tile, int numBanks, int stacksize,
     }
   }
   if (foundOverflow) {
-    InFlightDiagnostic error =
-        tile.emitOpError("allocated buffers exceeded available memory\n");
+    InFlightDiagnostic error = tile.emitOpError(
+        "allocated buffers exceeded available memory: Bank aware\n");
     auto &note = error.attachNote() << "Error in bank(s) : ";
     for (auto bank : overflow_banks)
       note << bank << " ";
@@ -274,6 +317,14 @@ LogicalResult checkAndPrintOverflow(TileOp tile, int numBanks, int stacksize,
   return success();
 }
 
+// Function to deallocate attributes of buffers in case of a failure
+void deAllocationBuffers(SmallVector<BufferOp, 4> &buffers) {
+  for (auto buffer : buffers) {
+    buffer->removeAttr("address");
+    buffer->removeAttr("mem_bank");
+  }
+}
+
 LogicalResult simpleBankAwareAllocation(TileOp tile) {
   auto device = tile->getParentOfType<AIE::DeviceOp>();
   if (!device)
@@ -294,7 +345,7 @@ LogicalResult simpleBankAwareAllocation(TileOp tile) {
   else
     maxDataMemorySize = targetModel.getLocalMemorySize();
 
-  int numBanks = getNumBanks(tile);
+  int numBanks = targetModel.getNumBanks(tile.getCol(), tile.getRow());
   int bankSize = maxDataMemorySize / numBanks;
 
   // Address range owned by the MemTile is 0x80000.
@@ -311,6 +362,7 @@ LogicalResult simpleBankAwareAllocation(TileOp tile) {
   fillBankLimits(numBanks, bankSize, bankLimits);
 
   SmallVector<BufferOp, 4> buffersToAlloc;
+  SmallVector<BufferOp, 4> preAllocatedBuffers;
   SmallVector<BufferOp, 4> allBuffers;
   // Collect all the buffers for this tile.
   device.walk<WalkOrder::PreOrder>([&](BufferOp buffer) {
@@ -330,6 +382,8 @@ LogicalResult simpleBankAwareAllocation(TileOp tile) {
                                                    nextAddrInBanks, bankLimits);
       if (!has_addr && !has_bank)
         buffersToAlloc.push_back(buffer);
+      else
+        preAllocatedBuffers.push_back(buffer);
     }
   }
 
@@ -340,10 +394,24 @@ LogicalResult simpleBankAwareAllocation(TileOp tile) {
             });
 
   // Set addresses for remaining buffers.
+  SmallVector<BufferOp, 4> allocatedBuffers;
   int bankIndex = 0;
-  for (auto buffer : buffersToAlloc)
-    bankIndex = setBufferAddress(buffer, numBanks, bankIndex, nextAddrInBanks,
-                                 bankLimits);
+  for (auto buffer : buffersToAlloc) {
+    // If the buffer doesn't fit in any of the bank space then
+    // it prints the current memory map of the banks,
+    // deallocates all the buffers, and
+    // returns a failure.
+    if (!setBufferAddress(buffer, numBanks, bankIndex, nextAddrInBanks,
+                          bankLimits)) {
+
+      printMemMap(tile, allocatedBuffers, preAllocatedBuffers, numBanks,
+                  bankLimits, stacksize);
+      deAllocationBuffers(allocatedBuffers);
+      return failure();
+    } else {
+      allocatedBuffers.push_back(buffer);
+    }
+  }
 
   // Sort by smallest address before printing memory map.
   std::sort(allBuffers.begin(), allBuffers.end(), [](BufferOp a, BufferOp b) {
@@ -379,16 +447,23 @@ struct AIEAssignBufferAddressesPass
     });
 
     // Select allocation scheme
-    if (clBasicAlloc) {
+    if (clAllocScheme == "basic-sequential") {
       for (auto tile : device.getOps<TileOp>()) {
         if (auto res = basicAllocation(tile); res.failed())
           return signalPassFailure();
       }
-    } else {
+    } else if (clAllocScheme == "bank-aware") {
       for (auto tile : device.getOps<TileOp>()) {
         if (auto res = simpleBankAwareAllocation(tile); res.failed())
           return signalPassFailure();
       }
+    } else {
+      for (auto tile : device.getOps<TileOp>()) {
+        if (auto res = simpleBankAwareAllocation(tile); res.failed()) {
+          if (auto res2 = basicAllocation(tile); res2.failed())
+            return signalPassFailure();
+        }
+      }
     }
   }
 };

lib/Targets/AIEVecToCpp/TranslateAIEVecToCpp.cpp

@@ -3444,4 +3444,4 @@ LogicalResult aievec::translateAIEVecToCpp(Operation *op, bool aie2,
                                            raw_ostream &os) {
   CppEmitter emitter(os, false, aie2);
   return emitter.emitOperation(*op, /*trailingSemicolon=*/false);
-}
+}

programming_examples/ml/bottleneck/Makefile

@@ -33,7 +33,6 @@ build/conv2dk1_skip.o: conv2dk1_skip.cc
 
 build/final.xclbin: build/${mlirFileName}.mlir build/conv2dk1.o build/conv2dk3.o build/conv2dk1_skip.o  
 	cd build && aiecc.py -v --aie-generate-cdo --aie-generate-npu --no-compile-host \
-		--basic-alloc-scheme \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt ${<F}
 
 clean:

programming_examples/ml/resnet/layers_conv2_x/Makefile

@@ -38,7 +38,7 @@ build/conv2dk1_skip.o: conv2dk1_skip.cc
 	xchesscc -d ${CHESSCC2_FLAGS} -DUINT8_ACT -c $< -o $@
 
 build/final.xclbin: build/${mlirFileName}.mlir build/conv2dk1_i8.o build/conv2dk1_skip_init.o build/conv2dk3.o build/conv2dk1_skip.o build/conv2dk1_ui8.o
-	cd build && aiecc.py --basic-alloc-scheme --aie-generate-cdo --aie-generate-npu --no-compile-host \
+	cd build && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt ${<F}
 clean:
 	rm -rf build log

programming_examples/vision/color_detect/Makefile

@@ -41,7 +41,7 @@ build/aie2_lineBased_8b_${COLORDETECT_WIDTH}.mlir: ${srcdir}/aie2_colorDetect.py
 
 build/final_${COLORDETECT_WIDTH}.xclbin: build/aie2_lineBased_8b_${COLORDETECT_WIDTH}.mlir build/rgba2hue.cc.o build/threshold.cc.o build/combined_bitwiseOR_gray2rgba_bitwiseAND.a
 	mkdir -p ${@D}
-	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host --basic-alloc-scheme \
+	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host --alloc-scheme=basic-sequential \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt $(<:%=../%)
 
 ${targetname}.exe: ${srcdir}/test.cpp

programming_examples/vision/color_threshold/Makefile

@@ -34,7 +34,7 @@ build/aie2_${COLORTHRESHOLD_WIDTH}.mlir: ${srcdir}/aie2_colorThreshold.py
 
 build/final_${COLORTHRESHOLD_WIDTH}.xclbin: build/aie2_${COLORTHRESHOLD_WIDTH}.mlir build/threshold.cc.o
 	mkdir -p ${@D}
-	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host --basic-alloc-scheme \
+	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host --alloc-scheme=basic-sequential \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt $(<:%=../%)
 
 ${targetname}.exe: ${srcdir}/test.cpp

programming_examples/vision/edge_detect/Makefile

@@ -41,7 +41,7 @@ build/aie2_lineBased_8b_${EDGEDETECT_WIDTH}.mlir: ${srcdir}/aie2_edgeDetect.py
 
 build/final_${EDGEDETECT_WIDTH}.xclbin: build/aie2_lineBased_8b_${EDGEDETECT_WIDTH}.mlir build/rgba2gray.cc.o build/gray2rgba.cc.o build/filter2d.cc.o build/threshold.cc.o build/addWeighted.cc.o build/combined_gray2rgba_addWeighted.a
 	mkdir -p ${@D}
-	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host --basic-alloc-scheme \
+	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host --alloc-scheme=basic-sequential \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt $(<:%=../%)
 
 ${targetname}.exe: ${srcdir}/test.cpp

programming_examples/vision/vision_passthrough/Makefile

@@ -34,7 +34,7 @@ build/passThrough.cc.o: passThrough.cc
 
 build/final_${PASSTHROUGH_WIDTH}.xclbin: build/aie2_lineBased_8b_${PASSTHROUGH_WIDTH}.mlir build/passThrough.cc.o
 	mkdir -p ${@D}
-	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host --basic-alloc-scheme \
+	cd ${@D} && aiecc.py --aie-generate-cdo --aie-generate-npu --no-compile-host --alloc-scheme=basic-sequential \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt $(<:%=../%)
 
 ${targetname}.exe: ${srcdir}/test.cpp

python/compiler/aiecc/cl_arguments.py

@@ -145,11 +145,10 @@ def parse_args(args=None):
         help="Disable linking of AIE code",
     )
     parser.add_argument(
-        "--basic-alloc-scheme",
-        dest="basic_alloc_scheme",
-        default=False,
-        action="store_true",
-        help="Use basic memory allocation scheme for AIE buffer address assignment",
+        "--alloc-scheme",
+        dest="alloc_scheme",
+        default=None,
+        help="Choose allocation scheme for AIE buffer address assignment; possibilities: basic-sequential, bank-aware. By default, bank-aware is tried followed by basic-sequential if it fails.",
     )
     parser.add_argument(
         "--aie-generate-airbin",