Committing some untested things

programming_examples/basic/passthrough_kernel/aie2.py

@@ -16,6 +16,7 @@
 from aie.api.kernels.binkernel import BinKernel
 from aie.api.phys.device import NPU1Col1
 from aie.api.program import MyProgram
+from aie.api.tensor import MyTensorType
 from aie.api.worker import MyWorker
 
 try:
@@ -28,19 +29,16 @@
 
 assert vector_size % 4 == 0
 line_size = vector_size // 4
-
-# TODO: clean up types
-inout_type = ((vector_size,), np.uint8)
-fifo_memref_type = ((line_size,), np.uint8)
+line_type = MyTensorType(np.uint8, (line_size,))
 
 # TODO: rely on depth inference
-of_in = MyObjectFifo(2, memref_type=fifo_memref_type)
-of_out = MyObjectFifo(2, memref_type=fifo_memref_type)
+of_in = MyObjectFifo(2, line_type)
+of_out = MyObjectFifo(2, line_type)
 
 passthrough_fn = BinKernel(
     "passThroughLine",
     "passThrough.cc.o",
-    [fifo_memref_type, fifo_memref_type, np.int32],
+    [line_type, line_type, np.int32],
 )
 
 

python/api/dataflow/endpoint.py

@@ -6,7 +6,7 @@
 * join
 """
 
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from ..phys.tile import MyTile
 from ..resolvable import Resolvable
 

python/api/dataflow/inout/inout.py

@@ -4,11 +4,12 @@
 """
 
 from abc import abstractmethod
+from typing import Callable
 
 from .... import ir
 from ....extras.util import np_dtype_to_mlir_type
 from ....dialects.aiex import runtime_sequence
-from ....dialects.memref import MemRefType
+from ...tensor import MyTensorType
 from ..endpoint import MyObjectFifoEndpoint
 from ..objectfifo import ObjectFifoHandle
 from ...phys.tile import MyTile
@@ -22,8 +23,8 @@ def get_fifos(self) -> list[ObjectFifoHandle]: ...
 class MyInOutProgram(InOutProgram):
     def __init__(
         self,
-        sequence_fn,  # TODO: needs a type
-        inout_types,  # TODO: needs a type
+        sequence_fn: Callable[..., None],
+        inout_types: list[MyTensorType],
         fifos=list[ObjectFifoHandle],
         coords: tuple[int, int] = (0, 0),  # TODO: how to get default
     ):
