[WIP] Enable init values for Object FIFO

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

@@ -1678,7 +1678,8 @@ def AIE_ObjectFifoCreateOp: AIE_Op<"objectfifo", [HasParent<"DeviceOp">, Symbol]
         // 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
+        OptionalAttr<AIEI32Attr>:$memtile_repeat,
+        OptionalAttr<ElementsAttr>:$init_values
   );
 
   let assemblyFormat = [{
@@ -1691,6 +1692,7 @@ def AIE_ObjectFifoCreateOp: AIE_Op<"objectfifo", [HasParent<"DeviceOp">, Symbol]
         `,`
         $elemNumber
     `)` attr-dict `:` $elemType
+    custom<ObjectFifoInitValues>(ref($elemNumber), ref($elemType), $init_values)
   }];
 
   let hasVerifier = 1;
@@ -1706,6 +1708,7 @@ def AIE_ObjectFifoCreateOp: AIE_Op<"objectfifo", [HasParent<"DeviceOp">, Symbol]
     }
 
     TileOp getProducerTileOp();
+    std::vector<std::vector<int>> getInitialValues();
 
     mlir::StringAttr name() {
       return getOperation()->getAttrOfType<mlir::StringAttr>(

lib/Dialect/AIE/IR/AIEDialect.cpp

@@ -507,6 +507,25 @@ TileOp ObjectFifoCreateOp::getProducerTileOp() {
   return cast<TileOp>(getProducerTile().getDefiningOp());
 }
 
+std::vector<std::vector<int32_t>> ObjectFifoCreateOp::getInitialValues() {
+  std::vector<std::vector<int32_t>> initValuesVector;
+  if (getInitValues().has_value()) {
+    auto initValues = llvm::cast<mlir::ElementsAttr>(getInitValues().value());
+    auto fifo = llvm::cast<AIEObjectFifoType>(getElemType());
+    auto elemType = llvm::cast<MemRefType>(fifo.getElementType());
+    int len = elemType.getNumElements();
+    for (int i = 0; i < size(); i++) {
+      std::vector<int> initBuffer;
+      for (int r = 0; r < len; r++) {
+        auto elem = initValues.getValues<mlir::ElementsAttr>()[i * len + r];
+        initBuffer.push_back(llvm::cast<IntegerAttr>(elem).getInt());
+      }
+      initValuesVector.push_back(initBuffer);
+    }
+  }
+  return initValuesVector;
+}
+
 namespace xilinx::AIE {
 
 ParseResult parseObjectFifoProducerTile(OpAsmParser &parser,
@@ -591,6 +610,46 @@ void printObjectFifoConsumerTiles(OpAsmPrinter &printer, Operation *op,
   }
 }
 
+static void printObjectFifoInitValues(OpAsmPrinter &p, ObjectFifoCreateOp op,
+                                      Attribute numElem, TypeAttr type, Attribute initValues) {
+  if (op.getInitValues()) {
+    p << "= ";
+    p.printAttributeWithoutType(initValues);
+  }
+}
+
+static ParseResult parseObjectFifoInitValues(OpAsmParser &parser, Attribute numElem,
+                                             TypeAttr type, Attribute &initValues) {
+  int64_t depth;
+  if (isa<ArrayAttr>(numElem)) {
+    depth = llvm::dyn_cast<mlir::IntegerAttr>(llvm::dyn_cast<mlir::ArrayAttr>(numElem)[0]).getInt();
+  } else {
+    depth = llvm::dyn_cast<mlir::IntegerAttr>(numElem).getInt();
+  }
+  auto objfifoType = llvm::cast<AIEObjectFifoType>(type.getValue());
+  auto memrefTypeNoDepth = llvm::cast<MemRefType>(objfifoType.getElementType());
+  ArrayRef shape = memrefTypeNoDepth.getShape();
+  std::vector<int64_t> newShape = {depth};
+  for (auto d : shape)
+    newShape.push_back(d);
+  auto memrefType = MemRefType::get(ArrayRef(newShape), memrefTypeNoDepth.getElementType());
+
+  if (!memrefType.hasStaticShape())
+    return parser.emitError(parser.getNameLoc())
+           << "type should be static shaped memref, but got " << memrefType;
+
+  if (parser.parseOptionalEqual())
+    return success();
+
+  Type tensorType = mlir::memref::getTensorTypeFromMemRefType(memrefType);
+  if (parser.parseAttribute(initValues, tensorType))
+    return failure();
+  if (!llvm::isa<ElementsAttr>(initValues))
+    return parser.emitError(parser.getNameLoc())
+           << "initial value should be an elements attribute";
+  return success();
+}
+
 } // namespace xilinx::AIE
 
 //===----------------------------------------------------------------------===//

lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp

@@ -310,10 +310,11 @@ struct AIEObjectFifoStatefulTransformPass
       int of_elem_index = 0; // used to give objectFifo elements a symbolic name
       for (int i = 0; i < numElem; i++) {
         // create corresponding aie1 locks
+        int initValue = op.getInitValues().has_value() ? 1 : 0;
         int lockID = lockAnalysis.getLockID(creation_tile);
         assert(lockID >= 0 && "No more locks to allocate!");
         auto lock = builder.create<LockOp>(builder.getUnknownLoc(),
-                                           creation_tile, lockID, 0);
+                                           creation_tile, lockID, initValue);
         lock.getOperation()->setAttr(
             SymbolTable::getSymbolAttrName(),
             builder.getStringAttr(op.name().str() + "_lock_" +
@@ -323,10 +324,11 @@ struct AIEObjectFifoStatefulTransformPass
       }
     } else {
       // create corresponding aie2 locks
+      auto initValues = op.getInitialValues();
       int prodLockID = lockAnalysis.getLockID(creation_tile);
       assert(prodLockID >= 0 && "No more locks to allocate!");
       auto prodLock = builder.create<LockOp>(
-          builder.getUnknownLoc(), creation_tile, prodLockID, numElem);
+          builder.getUnknownLoc(), creation_tile, prodLockID, numElem - initValues.size());
       prodLock.getOperation()->setAttr(
           SymbolTable::getSymbolAttrName(),
           builder.getStringAttr(op.name().str() + "_prod_lock"));
@@ -335,7 +337,7 @@ struct AIEObjectFifoStatefulTransformPass
       int consLockID = lockAnalysis.getLockID(creation_tile);
       assert(consLockID >= 0 && "No more locks to allocate!");
       auto consLock = builder.create<LockOp>(builder.getUnknownLoc(),
-                                             creation_tile, consLockID, 0);
+                                             creation_tile, consLockID, initValues.size());
       consLock.getOperation()->setAttr(
           SymbolTable::getSymbolAttrName(),
           builder.getStringAttr(op.name().str() + "_cons_lock"));
@@ -415,12 +417,16 @@ struct AIEObjectFifoStatefulTransformPass
     for (int i = 0; i < numElem; i++) {
       // if shimTile external buffers are collected from input code
       // create as many locks as there are external buffers
+      mlir::ElementsAttr initValues = nullptr;
       if (!creation_tile.isShimTile()) {
+        if (op.getInitValues().has_value()) {
+          initValues = builder.getI32VectorAttr(ArrayRef(op.getInitialValues()[i]));
+        }
         auto buff = builder.create<BufferOp>(
             builder.getUnknownLoc(), elemType, creation_tile,
             builder.getStringAttr(op.name().str() + "_buff_" +
                                   std::to_string(of_elem_index)),
-            /*address*/ nullptr, /*initial_value*/ nullptr,
+            /*address*/ nullptr, initValues,
             /*mem_bank*/ nullptr);
         buffers.push_back(buff);
       }
@@ -1249,9 +1255,11 @@ struct AIEObjectFifoStatefulTransformPass
           createOp.setElemNumberAttr(
               builder.getI32IntegerAttr(createOp.size()));
         else {
-          int prodMaxAcquire = findObjectFifoSize(
-              device, createOp.getProducerTileOp(), createOp);
-          createOp.setElemNumberAttr(builder.getI32IntegerAttr(prodMaxAcquire));
+          if (!createOp.getInitValues().has_value()) {
+            int prodMaxAcquire = findObjectFifoSize(
+                device, createOp.getProducerTileOp(), createOp);
+            createOp.setElemNumberAttr(builder.getI32IntegerAttr(prodMaxAcquire));
+          }
         }
         createObjectFifoElements(builder, lockAnalysis, createOp,
                                  share_direction);

test/objectFifo-stateful-transform/init_values_test.mlir

@@ -0,0 +1,53 @@
+//===- init_values_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 @init {
+// CHECK:   aie.device(xcve2302) {
+// CHECK:     memref.global "public" @of0 : memref<3xi32>
+// CHECK:     %tile_1_2 = aie.tile(1, 2)
+// CHECK:     %tile_1_3 = aie.tile(1, 3)
+// CHECK:     %of0_buff_0 = aie.buffer(%tile_1_2) {sym_name = "of0_buff_0"} : memref<3xi32> = dense<[0, 1, 2]>
+// CHECK:     %of0_buff_1 = aie.buffer(%tile_1_2) {sym_name = "of0_buff_1"} : memref<3xi32> = dense<[3, 4, 5]>
+// CHECK:     %of0_prod_lock = aie.lock(%tile_1_2, 0) {init = 0 : i32, sym_name = "of0_prod_lock"}
+// CHECK:     %of0_cons_lock = aie.lock(%tile_1_2, 1) {init = 2 : i32, sym_name = "of0_cons_lock"}
+// CHECK:   }
+// CHECK:   aie.device(xcvc1902) {
+// CHECK:     memref.global "public" @of3 : memref<2xi32>
+// CHECK:     %tile_1_2 = aie.tile(1, 2)
+// CHECK:     %tile_1_3 = aie.tile(1, 3)
+// CHECK:     %of2_buff_0 = aie.buffer(%tile_1_2) {sym_name = "of2_buff_0"} : memref<2xi32> = dense<[0, 1]>
+// CHECK:     %of2_buff_1 = aie.buffer(%tile_1_2) {sym_name = "of2_buff_1"} : memref<2xi32> = dense<[3, 4]>
+// CHECK:     %of2_lock_0 = aie.lock(%tile_1_2, 0) {init = 1 : i32, sym_name = "of2_lock_0"}
+// CHECK:     %of2_lock_1 = aie.lock(%tile_1_2, 1) {init = 1 : i32, sym_name = "of2_lock_1"}
+// CHECK:   }
+// CHECK: }
+
+module @init {
+ aie.device(xcve2302) {
+    %tile12 = aie.tile(1, 2)
+    %tile13 = aie.tile(1, 3)
+
+    aie.objectfifo @of0 (%tile12, {%tile13}, 2 : i32) : !aie.objectfifo<memref<3xi32>> = dense<[[0, 1, 2], [3, 4, 5]]>
+ }
+//  aie.device(xcve2302) {
+//     %tile12 = aie.tile(1, 2)
+//     %tile13 = aie.tile(1, 3)
+
+//     aie.objectfifo @of1 (%tile12, {%tile13}, 2 : i32) : !aie.objectfifo<memref<2x2xi32>> = dense<[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]>
+//  }
+ aie.device(xcvc1902) {
+    %tile12 = aie.tile(1, 2)
+    %tile13 = aie.tile(1, 3)
+
+    aie.objectfifo @of2 (%tile12, {%tile13}, 2 : i32) : !aie.objectfifo<memref<2xi32>> = dense<[[0, 1], [3, 4]]>
+ }
+}