diff --git a/python/aie-python-extras-req.txt b/python/aie-python-extras-req.txt index 36a31486ee..0a5ca35f57 100644 --- a/python/aie-python-extras-req.txt +++ b/python/aie-python-extras-req.txt @@ -1,2 +1,2 @@ # this is actually the aie branch of mlir-python-extras (see https://github.com/makslevental/mlir-python-extras/pull/33) -aie-python-extras @ https://github.com/makslevental/mlir-python-extras/archive/a0f48f0b67affc1c88d79fe94ef3fe5c27c73f80.zip +aie-python-extras @ https://github.com/makslevental/mlir-python-extras/archive/c7cf7ca8587b24e1f169199ddec5d4424f24c42e.zip diff --git a/python/dialects/aie.py b/python/dialects/aie.py index 8a90377057..8fc5306780 100644 --- a/python/dialects/aie.py +++ b/python/dialects/aie.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception import inspect +from collections import namedtuple +from dataclasses import dataclass from typing import List, Optional, Tuple, Union from ._aie_enum_gen import * @@ -20,6 +22,8 @@ _get_sym_name, get_user_code_loc, region_adder, + find_parent_of_type, + find_ops, ) from ..ir import ( ArrayAttr, @@ -31,6 +35,8 @@ IntegerAttr, IntegerType, TypeAttr, + DictAttr, + UnitAttr, _i32ArrayAttr, ) @@ -374,8 +380,10 @@ def buffer(buffer, tile, *, sym_name=None, address=None, loc=None, ip=None): _lock = lock -def lock(tile, *, lock_id=None, init=None, sym_name=None, loc=None, ip=None): - return _lock( +def lock( + tile, *, lock_id=None, init=None, sym_name=None, annot=None, loc=None, ip=None +): + l = _lock( tile, lock_id=lock_id, init=init, @@ -384,6 +392,9 @@ def lock(tile, *, lock_id=None, init=None, sym_name=None, loc=None, ip=None): loc=loc, ip=ip, ) + if annot is not None: + l.owner.attributes["annot"] = DictAttr.get({annot: UnitAttr.get()}) + return l @_cext.register_operation(_Dialect, replace=True) @@ -414,6 +425,101 @@ def flow( ) +def find_matching_flows( + tiles, + filter_source=False, + filter_dest=False, + source_annot=None, + dest_annot=None, + device=None, +): + assert not (filter_source and filter_dest), "Can only filter by source XOR dest" + if device is None: + device = find_parent_of_type(lambda op: isinstance(op, DeviceOp)) + + def _cb(op): + if isinstance(op, FlowOp): + if filter_source and op.source.owner.opview not in tiles: + return False + if filter_dest and op.dest.owner.opview not in tiles: + return False + + return ( + op.source.owner.opview in tiles + or op.dest.owner.opview in tiles + and ( + ( + "source_annot" in op.attributes + and source_annot in op.attributes["source_annot"] + ) + if source_annot is not None + else True + ) + and ( + ( + "dest_annot" in op.attributes + and dest_annot in op.attributes["dest_annot"] + ) + if dest_annot is not None + else True + ) + ) + + return find_ops(device, _cb) + + +def find_matching_locks(tiles, sym_name=None, annot=None, device=None): + if device is None: + device = find_parent_of_type(lambda op: isinstance(op, DeviceOp)) + + def _cb(op): + if isinstance(op, LockOp): + return ( + op.tile.owner.opview in tiles + and (sym_name == str(op.sym_name) if sym_name is not None else True) + and ( + ("annot" in op.attributes and annot in op.attributes["annot"]) + if annot is not None + else True + ) + ) + + return find_ops(device, _cb) + + +@dataclass +class Neighbors: + north: TileOp = None + west: TileOp = None + south: TileOp = None + + +def find_neighbors(tile, device=None): + if device is None: + device = find_parent_of_type(lambda op: isinstance(op, DeviceOp)) + + assert int(device.device) == int(AIEDevice.ipu), "only ipu supported" + + neighbors = {} + col, row = map(int, (tile.col, tile.row)) + if col > 0 and row > 0 and not (col, row) == (1, 1): + neighbors[col - 1, row] = "west" + if row > 1: + neighbors[col, row - 1] = "south" + if 0 < row < 5: + neighbors[col, row + 1] = "north" + + neighbors_ = {"north": None, "west": None, "south": None} + + for n in find_ops( + device, + lambda op: isinstance(op, TileOp) and (int(op.col), int(op.row)) in neighbors, + ): + neighbors_[neighbors[int(n.col), int(n.row)]] = n + + return Neighbors(**neighbors_) + + @_cext.register_operation(_Dialect, replace=True) class TileOp(TileOp): def __str__(self): @@ -435,6 +541,22 @@ def __eq__(self, other): def __hash__(self): return hash((self.col, self.row)) + def flows( + self, filter_source=False, filter_dest=False, source_annot=None, dest_annot=None + ): + return find_matching_flows( + [self], + filter_source=filter_source, + filter_dest=filter_dest, + source_annot=None, + dest_annot=None, + ) + + def locks(self, sym_name=None, annot=None, device=None): + return find_matching_locks( + [self], sym_name=sym_name, annot=annot, device=device + ) + def tile(col, row, *, loc=None, ip=None): return TileOp(col=col, row=row, loc=loc, ip=ip) diff --git a/python/dialects/aiex.py b/python/dialects/aiex.py index dbe14b9c63..401fc01e0e 100644 --- a/python/dialects/aiex.py +++ b/python/dialects/aiex.py @@ -12,12 +12,13 @@ from ._aiex_ops_gen import * from .aie import ( DMAChannelDir, - DeviceOp, - FlowOp, LockAction, TileOp, dma, dma_bd, + find_neighbors, + find_matching_flows, + find_matching_locks, flow, lock, tile, @@ -26,11 +27,8 @@ from .transform.structured import MixedValues, _dispatch_mixed_values from .._mlir_libs import get_dialect_registry from .._mlir_libs._aie import * -from ..extras.util import _get_previous_frame_idents, find_ops, find_parent_of_type from ..ir import DictAttr, IntegerAttr, UnitAttr -# Copyright (C) 2023, Advanced Micro Devices, Inc. -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # Comes from _aie register_dialect(get_dialect_registry()) @@ -373,26 +371,21 @@ def hold_lock(acq_lock, rel_lock, *, acq_val=None, rel_val=None): class TileArray: - def __init__(self, cols=5, rows=6, df=None, flows=None): + def __init__(self, cols=5, rows=6, df=None): if df is None: df = np.array( [[tile(c, r) for r in range(rows)] for c in range(cols)], ) self.df = df - if flows is None: - flows = Flows(self) - self._flows = flows def flow(self, other, *args, **kwargs): - return broadcast_flow(self.df, other.df, *args, **kwargs, flows=self.flows) + return broadcast_flow(self.df, other.df, *args, **kwargs) def __rshift__(self, other): - return broadcast_flow(self.df, other.df, flows=self.flows) + return broadcast_flow(self.df, other.df) def __lshift__(self, other): - r = np.frompyfunc(partial(broadcast_flow, flows=self.flows), 2, 1).outer( - other.df, self.df - ) + r = np.frompyfunc(partial(broadcast_flow), 2, 1).outer(other.df, self.df) if isinstance(r, np.ndarray): r = r.flatten().tolist() if len(r) == 1: @@ -429,9 +422,15 @@ def __contains__(self, item): assert isinstance(self.df, TileOp) return item == self.df + def flows(self, **kwargs): + return find_matching_flows(self, **kwargs) + + def locks(self, **kwargs): + return find_matching_locks(self, **kwargs) + @property - def flows(self): - return self._flows + def neighbors(self): + return np.vectorize(find_neighbors)(self.df) @property def shape(self): @@ -441,66 +440,6 @@ def __repr__(self): return f"<{self.__class__.__name__}: {self.df}>" -class Flows: - def __init__(self, tiles: TileArray): - self.tiles = tiles - - def find_matching_flows( - self, - tiles, - filter_source=False, - filter_dest=False, - source_annot=None, - dest_annot=None, - ): - assert not (filter_source and filter_dest), "Can only filter by source XOR dest" - device = find_parent_of_type(lambda op: isinstance(op, DeviceOp)) - - def _cb(op): - if isinstance(op, FlowOp): - if filter_source and op.source.owner.opview not in tiles: - return False - if filter_dest and op.dest.owner.opview not in tiles: - return False - - return ( - op.source.owner.opview in tiles - or op.dest.owner.opview in tiles - and ( - ( - "source_annot" in op.attributes - and source_annot in op.attributes["source_annot"] - ) - if source_annot is not None - else True - ) - and ( - ( - "dest_annot" in op.attributes - and dest_annot in op.attributes["dest_annot"] - ) - if dest_annot is not None - else True - ) - ) - - return find_ops(device, _cb) - - def __getitem__(self, item): - kwargs = {} - if len(item) > 2: - # not sure how but you don't need two backs here - # (the previous frame is the call site of TileArray().flows...) - previous_frame = inspect.currentframe().f_back - for kwarg in item[2:]: - k = _get_previous_frame_idents(kwarg, previous_frame) - assert len(k) == 1, f"{len(k)=}" - kwargs[k[0]] = kwarg - item = item[:2] - tiles = self.tiles[item] - return self.find_matching_flows(tiles, **kwargs) - - def broadcast_flow( source: Union[np.ndarray, TileOp], dest: Union[np.ndarray, TileOp], @@ -510,14 +449,13 @@ def broadcast_flow( dest_channel=None, source_annot=None, dest_annot=None, - flows: Flows = None, ): if isinstance(source, TileOp): source = np.asarray([source]) if isinstance(dest, TileOp): dest = np.asarray([dest]) for chan in [source_channel, dest_channel]: - assert (chan is None and flows is not None) or np.all( + assert chan is None or np.all( np.array(chan) != None ), "can't handle mixed auto channel assignment" @@ -534,19 +472,15 @@ def _find_next_channel(used_channels): if source_channel is None or np.all(np.array(source_channel) == None): source_channel = np.empty_like(source, dtype=None) for s, indices in zip(*map(list, np.unique(source, return_index=True))): - used_channels = set( - int(f.source_channel) - for f in flows.find_matching_flows([s], filter_source=True) - ) + matching_flows = find_matching_flows([s], filter_source=True) + used_channels = set(int(f.source_channel) for f in matching_flows) source_channel.flat[indices] = _find_next_channel(used_channels) if dest_channel is None or np.all(np.array(dest_channel) == None): used_channels = {} for d in np.unique(dest): - used_channels[d] = set( - int(f.dest_channel) - for f in flows.find_matching_flows([d], filter_dest=True) - ) + matching_flows = find_matching_flows([d], filter_dest=True) + used_channels[d] = set(int(f.dest_channel) for f in matching_flows) dest_channel = np.empty_like(dest, dtype=None) for idx, dst in np.ndenumerate(dest): dest_channel[idx] = _find_next_channel(used_channels[dst]) diff --git a/test/python/flow_dsl.py b/test/python/tile_array.py similarity index 78% rename from test/python/flow_dsl.py rename to test/python/tile_array.py index e102fc7ad5..4ac5d7284c 100644 --- a/test/python/flow_dsl.py +++ b/test/python/tile_array.py @@ -3,11 +3,20 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # # (c) Copyright 2023 AMD Inc. +import numpy as np # RUN: %python %s | FileCheck %s # REQUIRES: py310 -from aie.dialects.aie import AIEDevice, DMAChannelDir, LockAction, WireBundle, device +from aie.dialects.aie import ( + AIEDevice, + DMAChannelDir, + LockAction, + WireBundle, + device, + lock, + find_neighbors, +) from aie.dialects.aiex import TileArray from util import construct_and_print_module @@ -89,29 +98,29 @@ def ipu(): # CHECK: aie.flow(%tile_0_0, DMA : 3, %tile_2_2, DMA : 0) # CHECK: aie.flow(%tile_0_0, DMA : 4, %tile_2_2, DMA : 1) {dest_annot = {alice}, source_annot = {bob}} - for f in df.flows[2, 2]: + for f in df[2, 2].flows(): print(f) # CHECK: aie.flow(%tile_0_0, DMA : 4, %tile_2_2, DMA : 1) {dest_annot = {alice}, source_annot = {bob}} - for f in df.flows[2, 2, source_annot := "bob"]: + for f in df[2, 2].flows(source_annot="bob"): print(f) # CHECK: aie.flow(%tile_0_0, DMA : 4, %tile_2_2, DMA : 1) {dest_annot = {alice}, source_annot = {bob}} - for f in df.flows[2, 2, dest_annot := "alice"]: + for f in df[2, 2].flows(dest_annot="alice"): print(f) # CHECK: aie.flow(%tile_0_0, DMA : 4, %tile_2_2, DMA : 1) {dest_annot = {alice}, source_annot = {bob}} - for f in df.flows[2, 2, source_annot := "bob", dest_annot := "alice"]: + for f in df[2, 2].flows(source_annot="bob", dest_annot="alice"): print(f) - assert len(df.flows[0, 3, source_annot := "bob", dest_annot := "alice"]) == 0 + assert len(df[0, 3].flows(source_annot="bob", dest_annot="alice")) == 0 # CHECK: aie.flow(%tile_0_0, DMA : 1, %tile_0_3, DMA : 0) - for f in df.flows[0, 3]: + for f in df[0, 3].flows(): print(f) # CHECK: aie.flow(%tile_1_0, DMA : 0, %tile_1_3, DMA : 0) - for f in df.flows[1, 3, filter_dest := True]: + for f in df[1, 3].flows(filter_dest=True): print(f) # CHECK: module { @@ -198,3 +207,55 @@ def ipu(): # CHECK: aie.flow(%tile_0_3, DMA : 1, %tile_2_1, DMA : 3) for f in fls: print(f) + + +@construct_and_print_module +def locks(module): + @device(AIEDevice.ipu) + def ipu(): + tiles = TileArray() + + lock(tiles[0, 1].df) + # CHECK: %lock_0_1 = aie.lock(%tile_0_1) + for l in tiles[0, 1].locks(): + print(l) + + lock(tiles[0, 2].df) + lock(tiles[0, 2].df, annot="bob") + lock(tiles[0, 3].df) + lock(tiles[0, 3].df, annot="alice") + + # CHECK: %lock_0_2 = aie.lock(%tile_0_2) + # CHECK: %lock_0_2_0 = aie.lock(%tile_0_2) {annot = {bob}} + for l in tiles[0, 2].locks(): + print(l) + + # CHECK: %lock_0_2_0 = aie.lock(%tile_0_2) {annot = {bob}} + assert len(tiles[0, 2].locks(annot="bob")) + for l in tiles[0, 2].locks(annot="bob"): + print(l) + + assert len(tiles[0, 2].locks(annot="alice")) == 0 + + assert len(tiles[0, 3].locks(annot="alice")) == 1 + # CHECK: %lock_0_3_1 = aie.lock(%tile_0_3) {annot = {alice}} + for l in tiles[0, 3].locks(annot="alice"): + print(l) + + +@construct_and_print_module +def neighbors(module): + @device(AIEDevice.ipu) + def ipu(): + tiles = TileArray() + + # CHECK: Neighbors(north=%tile_2_3 = aie.tile(2, 3), west=%tile_1_2 = aie.tile(1, 2), south=%tile_2_1 = aie.tile(2, 1)) + print(find_neighbors(tiles[2, 2].df)) + + assert tiles[1:3, 1:3].neighbors.shape == (2, 2) + # CHECK: tile(col=1, row=1) Neighbors(north=%tile_1_2 = aie.tile(1, 2), west=None, south=None) + # CHECK: tile(col=1, row=2) Neighbors(north=%tile_1_3 = aie.tile(1, 3), west=%tile_0_2 = aie.tile(0, 2), south=%tile_1_1 = aie.tile(1, 1)) + # CHECK: tile(col=2, row=1) Neighbors(north=%tile_2_2 = aie.tile(2, 2), west=%tile_1_1 = aie.tile(1, 1), south=None) + # CHECK: tile(col=2, row=2) Neighbors(north=%tile_2_3 = aie.tile(2, 3), west=%tile_1_2 = aie.tile(1, 2), south=%tile_2_1 = aie.tile(2, 1)) + for idx, n in np.ndenumerate(tiles[1:3, 1:3].neighbors): + print(tiles[1:3, 1:3][idx].df, n)