Skip to content

Commit

Permalink
Cleanup wrappers for Object FIFO python bindings (#1091)
Browse files Browse the repository at this point in the history
  • Loading branch information
@AndraBisca
AndraBisca authored Mar 5, 2024
1 parent 0899ece commit c0a4daa
Show file tree
Hide file tree
Showing 20 changed files with 593 additions and 1,112 deletions.
78 changes: 46 additions & 32 deletions python/dialects/aie.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ def external_func(name, inputs, outputs=None, visibility="private"):


# Wrapper for func CallOp.
class Call(CallOp):
class call(CallOp):
"""Specialize CallOp class constructor to take python integers"""

def __init__(self, calleeOrResults, inputs=[], input_types=[]):
Expand Down Expand Up @@ -199,66 +199,80 @@ def __init__(self, shape, datatype, name=None, loc=None, ip=None):

# Create an aie objectFifo between specified tiles, with given depth and memref datatype.
# depth examples: 2, [2,2,7]
class OrderedObjectBuffer(ObjectFifoCreateOp):
class object_fifo(ObjectFifoCreateOp):
def __init__(
self,
name,
tile0,
tile1,
producerTile,
consumerTiles,
depth,
datatype,
dimensionsToStream=None,
dimensionsFromStreamPerConsumer=None,
):
self.datatype = datatype
if not isinstance(consumerTiles, List):
consumerTiles = [consumerTiles]
if dimensionsFromStreamPerConsumer is None:
dimensionsFromStreamPerConsumer = []
if dimensionsToStream is None:
dimensionsToStream = []
int_ty = IntegerType.get_signless(32)
if isinstance(depth, int):
int_depth = IntegerAttr.get(int_ty, depth)
else:
int_depths = []
for d in depth:
int_depths.append(IntegerAttr.get(int_ty, d))
int_depth = ArrayAttr.get(int_depths)
of_Ty = ObjectFifoType.get(datatype)
of_Ty = TypeAttr.get(ObjectFifoType.get(datatype))
super().__init__(
sym_name=name,
producerTile=tile0,
consumerTiles=tile1,
elemNumber=int_depth,
elem_type=TypeAttr.get(of_Ty),
producerTile=producerTile,
consumerTiles=consumerTiles,
elemNumber=depth,
elemType=of_Ty,
dimensionsToStream=dimensionsToStream,
dimensionsFromStreamPerConsumer=dimensionsFromStreamPerConsumer,
)

def acquire(self, port, num_elem):
subview_t = ObjectFifoSubviewType.get(self.datatype)
acq = ObjectFifoAcquireOp(subview_t, port, self.sym_name.value, num_elem)

# Create an aie objectFifo acquire op of given number of elements with given memref datatype,
# from objFifo with given name.
class ObjectFifoAcquireOp(ObjectFifoAcquireOp):
def __init__(self, port, of_name, num_elem, datatype):
subview_t = ObjectFifoSubviewType.get(datatype)
self.datatype = datatype
super().__init__(subview_t, port, of_name, num_elem)

def acquired_elem(self):
objects = []
if self.size.value == 1:
if acq.size.value == 1:
return ObjectFifoSubviewAccessOp(
self.datatype, self.subview, self.size.value - 1
self.datatype, acq.subview, acq.size.value - 1
)
for i in range(self.size.value):
objects.append(ObjectFifoSubviewAccessOp(self.datatype, self.subview, i))
for i in range(acq.size.value):
objects.append(ObjectFifoSubviewAccessOp(self.datatype, acq.subview, i))
return objects

def release(self, port, num_elem):
return objectfifo_release(port, self.sym_name.value, num_elem)

def acquire(port, of_name, num_elem, datatype):
return ObjectFifoAcquireOp(port, of_name, num_elem, datatype)

# Create an aie objectFifo_link between input and output objectFifos.
class object_fifo_link(ObjectFifoLinkOp):
"""Specialize ObjectFifoLinkOp class constructor to take python variables"""

def __init__(
self,
fifoIns,
fifoOuts,
):
if not isinstance(fifoIns, List):
fifoIns = [fifoIns]
if not isinstance(fifoOuts, List):
fifoOuts = [fifoOuts]
fifoInRefs = map(
lambda i: i if isinstance(i, str) else i.sym_name.value, fifoIns
)
fifoOutRefs = map(
lambda i: i if isinstance(i, str) else i.sym_name.value, fifoOuts
)
super().__init__(
fifoIns=fifoInRefs,
fifoOuts=fifoOutRefs,
)


# Create a packet flow between source and destination tile ports.
class PacketFlow(PacketFlowOp):
class packetflow(PacketFlowOp):
"""Specialize PacketFlowOp class constructor to take python integers"""

def __init__(
Expand Down
26 changes: 10 additions & 16 deletions reference_designs/ipu-xrt/add_one_objFifo/aie2.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,6 @@ def my_add_one_objFifo():
def device_body():
memRef_16_ty = T.memref(16, T.i32())
memRef_8_ty = T.memref(8, T.i32())
ofifo_memRef_16_ty = TypeAttr.get(ObjectFifoType.get(memRef_16_ty))
ofifo_memRef_8_ty = TypeAttr.get(ObjectFifoType.get(memRef_8_ty))

# Tile declarations
ShimTile = tile(0, 0)
Expand All @@ -29,14 +27,14 @@ def device_body():

# AIE-array data movement with object fifos
# Input
objectfifo("in0", ShimTile, [MemTile], 2, ofifo_memRef_16_ty, [], [])
objectfifo("in1", MemTile, [ComputeTile2], 2, ofifo_memRef_8_ty, [], [])
objectfifo_link(["in0"], ["in1"])
of_in0 = object_fifo("in0", ShimTile, MemTile, 2, memRef_16_ty)
of_in1 = object_fifo("in1", MemTile, ComputeTile2, 2, memRef_8_ty)
object_fifo_link(of_in0, of_in1)

# Output
objectfifo("out0", MemTile, [ShimTile], 2, ofifo_memRef_16_ty, [], [])
objectfifo("out1", ComputeTile2, [MemTile], 2, ofifo_memRef_8_ty, [], [])
objectfifo_link(["out1"], ["out0"])
of_out0 = object_fifo("out0", MemTile, ShimTile, 2, memRef_16_ty)
of_out1 = object_fifo("out1", ComputeTile2, MemTile, 2, memRef_8_ty)
object_fifo_link(of_out1, of_out0)

# Set up compute tiles

Expand All @@ -45,19 +43,15 @@ def device_body():
def core_body():
# Effective while(1)
for _ in for_(8):
elem_in = acquire(
ObjectFifoPort.Consume, "in1", 1, memRef_8_ty
).acquired_elem()
elem_out = acquire(
ObjectFifoPort.Produce, "out1", 1, memRef_8_ty
).acquired_elem()
elem_in = of_in1.acquire(ObjectFifoPort.Consume, 1)
elem_out = of_out1.acquire(ObjectFifoPort.Produce, 1)
for i in for_(8):
v0 = memref.load(elem_in, [i])
v1 = arith.addi(v0, arith.constant(1, T.i32()))
memref.store(v1, elem_out, [i])
yield_([])
objectfifo_release(ObjectFifoPort.Consume, "in1", 1)
objectfifo_release(ObjectFifoPort.Produce, "out1", 1)
of_in1.release(ObjectFifoPort.Consume, 1)
of_out1.release(ObjectFifoPort.Produce, 1)
yield_([])

# To/from AIE-array data movement
Expand Down
27 changes: 10 additions & 17 deletions reference_designs/ipu-xrt/log_hello_world/hello_world.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ def printf():
@device(AIEDevice.ipu)
def device_body():
memRef_ty = T.memref(N, T.i32())
ofifo_memRef_ty = TypeAttr.get(ObjectFifoType.get(memRef_ty))

# AIE Core Function declarations
kernel = external_func("kernel", inputs=[memRef_ty, memRef_ty, memRef_ty])
Expand All @@ -28,28 +27,22 @@ def device_body():
ComputeTile2 = tile(0, 2)

# AIE-array data movement with object fifos
objectfifo("inOF", ShimTile, [ComputeTile2], 2, ofifo_memRef_ty, [], [])
objectfifo("outOF", ComputeTile2, [ShimTile], 2, ofifo_memRef_ty, [], [])
objectfifo("logoutOF", ComputeTile2, [ShimTile], 2, ofifo_memRef_ty, [], [])
inOF = object_fifo("inOF", ShimTile, ComputeTile2, 2, memRef_ty)
outOF = object_fifo("outOF", ComputeTile2, ShimTile, 2, memRef_ty)
logoutOF = object_fifo("logoutOF", ComputeTile2, ShimTile, 2, memRef_ty)

# Set up compute tiles

# Compute tile 2
@core(ComputeTile2, "kernel.o")
def core_body():
elemOut = acquire(
ObjectFifoPort.Produce, "outOF", 1, memRef_ty
).acquired_elem()
elemIn = acquire(
ObjectFifoPort.Consume, "inOF", 1, memRef_ty
).acquired_elem()
elemLogout = acquire(
ObjectFifoPort.Produce, "logoutOF", 1, memRef_ty
).acquired_elem()
Call(kernel, [elemIn, elemOut, elemLogout])
objectfifo_release(ObjectFifoPort.Consume, "inOF", 1)
objectfifo_release(ObjectFifoPort.Produce, "outOF", 1)
objectfifo_release(ObjectFifoPort.Produce, "logoutOF", 1)
elemOut = outOF.acquire(ObjectFifoPort.Produce, 1)
elemIn = inOF.acquire(ObjectFifoPort.Consume, 1)
elemLogout = logoutOF.acquire(ObjectFifoPort.Produce, 1)
call(kernel, [elemIn, elemOut, elemLogout])
inOF.release(ObjectFifoPort.Consume, 1)
outOF.release(ObjectFifoPort.Produce, 1)
logoutOF.release(ObjectFifoPort.Produce, 1)

# To/from AIE-array data movement
@FuncOp.from_py_func(memRef_ty, memRef_ty, memRef_ty)
Expand Down
59 changes: 25 additions & 34 deletions reference_designs/ipu-xrt/matrix_multiplication/aie2.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,58 +83,55 @@ def device_body():

# AIE-array data movement with object fifos
# Input A
objectfifo("inA", shim_tile, [mem_tile], 2, ofifo_memref_a_ty, [], [])
objectfifo(
inA = object_fifo("inA", shim_tile, mem_tile, 2, memref_a_ty)
memA = object_fifo(
"memA",
mem_tile,
[compute_tile2],
compute_tile2,
2,
ofifo_memref_a_ty,
memref_a_ty,
[
(m // r, r * k * word_size_in // 4),
(k // s, s * word_size_in // 4),
(r, k * word_size_in // 4),
(s * word_size_in // 4, 1),
],
[],
)
objectfifo_link(["inA"], ["memA"])
object_fifo_link(inA, memA)

# Input B
objectfifo("inB", shim_tile, [mem_tile], 2, ofifo_memref_b_ty, [], [])
objectfifo(
inB = object_fifo("inB", shim_tile, mem_tile, 2, memref_b_ty)
memB = object_fifo(
"memB",
mem_tile,
[compute_tile2],
compute_tile2,
2,
ofifo_memref_b_ty,
memref_b_ty,
[
(k // s, s * n * word_size_in // 4),
(n // t, t * word_size_in // 4),
(s, n * word_size_in // 4),
(t * word_size_in // 4, 1),
],
[],
)
objectfifo_link(["inB"], ["memB"])
object_fifo_link(inB, memB)

# Output C
objectfifo("memC", compute_tile2, [mem_tile], 2, ofifo_memref_c_ty, [], [])
objectfifo(
memC = object_fifo("memC", compute_tile2, mem_tile, 2, memref_c_ty)
outC = object_fifo(
"outC",
mem_tile,
[shim_tile],
shim_tile,
2,
ofifo_memref_c_ty,
memref_c_ty,
[
(m // r, r * n * word_size_out // 4),
(r, t * word_size_out // 4),
(n // t, r * t * word_size_out // 4),
(t * word_size_out // 4, 1),
],
[],
)
objectfifo_link(["memC"], ["outC"])
object_fifo_link(memC, outC)

# Set up a circuit-switched flow from core to shim for tracing information
if enable_tracing:
Expand All @@ -147,30 +144,24 @@ def device_body():
def core_body():
for _ in for_(0xFFFFFFFF):
for _ in for_(tiles):
elem_out = acquire(
ObjectFifoPort.Produce, "memC", 1, memref_c_ty
).acquired_elem()
elem_out = memC.acquire(ObjectFifoPort.Produce, 1)
if vectorized:
Call(zero, [elem_out])
call(zero, [elem_out])
else:
Call(zero_scalar, [elem_out])
call(zero_scalar, [elem_out])

for _ in for_(K_div_k):
elem_in_a = acquire(
ObjectFifoPort.Consume, "memA", 1, memref_a_ty
).acquired_elem()
elem_in_b = acquire(
ObjectFifoPort.Consume, "memB", 1, memref_b_ty
).acquired_elem()
elem_in_a = memA.acquire(ObjectFifoPort.Consume, 1)
elem_in_b = memB.acquire(ObjectFifoPort.Consume, 1)
if vectorized:
Call(matmul, [elem_in_a, elem_in_b, elem_out])
call(matmul, [elem_in_a, elem_in_b, elem_out])
else:
Call(matmul_scalar, [elem_in_a, elem_in_b, elem_out])
objectfifo_release(ObjectFifoPort.Consume, "memA", 1)
objectfifo_release(ObjectFifoPort.Consume, "memB", 1)
call(matmul_scalar, [elem_in_a, elem_in_b, elem_out])
memA.release(ObjectFifoPort.Consume, 1)
memB.release(ObjectFifoPort.Consume, 1)
yield_([])

objectfifo_release(ObjectFifoPort.Produce, "memC", 1)
memC.release(ObjectFifoPort.Produce, 1)
yield_([])
yield_([])

Expand Down
Loading

0 comments on commit c0a4daa

Please sign in to comment.