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

@@ -181,6 +181,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,
@@ -268,6 +270,7 @@ class AIE1TargetModel : public AIETargetModel {
   }
   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;
@@ -342,6 +345,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=*/"",
+           "Select allocation scheme: basic-sequential or bank-aware. Default is bank-aware, falling back to basic-sequential if it fails.">,
   ];
 }
 

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

@@ -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,34 +187,86 @@ 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(
+      "Not all requested buffers fit in the available memory.\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
-// calls of this function to ensure a round-robin allocation scheme
-// over the available banks).
+// If it does not find one with enough space, it will throw an error.
+// 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,
                      std::vector<int64_t> &nextAddrInBanks,
                      std::vector<BankLimits> &bankLimits) {
   assert(startBankIndex < numBanks &&
          "Unexpected input value for startBankIndex");
   int bankIndex = startBankIndex;
+  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);
-      bankIndex = (bankIndex + 1) % numBanks;
+      allocated = true;
+      bankIndex++;
+      bankIndex %= numBanks;
       break;
     }
-    bankIndex = (bankIndex + 1) % numBanks;
+    // 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;
   }
-  return bankIndex;
+  return true;
 }
 
 LogicalResult checkAndPrintOverflow(TileOp tile, int numBanks, int stacksize,
@@ -275,6 +319,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)
@@ -295,7 +347,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.
@@ -312,6 +364,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) {
@@ -331,6 +384,8 @@ LogicalResult simpleBankAwareAllocation(TileOp tile) {
                                                    nextAddrInBanks, bankLimits);
       if (!has_addr && !has_bank)
         buffersToAlloc.push_back(buffer);
+      else
+        preAllocatedBuffers.push_back(buffer);
     }
   }
 
@@ -341,10 +396,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) {
@@ -380,16 +449,25 @@ 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>()) {
+        tile.emitWarning("Memory allocation scheme is either not provided or "
+                         "unrecognized. Defaulting to bank-aware allocation.");
+        if (auto res = simpleBankAwareAllocation(tile); res.failed()) {
+          if (auto res2 = basicAllocation(tile); res2.failed())
+            return signalPassFailure();
+        }
+      }
     }
   }
 };

lib/Targets/AIEVecToCpp/TranslateAIEVecToCpp.cpp

@@ -3470,4 +3470,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

@@ -37,7 +37,7 @@ build/conv2dk1_skip.o: conv2dk1_skip.cc
 build/final.xclbin: build/${mlirFileName}.mlir build/conv2dk1.o build/conv2dk3.o build/conv2dk1_skip.o  
 	mkdir -p ${@D}
 	cd ${@D} && aiecc.py -v --aie-generate-cdo --aie-generate-npu --no-compile-host \
-		--basic-alloc-scheme \
+		--alloc-scheme=basic-sequential \
 		--no-xchesscc --no-xbridge \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt $(<:%=../%) 
 

programming_examples/ml/resnet/layers_conv2_x/Makefile

@@ -44,9 +44,10 @@ build/conv2dk1_skip.o: conv2dk1_skip.cc
 
 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
 	mkdir -p ${@D}
-	cd ${@D} &&	aiecc.py --basic-alloc-scheme --aie-generate-cdo --aie-generate-npu --no-compile-host \
+	cd ${@D} &&	aiecc.py --alloc-scheme=basic-sequential --aie-generate-cdo --aie-generate-npu --no-compile-host \
 		--no-xchesscc --no-xbridge \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt $(<:%=../%)
+
 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 \
 		--no-xchesscc --no-xbridge \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt $(<:%=../%)
 

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 \
 		--no-xchesscc --no-xbridge \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt $(<:%=../%)
 

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 \
 		--no-xchesscc --no-xbridge \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt $(<:%=../%)
 

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 \
 		--no-xchesscc --no-xbridge \
 		--xclbin-name=${@F} --npu-insts-name=insts.txt $(<:%=../%)
 

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="Allocation scheme for AIE buffers: basic-sequential, bank-aware (default).",
     )
     parser.add_argument(
         "--generate-ctrl-pkt-overlay",

python/compiler/aiecc/main.py

@@ -35,7 +35,7 @@
 from aie.passmanager import PassManager
 
 INPUT_WITH_ADDRESSES_PIPELINE = (
-    lambda basic_alloc_scheme=False, dynamic_objFifos=False, ctrl_pkt_overlay=False: (
+    lambda scheme="", dynamic_objFifos=False, ctrl_pkt_overlay=False: (
         Pipeline()
         .lower_affine()
         .add_pass("aie-canonicalize-device")
@@ -56,7 +56,7 @@
                 "aie-generate-column-control-overlay",
                 route_shim_to_tile_ctrl=ctrl_pkt_overlay,
             )
-            .add_pass("aie-assign-buffer-addresses", basic_alloc=basic_alloc_scheme),
+            .add_pass("aie-assign-buffer-addresses", alloc_scheme=scheme),
         )
         .convert_scf_to_cf()
     )
@@ -1061,9 +1061,14 @@ async def run_flow(self):
             )
 
             file_with_addresses = self.prepend_tmp("input_with_addresses.mlir")
-            pass_pipeline = INPUT_WITH_ADDRESSES_PIPELINE(
-                opts.basic_alloc_scheme, opts.dynamic_objFifos, opts.ctrl_pkt_overlay
-            ).materialize(module=True)
+
+            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)
+
             run_passes(
                 pass_pipeline,
                 self.mlir_module_str,