Skip to content

Commit

Permalink
Attribute to force Object FIFO lowering via DMAs (#1444)
Browse files Browse the repository at this point in the history
  • Loading branch information
AndraBisca authored May 1, 2024
1 parent 6812f68 commit 7d48a21
Show file tree
Hide file tree
Showing 5 changed files with 87 additions and 5 deletions.
3 changes: 2 additions & 1 deletion include/aie/Dialect/AIE/IR/AIEOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -1609,7 +1609,8 @@ def AIE_ObjectFifoCreateOp: AIE_Op<"objectfifo", [HasParent<"DeviceOp">, Symbol]
AIE_ObjectFifo_Depth:$elemNumber,
TypeAttrOf<AIE_ObjectFifoType>:$elemType,
BDDimLayoutArrayAttr:$dimensionsToStream,
BDDimLayoutArrayArrayAttr:$dimensionsFromStreamPerConsumer
BDDimLayoutArrayArrayAttr:$dimensionsFromStreamPerConsumer,
DefaultValuedAttr<BoolAttr, "false">:$via_DMA
);

let assemblyFormat = [{
Expand Down
10 changes: 7 additions & 3 deletions lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -171,14 +171,18 @@ struct AIEObjectFifoStatefulTransformPass

// Return true if the objectFifo created by createOp requires a DMA to be set
// up. This is the case if the tiles are not adjacent (no shared memory), if
// the objectFifo broadcasts to multiple tiles, or if one of the consumers
// or the producer wants to use the multi-dimensional address generation
// features of the DMA.
// the objectFifo broadcasts to multiple tiles, if one of the consumers or
// the producer wants to use the multi-dimensional address generation
// features of the DMA, if the objectFifo is part of a LinkOp, or if the
// via_DMA attribute of the objectFifo is set.
bool requiresDMAs(ObjectFifoCreateOp createOp, int &share_direction) {
bool hasSharedMemory = false;
bool atLeastOneConsumerWantsTransform = false;
bool isUsedInLinkOp = false;

if (createOp.getVia_DMA())
return true;

if (createOp.getConsumerTiles().size() == 1 &&
createOp.getDimensionsToStream().empty()) {

Expand Down
2 changes: 2 additions & 0 deletions python/dialects/aie.py
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,7 @@ def __init__(
datatype,
dimensionsToStream=None,
dimensionsFromStreamPerConsumer=None,
via_DMA=None,
):
self.datatype = datatype
if not isinstance(consumerTiles, List):
Expand All @@ -227,6 +228,7 @@ def __init__(
elemType=of_Ty,
dimensionsToStream=dimensionsToStream,
dimensionsFromStreamPerConsumer=dimensionsFromStreamPerConsumer,
via_DMA=via_DMA,
)

def acquire(self, port, num_elem):
Expand Down
74 changes: 74 additions & 0 deletions test/objectFifo-stateful-transform/via_DMA_test.mlir
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
//===- via_DMA_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 @viaDMA {
// CHECK: aie.device(xcve2302) {
// CHECK: memref.global "public" @of_stream_cons : memref<16xi32>
// CHECK: memref.global "public" @of_stream : memref<16xi32>
// CHECK: memref.global "public" @of_shared : memref<16xi32>
// CHECK: %tile_1_2 = aie.tile(1, 2)
// CHECK: %tile_1_3 = aie.tile(1, 3)
// CHECK: %of_stream_cons_buff_0 = aie.buffer(%tile_1_3) {sym_name = "of_stream_cons_buff_0"} : memref<16xi32>
// CHECK: %of_stream_cons_buff_1 = aie.buffer(%tile_1_3) {sym_name = "of_stream_cons_buff_1"} : memref<16xi32>
// CHECK: %of_stream_cons_prod_lock = aie.lock(%tile_1_3, 0) {init = 2 : i32, sym_name = "of_stream_cons_prod_lock"}
// CHECK: %of_stream_cons_cons_lock = aie.lock(%tile_1_3, 1) {init = 0 : i32, sym_name = "of_stream_cons_cons_lock"}
// CHECK: %of_stream_buff_0 = aie.buffer(%tile_1_2) {sym_name = "of_stream_buff_0"} : memref<16xi32>
// CHECK: %of_stream_buff_1 = aie.buffer(%tile_1_2) {sym_name = "of_stream_buff_1"} : memref<16xi32>
// CHECK: %of_stream_prod_lock = aie.lock(%tile_1_2, 2) {init = 2 : i32, sym_name = "of_stream_prod_lock"}
// CHECK: %of_stream_cons_lock = aie.lock(%tile_1_2, 3) {init = 0 : i32, sym_name = "of_stream_cons_lock"}
// CHECK: %of_shared_buff_0 = aie.buffer(%tile_1_2) {sym_name = "of_shared_buff_0"} : memref<16xi32>
// CHECK: %of_shared_buff_1 = aie.buffer(%tile_1_2) {sym_name = "of_shared_buff_1"} : memref<16xi32>
// CHECK: %of_shared_prod_lock = aie.lock(%tile_1_2, 0) {init = 2 : i32, sym_name = "of_shared_prod_lock"}
// CHECK: %of_shared_cons_lock = aie.lock(%tile_1_2, 1) {init = 0 : i32, sym_name = "of_shared_cons_lock"}
// CHECK: aie.flow(%tile_1_2, DMA : 0, %tile_1_3, DMA : 0)
// CHECK: %mem_1_2 = aie.mem(%tile_1_2) {
// CHECK: %0 = aie.dma_start(MM2S, 0, ^bb1, ^bb3)
// CHECK: ^bb1: // 2 preds: ^bb0, ^bb2
// CHECK: aie.use_lock(%of_stream_cons_lock, AcquireGreaterEqual, 1)
// CHECK: aie.dma_bd(%of_stream_buff_0 : memref<16xi32>, 0, 16)
// CHECK: aie.use_lock(%of_stream_prod_lock, Release, 1)
// CHECK: aie.next_bd ^bb2
// CHECK: ^bb2: // pred: ^bb1
// CHECK: aie.use_lock(%of_stream_cons_lock, AcquireGreaterEqual, 1)
// CHECK: aie.dma_bd(%of_stream_buff_1 : memref<16xi32>, 0, 16)
// CHECK: aie.use_lock(%of_stream_prod_lock, Release, 1)
// CHECK: aie.next_bd ^bb1
// CHECK: ^bb3: // pred: ^bb0
// CHECK: aie.end
// CHECK: }
// CHECK: %mem_1_3 = aie.mem(%tile_1_3) {
// CHECK: %0 = aie.dma_start(S2MM, 0, ^bb1, ^bb3)
// CHECK: ^bb1: // 2 preds: ^bb0, ^bb2
// CHECK: aie.use_lock(%of_stream_cons_prod_lock, AcquireGreaterEqual, 1)
// CHECK: aie.dma_bd(%of_stream_cons_buff_0 : memref<16xi32>, 0, 16)
// CHECK: aie.use_lock(%of_stream_cons_cons_lock, Release, 1)
// CHECK: aie.next_bd ^bb2
// CHECK: ^bb2: // pred: ^bb1
// CHECK: aie.use_lock(%of_stream_cons_prod_lock, AcquireGreaterEqual, 1)
// CHECK: aie.dma_bd(%of_stream_cons_buff_1 : memref<16xi32>, 0, 16)
// CHECK: aie.use_lock(%of_stream_cons_cons_lock, Release, 1)
// CHECK: aie.next_bd ^bb1
// CHECK: ^bb3: // pred: ^bb0
// CHECK: aie.end
// CHECK: }
// CHECK: }
// CHECK: }

module @viaDMA {
aie.device(xcve2302) {
%tile12 = aie.tile(1, 2)
%tile13 = aie.tile(1, 3)

aie.objectfifo @of_shared (%tile12, {%tile13}, 2 : i32) : !aie.objectfifo<memref<16xi32>>
aie.objectfifo @of_stream (%tile12, {%tile13}, 2 : i32) {via_DMA = true} : !aie.objectfifo<memref<16xi32>>
}
}
3 changes: 2 additions & 1 deletion test/python/aie_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ def externalBufferOp():
# CHECK-LABEL: objFifo
# CHECK: %[[VAL0:.*]] = aie.tile(6, 6)
# CHECK: %[[VAL1:.*]] = aie.tile(2, 2)
# CHECK: aie.objectfifo @of0(%[[VAL0]] toStream [<size = 1, stride = 2>], {%[[VAL1]] fromStream [<size = 1, stride = 2>]}, 2 : i32) : !aie.objectfifo<memref<12xf16>>
# CHECK: aie.objectfifo @of0(%[[VAL0]] toStream [<size = 1, stride = 2>], {%[[VAL1]] fromStream [<size = 1, stride = 2>]}, 2 : i32) {via_DMA = true} : !aie.objectfifo<memref<12xf16>>
@construct_and_print_module
def objFifo():
dev = Device(AIEDevice.xcvc1902)
Expand All @@ -120,6 +120,7 @@ def objFifo():
T.memref(12, T.f16()),
[bd_dim_layout(size=1, stride=2)],
[[bd_dim_layout(size=1, stride=2)]],
via_DMA=True,
)
end()

Expand Down

0 comments on commit 7d48a21

Please sign in to comment.