diff --git a/include/aie/Dialect/AIE/IR/AIEOps.td b/include/aie/Dialect/AIE/IR/AIEOps.td index 59f7731d49..424f13146f 100644 --- a/include/aie/Dialect/AIE/IR/AIEOps.td +++ b/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:$via_shared_mem, // memtile_repeat==0 means "do it once" and don't repeat - OptionalAttr:$memtile_repeat + OptionalAttr:$memtile_repeat, + OptionalAttr:$init_values ); let assemblyFormat = [{ @@ -1691,6 +1692,7 @@ def AIE_ObjectFifoCreateOp: AIE_Op<"objectfifo", [HasParent<"DeviceOp">, Symbol] `,` $elemNumber `)` attr-dict `:` $elemType + custom(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> getInitialValues(); mlir::StringAttr name() { return getOperation()->getAttrOfType( diff --git a/lib/Dialect/AIE/IR/AIEDialect.cpp b/lib/Dialect/AIE/IR/AIEDialect.cpp index d7fafb311a..065250e323 100644 --- a/lib/Dialect/AIE/IR/AIEDialect.cpp +++ b/lib/Dialect/AIE/IR/AIEDialect.cpp @@ -507,6 +507,25 @@ TileOp ObjectFifoCreateOp::getProducerTileOp() { return cast(getProducerTile().getDefiningOp()); } +std::vector> ObjectFifoCreateOp::getInitialValues() { + std::vector> initValuesVector; + if (getInitValues().has_value()) { + auto initValues = llvm::cast(getInitValues().value()); + auto fifo = llvm::cast(getElemType()); + auto elemType = llvm::cast(fifo.getElementType()); + int len = elemType.getNumElements(); + for (int i = 0; i < size(); i++) { + std::vector initBuffer; + for (int r = 0; r < len; r++) { + auto elem = initValues.getValues()[i * len + r]; + initBuffer.push_back(llvm::cast(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(numElem)) { + depth = llvm::dyn_cast(llvm::dyn_cast(numElem)[0]).getInt(); + } else { + depth = llvm::dyn_cast(numElem).getInt(); + } + auto objfifoType = llvm::cast(type.getValue()); + auto memrefTypeNoDepth = llvm::cast(objfifoType.getElementType()); + ArrayRef shape = memrefTypeNoDepth.getShape(); + std::vector 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(initValues)) + return parser.emitError(parser.getNameLoc()) + << "initial value should be an elements attribute"; + return success(); +} + } // namespace xilinx::AIE //===----------------------------------------------------------------------===// diff --git a/lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp b/lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp index 388f14a8f6..964603d886 100644 --- a/lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp +++ b/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(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( - 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(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( 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); diff --git a/test/objectFifo-stateful-transform/init_values_test.mlir b/test/objectFifo-stateful-transform/init_values_test.mlir new file mode 100644 index 0000000000..2576246c3f --- /dev/null +++ b/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> = 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> = 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> = dense<[[0, 1], [3, 4]]> + } +}