Write doc strings for new IRON py API

python/iron/dataflow/endpoint.py

@@ -9,4 +9,6 @@
 
 
 class ObjectFifoEndpoint(Placeable):
+    """The endpoint of an ObjectFifo. Each ObjectFifoHandle has one ObjectFifoEndpoint"""
+
     pass

python/iron/device/device.py

@@ -6,26 +6,27 @@
 #
 # (c) Copyright 2024 Advanced Micro Devices, Inc.
 
-"""
-TODOs:
-* other devices types
-"""
-
 from abc import abstractmethod
 from ... import ir  # type: ignore
 from ...dialects.aie import AIEDevice, tile, TileOp  # type: ignore
 from ..resolvable import Resolvable
 from .tile import Tile
 
+# TODO: we need an NPU2 implementation.
+
 
 class Device(Resolvable):
     """
+    A base class for representations of a device of a specific type.
+
     Note: this class is abstract because it does not implement Resolve
     """
 
     class __DeviceTile(Resolvable):
         """
-        Interior class for tiles objects owned by a particular device
+        Interior class for tiles objects owned by a particular device.
+        This is needed to ensure we don't generate more than one MLIR operation corresponding
+        to the same logical tile within a device.
         """
 
         def __init__(self, col: int, row: int) -> None:
@@ -55,9 +56,19 @@ def op(self, op: TileOp):
             self._op = op
 
     def __init__(self, cols: int, rows: int) -> None:
+        """Initialize a representation of a device.
+
+        Args:
+            cols (int): Number of columns on the device
+            rows (int): Number of rows on the device.
+        """
         self._cols = cols
         self._rows = rows
         self._tiles: list[list[Device.__DeviceTile]] = []
+
+        # Create all "physical" tiles belonging to the device at initialization to
+        # ensure only one "physical" tile object is every created corresponding to the same
+        # coordinates.
         for c in range(self._cols):
             self._tiles.append([])
             for r in range(self._rows):
@@ -72,13 +83,32 @@ def cols(self) -> int:
         return self._cols
 
     @abstractmethod
-    def get_shim_tiles(self) -> list[Tile]: ...
+    def get_shim_tiles(self) -> list[Tile]:
+        """Returns a list of all shim tiles on the device.
+
+        Returns:
+            list[Tile]: A list of shim tiles.
+        """
+        ...
 
     @abstractmethod
-    def get_mem_tiles(self) -> list[Tile]: ...
+    def get_mem_tiles(self) -> list[Tile]:
+        """Returns a list of all mem tiles on the device.
+
+        Returns:
+            list[Tile]: A list of mem tiles.
+        """
+        ...
 
     @abstractmethod
