use WireBundle and ObjectFifoPort

python/dialects/_AIE_syntax.py

@@ -1,47 +1,11 @@
+from typing import List
+
 from ._AIE_ops_gen import *
 from ._AIE_enum_gen import *
 from ._AIE_util import *
 from .._mlir_libs._aieMlir import register_dialect, ObjectFifoType, ObjectFifoSubviewType
 
 
-# Helper function that returns the index of a named objectFifo port.
-def objectFifoPortToIndex(port):
-    if (port == "Produce"):
-        port_index = 0
-    elif (port == "Consume"):
-        port_index = 1
-    else:
-        port_index = -1
-    return port_index
-
-
-# Helper function that returns the index of a named WireBundle.
-def wireBundleToIndex(port):
-    if (port == "Core"):
-        port_index = 0
-    elif (port == "DMA"):
-        port_index = 1
-    elif (port == "FIFO"):
-        port_index = 2
-    elif (port == "South"):
-        port_index = 3
-    elif (port == "West"):
-        port_index = 4
-    elif (port == "North"):
-        port_index = 5
-    elif (port == "East"):
-        port_index = 6
-    elif (port == "PLIO"):
-        port_index = 7
-    elif (port == "NOC"):
-        port_index = 8
-    elif (port == "Trace"):
-        port_index = 9
-    else:
-        port_index = -1
-    return port_index
-
-
 def dim_tuple_attr_builder(wrap, stepsize):
     return Attribute.parse(f"#AIE.DimTuple<{wrap}, {stepsize}>")
 
@@ -124,7 +88,12 @@ def __init__(self, size, datatype, name=None):
 # depth examples: 2, [2,2,7]
 class OrderedObjectBuffer(ObjectFifoCreateOp):
     """Specialize ObjectFifoCreateOp class constructor to take python integers"""
-    def __init__(self, name, tile0, tile1, depth, datatype, dimensionsToStream = [], dimensionsFromStreamPerConsumer = []):
+    def __init__(self, name, tile0, tile1, depth, datatype, dimensionsToStream=None,
+                 dimensionsFromStreamPerConsumer=None):
+        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)
@@ -155,63 +124,42 @@ def __init__(self, ofIns, ofOuts):
             ofIns_sym.append(FlatSymbolRefAttr.get(o))
         for o in ofOuts:
             ofOuts_sym.append(FlatSymbolRefAttr.get(o))
-        ofIns_array = ArrayAttr.get(ofIns_sym)
-        ofOuts_array = ArrayAttr.get(ofOuts_sym)
-        super().__init__(fifoIns=ofIns_array, fifoOuts=ofOuts_array)
+        super().__init__(fifoIns=ofIns_sym, fifoOuts=ofOuts_sym)
 
 
 # Create an aie objectFifo acquire op of given number of elements with given memref datatype,
 # from objFifo with given name.
 class Acquire(ObjectFifoAcquireOp):
     """Specialize ObjectFifoAcquireOp class constructor to take python integers"""
     def __init__(self, of_name, port, num_elem, datatype):
-        of_sym = FlatSymbolRefAttr.get(of_name)
-        int_ty = IntegerType.get_signless(32)
-        int_port = IntegerAttr.get(int_ty, objectFifoPortToIndex(port))
         subview_Ty = ObjectFifoSubviewType.get(datatype)
         self.datatype = datatype
-        super().__init__(subview=subview_Ty, port=int_port, objFifo_name=of_sym, size=num_elem)
+        super().__init__(subview=subview_Ty, port=port, objFifo_name=of_name, size=num_elem)
 
     def acquiredElem(self):
         objects = []
         if self.size.value == 1:
-            return SubviewAccess(self.subview, self.size.value - 1, self.datatype)
+            return SubviewAccess(self.datatype, self.subview, self.size.value - 1)
         for i in range(self.size.value):
-            objects.append(SubviewAccess(self.subview, i, self.datatype))
+            objects.append(SubviewAccess(self.datatype, self.subview, i))
         return objects
 
 
-# Create an aie objectFifo access op on given subview with given memref datatype,
-# at given index.
-class SubviewAccess(ObjectFifoSubviewAccessOp):
-    """Rename ObjectFifoSubviewAccessOp class"""
-    def __init__(self, subview, index, datatype):
-        super().__init__(output=datatype, subview=subview, index=index)
+SubviewAccess = ObjectFifoSubviewAccessOp
 
-
-# Create an aie objectFifo release op of given number of elements from objFifo with given name.
-class Release(ObjectFifoReleaseOp):
-    """Specialize ObjectFifoReleaseOp class constructor to take python integers"""
-    def __init__(self, of_name, port, num_elem):
-        of_sym = FlatSymbolRefAttr.get(of_name)
-        int_ty = IntegerType.get_signless(32)
-        int_port = IntegerAttr.get(int_ty, objectFifoPortToIndex(port))
-        super().__init__(port=int_port, objFifo_name=of_sym, size=num_elem)
+Release = ObjectFifoReleaseOp
 
 
 # Create a flow between source and destination tile ports.
 class Flow(FlowOp):
     """Specialize FlowOp class constructor to take python integers"""
     def __init__(self, source, source_port, source_channel, dest, dest_port, dest_channel):
-        int_ty = IntegerType.get_signless(32)
-        sourceBundle = IntegerAttr.get(int_ty, wireBundleToIndex(source_port))
-        destBundle = IntegerAttr.get(int_ty, wireBundleToIndex(dest_port))
         super().__init__(
             source=source, 
-            sourceBundle=sourceBundle, 
+            sourceBundle=source_port,
             sourceChannel=source_channel, 
             dest=dest, 
-            destBundle=destBundle, 
+            destBundle=dest_port,
             destChannel=dest_channel
         )
 
