Clean up and speed up Shim DMA Allocation Retrieval (#2644)

Author: andrej

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

@@ -200,23 +200,6 @@ int32_t getBufferBaseAddress(mlir::Operation *bufOp);
 #define GET_OP_CLASSES
 #include "aie/Dialect/AIE/IR/AIEOps.h.inc"
 
-namespace xilinx::AIE {
-class DeviceOp;
-class ShimDMAAllocationOp;
-struct ShimDMAllocationGetter {
-public:
-  std::optional<AIE::ShimDMAAllocationOp> get(DeviceOp dev,
-                                              mlir::StringRef sym_name);
-
-private:
-  llvm::DenseMap<std::pair<DeviceOp, mlir::StringRef>,
-                 std::optional<AIE::ShimDMAAllocationOp>>
-      allocGetter;
-  std::optional<AIE::ShimDMAAllocationOp>
-  cachelessGet(DeviceOp dev, mlir::StringRef sym_name);
-};
-} // namespace xilinx::AIE
-
 namespace xilinx::AIE {
 
 void collectTiles(DeviceOp &device,

include/aie/Dialect/AIE/IR/AIEOps.td

@@ -1583,7 +1583,7 @@ def AIE_PutCascadeOp: AIE_Op<"put_cascade", []> {
   let assemblyFormat = [{ `(` $cascade_value `:` type($cascade_value) `)` attr-dict }];
 }
 
-def AIE_ShimDMAAllocationOp : AIE_Op<"shim_dma_allocation", [HasParent<"DeviceOp">]> {
+def AIE_ShimDMAAllocationOp : AIE_Op<"shim_dma_allocation", [HasParent<"DeviceOp">, Symbol]> {
   let summary = "Runtime allocation information for a single shim DMA";
   let description = [{
     This op exists for cases where shim_dma configuration is performed outside of MLIR-AIE
@@ -1607,7 +1607,7 @@ def AIE_ShimDMAAllocationOp : AIE_Op<"shim_dma_allocation", [HasParent<"DeviceOp
   }];
 
   let arguments = (
-    ins FlatSymbolRefAttr:$sym_name,
+    ins SymbolNameAttr:$sym_name,
         DMAChannelDir:$channel_dir,
         AIEI64Attr:$channel_index,
         AIEI64Attr:$col,
@@ -1626,15 +1626,6 @@ def AIE_ShimDMAAllocationOp : AIE_Op<"shim_dma_allocation", [HasParent<"DeviceOp
     static ::xilinx::AIE::ShimDMAAllocationOp getForSymbol(::xilinx::AIE::DeviceOp device, ::llvm::StringRef symbol);
   }];
 
-  let builders = [
-    OpBuilder<(ins "mlir::StringRef":$sym_name, "DMAChannelDir":$dir, "int":$channel_index,
-               "int":$col, "bool":$plio), [{
-      build($_builder, $_state, ::mlir::SymbolRefAttr::get(odsBuilder.getContext(), sym_name), 
-            DMAChannelDirAttr::get(odsBuilder.getContext(), dir),
-            $_builder.getI64IntegerAttr(channel_index), 
-            $_builder.getI64IntegerAttr(col), $_builder.getBoolAttr(plio), nullptr);
-    }]>
-  ];
 }
 
 def AIE_ObjectFifoCreateOp: AIE_Op<"objectfifo", [HasParent<"DeviceOp">, Symbol]> {

include/aie/Dialect/AIEX/IR/AIEX.td

@@ -690,7 +690,7 @@ def AIE_NpuDmaMemcpyNdOp: AIEX_Op<"npu.dma_memcpy_nd", [
         ConfinedAttr<DenseI64ArrayAttr, [DenseArrayCount<4>]>:$static_sizes,
         ConfinedAttr<DenseI64ArrayAttr, [DenseArrayCount<4>]>:$static_strides,
         OptionalAttr<PacketInfoAttr>:$packet,
-        FlatSymbolRefAttr:$metadata,
+        SymbolRefAttr:$metadata,
         I64Attr:$id,
         DefaultValuedOptionalAttr<BoolAttr, "false">:$issue_token,
         DefaultValuedOptionalAttr<I64Attr, "0">:$d0_zero_before,
@@ -1096,7 +1096,7 @@ def AIE_DMAConfigureTaskForOp : AIEX_Op<"dma_configure_task_for", [HasParent<"Ru
   let summary = "As dma_configure_task, but specify tile, direction and channel by reference to a Shim DMA allocation op";
 
   let arguments = (
-    ins FlatSymbolRefAttr:$alloc,
+    ins SymbolRefAttr:$alloc,
         DefaultValuedOptionalAttr<BoolAttr, "false">:$issue_token,
         DefaultValuedOptionalAttr<I32Attr, "0">:$repeat_count
   );

lib/Dialect/AIE/IR/AIEDialect.cpp

@@ -90,49 +90,6 @@ void AIEDialect::initialize() {
   addInterfaces<AIEInlinerInterface, AIEDialectFoldInterface>();
 }
 
-// Helper class to get a ShimDMAAllocationOp for a given <device, symbol name>
-//  pair. An object of this class is invalidated if, for any symbol_name, a
-//  ShimDMAAllocationOp that uses it changes, as the cache is not updated in
-//  this case.
-
-// Return the first ShimDMAAllocationOp nested inside the DeviceOp 'dev' that
-// uses the symbol 'sym_name'
-std::optional<xilinx::AIE::ShimDMAAllocationOp>
-xilinx::AIE::ShimDMAllocationGetter::get(DeviceOp dev, StringRef sym_name) {
-  auto key = std::make_pair(dev, sym_name);
-  auto it = allocGetter.find(key);
-  if (it != allocGetter.end()) {
-    return it->second;
-  }
-
-  // Call cachelessGet to look for the allocation operation
-  auto allocOp = cachelessGet(dev, sym_name);
-
-  // Only cache the value if it's not empty (i.e., not std::nullopt)
-  if (allocOp.has_value()) {
-    allocGetter[key] = allocOp; // Cache it
-  }
-
-  return allocOp; // Return the found or empty optional
-}
-
-// Finding the ShimDMAAllocationOp for a given <DeviceOp, symbol_name> pair
-// can be slow when the symbol is used in many places. This version of the
-// function is only called when the cache does not have a ShimDMAAllocationOp
-// stored from a previous lookup.
-std::optional<xilinx::AIE::ShimDMAAllocationOp>
-xilinx::AIE::ShimDMAllocationGetter::cachelessGet(DeviceOp dev,
-                                                  mlir::StringRef sym_name) {
-  auto *sym = dev.lookupSymbol(sym_name);
-  if (!sym)
-    return std::nullopt;
-  auto uses = SymbolTable::getSymbolUses(sym, dev);
-  for (auto use : *uses)
-    if (auto infoOp = dyn_cast<AIE::ShimDMAAllocationOp>(use.getUser()))
-      return infoOp;
-  return std::nullopt;
-}
-
 // Helper methods to retrieve the encoding associated to a burst length,
 // or to find the highest available burst length if the requested one is 0
 // (default value).
@@ -2367,11 +2324,10 @@ void BDChainOp::print(OpAsmPrinter &printer) {
 
 ShimDMAAllocationOp ShimDMAAllocationOp::getForSymbol(DeviceOp device,
                                                       llvm::StringRef symbol) {
-  auto alloc_ops = device.getOps<ShimDMAAllocationOp>();
-  for (auto it = alloc_ops.begin(); it != alloc_ops.end(); ++it) {
-    AIE::ShimDMAAllocationOp a = *it;
-    if (a.getSymName() == symbol) {
-      return a;
+  Operation *maybeOp = device.lookupSymbol(symbol);
+  if (maybeOp) {
+    if (ShimDMAAllocationOp op = dyn_cast<ShimDMAAllocationOp>(maybeOp)) {
+      return op;
     }
   }
   return nullptr;

lib/Dialect/AIE/Transforms/AIEGenerateColumnControlOverlay.cpp

@@ -363,13 +363,8 @@ struct AIEGenerateColumnControlOverlayPass
 
       builder.create<AIE::ShimDMAAllocationOp>(
           builder.getUnknownLoc(), StringRef(dma_name), dir,
-          rowToShimChanMap[tOp.rowIndex()], shimTile.colIndex(), false);
-      MemRefType memref_ty = MemRefType::get(
-          ArrayRef<int64_t>{2048}, IntegerType::get(builder.getContext(), 32),
-          nullptr, 0);
-      builder.create<memref::GlobalOp>(builder.getUnknownLoc(), dma_name,
-                                       builder.getStringAttr("public"),
-                                       memref_ty, nullptr, false, nullptr);
+          rowToShimChanMap[tOp.rowIndex()], shimTile.colIndex(), false,
+          nullptr);
     }
   }
 

lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp

@@ -1635,20 +1635,27 @@ struct AIEObjectFifoStatefulTransformPass
   /// shimDMAAllocationOp containing the channelDir, channelIndex and
   /// shimTile col assigned by the objectFifo lowering.
   void createObjectFifoAllocationInfo(OpBuilder &builder, MLIRContext *ctx,
-                                      FlatSymbolRefAttr obj_fifo, int colIndex,
-                                      DMAChannelDir channelDir,
+                                      ObjectFifoCreateOp &objFifoOp,
+                                      int colIndex, DMAChannelDir channelDir,
                                       int channelIndex, bool plio,
                                       std::optional<PacketInfoAttr> packet) {
     PacketInfoAttr packetInfo = nullptr;
     if (packet)
       packetInfo = *packet;
-    builder.create<ShimDMAAllocationOp>(builder.getUnknownLoc(), obj_fifo,
+    std::string alloc_name = getShimAllocationName(objFifoOp.getName());
+    // SymbolRefAttr::get(ctx, objFifoOp.getName())
+    builder.create<ShimDMAAllocationOp>(builder.getUnknownLoc(),
+                                        StringAttr::get(ctx, alloc_name),
                                         DMAChannelDirAttr::get(ctx, channelDir),
                                         builder.getI64IntegerAttr(channelIndex),
                                         builder.getI64IntegerAttr(colIndex),
                                         builder.getBoolAttr(plio), packetInfo);
   }
 
+  static std::string getShimAllocationName(llvm::StringRef objFifoName) {
+    return (objFifoName + "_shim_alloc").str();
+  }
+
   /// Function used to verify that an objectfifo is present in at most one
   /// ObjectFifoLinkOp.
   void verifyObjectFifoLinks(DeviceOp &device) {
@@ -1908,9 +1915,9 @@ struct AIEObjectFifoStatefulTransformPass
 
       if (producer.getProducerTileOp().isShimTile())
         createObjectFifoAllocationInfo(
-            builder, ctx, SymbolRefAttr::get(ctx, producer.getName()),
-            producer.getProducerTileOp().colIndex(), producerChan.direction,
-            producerChan.channel, producer.getPlio(), bdPacket);
+            builder, ctx, producer, producer.getProducerTileOp().colIndex(),
+            producerChan.direction, producerChan.channel, producer.getPlio(),
+            bdPacket);
 
       PacketFlowOp packetflow;
       if (clPacketSwObjectFifos) {
@@ -1963,9 +1970,9 @@ struct AIEObjectFifoStatefulTransformPass
 
         if (consumer.getProducerTileOp().isShimTile())
           createObjectFifoAllocationInfo(
-              builder, ctx, SymbolRefAttr::get(ctx, producer.getName()),
-              consumer.getProducerTileOp().colIndex(), consumerChan.direction,
-              consumerChan.channel, producer.getPlio(), {});
+              builder, ctx, producer, consumer.getProducerTileOp().colIndex(),
+              consumerChan.direction, consumerChan.channel, producer.getPlio(),
+              {});
 
         if (!clPacketSwObjectFifos) {
           // create flow
@@ -2233,34 +2240,41 @@ struct AIEObjectFifoStatefulTransformPass
       if (res.wasInterrupted())
         return signalPassFailure();
     }
-    // make global symbols to replace the to be erased ObjectFifoCreateOps
-    for (auto createOp : device.getOps<ObjectFifoCreateOp>()) {
-      builder.setInsertionPointToStart(device.getBody());
-      auto sym_name = createOp.getName();
-      createOp->setAttr(SymbolTable::getSymbolAttrName(),
-                        builder.getStringAttr("__erase_" + sym_name));
-      auto memrefType = llvm::cast<AIEObjectFifoType>(createOp.getElemType())
-                            .getElementType();
-      builder.create<memref::GlobalOp>(builder.getUnknownLoc(), sym_name,
-                                       builder.getStringAttr("public"),
-                                       memrefType, nullptr, false, nullptr);
-    }
 
     //===------------------------------------------------------------------===//
     // Remove old ops
     //===------------------------------------------------------------------===//
     SetVector<Operation *> opsToErase;
     device.walk([&](Operation *op) {
-      if (isa<ObjectFifoCreateOp, ObjectFifoLinkOp,
-              ObjectFifoRegisterExternalBuffersOp, ObjectFifoAcquireOp,
-              ObjectFifoSubviewAccessOp, ObjectFifoReleaseOp,
-              ObjectFifoAllocateOp>(op))
+      if (isa<ObjectFifoLinkOp, ObjectFifoRegisterExternalBuffersOp,
+              ObjectFifoAcquireOp, ObjectFifoSubviewAccessOp,
+              ObjectFifoReleaseOp, ObjectFifoAllocateOp>(op))
         opsToErase.insert(op);
     });
     SmallVector<Operation *> sorted{opsToErase.begin(), opsToErase.end()};
     computeTopologicalSorting(sorted);
     for (auto *op : llvm::reverse(sorted))
       op->erase();
+
+    //===------------------------------------------------------------------===//
+    // Replace any remaining uses of object fifo symbol with symbol of its shim
+    // dma allocation.
+    //===------------------------------------------------------------------===//
+    opsToErase.clear();
+    for (auto createOp : device.getOps<ObjectFifoCreateOp>()) {
+      std::string shimAllocName = getShimAllocationName(createOp.getName());
+      if (failed(SymbolTable::replaceAllSymbolUses(
+              createOp.getNameAttr(), builder.getStringAttr(shimAllocName),
+              device))) {
+        createOp.emitError(
+            "failed to replace symbol uses with shim allocation");
+        return signalPassFailure();
+      }
+      opsToErase.insert(createOp);
+    }
+    for (auto *op : opsToErase) {
+      op->erase();
+    }
   }
 };
 

lib/Dialect/AIEX/IR/AIEXDialect.cpp

@@ -453,10 +453,10 @@ LogicalResult AIEX::NpuDmaMemcpyNdOp::verify() {
   // even if it exceeds the maximum stride/wrap size of any one dimension,
   // and simply do not lower any data layout transformations, since there is
   // no other way to express this at the dma_memcpy_nd interface otherwise.
-  AIE::ShimDMAllocationGetter allocGetter;
   AIE::DeviceOp dev = getOperation()->getParentOfType<AIE::DeviceOp>();
-  if (auto allocOp = allocGetter.get(dev, getMetadata())) {
-    int col = allocOp->getCol();
+  if (auto allocOp = AIE::ShimDMAAllocationOp::getForSymbol(
+          dev, getMetadata().getRootReference())) {
+    int col = allocOp.getCol();
     bool skipTransformationChecks = isLinearTransferWithoutTransformation();
     if (failed(verifyStridesWraps(*this, buffer, col, 0, inputSizes,
                                   inputStrides, hardwareSizes, hardwareStrides,

lib/Dialect/AIEX/Transforms/AIECtrlPacketToDma.cpp

@@ -160,7 +160,8 @@ struct AIECtrlPacketToDmaPass : AIECtrlPacketToDmaBase<AIECtrlPacketToDmaPass> {
         const std::vector<int64_t> staticSizes = {1, 1, 1, batchIt->totalSize};
         const std::vector<int64_t> staticStrides = {0, 0, 0, 1};
 
-        StringRef metadata = builder.getStringAttr(batchIt->shimDmaAllocName);
+        SymbolRefAttr metadata =
+            SymbolRefAttr::get(builder.getContext(), batchIt->shimDmaAllocName);
         builder.create<NpuDmaMemcpyNdOp>(
             builder.getUnknownLoc(), newBlockArg, SmallVector<Value>{},
             SmallVector<Value>{}, SmallVector<Value>{}, ArrayRef(staticOffsets),