Nov 21, 2025 · Nov 21, 2025 · Nov 21, 2025 · Nov 21, 2025 · Nov 21, 2025 · Nov 21, 2025
diff --git a/python/iron/device/device.py b/python/iron/device/device.py

 from ... import ir  # type: ignore
 from ...dialects._aie_enum_gen import WireBundle  # type: ignore
 from ...dialects.aie import AIEDevice, tile,TileOp,get_target_model  # type: ignore
 from ...dialects.aie import AIEDevice, tile, get_target_model  # type: ignore
 from ..resolvable import Resolvable
 from .tile import Tile

 import re


 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):
 class DeviceLike(Resolvable):
    class __DeviceTile(Tile, Resolvable):
        """
        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:
            self._col: int = col
            self._row: int = row
            self._op: TileOp | None = None
            super().__init__()
            Tile.__init__(self, col, row)
            Resolvable.__init__(self)

        def resolve(
            self,
        ) -> None:
            if self._op == None:
                self._op = tile(
                    self._col,
                    self._row,
                    self.col,
                    self.row,
                    loc=loc,
                    ip=ip,
                    allocation_scheme=allocation_scheme,
                )

        @property
        def op(self) -> TileOp:
            if not self._op:
                raise ValueError("Cannot get operation before it is set.")
            return self._op

        @op.setter
        def op(self, op: TileOp):
            if self._op:
                raise ValueError("Cannot set operation more than once.")
            self._op = op

    def __init__(self, device: AIEDevice) -> None:
        """Initialize a representation of a device.

        Args:
            device (AIEDevice): aie device
        """
    def __init__(
        self, device: AIEDevice, tiles: list[list[Tile]] | None = None
    ) -> None:
        self._device = device
        self._tiles: list[list[Device.__DeviceTile]] = []
        self._tm = get_target_model(device)
        for c in range(self._tm.columns()):
            self._tiles.append([])
            for r in range(self._tm.rows()):
                self._tiles[c].append(Device.__DeviceTile(c, r))

    def tile_iterator(self) -> Generator[Tile, None, None]:
        """
        Iterates over the device tiles deterministically
        """
        for c in range(self._tm.columns()):
            for r in range(self._tm.rows()):
                yield self._tiles[c][r]
        return None

    @property
    def rows(self) -> int:
        return self._tm.rows()

    @property
    def cols(self) -> int:
        return self._tm.columns()
        if tiles is None:
            self._tiles: list[list[DeviceLike.__DeviceTile]] = []
            for c in range(self._tm.columns()):
                self._tiles.append([])
                for r in range(self._tm.rows()):
                    self._tiles[c].append(DeviceLike.__DeviceTile(c, r))
        else:
            self._tiles = tiles
        self.ncols = len(self._tiles)
        self.nrows = len(self._tiles[0]) if self.ncols else 0

    def get_shim_tiles(self) -> list[Tile]:
        """Returns a list of all shim tiles on the device.
            list[Tile]: A list of shim tiles.
        """
        return [
 Tile(t._col, t._row)
 t
            for t in self.tile_iterator()
            if self._tm.is_shim_noc_or_pl_tile(t._col, t._row)
            if self._tm.is_shim_noc_or_pl_tile(t.col, t.row)
        ]

    def get_mem_tiles(self) -> list[Tile]:
        Returns:
            list[Tile]: A list of mem tiles.
        """
        return [
            Tile(t._col, t._row)
            for t in self.tile_iterator()
            if self._tm.is_mem_tile(t._col, t._row)
        ]
        return [t for t in self.tile_iterator() if self._tm.is_mem_tile(t.col, t.row)]

    def get_compute_tiles(self) -> list[Tile]:
        """Returns a list of all compute tiles on the device.
            list[Tile]: A list of compute tiles.
        """
        return [
            Tile(t._col, t._row)
            Tile(t.col, t.row)
            for t in self.tile_iterator()
            if self._tm.is_core_tile(t._col, t._row)
            if self._tm.is_core_tile(t.col, t.row)
        ]

    def get_num_source_switchbox_connections(self, t: Tile) -> int:
            int: Number of connections (channels) available on the tile
        """
        return self._tm.is_legal_mem_affinity(
            src_tile.col, src_tile.row, dst_tile.col, dst_tile.rol
            src_tile.col, src_tile.row, dst_tile.col, dst_tile.row
        )

    def resolve(
        self,
        loc: ir.Location | None = None,
        ip: ir.InsertionPoint | None = None,
    ) -> None:
        return self._device

    def resolve_tile(
        self,
        tile: Tile,
        loc: ir.Location | None = None,
        ip: ir.InsertionPoint | None = None,
    ) -> None:
        self._tiles[tile.col][tile.row].resolve(loc, ip, tile.allocation_scheme)
    ):
        self._tiles[tile.col][tile.row].resolve(
            loc, ip, getattr(tile, "allocation_scheme", None)
        )
        tile.op = self._tiles[tile.col][tile.row].op

    def tile_iterator(self) -> Generator[Tile, None, None]:
        """
        Iterates over the available device tiles deterministically
        """
        for c in range(self._tm.columns()):
            for r in range(self._tm.rows()):
                yield self._tiles[c][r]


 class DeviceView(DeviceLike):
    def __init__(self, device: "Device", tiles: list[Tile]):
        super().__init__(device=device._device, tiles=tiles)
        self._device_instance = device
        self._coords = set()
        for col in tiles:
            for t in col:
                self._coords.add((t.col, t.row))

    def tile_iterator(self) -> Generator[Tile, None, None]:
        # Keep ordering consistent from the device we sliced from.
        for c in range(len(self._tiles)):
            for r in range(len(self._tiles[0])):
                yield self._tiles[c][r]


 class Device(DeviceLike):
    """
    A base class for representations of a device of a specific type.
    """

    def __init__(self, device: AIEDevice) -> None:
        """Initialize a representation of a device.

        Args:
            device (AIEDevice): aie device
        """
        super().__init__(device=device)
        self._claimed_tiles = set()

    def __getitem__(self, key):
        if isinstance(key, tuple):
            if len(key) > 2:
                raise IndexError("Only 2D slicing is supported for devices.")
            if len(key) == 2:
                col_slice, row_slice = key
            else:
                col_slice, row_slice = key[0], slice(None, None, None)
        elif isinstance(key, (int, slice)):
            col_slice, row_slice = key, slice(None, None, None)
        else:
            raise IndexError(
                "Device indices must be integers, slices, or a 2-tuple of those."
            )

        if isinstance(col_slice, int) and isinstance(row_slice, int):
            if col_slice >= self._tm.columns() or row_slice >= self._tm.rows():
                raise IndexError("Tile index out of range.")

        # Handle slices and integers for cols
        if isinstance(col_slice, int):
            cols = [col_slice]
        else:
            cols = range(self._tm.columns())[col_slice]

        # Handle slices and integers for rows
        if isinstance(row_slice, int):
            rows = [row_slice]
        else:
            rows = range(self._tm.rows())[row_slice]

        if not cols or not rows:
            return DeviceView(self, [])

        tiles_to_claim = []
        coords_to_claim = set()
        for c in cols:
            tiles_to_claim.append([])
            for r in rows:
                if (c, r) in self._claimed_tiles:
                    raise ValueError(f"Tile ({c}, {r}) has already been claimed.")
                coords_to_claim.add((c, r))
                tiles_to_claim[-1].append(self._tiles[c][r])

        self._claimed_tiles.update(coords_to_claim)
        return DeviceView(self, tiles_to_claim)

    def tile_iterator(self) -> Generator[Tile, None, None]:
        for t in super().tile_iterator():
            if (t.col, t.row) not in self._claimed_tiles:
                yield self._tiles[t.col][t.row]


 def create_class(class_name, device):

    def _device__init__(self) -> None:
        super(globals()[class_name], self).__init__(device=device)

    def _device_resolve(
        self,
        loc: ir.Location | None = None,
        ip: ir.InsertionPoint | None = None,
    ) -> None:
        return device

    globals()[class_name] = type(
        class_name,
        (Device,),
        {
            "__init__": _device__init__,
            "resolve": _device_resolve,
            "__doc__": f"A representation of a device that resolves to {device}",
        },
    )
diff --git a/python/iron/hostruntime/config.py b/python/iron/hostruntime/config.py
                "RyzenAI-npu6",
            ]
        ):
            return NPU2()
            return NPU2
        elif any(
            keyword.lower() in output.lower()
            for keyword in [
                "RyzenAI-npu1",
            ]
        ):
            return NPU1()
            return NPU1
        else:
            raise RuntimeError("No supported NPU device found.")

    global config
    if "device" not in config:
        config["device"] = detect_npu_device()

    return config["device"]
    return config["device"]() if config["device"] else None
diff --git a/python/iron/localbuffer.py b/python/iron/localbuffer.py
 from collections import defaultdict
 import numpy as np

 from .. import ir
 from ..ir import Module, Context, InsertionPoint, Location, Operation
 from ..dialects.aie import buffer
 from .worker import Worker
 from .program import CurrentDeviceOp, CurrentModule
 from .device import Tile


 class LocalBuffer(buffer):
            idx = self.__get_index(current_core_placement)
            name = f"buf_{col}_{row}_{idx}"

        # if initial_value is None:
        device_op = get_device_op_from_module()
        block = device_op.regions[0].blocks[0]

        # Count number of tile ops and track last one
        tile_count = 0
        last_tile_op = None
        for op in block.operations:
            if op.name == "aie.tile":
                tile_count += 1
                last_tile_op = op

        ip = None if tile_count == 1 else InsertionPoint(last_tile_op)

        super().__init__(
            tile=current_core_placement,
            datatype=type,
            name=name,
            initial_value=initial_value,
            ip=ip,
        )

    @classmethod
    def __get_index(cls, placement: Tile) -> int:
    def __get_index(cls, placement):
        idx = cls.__buf_tile_index[placement]
        cls.__buf_tile_index[placement] += 1
        return idx


 def get_device_op_from_module():
    module = CurrentModule.get()
    if module is None:
        raise RuntimeError("No MLIR module context available")
    # The device op should be the first (or only) op in the module's body
    for op in module.body.operations:
        if op.name == "aie.device":
            return op
    raise RuntimeError("No aie.device op found in module")
Original file line number	Diff line number	Diff line change
Expand Up		@@ -11,32 +11,24 @@

		from ... import ir # type: ignore
		from ...dialects._aie_enum_gen import WireBundle # type: ignore
		from ...dialects.aie import AIEDevice, tile,TileOp,get_target_model # type: ignore
		from ...dialects.aie import AIEDevice, tile, get_target_model # type: ignore
		from ..resolvable import Resolvable
		from .tile import Tile

		import re


		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):
		class DeviceLike(Resolvable):
		class __DeviceTile(Tile, Resolvable):
		"""
		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:
		self._col: int = col
		self._row: int = row
		self._op: TileOp \| None = None
		super().__init__()
		Tile.__init__(self, col, row)
		Resolvable.__init__(self)

		def resolve(
		self,
Expand All		@@ -46,55 +38,28 @@ def resolve(
		) -> None:
		if self._op == None:
		self._op = tile(
		self._col,
		self._row,
		self.col,
		self.row,
		loc=loc,
		ip=ip,
		allocation_scheme=allocation_scheme,
		)

		@property
		def op(self) -> TileOp:
		if not self._op:
		raise ValueError("Cannot get operation before it is set.")
		return self._op

		@op.setter
		def op(self, op: TileOp):
		if self._op:
		raise ValueError("Cannot set operation more than once.")
		self._op = op

		def __init__(self, device: AIEDevice) -> None:
		"""Initialize a representation of a device.

		Args:
		device (AIEDevice): aie device
		"""
		def __init__(
		self, device: AIEDevice, tiles: list[list[Tile]] \| None = None
		) -> None:
		self._device = device
		self._tiles: list[list[Device.__DeviceTile]] = []
		self._tm = get_target_model(device)
		for c in range(self._tm.columns()):
		self._tiles.append([])
		for r in range(self._tm.rows()):
		self._tiles[c].append(Device.__DeviceTile(c, r))

		def tile_iterator(self) -> Generator[Tile, None, None]:
		"""
		Iterates over the device tiles deterministically
		"""
		for c in range(self._tm.columns()):
		for r in range(self._tm.rows()):
		yield self._tiles[c][r]
		return None

		@property
		def rows(self) -> int:
		return self._tm.rows()

		@property
		def cols(self) -> int:
		return self._tm.columns()
		if tiles is None:
		self._tiles: list[list[DeviceLike.__DeviceTile]] = []
		for c in range(self._tm.columns()):
		self._tiles.append([])
		for r in range(self._tm.rows()):
		self._tiles[c].append(DeviceLike.__DeviceTile(c, r))
		else:
		self._tiles = tiles
		self.ncols = len(self._tiles)
		self.nrows = len(self._tiles[0]) if self.ncols else 0

		def get_shim_tiles(self) -> list[Tile]:
		"""Returns a list of all shim tiles on the device.
Expand All		@@ -103,9 +68,9 @@ def get_shim_tiles(self) -> list[Tile]:
		list[Tile]: A list of shim tiles.
		"""
		return [
		Tile(t._col, t._row)
		t
		for t in self.tile_iterator()
		if self._tm.is_shim_noc_or_pl_tile(t._col, t._row)
		if self._tm.is_shim_noc_or_pl_tile(t.col, t.row)
		]

		def get_mem_tiles(self) -> list[Tile]:
Expand All		@@ -114,11 +79,7 @@ def get_mem_tiles(self) -> list[Tile]:
		Returns:
		list[Tile]: A list of mem tiles.
		"""
		return [
		Tile(t._col, t._row)
		for t in self.tile_iterator()
		if self._tm.is_mem_tile(t._col, t._row)
		]
		return [t for t in self.tile_iterator() if self._tm.is_mem_tile(t.col, t.row)]

		def get_compute_tiles(self) -> list[Tile]:
		"""Returns a list of all compute tiles on the device.
Expand All		@@ -127,9 +88,9 @@ def get_compute_tiles(self) -> list[Tile]:
		list[Tile]: A list of compute tiles.
		"""
		return [
		Tile(t._col, t._row)
		Tile(t.col, t.row)
hunhoffe marked this conversation as resolved. Show resolvedHide resolved
		for t in self.tile_iterator()
		if self._tm.is_core_tile(t._col, t._row)
		if self._tm.is_core_tile(t.col, t.row)
		]

		def get_num_source_switchbox_connections(self, t: Tile) -> int:
Expand DownExpand Up		@@ -197,37 +158,129 @@ def is_legal_mem_affinity(self, src_tile: Tile, dst_tile: Tile) -> bool:
		int: Number of connections (channels) available on the tile
		"""
		return self._tm.is_legal_mem_affinity(
		src_tile.col, src_tile.row, dst_tile.col, dst_tile.rol
		src_tile.col, src_tile.row, dst_tile.col, dst_tile.row
		)

		def resolve(
		self,
		loc: ir.Location \| None = None,
		ip: ir.InsertionPoint \| None = None,
		) -> None:
		return self._device

		def resolve_tile(
		self,
		tile: Tile,
		loc: ir.Location \| None = None,
		ip: ir.InsertionPoint \| None = None,
		) -> None:
		self._tiles[tile.col][tile.row].resolve(loc, ip, tile.allocation_scheme)
		):
		self._tiles[tile.col][tile.row].resolve(
		loc, ip, getattr(tile, "allocation_scheme", None)
		)
		tile.op = self._tiles[tile.col][tile.row].op

		def tile_iterator(self) -> Generator[Tile, None, None]:
		"""
		Iterates over the available device tiles deterministically
		"""
		for c in range(self._tm.columns()):
		for r in range(self._tm.rows()):
		yield self._tiles[c][r]


		class DeviceView(DeviceLike):
		def __init__(self, device: "Device", tiles: list[Tile]):
		super().__init__(device=device._device, tiles=tiles)
		self._device_instance = device
		self._coords = set()
		for col in tiles:
		for t in col:
