Small change

Xilinx · Sep 13, 2024 · ea55152 · ea55152
1 parent 9197287
commit ea55152
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 4 deletions.
diff --git a/python/api/dataflow/objectfifo.py b/python/api/dataflow/objectfifo.py
@@ -7,9 +7,13 @@
 """
 
 from ... import ir
+from ..._mlir_libs._aie import ObjectFifoSubviewType
 from ...extras.util import np_dtype_to_mlir_type
-from ...dialects._aie_enum_gen import ObjectFifoPort
+
+from ...dialects._aie_enum_gen import ObjectFifoPort 
+from ...dialects._aie_ops_gen import ObjectFifoSubviewAccessOp, ObjectFifoAcquireOp, objectfifo_release
 from ...dialects.aie import object_fifo
+
 from ...dialects.memref import MemRefType
 from ..resolvable import Resolvable
 from .endpoint import MyObjectFifoEndpoint
@@ -70,6 +74,8 @@ def resolve(
             self.__end2.get_tile(),
             self.__depth,
             memRef_ty,
+            loc=loc,
+            ip=ip,
         )
 
     def set_endpoint(self, endpoint, first=True):
@@ -99,7 +105,20 @@ def _acquire(
         assert (
             num_elem <= self.__depth
         ), "Cannot consume elements to exceed ObjectFifo depth"
-        return self.__op.acquire(port, num_elem)
+        dtype = np_dtype_to_mlir_type(self.__memref_type[1])
+        assert dtype != None
+        memRef_ty = MemRefType.get(shape=self.__memref_type[0], element_type=dtype)
+        subview_t = ObjectFifoSubviewType.get(memRef_ty)
+        acq = ObjectFifoAcquireOp(subview_t, port, self.name, num_elem)
+
+        objects = []
+        if acq.size.value == 1:
+            return ObjectFifoSubviewAccessOp(
+                memRef_ty, acq.subview, acq.size.value - 1
+            )
+        for i in range(acq.size.value):
+            objects.append(ObjectFifoSubviewAccessOp(memRef_ty, acq.subview, i))
+        return objects
 
     def _release(
         self, port: ObjectFifoPort, num_elem: int, loc=None, ip=None, context=None
@@ -108,4 +127,4 @@ def _release(
         assert (
             num_elem <= self.__depth
         ), "Cannot consume elements to exceed ObjectFifo depth"
-        self.op.release(port, num_elem)
+        objectfifo_release(port, self.name, num_elem, loc=loc, ip=ip)
diff --git a/python/dialects/aie.py b/python/dialects/aie.py
@@ -279,6 +279,8 @@ def __init__(
         dimensionsFromStreamPerConsumer=None,
         via_DMA=None,
         plio=None,
+        ip=None,
+        loc=None,
     ):
         self.datatype = datatype
         if not isinstance(consumerTiles, List):
@@ -287,7 +289,6 @@ def __init__(
             dimensionsFromStreamPerConsumer = []
         if dimensionsToStream is None:
             dimensionsToStream = []
-        int_ty = IntegerType.get_signless(32)
         of_Ty = TypeAttr.get(ObjectFifoType.get(datatype))
         super().__init__(
             sym_name=name,
@@ -299,6 +300,8 @@ def __init__(
             dimensionsFromStreamPerConsumer=dimensionsFromStreamPerConsumer,
             via_DMA=via_DMA,
             plio=plio,
+            ip=ip,
+            loc=loc,
         )
 
     def acquire(self, port, num_elem):