@@ -220,20 +168,15 @@ def __init__(self, source, source_port, source_channel, dest, dest_port, dest_ch
 class PacketFlow(PacketFlowOp):
     """Specialize PacketFlowOp class constructor to take python integers"""
     def __init__(self, pkt_id, source, source_port, source_channel, dest, dest_port, dest_channel):
-        int8_ty = IntegerType.get_signless(8)
-        int_ty = IntegerType.get_signless(32)
-        int8_pkt_id = IntegerAttr.get(int8_ty, pkt_id)
-        sourceBundle = IntegerAttr.get(int_ty, wireBundleToIndex(source_port))
-        destBundle = IntegerAttr.get(int_ty, wireBundleToIndex(dest_port))
-        super().__init__(ID=int8_pkt_id)
+        super().__init__(ID=pkt_id)
         bb = Block.create_at_start(self.ports)
         with InsertionPoint(bb):
-            src = PacketSourceOp(source, sourceBundle, source_channel)
-            dest = PacketDestOp(dest, destBundle, dest_channel)
+            src = PacketSourceOp(source, source_port, source_channel)
+            dest = PacketDestOp(dest, dest_port, dest_channel)
             end = EndOp()
 
 
 #### Global Wrappers ####
 core = region_op(Core, terminator=lambda *args: EndOp())
 device = region_op(Device)
-forLoop = region_op(For, terminator=lambda *args: YieldOp([]))
+forLoop = region_op(For, terminator=YieldOp)

test/python/aie_ops.py

@@ -141,7 +141,7 @@ def objFifoAcquire():
         C = Core(tile1)
         bb = Block.create_at_start(C.body)
         with InsertionPoint(bb):    
-            acq = Acquire(of_name="of0", port="Consume", num_elem=1, datatype=memTy)
+            acq = Acquire(of_name="of0", port=ObjectFifoPort.Consume, num_elem=1, datatype=memTy)
             EndOp()
 
 
@@ -164,8 +164,8 @@ def objFifoSubviewAccess():
         C = Core(tile1)
         bb = Block.create_at_start(C.body)
         with InsertionPoint(bb):
-            acq = Acquire(of_name="of0", port="Consume", num_elem=1, datatype=memTy)
-            subview = SubviewAccess(subview=acq, index=0, datatype=memTy)
+            acq = Acquire(of_name="of0", port=ObjectFifoPort.Consume, num_elem=1, datatype=memTy)
+            subview = SubviewAccess(memTy, subview=acq, index=0)
             EndOp()
 
 
@@ -187,7 +187,7 @@ def objFifoRelease():
         C = Core(tile0)
         bb = Block.create_at_start(C.body)
         with InsertionPoint(bb):
-            acq = Release(of_name="of0", port="Produce", num_elem=1)
+            acq = Release(ObjectFifoPort.Produce, "of0", 1)
             EndOp()
 
 
@@ -199,7 +199,7 @@ def objFifoRelease():
 def flowOp():
     S = Tile(0, 0)
     T = Tile(0, 2)
-    Flow(T, "Trace", 0, S, "DMA", 1)
+    Flow(T, WireBundle.Trace, 0, S, WireBundle.DMA, 1)
 
 
 # CHECK-LABEL: packetFlowOp
@@ -213,4 +213,4 @@ def flowOp():
 def packetFlowOp():
     S = Tile(0, 0)
     T = Tile(0, 2)
-    PacketFlow(0, T, "Trace", 0, S, "DMA", 1)
+    PacketFlow(0, T, WireBundle.Trace, 0, S, WireBundle.DMA, 1)

test/python/codeRegion.py

@@ -51,6 +51,6 @@ def deviceBody():
         def coreBody():
             @forLoop(lowerBound = 0, upperBound = 10, step = 1)
             def loopBody():
-                elem0 = Acquire("of1", "Consume", 1, memRef_64_ty).acquiredElem()
+                elem0 = Acquire("of1", ObjectFifoPort.Consume, 1, memRef_64_ty).acquiredElem()
                 res = Call("test_func", [elem0], [int_ty])
-                Release("of1", "Consume", 1)
+                Release(ObjectFifoPort.Consume, "of1", 1)

test/python/core_ext_kernel.py

@@ -53,11 +53,11 @@ def core_ext_kernel():
 
         C = Core(T, "test.o")
         bb = Block.create_at_start(C.body)
-        with InsertionPoint(bb):    
+        with InsertionPoint(bb):
             loop = For(lowerBound = 0, upperBound = 10, step = 1)
             with InsertionPoint(loop.body):
-                elem0 = Acquire("of1", "Consume", 1, memRef_64_ty).acquiredElem()
+                elem0 = Acquire("of1", ObjectFifoPort.Consume, 1, memRef_64_ty).acquiredElem()
                 res = Call("test_func", [elem0, integerConstant(4)], [int_ty])
-                Release("of1", "Consume", 1)
+                Release(ObjectFifoPort.Consume, "of1", 1)
                 YieldOp([])
             EndOp()          

test/python/objFifo.py

@@ -40,7 +40,7 @@ def objFifo_example():
         C = Core(T)
         bb = Block.create_at_start(C.body)
         with InsertionPoint(bb):    
-            elem0 = Acquire("of0", "Consume", 1, memRef_ty).acquiredElem()
+            elem0 = Acquire("of0", ObjectFifoPort.Consume, 1, memRef_ty).acquiredElem()
             Store(10, elem0, 0)
-            Release("of0", "Consume", 1)
+            Release(ObjectFifoPort.Consume, "of0", 1)
             EndOp()