hunhoffe marked this conversation as resolved. OutdatedShow resolvedHide resolved
		self._coords.add((t.col, t.row))

		def tile_iterator(self) -> Generator[Tile, None, None]:
		# Keep ordering consistent from the device we sliced from.
		for c in range(len(self._tiles)):
		for r in range(len(self._tiles[0])):
		yield self._tiles[c][r]


		class Device(DeviceLike):
		"""
		A base class for representations of a device of a specific type.
		"""

		def __init__(self, device: AIEDevice) -> None:
		"""Initialize a representation of a device.

		Args:
		device (AIEDevice): aie device
		"""
		super().__init__(device=device)
		self._claimed_tiles = set()

		def __getitem__(self, key):
		if isinstance(key, tuple):
		if len(key) > 2:
		raise IndexError("Only 2D slicing is supported for devices.")
		if len(key) == 2:
		col_slice, row_slice = key
		else:
		col_slice, row_slice = key[0], slice(None, None, None)
		elif isinstance(key, (int, slice)):
		col_slice, row_slice = key, slice(None, None, None)
		else:
		raise IndexError(
		"Device indices must be integers, slices, or a 2-tuple of those."
		)

		if isinstance(col_slice, int) and isinstance(row_slice, int):
		if col_slice >= self._tm.columns() or row_slice >= self._tm.rows():
		raise IndexError("Tile index out of range.")

		# Handle slices and integers for cols
		if isinstance(col_slice, int):
		cols = [col_slice]
		else:
		cols = range(self._tm.columns())[col_slice]

		# Handle slices and integers for rows
		if isinstance(row_slice, int):
		rows = [row_slice]
		else:
		rows = range(self._tm.rows())[row_slice]

		if not cols or not rows:
		return DeviceView(self, [])

		tiles_to_claim = []
		coords_to_claim = set()
		for c in cols:
		tiles_to_claim.append([])
		for r in rows:
		if (c, r) in self._claimed_tiles:
		raise ValueError(f"Tile ({c}, {r}) has already been claimed.")
		coords_to_claim.add((c, r))
		tiles_to_claim[-1].append(self._tiles[c][r])

		self._claimed_tiles.update(coords_to_claim)
		return DeviceView(self, tiles_to_claim)

		def tile_iterator(self) -> Generator[Tile, None, None]:
		for t in super().tile_iterator():
		if (t.col, t.row) not in self._claimed_tiles:
		yield self._tiles[t.col][t.row]


		def create_class(class_name, device):

		def _device__init__(self) -> None:
		super(globals()[class_name], self).__init__(device=device)

		def _device_resolve(
		self,
		loc: ir.Location \| None = None,
		ip: ir.InsertionPoint \| None = None,
		) -> None:
		return device

		globals()[class_name] = type(
		class_name,
		(Device,),
		{
		"__init__": _device__init__,
		"resolve": _device_resolve,
		"__doc__": f"A representation of a device that resolves to {device}",
		},
		)
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -43,7 +43,7 @@ def detect_npu_device():
		"RyzenAI-npu6",
		]
		):
		return NPU2()
		return NPU2
		elif any(
		keyword.lower() in output.lower()
		for keyword in [
Expand All		@@ -52,7 +52,7 @@ def detect_npu_device():
		"RyzenAI-npu1",
		]
		):
		return NPU1()
		return NPU1
		else:
		raise RuntimeError("No supported NPU device found.")

