[WIP] experiment, please ignore

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

@@ -1681,9 +1681,6 @@ def AIE_ObjectFifoCreateOp: AIE_Op<"objectfifo", [HasParent<"DeviceOp">, Symbol]
         // disable_synchronization==true will skip lock generation for 
         // objectfifo synchronous accesses
         DefaultValuedAttr<BoolAttr, "false">:$disable_synchronization,
-        // via_shared_mem==0 means use producer tile's memory module
-        // via_shared_mem==1 means use consumer tile's memory module
-        OptionalAttr<AIEI32Attr>:$via_shared_mem,
         // memtile_repeat==0 means "do it once" and don't repeat
         OptionalAttr<AIEI32Attr>:$memtile_repeat
   );
@@ -1738,6 +1735,31 @@ def AIE_ObjectFifoCreateOp: AIE_Op<"objectfifo", [HasParent<"DeviceOp">, Symbol]
   ];
 }
 
+def AIE_ObjectFifoAllocateOp: AIE_Op<"objectfifo.allocate", [HasParent<"DeviceOp">]> {
+  let summary = "Sets a tile's memory module as the target for an objectfifo's memory allocation";
+  let description = [{
+    TODO
+  }];
+
+  let arguments = (
+    ins FlatSymbolRefAttr:$objFifo_name,
+        Index:$delegateTile
+  );
+
+  let hasCustomAssemblyFormat = 1;
+
+  let assemblyFormat = [{
+    $objFifo_name `(` $delegateTile `)` attr-dict
+  }];
+
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+    ObjectFifoCreateOp getObjectFifo();
+    TileOp getDelegateTileOp();
+  }];
+}
+
 def AIE_ObjectFifoLinkOp: AIE_Op<"objectfifo.link", [HasParent<"DeviceOp">]> {
   let summary = "Links two objectFifos through an intermediary tile's DMA";
   let description = [{

lib/Dialect/AIE/IR/AIEDialect.cpp

@@ -488,12 +488,6 @@ LogicalResult ObjectFifoCreateOp::verify() {
         "on shim tile producers");
   }
 
-  if (getViaSharedMem().has_value()) {
-    if (getConsumerTiles().size() > 1)
-      return emitError(
-          "`via_shared_mem` can only be used in 1-to-1 object FIFOs");
-  }
-
   if (getMemtileRepeat().has_value()) {
     if (!getProducerTileOp().isMemTile())
       return emitError("`memtile_repeat` can only be used with a mem tile "
@@ -593,6 +587,35 @@ void printObjectFifoConsumerTiles(OpAsmPrinter &printer, Operation *op,
 
 } // namespace xilinx::AIE
 
+//===----------------------------------------------------------------------===//
+// ObjectFifoAllocateOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ObjectFifoAllocateOp::verify() {
+  ObjectFifoCreateOp objFifo = getObjectFifo();
+  if (!objFifo)
+    return emitError("cannot retrieve associated object FIFO");
+  if (objFifo.getConsumerTiles().size() != 1)
+    return emitError("can only be used in 1-to-1 object FIFOs");
+  return success();
+}
+
+TileOp ObjectFifoAllocateOp::getDelegateTileOp() {
+  return cast<TileOp>(getDelegateTile().getDefiningOp());
+}
+
+ObjectFifoCreateOp ObjectFifoAllocateOp::getObjectFifo() {
+  Operation *parent = getOperation();
+  while ((parent = parent->getParentOp())) {
+    if (parent->hasTrait<OpTrait::SymbolTable>()) {
+      if (auto *st = SymbolTable::lookupSymbolIn(parent, getObjFifoName());
+          isa_and_nonnull<ObjectFifoCreateOp>(st))
+        return dyn_cast<ObjectFifoCreateOp>(st);
+    }
+  }
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // ObjectFifoLinkOp
 //===----------------------------------------------------------------------===//
@@ -807,6 +830,8 @@ LogicalResult ObjectFifoAcquireOp::verify() {
 
   auto coreTile = parent.getTile();
   auto objFifo = getObjectFifo();
+  if (!objFifo)
+    return emitError("cannot retrieve associated object FIFO");
   if (getPort() == ObjectFifoPort::Produce) {
     if (coreTile != objFifo.getProducerTile())
       return parent.emitOpError(
@@ -865,6 +890,8 @@ LogicalResult ObjectFifoReleaseOp::verify() {
 
   auto coreTile = parent.getTile();
   auto objFifo = getObjectFifo();
+  if (!objFifo)
+    return emitError("cannot retrieve associated object FIFO");
   if (getPort() == ObjectFifoPort::Produce) {
     if (coreTile != objFifo.getProducerTile())
       return parent.emitOpError(

lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp

@@ -230,42 +230,6 @@ struct AIEObjectFifoStatefulTransformPass
            isUsedInLinkOp;
   }
 
-  // Checks if via_shared_mem attribute of the objectfifo is set and if so
-  // tries to apply it. If the desired shared memory module is available to
-  // both producer and consumer then it will be used, otherwise a warning is
-  // emitted and the original shared memory module is used instead.
-  void checkAndApplyViaSharedMemAttribute(ObjectFifoCreateOp createOp,
-                                          int &share_direction) {
-    if (createOp.getViaSharedMem().has_value()) {
-      int desiredSharedTile = createOp.getViaSharedMem().value();
-      int desiredSharedModule = 1;
-      if (desiredSharedTile == 0)
-        desiredSharedModule = -1;
-      if (share_direction != desiredSharedModule) {
-        bool desiredSharedModuleIsShared = false;
-        int newShareDirection = 0;
-        for (auto consumerTile : createOp.getConsumerTiles()) {
-          if (auto consumerTileOp =
-                  dyn_cast<TileOp>(consumerTile.getDefiningOp()))
-            if (share_direction == -1)
-              ///   * -1 if the shared memory module is that of the first input
-              ///   tile,
-              ///   * 1 if it is that of the second input tile
-              desiredSharedModuleIsShared =
-                  isSharedMemory(consumerTileOp, createOp.getProducerTileOp(),
-                                 &newShareDirection);
-        }
-        if (desiredSharedModuleIsShared) {
-          if (share_direction == newShareDirection)
-            share_direction = (share_direction == -1) ? 1 : -1;
-          else
-            createOp->emitWarning("Memory module specified by `via_shared_mem` "
-                                  "is not available as shared memory module");
-        }
-      }
-    }
-  }
-
   /// Function to retrieve ObjectFifoLinkOp of ObjectFifoCreateOp,
   /// if it belongs to one.
   std::optional<ObjectFifoLinkOp> getOptionalLinkOp(ObjectFifoCreateOp op) {
@@ -400,12 +364,30 @@ struct AIEObjectFifoStatefulTransformPass
     }
 
     TileOp creation_tile;
+    auto consumerTileOp =
+        dyn_cast<TileOp>(op.getConsumerTiles()[0].getDefiningOp());
     if (share_direction == 0 || share_direction == -1)
       creation_tile = op.getProducerTileOp();
-    else {
-      auto consumerTileOp =
-          dyn_cast<TileOp>(op.getConsumerTiles()[0].getDefiningOp());
+    else
       creation_tile = consumerTileOp;
+
+    ObjectFifoAllocateOp opAlloc;
+    auto device = op->getParentOfType<DeviceOp>();
+    for (ObjectFifoAllocateOp alloc : device.getOps<ObjectFifoAllocateOp>()) {
+      if (alloc.getObjectFifo() == op)
+        opAlloc = alloc;
+    }
+    if (opAlloc) {
+      TileOp delegate = opAlloc.getDelegateTileOp();
+      int prodShareDir;
+      int consShareDir;
+      isSharedMemory(delegate, op.getProducerTileOp(), &prodShareDir);
+      isSharedMemory(delegate, consumerTileOp, &consShareDir);
+      if (prodShareDir == -1 && consShareDir == -1)
+        creation_tile = delegate;
+      else
+        opAlloc.emitOpError("objectfifo has no shared memory access to "
+                            "delegate tile's memory module");
     }
 
     // Reset opbuilder location to after the last tile declaration
@@ -1415,14 +1397,9 @@ struct AIEObjectFifoStatefulTransformPass
 
       // if split, the necessary size for producer fifo might change
       if (shared) {
-        checkAndApplyViaSharedMemAttribute(createOp, share_direction);
         createObjectFifoElements(builder, lockAnalysis, createOp,
                                  share_direction);
       } else {
-        if (createOp.getViaSharedMem().has_value())
-          createOp->emitWarning("No access to shared memory module; ignoring "
-                                "`via_shared_mem`");
-
         if (isa<ArrayAttr>(createOp.getElemNumber()))
           createOp.setElemNumberAttr(
               builder.getI32IntegerAttr(createOp.size()));
@@ -1721,7 +1698,8 @@ struct AIEObjectFifoStatefulTransformPass
     device.walk([&](Operation *op) {
       if (isa<ObjectFifoCreateOp, ObjectFifoLinkOp,
               ObjectFifoRegisterExternalBuffersOp, ObjectFifoAcquireOp,
-              ObjectFifoSubviewAccessOp, ObjectFifoReleaseOp>(op))
+              ObjectFifoSubviewAccessOp, ObjectFifoReleaseOp,
+              ObjectFifoAllocateOp>(op))
         opsToErase.insert(op);
     });
     topologicalSort(opsToErase);

python/dialects/aie.py

@@ -419,14 +419,8 @@ def acquire(self, port, num_elem):
     def release(self, port, num_elem):
         return objectfifo_release(port, self.sym_name.value, num_elem)
 
-    def set_via_shared_mem(self, port):
-        num = 0
-        if port == ObjectFifoPort.Produce:
-            num = 0
-        elif port == ObjectFifoPort.Consume:
-            num = 1
-        int_num = IntegerAttr.get(T.i32(), num)
-        self.attributes["via_shared_mem"] = int_num
+    def allocate(self, tile):
+        return objectfifo_allocate(self.sym_name.value, tile)
 
     def set_memtile_repeat(self, num):
         int_num = IntegerAttr.get(T.i32(), num)

test/objectFifo-stateful-transform/allocate_test.mlir

@@ -0,0 +1,46 @@
+//===- allocate_test.mlir ---------------------------------------*- MLIR -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// Copyright (C) 2024, Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: aie-opt --aie-objectFifo-stateful-transform %s | FileCheck %s
+
+// CHECK-LABEL:   aie.device(xcve2302) {
+// CHECK:           memref.global "public" @of2 : memref<16xi32>
+// CHECK:           memref.global "public" @of1 : memref<16xi32>
+// CHECK:           memref.global "public" @of0 : memref<16xi32>
+// CHECK:           %[[VAL_0:.*]] = aie.tile(1, 2)
+// CHECK:           %[[VAL_1:.*]] = aie.tile(1, 3)
+// CHECK:           %[[VAL_2:.*]] = aie.tile(2, 2)
+// CHECK:           %of2_buff_0 = aie.buffer(%tile_1_2) {sym_name = "of2_buff_0"} : memref<16xi32> 
+// CHECK:           %of2_prod_lock = aie.lock(%tile_1_2, 0) {init = 1 : i32, sym_name = "of2_prod_lock"}
+// CHECK:           %of2_cons_lock = aie.lock(%tile_1_2, 1) {init = 0 : i32, sym_name = "of2_cons_lock"}
+// CHECK:           %of1_buff_0 = aie.buffer(%tile_1_3) {sym_name = "of1_buff_0"} : memref<16xi32> 
+// CHECK:           %of1_prod_lock = aie.lock(%tile_1_3, 2) {init = 1 : i32, sym_name = "of1_prod_lock"}
+// CHECK:           %of1_cons_lock = aie.lock(%tile_1_3, 3) {init = 0 : i32, sym_name = "of1_cons_lock"}
+// CHECK:           %of0_buff_0 = aie.buffer(%tile_1_3) {sym_name = "of0_buff_0"} : memref<16xi32> 
+// CHECK:           %of0_prod_lock = aie.lock(%tile_1_3, 0) {init = 1 : i32, sym_name = "of0_prod_lock"}
+// CHECK:           %of0_cons_lock = aie.lock(%tile_1_3, 1) {init = 0 : i32, sym_name = "of0_cons_lock"}
+// CHECK:         }
+
+module @same_core {
+    aie.device(xcve2302) {
+        %tile12 = aie.tile(1, 2)
+        %tile13 = aie.tile(1, 3)
+        %tile22 = aie.tile(2, 2)
+
+        aie.objectfifo @of0 (%tile12, {%tile12}, 1 : i32) : !aie.objectfifo<memref<16xi32>>
+        aie.objectfifo.allocate @of0 (%tile13)
+
+        aie.objectfifo @of1 (%tile12, {%tile13}, 1 : i32) : !aie.objectfifo<memref<16xi32>>
+        aie.objectfifo.allocate @of1 (%tile13)
+
+        aie.objectfifo @of2 (%tile12, {%tile22}, 1 : i32) : !aie.objectfifo<memref<16xi32>>
+        aie.objectfifo.allocate @of2 (%tile12)
+    }
+}

test/python/objFifo.py

@@ -25,7 +25,8 @@
 # CHECK:      %tile_1_3 = aie.tile(1, 3)
 # CHECK:      aie.objectfifo @of0(%tile_0_0, {%tile_1_2}, 2 : i32) : !aie.objectfifo<memref<256xi32>>
 # CHECK:      aie.objectfifo @of1(%tile_0_1, {%tile_1_2}, 2 : i32) {memtile_repeat = 4 : i32} : !aie.objectfifo<memref<256xi32>>
-# CHECK:      aie.objectfifo @of2(%tile_1_2, {%tile_1_3}, 2 : i32) {via_shared_mem = 1 : i32} : !aie.objectfifo<memref<256xi32>>
+# CHECK:      aie.objectfifo @of2(%tile_1_2, {%tile_1_3}, 2 : i32) : !aie.objectfifo<memref<256xi32>>
+# CHECK:      aie.objectfifo.allocate @of2(%tile_1_3)
 # CHECK:      %core_1_2 = aie.core(%tile_1_2) {
 # CHECK:        %0 = aie.objectfifo.acquire @of0(Consume, 1) : !aie.objectfifosubview<memref<256xi32>>
 # CHECK:        %1 = aie.objectfifo.subview.access %0[0] : !aie.objectfifosubview<memref<256xi32>> -> memref<256xi32>
@@ -53,7 +54,7 @@ def objFifo_example():
         of1 = object_fifo("of1", M, T_, 2, np.ndarray[(256,), np.dtype[np.int32]])
         of1.set_memtile_repeat(4)
         of2 = object_fifo("of2", T_, C_, 2, np.ndarray[(256,), np.dtype[np.int32]])
-        of2.set_via_shared_mem(ObjectFifoPort.Consume)
+        of2.allocate(C_)
 
         C = Core(T_)
         bb = Block.create_at_start(C.body)