-    def get_compute_tiles(self) -> list[Tile]: ...
+    def get_compute_tiles(self) -> list[Tile]:
+        """Returns a list of all compute tiles on the device.
+
+        Returns:
+            list[Tile]: A list of compute tiles.
+        """
+        # TODO: should this be shaped?
+        ...
 
     def resolve_tile(
         self,
@@ -91,6 +121,8 @@ def resolve_tile(
 
 
 class NPU1Col1(Device):
+    """A representation of a device that resolves to AIEDevice.npu1_1col"""
+
     def __init__(self) -> None:
         super().__init__(cols=1, rows=6)
 
@@ -112,6 +144,8 @@ def resolve(
 
 
 class NPU1Col2(Device):
+    """A representation of a device that resolves to AIEDevice.npu1_2col"""
+
     def __init__(self) -> None:
         super().__init__(cols=2, rows=6)
 
@@ -143,6 +177,8 @@ def resolve(
 
 
 class NPU1Col3(Device):
+    """A representation of a device that resolves to AIEDevice.npu1_3col"""
+
     def __init__(self) -> None:
         super().__init__(cols=3, rows=6)
 
@@ -174,6 +210,8 @@ def resolve(
 
 
 class NPU1Col4(Device):
+    """A representation of a device that resolves to AIEDevice.npu1_4col"""
+
     def __init__(self) -> None:
         super().__init__(cols=4, rows=6)
 
@@ -205,6 +243,10 @@ def resolve(
 
 
 class XCVC1902(Device):
+    """A placeholder representation of a device that resolves to IEDevice.xcvc1902
+    TODO: this needs to be implemented.
+    """
+
     def __init__(self) -> None:
         raise NotImplementedError("This device type is not yet implementated")
 

python/iron/device/tile.py

@@ -6,21 +6,17 @@
 #
 # (c) Copyright 2024 Advanced Micro Devices, Inc.
 
-"""
-TODOs:
-* docs
-* error handling
-* tile types
-"""
-
 from ...dialects.aie import TileOp
 
 
 class Tile:
+    """An object representing a single component denoted by coordinates on a device."""
+
     def __init__(self, col: int, row: int) -> None:
         self.col: int = col
         self.row: int = row
         self._op: TileOp | None = None
+        # TODO: each tile should probably have a type, e.g., Shim or Mem or Compute
 
     @property
     def op(self) -> TileOp:
@@ -47,14 +43,20 @@ def __hash__(self):
 
 
 class AnyShimTile:
+    """A placeholder that should be replaced with a concrete Tile() representing a Shim tile on a device."""
+
     pass
 
 
 class AnyMemTile:
+    """A placeholder that should be replaced with a concrete Tile() representing a Mem tile on a device."""
+
     pass
 
 
 class AnyComputeTile:
+    """A placeholder that should be replaced with a concrete Tile() representing a Compute tile on a device."""
+
     pass
 
 

python/iron/globalbuffer.py

@@ -20,6 +20,11 @@
 
 
 class GlobalBuffer(Resolvable, Placeable):
+    """A buffer that is available both to Workers and to the Runtime for operations.
+    This is often used for Runtime Parameters.
+    """
+
+    """This is used to generate unique names if none is given during construction"""
     __gbuf_index = 0
 
     def __init__(
@@ -30,6 +35,19 @@ def __init__(
         placement: PlacementTile | None = None,
         use_write_rtp: bool = False,
     ):
+        """A GlobalBuffer is a memory region declared at the top-level of the design, allowing it to
+        be accessed by both Workers and the Runtime.
+
+        Args:
+            type (type[np.ndarray] | None, optional): The type of the buffer. Defaults to None.
+            initial_value (np.ndarray | None, optional): An initial value to set the buffer to. Should be of same datatype and shape as the buffer. Defaults to None.
+            name (str | None, optional): The name of the buffer. If none is given, a unique name will be generated. Defaults to None.
+            placement (PlacementTile | None, optional): A placement location for the buffer. Defaults to None.
+            use_write_rtp (bool, optional): If use_write_rtp, write_rtp/read_rtp operations will be generated. Otherwise, traditional write/read operations will be used. Defaults to False.
+
+        Raises:
+            ValueError: Arguments are validated.
+        """
         if type is None and initial_value is None:
             raise ValueError("Must provide either type, initial value, or both.")
         if type is None:
@@ -51,10 +69,12 @@ def __get_index(cls) -> int:
 
     @property
     def shape(self) -> Sequence[int]:
+        """The shape of the buffer"""
         return np_ndarray_type_get_shape(self._obj_type)
 
     @property
     def dtype(self) -> np.dtype:
+        """The per-element datatype of the buffer."""
         return np_ndarray_type_get_dtype(self._obj_type)
 
     def __getitem__(self, idx):

python/iron/kernel.py

@@ -22,6 +22,14 @@ def __init__(
         bin_name: str,
         arg_types: list[type[np.ndarray] | np.dtype] = [],
     ) -> None:
+        """A Kernel is an externally defined function that eventually resolves to a FuncOp. If it is called,
+        a CallOp will be generated.
+
+        Args:
+            name (str): The name of the function
+            bin_name (str): The name of the binary (used for linking to a compute core)
+            arg_types (list[type[np.ndarray]  |  np.dtype], optional): The type signature of the function. Defaults to [].
+        """
         self._name = name
         self._bin_name = bin_name
         self._arg_types = arg_types

python/iron/localbuffer.py

@@ -14,6 +14,11 @@
 
 
 class LocalBuffer(buffer):
+    """A LocalBuffer will generate a BufferOp; it is called a LocalBuffer because it should only
+    be initialized within a function given to a Worker.
+    """
+
+    """This is used to ensure default names are uniquely given within a design"""
     __buf_tile_index = defaultdict(int)
 
     def __init__(
@@ -22,6 +27,16 @@ def __init__(
         initial_value: np.ndarray | None = None,
         name: str | None = None,
     ):
+        """Initialize a LocalBuffer. Placement will be set to the current core.
+
+        Args:
+            type (type[np.ndarray] | None, optional): The tensor type representing the buffer to allocate. Defaults to None.
+            initial_value (np.ndarray | None, optional): An initial value of the buffer. The type of this should match the given type. Defaults to None.
+            name (str | None, optional): The name of the buffer. If None is given, one will be autogenerated. Defaults to None.
+
+        Raises:
+            ValueError: Arguments are validated.
+        """
         if type is None and initial_value is None:
             raise ValueError("Must provide either type, initial value, or both.")
         if type is None:

python/iron/placeable.py

@@ -9,21 +9,49 @@
 
 
 class Placeable:
-    def __init__(self, tile: PlacementTile):
+    """Placeable is a base class of an object that might be Placed by a Placer."""
+
+    def __init__(self, tile: PlacementTile | None):
+        """Initialize a Placeable object.
+
+        Args:
+            tile (PlacementTile): A placeable object has a tile. This may be None during construction.
+        """
         self._tile = tile
 
     def place(self, tile: Tile) -> None:
+        """Place the object by assigning the object to a Tile.
+
+        Args:
+            tile (Tile): The placement tile.
+
+        Raises:
+            AlreadyPlacedError: If the object's tile is already set to a Tile object.
+        """
         if isinstance(self._tile, Tile):
             raise AlreadyPlacedError(self.__class__, self._tile, tile)
         self._tile = tile
 
     @property
-    def tile(self) -> PlacementTile:
+    def tile(self) -> PlacementTile | None:
+        """Return the tile of the placeable object.
+
+        Returns:
+            PlacementTile: The current placement of the object.
+        """
         return self._tile
 
 
 class AlreadyPlacedError(Exception):
+    """Placeable objects may raise this error if one attempts to assign them to a Tile more than once."""
+
     def __init__(self, cls, current_tile: Tile, new_tile: Tile):
+        """Create an AlreadyPlacedError
+
+        Args:
+            current_tile (Tile): The current placement tile
+            new_tile (Tile): The placement tile given for the second attempt to place the object.
+        """
         self.message = (
             f"{cls} already placed at {current_tile}; cannot place at {new_tile}"
         )

python/iron/placers.py

@@ -17,6 +17,9 @@
 
 
 class Placer(metaclass=ABCMeta):
+    """Placer is an abstract class to define the interface between the Program
+    and the Placer.
+    """
 
     @abstractmethod
     def make_placement(
@@ -25,10 +28,27 @@ def make_placement(
         rt: Runtime,
         workers: list[Worker],
         object_fifos: list[ObjectFifoHandle],
-    ): ...
+    ):
+        """Assign placement informatio to a program.
+
+        Args:
+            device (Device): The device to use for placement.
+            rt (Runtime): The runtime information for the program.
+            workers (list[Worker]): The workers included in the program.
+            object_fifos (list[ObjectFifoHandle]): The object fifos used by the program.
+        """
+        ...
 
 
 class SequentialPlacer(Placer):
+    """SequentialPlacer is a simple implementation of a placer. The SequentialPlacer is to named
+    because it will sequentially place workers to Compute Tiles. After workers are placed, Memory Tiles and
+    Shim Tiles are placed so as to match the column of the given compute tile.
+
+    The SequentialPlacer does not do any validation of placement and can often yield invalid placements
+    that exceed resource limits for channels, memory, etc. For complex or resource sensitive designs,
+    a more complex placer or manual placement is required.
+    """
 
     def __init__(self):
         super().__init__()
@@ -93,7 +113,7 @@ def make_placement(
 
     def _get_common_col(self, tiles: list[Tile]) -> int:
         """
-        This is a simplistic utility function that calculates a column that is "close" or "common"
+        A utility function that calculates a column that is "close" or "common"
         to a set of tiles. It is a simple heuristic using the average to represent "distance"
         """
         cols = [t.col for t in tiles if isinstance(t, Tile)]
@@ -103,6 +123,9 @@ def _get_common_col(self, tiles: list[Tile]) -> int:
         return avg_col
 
     def _find_col_match(self, col: int, tiles: list[Tile]) -> Tile:
+        """
+        A utility function that sequentially searches a list of tiles to find one with a matching column.
+        """
         for t in tiles:
             if t.col == col:
                 return t