Expand DownExpand Up		@@ -89,5 +89,4 @@ def get_current_device():
		global config
		if "device" not in config:
		config["device"] = detect_npu_device()

		return config["device"]
		return config["device"]() if config["device"] else None
Original file line number	Diff line number	Diff line change
Expand Up		@@ -8,12 +8,8 @@
		from collections import defaultdict
		import numpy as np

		from .. import ir
		from ..ir import Module, Context, InsertionPoint, Location, Operation
		from ..dialects.aie import buffer
		from .worker import Worker
		from .program import CurrentDeviceOp, CurrentModule
		from .device import Tile


		class LocalBuffer(buffer):
Expand DownExpand Up		@@ -60,41 +56,15 @@ def __init__(
		idx = self.__get_index(current_core_placement)
		name = f"buf_{col}_{row}_{idx}"

		# if initial_value is None:
		device_op = get_device_op_from_module()
		block = device_op.regions[0].blocks[0]

		# Count number of tile ops and track last one
		tile_count = 0
		last_tile_op = None
		for op in block.operations:
		if op.name == "aie.tile":
		tile_count += 1
		last_tile_op = op

		ip = None if tile_count == 1 else InsertionPoint(last_tile_op)

		super().__init__(
		tile=current_core_placement,
		datatype=type,
		name=name,
		initial_value=initial_value,
		ip=ip,
		)

		@classmethod
		def __get_index(cls, placement: Tile) -> int:
		def __get_index(cls, placement):
		idx = cls.__buf_tile_index[placement]
		cls.__buf_tile_index[placement] += 1
		return idx


		def get_device_op_from_module():
		module = CurrentModule.get()
		if module is None:
		raise RuntimeError("No MLIR module context available")
		# The device op should be the first (or only) op in the module's body
		for op in module.body.operations:
		if op.name == "aie.device":
		return op
		raise RuntimeError("No aie.device op found in module")