From 4629dcd3894339da7ed050d9f91d81a1da5c981e Mon Sep 17 00:00:00 2001 From: Andra Bisca Date: Tue, 20 Aug 2024 12:29:10 +0200 Subject: [PATCH] Attribute to choose shared memory module in ObjFIFO (#1529) Co-authored-by: AndraBisca Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- include/aie/Dialect/AIE/IR/AIEOps.td | 5 ++- lib/Dialect/AIE/IR/AIEDialect.cpp | 6 +++ .../AIEObjectFifoStatefulTransform.cpp | 45 ++++++++++++++++++- python/dialects/aie.py | 9 ++++ .../via_shared_mem_test.mlir | 36 +++++++++++++++ test/python/objFifo.py | 11 +++-- 6 files changed, 106 insertions(+), 6 deletions(-) create mode 100644 test/objectFifo-stateful-transform/via_shared_mem_test.mlir diff --git a/include/aie/Dialect/AIE/IR/AIEOps.td b/include/aie/Dialect/AIE/IR/AIEOps.td index d872e1640b..d04bcd7ded 100644 --- a/include/aie/Dialect/AIE/IR/AIEOps.td +++ b/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:$via_DMA, - DefaultValuedAttr:$plio + DefaultValuedAttr:$plio, + // via_shared_mem==0 means use producer tile's memory module + // via_shared_mem==1 means use consumer tile's memory module + OptionalAttr:$via_shared_mem ); let assemblyFormat = [{ diff --git a/lib/Dialect/AIE/IR/AIEDialect.cpp b/lib/Dialect/AIE/IR/AIEDialect.cpp index e5c9f5946a..eba7604370 100644 --- a/lib/Dialect/AIE/IR/AIEDialect.cpp +++ b/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(); } diff --git a/lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp b/lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp index b787e99e3f..90b915c9cd 100644 --- a/lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp +++ b/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(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 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(createOp.getElemNumber())) createOp.setElemNumberAttr( builder.getI32IntegerAttr(createOp.size())); diff --git a/python/dialects/aie.py b/python/dialects/aie.py index 52b3e90711..2c08a5a78e 100644 --- a/python/dialects/aie.py +++ b/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): diff --git a/test/objectFifo-stateful-transform/via_shared_mem_test.mlir b/test/objectFifo-stateful-transform/via_shared_mem_test.mlir new file mode 100644 index 0000000000..19140c3f26 --- /dev/null +++ b/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> + aie.objectfifo @of1 (%tile12, {%tile13}, 1 : i32) {via_shared_mem = 1 : i32} : !aie.objectfifo> + } +} diff --git a/test/python/objFifo.py b/test/python/objFifo.py index 6bf5baaee1..586a3929d3 100644 --- a/test/python/objFifo.py +++ b/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> +# CHECK: %tile_1_3 = aie.tile(1, 3) +# CHECK: aie.objectfifo @of0(%tile_0_0, {%tile_1_2}, 2 : i32) : !aie.objectfifo> +# CHECK: aie.objectfifo @of1(%tile_1_2, {%tile_1_3}, 2 : i32) {via_shared_mem = 1 : i32} : !aie.objectfifo> # CHECK: %core_1_2 = aie.core(%tile_1_2) { # CHECK: %0 = aie.objectfifo.acquire @of0(Consume, 1) : !aie.objectfifosubview> # CHECK: %1 = aie.objectfifo.subview.access %0[0] : !aie.objectfifosubview> -> 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)