Attribute to choose shared memory module in ObjFIFO (Xilinx#1529)

Co-authored-by: AndraBisca <andrab@amd.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

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

@@ -1665,7 +1665,10 @@ def AIE_ObjectFifoCreateOp: AIE_Op<"objectfifo", [HasParent<"DeviceOp">, Symbol]
         BDDimLayoutArrayAttr:$dimensionsToStream,
         BDDimLayoutArrayArrayAttr:$dimensionsFromStreamPerConsumer,
         DefaultValuedAttr<BoolAttr, "false">:$via_DMA,
-        DefaultValuedAttr<BoolAttr, "false">:$plio
+        DefaultValuedAttr<BoolAttr, "false">:$plio,
+        // 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
   );
 
   let assemblyFormat = [{

lib/Dialect/AIE/IR/AIEDialect.cpp

@@ -487,6 +487,12 @@ 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");
+  }
+
   return success();
 }
 

lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp

@@ -229,6 +229,42 @@ 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) {
@@ -1108,10 +1144,15 @@ struct AIEObjectFifoStatefulTransformPass
                               createOp.getProducerTile());
 
       // if split, the necessary size for producer fifo might change
-      if (shared)
+      if (shared) {
+        checkAndApplyViaSharedMemAttribute(createOp, share_direction);
         createObjectFifoElements(builder, lockAnalysis, createOp,
                                  share_direction);
-      else {
+      } 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()));

python/dialects/aie.py

@@ -315,6 +315,15 @@ 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
+
 
 # Create an aie objectFifo_link between input and output objectFifos.
 class object_fifo_link(ObjectFifoLinkOp):

test/objectFifo-stateful-transform/via_shared_mem_test.mlir

@@ -0,0 +1,36 @@
+//===- via_shared_mem_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: module @viaSharedMem {
+// CHECK:   aie.device(npu1) {
+// CHECK:     memref.global "public" @of1 : memref<16xi32>
+// CHECK:     memref.global "public" @of0 : memref<16xi32>
+// CHECK:     %tile_1_2 = aie.tile(1, 2)
+// CHECK:     %tile_1_3 = aie.tile(1, 3)
+// 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, 0) {init = 1 : i32, sym_name = "of1_prod_lock"}
+// CHECK:     %of1_cons_lock = aie.lock(%tile_1_3, 1) {init = 0 : i32, sym_name = "of1_cons_lock"}
+// CHECK:     %of0_buff_0 = aie.buffer(%tile_1_2) {sym_name = "of0_buff_0"} : memref<16xi32> 
+// CHECK:     %of0_prod_lock = aie.lock(%tile_1_2, 0) {init = 1 : i32, sym_name = "of0_prod_lock"}
+// CHECK:     %of0_cons_lock = aie.lock(%tile_1_2, 1) {init = 0 : i32, sym_name = "of0_cons_lock"}
+// CHECK:   }
+// CHECK: }
+
+module @viaSharedMem {
+ aie.device(npu1) {
+    %tile12 = aie.tile(1, 2)
+    %tile13 = aie.tile(1, 3)
+
+    aie.objectfifo @of0 (%tile12, {%tile13}, 1 : i32) : !aie.objectfifo<memref<16xi32>>
+    aie.objectfifo @of1 (%tile12, {%tile13}, 1 : i32) {via_shared_mem = 1 : i32} : !aie.objectfifo<memref<16xi32>>
+ }
+}

test/python/objFifo.py

@@ -21,9 +21,11 @@
 
 # CHECK:  module {
 # CHECK:    aie.device(xcve2302) {
-# CHECK:      %tile_0_2 = aie.tile(0, 2)
+# CHECK:      %tile_0_0 = aie.tile(0, 0)
 # CHECK:      %tile_1_2 = aie.tile(1, 2)
-# CHECK:      aie.objectfifo @of0(%tile_0_2, {%tile_1_2}, 2 : i32) : !aie.objectfifo<memref<256xi32>>
+# 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_1_2, {%tile_1_3}, 2 : i32) {via_shared_mem = 1 : i32} : !aie.objectfifo<memref<256xi32>>
 # 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>
@@ -40,10 +42,13 @@ def objFifo_example():
     dev = Device(AIEDevice.xcve2302)
     dev_block = Block.create_at_start(dev.body_region)
     with InsertionPoint(dev_block):
-        S = tile(0, 2)
+        S = tile(0, 0)
         T_ = tile(1, 2)
+        C_ = tile(1, 3)
 
         of0 = object_fifo("of0", S, T_, 2, T.memref(256, T.i32()))
+        of1 = object_fifo("of1", T_, C_, 2, T.memref(256, T.i32()))
+        of1.set_via_shared_mem(ObjectFifoPort.Consume)
 
         C = Core(T_)
         bb = Block.create_at_start(C.body)