@@ -36,7 +37,7 @@ def __init__(
 
     def get_tile(self) -> MyTile:
         assert self.tile != None
-        return self.tile.op
+        return self.tile
 
     def get_fifos(self) -> list[ObjectFifoHandle]:
         return self.fifos
@@ -45,14 +46,8 @@ def resolve(
         self,
         loc: ir.Location = None,
         ip: ir.InsertionPoint = None,
-        context: ir.Context = None,
     ) -> None:
-        # TODO: cleanup types
-        my_memref_types = [
-            MemRefType.get(shape=t[0], element_type=np_dtype_to_mlir_type(t[1]))
-            for t in self.inout_types
-        ]
-
+        my_memref_types = [t.memref_type for t in self.inout_types]
         @runtime_sequence(*my_memref_types)
         def sequence(*args, **kwargs):
             self.sequence_fn(*args, *self.fifos, **kwargs)

python/api/dataflow/inout/simplefifoinout.py

@@ -5,13 +5,14 @@
 """
 
 import numpy as np
+from typing import Optional
 
 from .... import ir
-from ....extras.util import np_dtype_to_mlir_type
 from ....dialects.aiex import runtime_sequence, npu_dma_memcpy_nd, npu_sync, T
 from .inout import InOutProgram
 from ...phys.tile import MyTile
 from ..objectfifo import ObjectFifoHandle
+from ...tensor import MyTensorType
 
 
 class SimpleFifoInOutProgram(InOutProgram):
@@ -21,14 +22,17 @@ def __init__(
         bytes_in: int,
         fifo_out: ObjectFifoHandle,
         bytes_out: int,
-        in_sizes: list[int] = None,
-        in_strides: list[int] = None,
-        out_sizes: list[int] = None,
-        out_strides: list[int] = None,
-        dtype=np.uint8,  # TODO: needs type
+        in_sizes: Optional[list[int]] = None,
+        in_strides: Optional[list[int]] = None,
+        out_sizes: Optional[list[int]] = None,
+        out_strides: Optional[list[int]] = None,
+        dtype: np.generic = np.uint8,
     ):
         assert bytes_in % np.prod(fifo_in.obj_type[0]) == 0
         assert bytes_out % np.prod(fifo_out.obj_type[0]) == 0
+        assert bytes_in > 0
+        assert bytes_out > 0
+
         self.fifo_in = fifo_in
         self.fifo_out = fifo_out
         self.bytes_in = bytes_in
@@ -77,7 +81,7 @@ def __init__(
 
     def get_tile(self) -> MyTile:
         assert self.tile != None
-        return self.tile.op
+        return self.tile
 
     def get_fifos(self) -> list[ObjectFifoHandle]:
         return [self.fifo_in, self.fifo_out]
@@ -86,10 +90,9 @@ def resolve(
         self,
         loc: ir.Location = None,
         ip: ir.InsertionPoint = None,
-        context: ir.Context = None,
     ) -> None:
-        tensor_in_ty = T.memref(self.bytes_in, np_dtype_to_mlir_type(self.dtype))
-        tensor_out_ty = T.memref(self.bytes_out, np_dtype_to_mlir_type(self.dtype))
+        tensor_in_ty = MyTensorType(self.bytes_in, self.dtype).memref_type
+        tensor_out_ty = MyTensorType(self.bytes_out, self.dtype).memref_type
 
         @runtime_sequence(tensor_in_ty, tensor_out_ty)
         def sequence(inTensor, outTensor):

python/api/dataflow/objectfifo.py

@@ -7,6 +7,7 @@
 
 # Address circular dependency between MyObjectFifo and ObjectFifoHandle
 from __future__ import annotations
+from typing import Optional
 
 from ... import ir
 from ..._mlir_libs._aie import ObjectFifoSubviewType
@@ -21,26 +22,25 @@
 )
 from ...dialects.aie import object_fifo
 
-from ...dialects.memref import MemRefType
 from ..resolvable import Resolvable
 from .endpoint import MyObjectFifoEndpoint
-
+from ..tensor import MyTensorType
 
 class MyObjectFifo(Resolvable):
     __of_index = 0
 
     def __init__(
         self,
         depth: int = 1,
-        memref_type=None,
+        obj_type: MyTensorType = None,
         name: str = None,
         end1: MyObjectFifoEndpoint = None,
         end2: MyObjectFifoEndpoint = None,
-        dimensionsToStream=None,  # TODO(erika): needs a type
-        dimensionsFromStreamPerConsumer=None,  # TODO(erika): needs a type
+        dimensionsToStream=list[list[int]],                 # TODO(erika): needs a type
+        dimensionsFromStreamPerConsumer=list[list[int]],    # TODO(erika): needs a type
     ):
-        self.__depth: int = depth
-        self.__memref_type = memref_type
+        self.__depth = depth
+        self.__obj_type = obj_type
         self.end1: MyObjectFifoEndpoint = end1
         self.end2: MyObjectFifoEndpoint = end2
         self.dimensionToStream = dimensionsToStream
@@ -50,9 +50,9 @@ def __init__(
             self.name = f"myof{MyObjectFifo.__get_index()}"
         else:
             self.name = name
-        self.__op = None
-        self.__first = None
-        self.__second = None
+        self.__op: Optional[ObjectFifoCreateOp] = None
+        self.__first: ObjectFifoHandle = ObjectFifoHandle(self, True)
+        self.__second: ObjectFifoHandle = ObjectFifoHandle(self, False)
 
     @classmethod
     def __get_index(cls) -> int:
@@ -67,100 +67,79 @@ def op(self) -> ObjectFifoCreateOp:
 
     @property
     def first(self) -> ObjectFifoHandle:
-        if self.__first == None:
-            self.__first = ObjectFifoHandle(self, True)
         return self.__first
 
-    # TODO: type this
-    @property
-    def obj_type(self):
-        return self.__memref_type
-
     @property
     def second(self) -> ObjectFifoHandle:
-        if self.__second == None:
-            self.__second = ObjectFifoHandle(self, False)
         return self.__second
 
+    @property
+    def obj_type(self) -> MyTensorType:
+        return self.__obj_type
+
     def resolve(
         self,
         loc: ir.Location = None,
         ip: ir.InsertionPoint = None,
-        context: ir.Context = None,
     ) -> None:
-        if self.__op != None:
-            return
-        assert self.end1 != None, "ObjectFifo missing endpoint 1"
-        assert self.end2 != None, "ObjectFifo missing endpoint 2"
-        assert self.__memref_type != None, "ObjectFifo missing memref_type"
-        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)
-        self.__op = object_fifo(
-            self.name,
-            self.end1.get_tile(),
-            self.end2.get_tile(),
-            self.__depth,
-            memRef_ty,
-            dimensionsToStream=self.dimensionToStream,
-            dimensionsFromStreamPerConsumer=self.dimensionsFromStreamPerConsumer,
-            loc=loc,
-            ip=ip,
-        )
-
-    def _set_endpoint(self, endpoint, first=True):
+        if self.__op == None:
+            assert self.end1 != None, "ObjectFifo missing endpoint 1"
+            assert self.end2 != None, "ObjectFifo missing endpoint 2"
+            assert self.__memref_type != None, "ObjectFifo missing memref_type"
+            self.__op = object_fifo(
+                self.name,
+                self.end1.get_tile().op,
+                self.end2.get_tile().op,
+                self.__depth,
+                self.__obj_type.memref_type,
+                dimensionsToStream=self.dimensionToStream,
+                dimensionsFromStreamPerConsumer=self.dimensionsFromStreamPerConsumer,
+                loc=loc,
+                ip=ip,
+            )
+
+    def _set_endpoint(self, endpoint: MyObjectFifoEndpoint, first: bool = True) -> None:
         if first:
             assert self.end1 == None, "ObjectFifo already assigned endpoint 1"
             self.end1 = endpoint
         else:
             assert self.end2 == None, "ObjectFifo already assigned endpoint 2"
             self.end2 = endpoint
 
-    def _acquire(
-        self, port: ObjectFifoPort, num_elem: int, loc=None, ip=None, context=None
-    ):
+    def _acquire(self, port: ObjectFifoPort, num_elem: int, loc: ir.Location = None, ip: ir.InsertionPoint = None):
         assert num_elem > 0, "Must consume at least one element"
-        assert (
-            num_elem <= self.__depth
-        ), "Cannot consume elements to exceed ObjectFifo depth"
-        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)
+        assert num_elem <= self.__depth, "Cannot consume elements to exceed ObjectFifo depth"
+
+        subview_t = ObjectFifoSubviewType.get(self.__obj_type.memref_type)
+        acq = ObjectFifoAcquireOp(subview_t, port, self.name, num_elem, loc=loc, ip=ip)
 
         objects = []
         if acq.size.value == 1:
-            return ObjectFifoSubviewAccessOp(memRef_ty, acq.subview, acq.size.value - 1)
+            return ObjectFifoSubviewAccessOp(self.__obj_type.memref_type, acq.subview, acq.size.value - 1, loc=loc, ip=ip)
         for i in range(acq.size.value):
-            objects.append(ObjectFifoSubviewAccessOp(memRef_ty, acq.subview, i))
+            objects.append(ObjectFifoSubviewAccessOp(self.__obj_type.memref_type, acq.subview, i, loc=loc, ip=ip))
         return objects
 
-    def _release(
-        self, port: ObjectFifoPort, num_elem: int, loc=None, ip=None, context=None
-    ):
+    def _release(self, port: ObjectFifoPort, num_elem: int, loc: ir.Location = None, ip: ir.InsertionPoint = None):
         assert num_elem > 0, "Must consume at least one element"
-        assert (
-            num_elem <= self.__depth
-        ), "Cannot consume elements to exceed ObjectFifo depth"
+        assert num_elem <= self.__depth, "Cannot consume elements to exceed ObjectFifo depth"
         objectfifo_release(port, self.name, num_elem, loc=loc, ip=ip)
 
 
 class ObjectFifoHandle(Resolvable):
     def __init__(self, of: MyObjectFifo, is_first: bool):
-        self.__port = ObjectFifoPort.Produce if is_first else ObjectFifoPort.Consume
+        self.__port: ObjectFifoPort = ObjectFifoPort.Produce if is_first else ObjectFifoPort.Consume
         self.__is_first = is_first
         self.__object_fifo = of
 
-    def acquire(self, num_elem: int, loc=None, ip=None, context=None):
-        return self.__object_fifo._acquire(self.__port, num_elem)
+    def acquire(self, num_elem: int, loc: ir.Location = None, ip: ir.InsertionPoint = None):
+        return self.__object_fifo._acquire(self.__port, num_elem, loc=loc, ip=ip)
 
-    def release(self, num_elem: int, loc=None, ip=None, context=None):
-        return self.__object_fifo._release(self.__port, num_elem)
+    def release(self, num_elem: int, loc: ir.Location = None, ip: ir.InsertionPoint = None):
+        return self.__object_fifo._release(self.__port, num_elem, loc=loc, ip=ip)
 
-    # TODO: type this
     @property
-    def obj_type(self):
+    def obj_type(self) -> MyTensorType:
         return self.__object_fifo.obj_type
 
     @property
@@ -178,6 +157,5 @@ def resolve(
         self,
         loc: ir.Location = None,
         ip: ir.InsertionPoint = None,
-        context: ir.Context = None,
     ) -> None:
-        return self.__object_fifo.resolve(loc, ip, context)
+        return self.__object_fifo.resolve(loc=loc, ip=ip)