diff --git a/.github/workflows/lintAndFormat.yml b/.github/workflows/lintAndFormat.yml index 05c6d4877ab..6665938dd61 100644 --- a/.github/workflows/lintAndFormat.yml +++ b/.github/workflows/lintAndFormat.yml @@ -129,6 +129,9 @@ jobs: ln -sfn "$PWD/build/python" install/python pyright + - name: Lint Python (ruff) + run: ruff check + - name: Analyze id: clang-tidy-fixes run: | diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b1caf4b5aad..06bbbfe3c89 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -36,6 +36,12 @@ repos: ironenv/.* )$ +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: cb8c523fd4835aba42af70f4cad5568db4df0b6c # frozen: v0.16.0 + hooks: + - id: ruff-check + stages: [pre-push] + # Baseline hygiene: validators run on every commit, fixers gated to pre-push - repo: https://github.com/pre-commit/pre-commit-hooks rev: 3e8a8703264a2f4a69428a0aa4dcb512790b2c8c # frozen: v6.0.0 diff --git a/.pylintrc b/.pylintrc deleted file mode 100644 index adba8f77a13..00000000000 --- a/.pylintrc +++ /dev/null @@ -1,66 +0,0 @@ -# Copyright (C) 2023 Advanced Micro Devices, Inc. -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -[MAIN] - -py-version=3.9 -suggestion-mode=yes - -[BASIC] - -argument-naming-style=snake_case -attr-naming-style=snake_case -bad-names=foo, - bar, - baz, - toto, - tutu, - tata - -class-attribute-naming-style=snake_case -class-const-naming-style=UPPER_CASE -class-naming-style=PascalCase -const-naming-style=UPPER_CASE -docstring-min-length=-1 -function-naming-style=snake_case -include-naming-hint=yes -inlinevar-naming-style=snake_case -method-naming-style=snake_case -module-naming-style=snake_case -variable-naming-style=snake_case - -[CLASSES] - -valid-classmethod-first-arg=cls -valid-metaclass-classmethod-first-arg=mcs - -[IMPORTS] - -allow-wildcard-with-all=no - -[SIMILARITIES] - -# Minimum lines number of a similarity. -min-similarity-lines=4 - -[STRING] - -# This flag controls whether inconsistent-quotes generates a warning when the -# character used as a quote delimiter is used inconsistently within a module. -check-quote-consistency=yes - -# This flag controls whether the implicit-str-concat should generate a warning -# on implicit string concatenation in sequences defined over several lines. -check-str-concat-over-line-jumps=yes - -[MESSAGES CONTROL] - -# C0116: Missing function or method docstring (missing-function-docstring) -# C0114: Missing module docstring (missing-module-docstring) -# E0402: Attempted relative import beyond top-level package (relative-beyond-top-level) -# C0115: Missing class docstring (missing-class-docstring) -# C0415: Import outside toplevel -# R0913: Too many arguments -# R0914: Too many local variables -# W1514: Using open without explicitly specifying an encoding (unspecified-encoding) -disable=C0116,C0114,E0402,C0115,C0415,R0913,R0914,W1514 diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 9af332d1ea8..82f06f2bda6 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -38,8 +38,10 @@ declare_mlir_python_sources(AIEPythonSources.Utils ADD_TO_PARENT AIEPythonSources SOURCES utils/__init__.py + utils/_log_setup.py utils/benchmark.py utils/config.py + utils/tensor_factory.py utils/test.py utils/ml.py utils/npukernel.py diff --git a/python/compiler/txn2mlir/main.py b/python/compiler/txn2mlir/main.py index 36f02744608..04ff06465f8 100755 --- a/python/compiler/txn2mlir/main.py +++ b/python/compiler/txn2mlir/main.py @@ -4,11 +4,15 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # -from aie.ir import * # pyright: ignore[reportMissingImports, reportAttributeAccessIssue] -from aie.dialects.aie import * # pyright: ignore[reportMissingImports] - import argparse +from aie.dialects.aie import ( + transaction_binary_to_mlir, # pyright: ignore[reportAttributeAccessIssue] +) +from aie.ir import ( # pyright: ignore[reportMissingImports] + Context, +) + def main(): # Parse arguments @@ -20,7 +24,7 @@ def main(): # Read the data from the file data = args.file.read() - with Context() as ctx: # pyright: ignore[reportUndefinedVariable] + with Context() as ctx: module = transaction_binary_to_mlir(ctx, data) print(str(module)) diff --git a/python/helpers/dialects/func.py b/python/helpers/dialects/func.py index f55283ac039..75d3b251ca6 100644 --- a/python/helpers/dialects/func.py +++ b/python/helpers/dialects/func.py @@ -2,46 +2,43 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception import inspect +from functools import update_wrapper +from typing import Any, List, get_args, get_origin + import numpy as np -from functools import lru_cache, update_wrapper -import sys -from typing import get_args, get_origin, Any, List -from ...extras.meta import op_region_builder # pyright: ignore[reportMissingImports] -from ...extras.util import ( # pyright: ignore[reportMissingImports] - get_user_code_loc, - make_maybe_no_args_decorator, -) -from ..util import get_arg_types, NpuDType, try_convert_np_type_to_mlir_type from ...dialects._ods_common import ( # pyright: ignore[reportMissingImports] get_op_result_or_op_results, ) -from ...dialects.func import * # pyright: ignore[reportMissingImports] from ...dialects.func import ( # pyright: ignore[reportMissingImports] - FuncOp, CallOp, + FuncOp, ReturnOp, ) -from ...ir import ( # pyright: ignore[reportMissingImports] +from ...extras.dialects.arith import ( # pyright: ignore[reportMissingImports] + ScalarValue, + index_cast, +) +from ...extras.meta import op_region_builder # pyright: ignore[reportMissingImports] +from ...extras.util import ( # pyright: ignore[reportMissingImports] + get_user_code_loc, + make_maybe_no_args_decorator, +) +from ...ir import ( # pyright: ignore[reportMissingImports] # pyright: ignore[reportMissingImports] Context, FlatSymbolRefAttr, FunctionType, + IndexType, InsertionPoint, - OpView, + IntegerType, Operation, + OpResult, + OpView, Type, TypeAttr, Value, ) -from ...extras.dialects.arith import ( # pyright: ignore[reportMissingImports] - ScalarValue, - index_cast, -) -from ...ir import ( # pyright: ignore[reportMissingImports] - IndexType, - IntegerType, - OpResult, -) +from ..util import NpuDType, get_arg_types, try_convert_np_type_to_mlir_type def call( @@ -139,17 +136,20 @@ def call( def isalambda(v): - LAMBDA = lambda: 0 + # A `def` here would defeat the check: Python only assigns the literal + # name "" to actual lambda expressions, not to `def`s. + LAMBDA = lambda: 0 # noqa: E731 + return isinstance(v, type(LAMBDA)) and v.__name__ == LAMBDA.__name__ def prep_func_types(sig, return_types): assert not ( - not sig.return_annotation is inspect.Signature.empty and len(return_types) > 0 - ), f"func can use return annotation or explicit return_types but not both" + sig.return_annotation is not inspect.Signature.empty and len(return_types) > 0 + ), "func can use return annotation or explicit return_types but not both" return_types = ( sig.return_annotation - if not sig.return_annotation is inspect.Signature.empty + if sig.return_annotation is not inspect.Signature.empty else return_types ) if not isinstance(return_types, (tuple, list)): @@ -162,7 +162,7 @@ def prep_func_types(sig, return_types): input_types = [ p.annotation for p in sig.parameters.values() - if not p.annotation is inspect.Signature.empty + if p.annotation is not inspect.Signature.empty ] # convert ndarray types to memref types assert all( @@ -233,7 +233,7 @@ def __init__( if self._is_decl(): assert len(self.input_types) == len( sig.parameters - ), f"func decl needs all input types annotated" + ), "func decl needs all input types annotated" self.sym_visibility = "private" self.emit() diff --git a/python/helpers/dialects/scf.py b/python/helpers/dialects/scf.py index e61b489737b..564bbf1bc1a 100644 --- a/python/helpers/dialects/scf.py +++ b/python/helpers/dialects/scf.py @@ -1,19 +1,17 @@ # Copyright (C) 2024-2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -import numpy as np +from contextlib import contextmanager from typing import Sequence -from ...ir import ( # pyright: ignore[reportMissingImports] - IndexType, - InsertionPoint, - Value, -) +import numpy as np + from ...dialects.scf import ( # pyright: ignore[reportMissingImports] - IfOp, ForOp, + IfOp, yield_, ) +from ...extras import types as T # pyright: ignore[reportMissingImports] from ...extras.dialects.arith import ( # pyright: ignore[reportMissingImports] constant, index_cast, @@ -21,8 +19,11 @@ from ...extras.util import ( # pyright: ignore[reportMissingImports] get_user_code_loc, ) -from contextlib import contextmanager -from ...extras import types as T # pyright: ignore[reportMissingImports] +from ...ir import ( # pyright: ignore[reportMissingImports] + IndexType, + InsertionPoint, + Value, +) def _for( diff --git a/python/helpers/taplib/__init__.py b/python/helpers/taplib/__init__.py index 69c3edfa1fa..9257b9f6004 100644 --- a/python/helpers/taplib/__init__.py +++ b/python/helpers/taplib/__init__.py @@ -6,3 +6,9 @@ TensorAccessSequence, ) from .tensortiler2d import TensorTiler2D + +__all__ = [ + "TensorAccessPattern", + "TensorAccessSequence", + "TensorTiler2D", +] diff --git a/python/helpers/taplib/tap.py b/python/helpers/taplib/tap.py index c1b48f5aaaa..2a1e0872fd5 100644 --- a/python/helpers/taplib/tap.py +++ b/python/helpers/taplib/tap.py @@ -3,10 +3,11 @@ from __future__ import annotations +import itertools from copy import deepcopy +from typing import Generator, Sequence + import numpy as np -import itertools -from typing import Sequence, Generator from .utils import ( validate_and_clean_sizes_strides, diff --git a/python/helpers/taplib/tas.py b/python/helpers/taplib/tas.py index f01b2360a22..568c1b29fb1 100644 --- a/python/helpers/taplib/tas.py +++ b/python/helpers/taplib/tas.py @@ -2,10 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception from __future__ import annotations + from collections import abc from copy import deepcopy +from typing import TYPE_CHECKING, Any, Callable, Sequence + import numpy as np -from typing import Any, Callable, Sequence, TYPE_CHECKING if TYPE_CHECKING: from matplotlib.animation import FuncAnimation @@ -18,6 +20,13 @@ ) +def _constant_fn(value): + def _fn(_step, _prev): + return value + + return _fn + + class TensorAccessSequence(abc.MutableSequence, abc.Iterable): """ TensorAccessSequence is a MutableSequence and an Iterable. Generally, it is a thin wrapper around a list[TensorAccessPattern]. @@ -63,7 +72,7 @@ def __init__( # Check tensor dims, offset, sizes, strides self._tensor_dims = validate_tensor_dims(tensor_dims) - if not (offset is None): + if offset is not None: offset = validate_offset(offset, self._tensor_dims) sizes, strides = validate_and_clean_sizes_strides( sizes, strides, allow_none=True @@ -75,34 +84,39 @@ def __init__( if num_steps == 0: if ( - offset != None - or sizes != None - or strides != None - or offset_fn != None - or sizes_fn != None - or strides_fn != None + offset is not None + or sizes is not None + or strides is not None + or offset_fn is not None + or sizes_fn is not None + or strides_fn is not None ): raise ValueError( - f"If num_steps=0, no sizes/strides/offset information may be specified" + "If num_steps=0, no sizes/strides/offset information may be specified" ) self._taps = [] else: # Make sure values or not None if iteration functions are None; also set default iter fn - if offset_fn is None: + if offset_fn is not None: + resolved_offset_fn = offset_fn + else: if offset is None: raise ValueError("Offset must be provided if offset_fn is None") - const_offset = offset - offset_fn = lambda _step, _prev_offset: const_offset - if sizes_fn is None: + resolved_offset_fn = _constant_fn(offset) + + if sizes_fn is not None: + resolved_sizes_fn = sizes_fn + else: if sizes is None: raise ValueError("Sizes must be provided if size_fn is None") - const_sizes = sizes - sizes_fn = lambda _step, _prev_sizes: const_sizes - if strides_fn is None: + resolved_sizes_fn = _constant_fn(sizes) + + if strides_fn is not None: + resolved_strides_fn = strides_fn + else: if strides is None: raise ValueError("Strides must be provided if stride_fn is None") - const_strides = strides - strides_fn = lambda _step, _prev_strides: const_strides + resolved_strides_fn = _constant_fn(strides) # Pre-calculate taps, because better for error handling up-front (and for visualizing full iter) # This is somewhat against the mentality behind iterations, but should be okay at the scale this @@ -112,9 +126,9 @@ def __init__( cur_sizes: Any = sizes cur_strides: Any = strides for step in range(num_steps): - cur_offset = offset_fn(step, cur_offset) - cur_sizes = sizes_fn(step, cur_sizes) - cur_strides = strides_fn(step, cur_strides) + cur_offset = resolved_offset_fn(step, cur_offset) + cur_sizes = resolved_sizes_fn(step, cur_sizes) + cur_strides = resolved_strides_fn(step, cur_strides) self._taps.append( TensorAccessPattern( diff --git a/python/helpers/taplib/tensortiler2d.py b/python/helpers/taplib/tensortiler2d.py index 04611432da4..da663bdc4d8 100644 --- a/python/helpers/taplib/tensortiler2d.py +++ b/python/helpers/taplib/tensortiler2d.py @@ -3,12 +3,12 @@ from copy import deepcopy from functools import partial -import numpy as np from typing import Sequence -from .tas import TensorAccessSequence -from aie.utils import ceildiv +import numpy as np +from aie.utils.tensor_factory import ceildiv +from .tas import TensorAccessSequence from .utils import validate_and_clean_sizes_strides, validate_tensor_dims @@ -138,7 +138,11 @@ def step_tiler( tile_col_major (bool, optional): Iterate column major within each tile. Defaults to False. tile_group_col_major (bool, optional): Iterate column major between tiles in a group within a TensorAccessSequence. Defaults to False. iter_col_major (bool, optional): Iterate column major over tiles within the TensorAccessSequence. Defaults to False. - allow_partial (bool, optional): Whether to allow partial tile groups. A tensor always decomposes into tiles evenly, but it may not decompose into tile *groups* evenly. When False, a partial tile grouping raises a ValueError; when True, partial groupings at the tensor edges are permitted. Defaults to False. + allow_partial (bool, optional): Whether to allow partial tile groups. A tensor + always decomposes into tiles evenly, but it may not decompose into tile + *groups* evenly. When False, a partial tile grouping raises a ValueError; + when True, partial groupings at the tensor edges are permitted. Defaults + to False. pattern_repeat (int, optional): Apply the access pattern n times within a single TensorAccessPattern. Defaults to 1. prune_step (bool, optional): Prune the iteration steps in the tiling process. Defaults to True. diff --git a/python/helpers/taplib/utils.py b/python/helpers/taplib/utils.py index a2a8a1feb3b..0de9e7c005c 100644 --- a/python/helpers/taplib/utils.py +++ b/python/helpers/taplib/utils.py @@ -2,9 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception from copy import deepcopy -import numpy as np from typing import Sequence +import numpy as np + def validate_and_clean_sizes_strides( sizes: Sequence[int] | None, @@ -35,7 +36,7 @@ def validate_and_clean_sizes_strides( raise ValueError("Strides is None, but expected Sequence[int]") # After this point can assume None is ok for sizes/strides - if not (expected_dims is None): + if expected_dims is not None: if expected_dims < 1: raise ValueError(f"Expected dimensions ({expected_dims}) should be >= 1") @@ -44,9 +45,9 @@ def validate_and_clean_sizes_strides( return None, None # Validate dimensions - if (not (sizes is None)) and len(sizes) == 0: + if (sizes is not None) and len(sizes) == 0: raise ValueError("len(sizes) must be >0") - if (not (strides is None)) and len(strides) == 0: + if (strides is not None) and len(strides) == 0: raise ValueError("len(strides) must be >0") if sizes and strides: @@ -110,7 +111,7 @@ def validate_tensor_dims( Returns: Sequence[int]: The validated tensor dimensions. """ - if not (expected_dims is None): + if expected_dims is not None: if expected_dims < 1: raise ValueError(f"Expected dimensions ({expected_dims}) should be >= 1") tensor_dims = deepcopy(tensor_dims) @@ -130,7 +131,7 @@ def validate_tensor_dims( if len(tensor_dims) == 1: tensor_dims = [1, tensor_dims[0]] - if not (expected_dims is None) and len(tensor_dims) != expected_dims: + if expected_dims is not None and len(tensor_dims) != expected_dims: raise ValueError( f"Tensor dimension ({tensor_dims}) does not match expected dimension ({expected_dims})" ) diff --git a/python/helpers/taplib/visualization2d.py b/python/helpers/taplib/visualization2d.py index d94f55e164d..f07cb7265e9 100644 --- a/python/helpers/taplib/visualization2d.py +++ b/python/helpers/taplib/visualization2d.py @@ -1,14 +1,14 @@ # Copyright (C) 2024-2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +import os +import warnings + import matplotlib.animation as animation import matplotlib.patheffects as pe import matplotlib.pyplot as plt import numpy as np -import os -import warnings - -from aie.utils import ceildiv +from aie.utils.tensor_factory import ceildiv def animate_from_accesses( @@ -18,7 +18,7 @@ def animate_from_accesses( ) -> animation.FuncAnimation: if len(access_order_tensors) < 1: raise ValueError("At least one access order tensor is required.") - if not (access_count_tensors is None): + if access_count_tensors is not None: if len(access_count_tensors) < 1: raise ValueError( "access_count_tensor should be None or requires at least one tensor" @@ -36,7 +36,7 @@ def animate_from_accesses( fig_height = min(fig_width, fig_width * height_width_ratio) ax_count = None - if not (access_count_tensors is None): + if access_count_tensors is not None: fig_height *= 2 fig, (ax_order, ax_count) = plt.subplots(2, 1) else: @@ -55,7 +55,7 @@ def animate_from_accesses( if ax_count is not None: ax_count.xaxis.tick_top() ax_count.invert_yaxis() - ax_count.set_title(f"Access Counts") + ax_count.set_title("Access Counts") def animate_order(i): access_heatmap = ax_order.pcolormesh(access_order_tensors[i]) @@ -110,7 +110,7 @@ def visualize_from_accesses( fig_height = min(fig_width, fig_width * height_width_ratio) ax_count = None - if not (access_count_tensor is None): + if access_count_tensor is not None: fig_height *= 2 fig, (ax_order, ax_count) = plt.subplots(2, 1) else: diff --git a/python/helpers/util.py b/python/helpers/util.py index abd23e6e998..dad16310f47 100644 --- a/python/helpers/util.py +++ b/python/helpers/util.py @@ -1,11 +1,13 @@ # Copyright (C) 2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception from collections import defaultdict -import numpy as np from typing import Any, Sequence, TypeVar, get_args, get_origin + +import numpy as np from aie._mlir_libs import ( _aie as CustomTypes, # pyright: ignore[reportAttributeAccessIssue] ) +from ml_dtypes import bfloat16 from ..dialects.arith import ConstantOp # pyright: ignore[reportMissingImports] from ..extras import types as T # pyright: ignore[reportMissingImports] @@ -19,7 +21,6 @@ Value, VectorType, ) -from ml_dtypes import bfloat16 # Custom types @@ -100,9 +101,9 @@ def get(): | v16bfp16ebs16 ) -_mlir_type_ctor_to_np_dtype = lambda: { - v: k for k, v in _np_dtype_to_mlir_type_ctor.items() -} + +def _mlir_type_ctor_to_np_dtype(): + return {v: k for k, v in _np_dtype_to_mlir_type_ctor.items()} def np_dtype_to_mlir_type(np_dtype): diff --git a/python/iron/__init__.py b/python/iron/__init__.py index 02d56b6988b..d7a8c09a64b 100644 --- a/python/iron/__init__.py +++ b/python/iron/__init__.py @@ -34,13 +34,36 @@ surface at MLIR verification time. """ +from aie.utils import ( + arange, + ceildiv, + ensure_current_device, + full, + get_current_device, + ones, + rand, + randint, + set_current_device, + set_tensor_class, + tensor, + zeros, + zeros_like, +) +from aie.utils.callabledesign import CallableDesign +from aie.utils.compile.jit import ( + CompilableDesign, + CompileTime, + In, + InOut, + Out, + compile_context, + compileconfig, + get_compile_arg, +) +from aie.utils.jit import jit + +from . import algorithms, kernels from .buffer import Buffer -from .kernel import ExternalFunction, Kernel -from .lock import Lock -from .scratchpad_parameter import ScratchpadParameter -from .program import Program -from .worker import Worker, WorkerRuntimeBarrier -from .runtime import Runtime, TaskGroup, RuntimeData, Task, sync_parameters from .dataflow import ( Acquire, Bd, @@ -55,36 +78,13 @@ StreamDims, TileDma, ) -from .dtype import str_to_dtype, dtype_to_str -from aie.utils.compile.jit import ( - CompilableDesign, - compile_context, - CompileTime, - In, - InOut, - Out, - compileconfig, - get_compile_arg, -) -from aie.utils.jit import jit -from aie.utils.callabledesign import CallableDesign -from . import kernels -from . import algorithms -from aie.utils import ( - tensor, - ones, - zeros, - full, - randint, - rand, - arange, - zeros_like, - ceildiv, - set_tensor_class, - set_current_device, - get_current_device, - ensure_current_device, -) +from .dtype import dtype_to_str, str_to_dtype +from .kernel import ExternalFunction, Kernel +from .lock import Lock +from .program import Program +from .runtime import Runtime, RuntimeData, Task, TaskGroup, sync_parameters +from .scratchpad_parameter import ScratchpadParameter +from .worker import Worker, WorkerRuntimeBarrier __all__ = [ # Core design abstractions diff --git a/python/iron/algorithms/__init__.py b/python/iron/algorithms/__init__.py index 13cfe533f20..fbd1839f91e 100644 --- a/python/iron/algorithms/__init__.py +++ b/python/iron/algorithms/__init__.py @@ -5,6 +5,12 @@ # """High-level algorithm templates built on IRON (transform, for_each, etc.).""" +from ._transform import ( + transform, + transform_binary, + transform_parallel, + transform_parallel_binary, +) from .conv_pipeline import ( row_at_a_time, row_at_a_time_tiled, @@ -13,9 +19,16 @@ ) from .for_each import for_each from .reduce import reduce -from ._transform import ( - transform, - transform_binary, - transform_parallel, - transform_parallel_binary, -) + +__all__ = [ + "row_at_a_time", + "row_at_a_time_tiled", + "row_at_a_time_with_skip", + "sliding_3row", + "for_each", + "reduce", + "transform", + "transform_binary", + "transform_parallel", + "transform_parallel_binary", +] diff --git a/python/iron/algorithms/_transform.py b/python/iron/algorithms/_transform.py index 32b5f7cd608..aa2056b18a0 100644 --- a/python/iron/algorithms/_transform.py +++ b/python/iron/algorithms/_transform.py @@ -6,11 +6,14 @@ """Tiled transform algorithms (unary/binary, single-core/parallel) built on IRON.""" import numpy as np - -from aie.iron import ObjectFifo, Program, Runtime, TaskGroup, Worker from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ -import aie.iron as iron +from aie.iron.dataflow import ObjectFifo +from aie.iron.kernel import ExternalFunction +from aie.iron.program import Program +from aie.iron.runtime import Runtime, TaskGroup +from aie.iron.worker import Worker +from aie.utils import get_current_device def _check_num_channels(num_channels: int) -> None: @@ -45,7 +48,7 @@ def _transform_gen(func, inputs: list, output, *params, tile_size=16, trace_size (or lambda) is expected to emit event0()/event1() markers; the trace shim records cycles between them. """ - is_external_func = isinstance(func, iron.ExternalFunction) + is_external_func = isinstance(func, ExternalFunction) # Validate tile_size matches ExternalFunction's tile_size() if defined if is_external_func and func.tile_size() != tile_size: @@ -201,7 +204,7 @@ def sequence(*args): ) # Place program components and generate an MLIR module - device = iron.get_current_device() + device = get_current_device() if device is None: raise RuntimeError( "iron.algorithms.transform requires an active NPU device. " @@ -262,7 +265,7 @@ def _transform_parallel_gen( "num_channels=2 is not supported together with shared *params; " "use num_channels=1 instead." ) - is_external_func = isinstance(func, iron.ExternalFunction) + is_external_func = isinstance(func, ExternalFunction) # Validate tile_size matches ExternalFunction arg_type if is_external_func and func.tile_size() != tile_size: @@ -290,7 +293,7 @@ def _transform_parallel_gen( dtype = ref_dtype # Determine number of columns based on device - device = iron.get_current_device() + device = get_current_device() if device is None: raise RuntimeError( "iron.algorithms.transform_parallel requires an active NPU device. " diff --git a/python/iron/algorithms/for_each.py b/python/iron/algorithms/for_each.py index 1e0d88b10c2..d8b840de22b 100644 --- a/python/iron/algorithms/for_each.py +++ b/python/iron/algorithms/for_each.py @@ -6,10 +6,13 @@ """``for_each``: apply a function in-place over a tiled tensor on an AIE core.""" import numpy as np - -from aie.iron import ObjectFifo, Program, Runtime, Worker from aie.iron.controlflow import range_ -import aie.iron as iron +from aie.iron.dataflow import ObjectFifo +from aie.iron.kernel import ExternalFunction +from aie.iron.program import Program +from aie.iron.runtime import Runtime +from aie.iron.worker import Worker +from aie.utils import get_current_device def for_each(func, tensor_ty, tile_size=16): @@ -87,7 +90,7 @@ def _for_each_real(func, tensor, *params, tile_size=16): Returns: mlir.ir.Module: The compiled MLIR module ready for execution. """ - is_external_func = isinstance(func, iron.ExternalFunction) + is_external_func = isinstance(func, ExternalFunction) num_elements = np.size(tensor) # Validate tile_size matches ExternalFunction's tile_size() if defined @@ -217,7 +220,7 @@ def sequence(*args): ) # Place program components and generate an MLIR module - device = iron.get_current_device() + device = get_current_device() if device is None: raise RuntimeError( "iron.algorithms.for_each requires an active NPU device. " diff --git a/python/iron/algorithms/reduce.py b/python/iron/algorithms/reduce.py index 0cac6c4a433..8685e69a69f 100644 --- a/python/iron/algorithms/reduce.py +++ b/python/iron/algorithms/reduce.py @@ -18,9 +18,12 @@ """ import numpy as np - -from aie.iron import ObjectFifo, Program, Runtime, Worker -import aie.iron as iron +from aie.iron.dataflow import ObjectFifo +from aie.iron.kernel import ExternalFunction +from aie.iron.program import Program +from aie.iron.runtime import Runtime +from aie.iron.worker import Worker +from aie.utils import get_current_device from ._transform import make_param_descriptor @@ -42,7 +45,7 @@ def _reduce_gen(func, input_desc, output_desc, *, trace_size=0): ``trace_size``-byte runtime trace buffer (default: 0). Kernel is expected to emit ``event0()``/``event1()`` markers. """ - if not isinstance(func, iron.ExternalFunction): + if not isinstance(func, ExternalFunction): raise TypeError( "_reduce_gen requires an ExternalFunction; reductions need " "accumulator state across the input elements which a per-element " @@ -75,7 +78,7 @@ def sequence(a_in, c_out, in_h, out_h): rt = Runtime(sequence, [in_ty, out_ty, of_in.prod(), of_out.cons()]) - device = iron.get_current_device() + device = get_current_device() if device is None: raise RuntimeError( "iron.algorithms.reduce requires an active NPU device. " diff --git a/python/iron/buffer.py b/python/iron/buffer.py index 03a48a70c3a..6be89470a33 100644 --- a/python/iron/buffer.py +++ b/python/iron/buffer.py @@ -6,8 +6,9 @@ """Named memory region accessible by both Workers and the Runtime.""" import itertools +from typing import TYPE_CHECKING, Sequence + import numpy as np -from typing import Sequence, TYPE_CHECKING from .. import ir # pyright: ignore[reportMissingImports, reportAttributeAccessIssue] from ..dialects.aie import buffer @@ -17,7 +18,7 @@ np_ndarray_type_get_shape, ) from .device import Tile -from .resolvable import Resolvable, NotResolvedError +from .resolvable import NotResolvedError, Resolvable if TYPE_CHECKING: from .worker import Worker @@ -45,11 +46,19 @@ def __init__( 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. - tile (Tile | None, optional): The tile for the buffer. Automatically set to the Worker's tile when the buffer is passed in the Worker's fn_args list. 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. - address (int | None, optional): Pin the buffer to a fixed L1 address. Needed for host-written RTP buffers the runtime pokes at a hardcoded address. Defaults to None (compiler-assigned). + 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. + tile (Tile | None, optional): The tile for the buffer. Automatically set to the + Worker's tile when the buffer is passed in the Worker's fn_args list. + 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. + address (int | None, optional): Pin the buffer to a fixed L1 address. Needed + for host-written RTP buffers the runtime pokes at a hardcoded address. + Defaults to None (compiler-assigned). Raises: ValueError: If neither ``type`` nor ``initial_value`` is provided. diff --git a/python/iron/compile/__init__.py b/python/iron/compile/__init__.py index 6958e35914f..0994f9b9360 100644 --- a/python/iron/compile/__init__.py +++ b/python/iron/compile/__init__.py @@ -2,10 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception """Backwards-compatible re-export from aie.utils.compile.jit.""" -from aie.utils.compile.jit.context import compile_context, get_compile_arg -from aie.utils.compile.jit.markers import CompileTime, In, InOut, Out from aie.utils.compile.jit.compilabledesign import CompilableDesign from aie.utils.compile.jit.compileconfig import compileconfig +from aie.utils.compile.jit.context import compile_context, get_compile_arg +from aie.utils.compile.jit.markers import CompileTime, In, InOut, Out __all__ = [ "CompilableDesign", diff --git a/python/iron/controlflow.py b/python/iron/controlflow.py index ec89ccee24d..e30c99d3147 100644 --- a/python/iron/controlflow.py +++ b/python/iron/controlflow.py @@ -3,6 +3,8 @@ from aie.helpers.dialects.scf import ( _for, +) +from aie.helpers.dialects.scf import ( yield_ as _yield_, # pyright: ignore[reportAttributeAccessIssue] ) from aie.iron.runtime.dmataskhandle import Task diff --git a/python/iron/dataflow/__init__.py b/python/iron/dataflow/__init__.py index c8361bfa349..254ab8fac56 100644 --- a/python/iron/dataflow/__init__.py +++ b/python/iron/dataflow/__init__.py @@ -15,14 +15,32 @@ [`Acquire`][iron.Acquire], [`Release`][iron.Release] """ +from .cascadeflow import CascadeFlow +from .flow import Flow, PacketDest, PacketFlow from .objectfifo import ( ObjectFifo, + ObjectFifoEndpoint, ObjectFifoHandle, ObjectFifoLink, - ObjectFifoEndpoint, PadDims, StreamDims, ) -from .cascadeflow import CascadeFlow -from .flow import Flow, PacketDest, PacketFlow from .tile_dma import Acquire, Bd, DmaChannel, Release, TileDma + +__all__ = [ + "ObjectFifo", + "ObjectFifoHandle", + "ObjectFifoLink", + "ObjectFifoEndpoint", + "PadDims", + "StreamDims", + "CascadeFlow", + "Flow", + "PacketDest", + "PacketFlow", + "Acquire", + "Bd", + "DmaChannel", + "Release", + "TileDma", +] diff --git a/python/iron/dataflow/cascadeflow.py b/python/iron/dataflow/cascadeflow.py index 7a2ff4cdd9b..09aab640f59 100644 --- a/python/iron/dataflow/cascadeflow.py +++ b/python/iron/dataflow/cascadeflow.py @@ -6,6 +6,7 @@ """CascadeFlow: a directed cascade stream connection between two Workers.""" from __future__ import annotations + from typing import TYPE_CHECKING from ...dialects.aie import ( diff --git a/python/iron/dataflow/flow.py b/python/iron/dataflow/flow.py index b812cd72ee7..2dfc7f12dec 100644 --- a/python/iron/dataflow/flow.py +++ b/python/iron/dataflow/flow.py @@ -32,7 +32,11 @@ ) from ...dialects.aie import ( flow as _flow_op, +) +from ...dialects.aie import ( packetflow as _packetflow_op, +) +from ...dialects.aie import ( shim_dma_allocation, # pyright: ignore[reportAttributeAccessIssue] ) from ..device import Tile # noqa: F401 (re-exported via package) diff --git a/python/iron/dataflow/objectfifo.py b/python/iron/dataflow/objectfifo.py index 171b6ea2649..d905e526771 100644 --- a/python/iron/dataflow/objectfifo.py +++ b/python/iron/dataflow/objectfifo.py @@ -4,17 +4,11 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # from __future__ import annotations + import itertools -import numpy as np from typing import Sequence, TypeAlias -# Named aliases for the (size, stride) pair-lists used by DMA stream -# layout transforms and pad-then-stream descriptors. Both are list[(size, -# stride)] from highest to lowest dimension; the names exist purely to -# make signatures and docs self-documenting — they are not validated -# types at runtime. -StreamDims: TypeAlias = list[Sequence[int]] -PadDims: TypeAlias = list[Sequence[int]] +import numpy as np from ... import ir # pyright: ignore[reportMissingImports, reportAttributeAccessIssue] from ...dialects._aie_enum_gen import ( # pyright: ignore[reportMissingImports] @@ -27,15 +21,22 @@ from ...dialects.aie import object_fifo, object_fifo_link from ...helpers.util import ( NpuDType, - np_ndarray_type_to_memref_type, np_ndarray_type_get_dtype, np_ndarray_type_get_shape, + np_ndarray_type_to_memref_type, single_elem_or_list_to_list, ) - -from ..resolvable import Resolvable, NotResolvedError +from ..device import AnyMemTile, Tile +from ..resolvable import NotResolvedError, Resolvable from .endpoint import ObjectFifoEndpoint -from ..device import Tile, AnyMemTile + +# Named aliases for the (size, stride) pair-lists used by DMA stream +# layout transforms and pad-then-stream descriptors. Both are list[(size, +# stride)] from highest to lowest dimension; the names exist purely to +# make signatures and docs self-documenting — they are not validated +# types at runtime. +StreamDims: TypeAlias = list[Sequence[int]] +PadDims: TypeAlias = list[Sequence[int]] def _same_shim_pin(a: "Tile | None", b: "Tile | None") -> bool: @@ -93,15 +94,43 @@ def __init__( obj_type (type[np.ndarray]): The type of each buffer in the ObjectFifo depth (int | None, optional): The default depth of the ObjectFifo endpoints. Defaults to 2. name (str | None, optional): The name of the ObjectFifo. If None is given, a unique name will be generated. Defaults to None. - dims_to_stream (StreamDims | None, optional): Data layout transformations applied when data is pushed onto the AXI stream, described as pairs of (size, stride) from highest to lowest dimension. Defaults to None. - dims_from_stream_per_cons (StreamDims | None, optional): List of data layout transformations applied by each consumer when data is read from the AXI stream, described as pairs of (size, stride) from highest to lowest dimension. Defaults to None. + dims_to_stream (StreamDims | None, optional): Data layout transformations applied + when data is pushed onto the AXI stream, described as pairs of (size, stride) + from highest to lowest dimension. Defaults to None. + dims_from_stream_per_cons (StreamDims | None, optional): List of data layout + transformations applied by each consumer when data is read from the AXI + stream, described as pairs of (size, stride) from highest to lowest + dimension. Defaults to None. plio (bool, optional): Whether the ObjectFifo uses PLIO connections. Defaults to False. - disable_synchronization (bool, optional): When True, disables lock-based synchronization on the ObjectFifo. Defaults to False. - repeat_count (int | None, optional): If set, causes the MemTile DMA to replay the buffer descriptor this many times without a new DMA transfer from L3. Distinct from ``iter_count`` (BD-chain iteration count). Defaults to None. - delegate_tile (Tile | None, optional): Shared-memory delegate tile. When set, the ObjectFifo's underlying buffer pool is allocated on this tile's memory module instead of the default placement. Lowers to ``aie.objectfifo.allocate``. *Only valid when both producer and consumer have shared-memory access to the delegate tile* (e.g. self-loop fifos where prod == cons, or fifos between adjacent tiles spilling to a neighboring MemTile). The delegate is the storage location, not a producer- or consumer-side concept; the underlying op verifier rejects this if either endpoint cannot share memory with the delegate. Defaults to None. - init_values (list[np.ndarray] | None, optional): Per-buffer static initial values for the producer endpoint. One ndarray per producer-side buffer; the producer tile must be able to hold static data at design startup (e.g. a MemTile). Lowers to the ``initValues`` attribute on the underlying ``aie.objectfifo`` op. Defaults to None. - consumer_obj_type (type[np.ndarray] | None, optional): Consumer element type for asymmetric transfer granularity. When set, the producer sends obj_type-sized transfers and the consumer receives consumer_obj_type-sized transfers. Producer element count must be an integer multiple of consumer element count. Defaults to None. - aie_stream (tuple[int, int] | None, optional): Mark the fifo as a direct AIE-stream connection by stamping the ``aie_stream`` / ``aie_stream_port`` attributes ``(end, port)`` on the underlying ``aie.objectfifo`` op. Use with kernels that emit on the wire via ``put_ms()`` instead of going through an L1 buffer. Defaults to None. + disable_synchronization (bool, optional): When True, disables lock-based + synchronization on the ObjectFifo. Defaults to False. + repeat_count (int | None, optional): If set, causes the MemTile DMA to replay the + buffer descriptor this many times without a new DMA transfer from L3. + Distinct from ``iter_count`` (BD-chain iteration count). Defaults to None. + delegate_tile (Tile | None, optional): Shared-memory delegate tile. When set, the + ObjectFifo's underlying buffer pool is allocated on this tile's memory module + instead of the default placement. Lowers to ``aie.objectfifo.allocate``. *Only + valid when both producer and consumer have shared-memory access to the + delegate tile* (e.g. self-loop fifos where prod == cons, or fifos between + adjacent tiles spilling to a neighboring MemTile). The delegate is the storage + location, not a producer- or consumer-side concept; the underlying op verifier + rejects this if either endpoint cannot share memory with the delegate. + Defaults to None. + init_values (list[np.ndarray] | None, optional): Per-buffer static initial values + for the producer endpoint. One ndarray per producer-side buffer; the producer + tile must be able to hold static data at design startup (e.g. a MemTile). + Lowers to the ``initValues`` attribute on the underlying ``aie.objectfifo`` + op. Defaults to None. + consumer_obj_type (type[np.ndarray] | None, optional): Consumer element type for + asymmetric transfer granularity. When set, the producer sends obj_type-sized + transfers and the consumer receives consumer_obj_type-sized transfers. + Producer element count must be an integer multiple of consumer element count. + Defaults to None. + aie_stream (tuple[int, int] | None, optional): Mark the fifo as a direct + AIE-stream connection by stamping the ``aie_stream`` / ``aie_stream_port`` + attributes ``(end, port)`` on the underlying ``aie.objectfifo`` op. Use with + kernels that emit on the wire via ``put_ms()`` instead of going through an L1 + buffer. Defaults to None. Raises: ValueError: If ``depth`` is provided and is less than 1. @@ -223,7 +252,7 @@ def prod( if self._prod: if depth is None: if self._depth is None: - raise ValueError(f"If depth is None, then depth must be specified.") + raise ValueError("If depth is None, then depth must be specified.") else: depth = self._depth elif depth < 1: @@ -268,7 +297,7 @@ def cons( """ if depth is None: if self._depth is None: - raise ValueError(f"If depth is None, then depth must be specified.") + raise ValueError("If depth is None, then depth must be specified.") else: depth = self._depth @@ -298,7 +327,7 @@ def tiles(self, cons_only: bool = False) -> list[Tile]: """ tiles = [] if not cons_only: - if self._prod == None: + if self._prod is None: raise ValueError( "Cannot return prod.tile.op because prod was not created." ) @@ -316,7 +345,7 @@ def tiles(self, cons_only: bool = False) -> list[Tile]: return tiles def _prod_tile_op(self) -> Tile: - if self._prod == None: + if self._prod is None: raise ValueError( f"Cannot return prod.tile.op for ObjectFifo {self.name} because prod was not created." ) @@ -519,7 +548,9 @@ def acquire( """Acquire access to some elements of the ObjectFifo, using ObjectFifo synchronization to moderate access. Args: - num_elem (int): Number of elements to acquire. If some elements are already acquired, only the additional elements needed to reach a total of ``num_elem`` are acquired. + num_elem (int): Number of elements to acquire. If some elements are already + acquired, only the additional elements needed to reach a total of + ``num_elem`` are acquired. Raises: ValueError: Number of elements cannot exceed ObjectFifo depth. @@ -654,11 +685,11 @@ def _emit_transfer( Lazy imports break the runtime<->dataflow import cycle. """ + from ..runtime._context import active_sequence from ..runtime.data import RuntimeData from ..runtime.dmatask import DMATask from ..runtime.dmataskhandle import Task from ..runtime.endpoint import RuntimeEndpoint - from ..runtime._context import active_sequence from ..scratchpad_parameter import ScratchpadParameter active = active_sequence() @@ -825,7 +856,9 @@ def join( obj_types (list[type[np.ndarray]], optional): The type of the buffers corresponding to each new ObjectFifo. Defaults to None. names (list[str] | None, optional): The name of each new ObjectFifo. If not given, unique names will be generated. Defaults to None. dims_to_stream (list[list[Sequence[int] | None]] | None, optional): The dimensionsToStream to assign to each new ObjectFifo. Defaults to None. - dims_from_stream (list[list[Sequence[int] | None]] | None, optional): The dimensionsFromStream to assign to each new ObjectFifo consumer. Defaults to None. + dims_from_stream (list[list[Sequence[int] | None]] | None, optional): The + dimensionsFromStream to assign to each new ObjectFifo consumer. Defaults + to None. plio (bool, optional): Set plio on each new ObjectFifo. Defaults to False. repeat_counts (list[int | None] | None, optional): Per-sub-fifo MemTile DMA repeat count (see ObjectFifo.repeat_count). Defaults to None. @@ -1058,8 +1091,10 @@ def __init__( srcs (list[ObjectFifoHandle] | ObjectFifoHandle): A list of consumer ObjectFifoHandles to link. dsts (list[ObjectFifoHandle] | ObjectFifoHandle): A list of producer ObjectFifoHandles to link. tile (Tile, optional): The tile where the link occurs. Also accepts None (treated as AnyMemTile). Defaults to AnyMemTile. - src_offsets (list[int] | None, optional): If many sources, one offset per source is required to split the destination. Defaults to None (empty list). - dst_offsets (list[int] | None, optional): If many destinations, one offset per destination is required to split the source. Defaults to None (empty list). + src_offsets (list[int] | None, optional): If many sources, one offset per source + is required to split the destination. Defaults to None (empty list). + dst_offsets (list[int] | None, optional): If many destinations, one offset per + destination is required to split the source. Defaults to None (empty list). Raises: ValueError: Arguments are validated. diff --git a/python/iron/device/__init__.py b/python/iron/device/__init__.py index d87ea0523b3..08ce45b9ea6 100644 --- a/python/iron/device/__init__.py +++ b/python/iron/device/__init__.py @@ -7,7 +7,16 @@ from . import device as _device_module from .device import Device -from .tile import AnyShimTile, AnyMemTile, AnyComputeTile, Tile +from .tile import AnyComputeTile, AnyMemTile, AnyShimTile, Tile + +__all__ = [ + "Device", + "AnyShimTile", + "AnyMemTile", + "AnyComputeTile", + "Tile", + "from_name", +] def __getattr__(name: str) -> type[Device]: diff --git a/python/iron/device/device.py b/python/iron/device/device.py index 7584edd00a0..e68e2dece06 100644 --- a/python/iron/device/device.py +++ b/python/iron/device/device.py @@ -4,6 +4,8 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # +import re + from ... import ir # pyright: ignore[reportMissingImports, reportAttributeAccessIssue] from ...dialects._aie_enum_gen import ( # pyright: ignore[reportMissingImports] AIEArch, @@ -11,15 +13,13 @@ ) from ...dialects.aie import ( AIEDevice, # pyright: ignore[reportAttributeAccessIssue] - logical_tile, LogicalTileOp, get_target_model, # pyright: ignore[reportAttributeAccessIssue] + logical_tile, ) from ..resolvable import Resolvable from .tile import Tile -import re - class Device(Resolvable): """A representation of a device of a specific type. diff --git a/python/iron/kernel.py b/python/iron/kernel.py index 1c41ab772bb..0725087732c 100644 --- a/python/iron/kernel.py +++ b/python/iron/kernel.py @@ -11,15 +11,15 @@ import numpy as np -logger = logging.getLogger(__name__) - from .. import ir # pyright: ignore[reportMissingImports, reportAttributeAccessIssue] from ..dialects import memref # pyright: ignore[reportAttributeAccessIssue] +from ..dialects.aie import external_func from ..extras.dialects.func import FuncOp # pyright: ignore[reportMissingImports] from ..helpers.dialects.func import call -from ..dialects.aie import external_func -from .resolvable import Resolvable from .buffer import Buffer +from .resolvable import Resolvable + +logger = logging.getLogger(__name__) def _is_contiguous_row_major(mr): diff --git a/python/iron/kernels/__init__.py b/python/iron/kernels/__init__.py index b55e05aff6a..f781566672d 100644 --- a/python/iron/kernels/__init__.py +++ b/python/iron/kernels/__init__.py @@ -14,50 +14,50 @@ - `conv` — conv2dk1, conv2dk3, conv2dk1_skip, conv2dk1_i8, conv2dk14, conv2dk1_skip_init, bn_* """ -from .eltwise import passthrough, scale, add, mul, relu -from .reduce import reduce_add, reduce_min, reduce_max, compute_max -from .vision import ( - rgba2hue, - threshold, - bitwise_or, - bitwise_and, - gray2rgba, - rgba2gray, - filter2d, - add_weighted, -) from .activation import ( - softmax, - gelu, - silu, - swiglu, bf16_exp, + bf16_exp_ref, + gelu, + gelu_ref, relu_ref, + silu, silu_ref, - gelu_ref, - bf16_exp_ref, + softmax, softmax_ref, + swiglu, ) -from .linalg import mm, mv, cascade_mm from .conv import ( - conv2dk1, - conv2dk3, - conv2dk1_skip, - conv2dk1_i8, - conv2dk14, - conv2dk1_skip_init, - bn_conv2dk1_relu, - bn_conv2dk3, bn_conv2dk1_i8, + bn_conv2dk1_input_split_partial_put_ui8, + bn_conv2dk1_input_split_partial_skip_get, + bn_conv2dk1_partial_get_relu_i8, + bn_conv2dk1_partial_put_i8, + bn_conv2dk1_relu, + bn_conv2dk1_relu_xy_pool_padded, bn_conv2dk1_skip, + bn_conv2dk3, bn_conv2dk3_dw, - bn_conv2dk1_relu_xy_pool_padded, - bn_fc_relu_ui16_pad, - bn_conv2dk1_partial_put_i8, - bn_conv2dk1_partial_get_relu_i8, bn_conv2dk3_dw_out_split, - bn_conv2dk1_input_split_partial_put_ui8, - bn_conv2dk1_input_split_partial_skip_get, + bn_fc_relu_ui16_pad, + conv2dk1, + conv2dk1_i8, + conv2dk1_skip, + conv2dk1_skip_init, + conv2dk3, + conv2dk14, +) +from .eltwise import add, mul, passthrough, relu, scale +from .linalg import cascade_mm, mm, mv +from .reduce import compute_max, reduce_add, reduce_max, reduce_min +from .vision import ( + add_weighted, + bitwise_and, + bitwise_or, + filter2d, + gray2rgba, + rgba2gray, + rgba2hue, + threshold, ) __all__ = [ diff --git a/python/iron/kernels/_common.py b/python/iron/kernels/_common.py index 26af6c4707c..7160b93271c 100644 --- a/python/iron/kernels/_common.py +++ b/python/iron/kernels/_common.py @@ -8,9 +8,8 @@ import hashlib import logging from pathlib import Path -import numpy as np -from ml_dtypes import bfloat16 +import numpy as np from aie.iron.kernel import ExternalFunction _log = logging.getLogger(__name__) diff --git a/python/iron/kernels/activation.py b/python/iron/kernels/activation.py index d1748211c50..351f66bdc5a 100644 --- a/python/iron/kernels/activation.py +++ b/python/iron/kernels/activation.py @@ -19,10 +19,10 @@ """ from pathlib import Path -import numpy as np -from ml_dtypes import bfloat16 +import numpy as np from aie.iron.kernel import ExternalFunction +from ml_dtypes import bfloat16 from ._common import ( _detect_arch, diff --git a/python/iron/kernels/conv.py b/python/iron/kernels/conv.py index fb02b31ead9..88ec2f79b84 100644 --- a/python/iron/kernels/conv.py +++ b/python/iron/kernels/conv.py @@ -6,7 +6,6 @@ """Convolution kernel factories: conv2dk1/3/14, bottleneck (bn_*) variants.""" import numpy as np - from aie.iron.kernel import ExternalFunction from ._common import ( diff --git a/python/iron/kernels/eltwise.py b/python/iron/kernels/eltwise.py index f912d8ec125..9dfe9e34e40 100644 --- a/python/iron/kernels/eltwise.py +++ b/python/iron/kernels/eltwise.py @@ -6,9 +6,8 @@ """Element-wise kernel factories: passthrough, scale, add, mul, relu.""" import numpy as np -from ml_dtypes import bfloat16 - from aie.iron.kernel import ExternalFunction +from ml_dtypes import bfloat16 from ._common import ( _default_source_path, diff --git a/python/iron/kernels/linalg.py b/python/iron/kernels/linalg.py index 9b92808b885..098359e2959 100644 --- a/python/iron/kernels/linalg.py +++ b/python/iron/kernels/linalg.py @@ -6,9 +6,8 @@ """Linear algebra kernel factories: mm, mv, cascade_mm.""" import numpy as np -from ml_dtypes import bfloat16 - from aie.iron.kernel import ExternalFunction, Kernel +from ml_dtypes import bfloat16 from ._common import _default_source_path, _detect_arch, _make_extern diff --git a/python/iron/kernels/reduce.py b/python/iron/kernels/reduce.py index 53102015679..0827926fb20 100644 --- a/python/iron/kernels/reduce.py +++ b/python/iron/kernels/reduce.py @@ -6,9 +6,8 @@ """Reduction kernel factories: reduce_add, reduce_min, reduce_max, compute_max.""" import numpy as np -from ml_dtypes import bfloat16 - from aie.iron.kernel import ExternalFunction +from ml_dtypes import bfloat16 from ._common import _default_source_path, _make_extern, _min_dma_aligned_elems diff --git a/python/iron/kernels/vision.py b/python/iron/kernels/vision.py index 699dc360745..71f849292b7 100644 --- a/python/iron/kernels/vision.py +++ b/python/iron/kernels/vision.py @@ -6,7 +6,6 @@ """Vision kernel factories: color conversion, threshold, filter2d, add_weighted.""" import numpy as np - from aie.iron.kernel import ExternalFunction from ._common import ( diff --git a/python/iron/lock.py b/python/iron/lock.py index 76e7ed4d56c..a4d44cd0892 100644 --- a/python/iron/lock.py +++ b/python/iron/lock.py @@ -17,6 +17,8 @@ from ..dialects._aie_enum_gen import LockAction # pyright: ignore[reportMissingImports] from ..dialects.aie import ( lock as _lock_op, +) +from ..dialects.aie import ( use_lock as _use_lock, # pyright: ignore[reportAttributeAccessIssue] ) from .device import Tile diff --git a/python/iron/program.py b/python/iron/program.py index 7d26729c395..72046156dce 100644 --- a/python/iron/program.py +++ b/python/iron/program.py @@ -6,18 +6,17 @@ import logging -logger = logging.getLogger(__name__) - +from ..dialects.aie import device from ..extras.context import mlir_mod_ctx # pyright: ignore[reportMissingImports] from ..helpers.dialects.func import FuncBase -from ..dialects.aie import device - +from ..utils import trace as trace_utils +from ..utils.compile.jit.context import get_compile_arg from .device import Device +from .resolvable import Resolvable from .runtime import Runtime from .scratchpad_parameter import ScratchpadParameter -from .resolvable import Resolvable -from ..utils import trace as trace_utils -from ..utils.compile.jit.context import get_compile_arg + +logger = logging.getLogger(__name__) class Program: @@ -276,5 +275,5 @@ def device_body(): def _print_verify(self, ctx): verify = ctx.module.operation.verify() - if verify != True: + if not verify: raise RuntimeError(f"MLIR module failed verification: {verify}") diff --git a/python/iron/runtime/__init__.py b/python/iron/runtime/__init__.py index d13a915b2db..eeddfcec263 100644 --- a/python/iron/runtime/__init__.py +++ b/python/iron/runtime/__init__.py @@ -5,7 +5,9 @@ # """Runtime: host-side data movement and worker execution orchestration.""" -from .runtime import Runtime, sync_parameters -from .taskgroup import TaskGroup from .data import RuntimeData from .dmataskhandle import Task +from .runtime import Runtime, sync_parameters +from .taskgroup import TaskGroup + +__all__ = ["Runtime", "RuntimeData", "Task", "TaskGroup", "sync_parameters"] diff --git a/python/iron/runtime/_context.py b/python/iron/runtime/_context.py index 515ed022311..18e6ce8c0bf 100644 --- a/python/iron/runtime/_context.py +++ b/python/iron/runtime/_context.py @@ -14,6 +14,7 @@ """ from __future__ import annotations + from contextlib import contextmanager from contextvars import ContextVar from typing import TYPE_CHECKING, Iterator diff --git a/python/iron/runtime/data.py b/python/iron/runtime/data.py index d447285107e..74369179295 100644 --- a/python/iron/runtime/data.py +++ b/python/iron/runtime/data.py @@ -4,18 +4,19 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # +from typing import Sequence, get_origin + import numpy as np -from typing import get_origin, Sequence from ...extras.dialects.memref import ( # pyright: ignore[reportMissingImports] MemRefValue, ) +from ...helpers.taplib import TensorAccessPattern, TensorTiler2D from ...helpers.util import ( NpuDType, np_ndarray_type_get_dtype, np_ndarray_type_get_shape, ) -from ...helpers.taplib import TensorAccessPattern, TensorTiler2D class RuntimeData: diff --git a/python/iron/runtime/dmatask.py b/python/iron/runtime/dmatask.py index 921b4ae8460..ae0d7efc806 100644 --- a/python/iron/runtime/dmatask.py +++ b/python/iron/runtime/dmatask.py @@ -6,14 +6,13 @@ """DMATask: a RuntimeTask that generates a shim DMA transfer operation.""" from ... import ir # pyright: ignore[reportMissingImports, reportAttributeAccessIssue] - from ...dialects._aiex_ops_gen import ( # pyright: ignore[reportMissingImports] dma_start_task, ) from ...dialects.aiex import shim_dma_single_bd_task +from ...helpers.taplib import TensorAccessPattern from ..dataflow import ObjectFifoHandle from .data import RuntimeData -from ...helpers.taplib import TensorAccessPattern from .task import RuntimeTask from .taskgroup import TaskGroup @@ -46,7 +45,8 @@ def __init__( tap (TensorAccessPattern | None, optional): The static access pattern. Mutually exclusive with sizes/strides/offset/transfer_len. task_group (TaskGroup | None, optional): The task group associated with the operation. Defaults to None. wait (bool, optional): Whether this task should conclude with a call to await or a call to free. Defaults to False. - offset_parameter (str | None, optional): Name of a ScratchpadParameter whose value is used as the element offset for this DMA transfer. Defaults to None. + offset_parameter (str | None, optional): Name of a ScratchpadParameter whose + value is used as the element offset for this DMA transfer. Defaults to None. packet (tuple[int, int] | None, optional): Stamp the shim DMA's BD with a packet header `(pkt_type, pkt_id)`. Pairs with downstream packet-switched routing (e.g. ObjectFifos diff --git a/python/iron/runtime/endpoint.py b/python/iron/runtime/endpoint.py index c3c28801051..a3e4707cd71 100644 --- a/python/iron/runtime/endpoint.py +++ b/python/iron/runtime/endpoint.py @@ -6,11 +6,11 @@ from __future__ import annotations -from ..dataflow.endpoint import ObjectFifoEndpoint -from ..device import Tile, AnyShimTile from ...dialects._aie_enum_gen import ( # pyright: ignore[reportMissingImports] AIETileType, ) +from ..dataflow.endpoint import ObjectFifoEndpoint +from ..device import AnyShimTile, Tile class RuntimeEndpoint(ObjectFifoEndpoint): diff --git a/python/iron/runtime/runtime.py b/python/iron/runtime/runtime.py index 500ccff6617..a811d96d17c 100644 --- a/python/iron/runtime/runtime.py +++ b/python/iron/runtime/runtime.py @@ -16,40 +16,40 @@ """ from __future__ import annotations + import itertools import logging -import numpy as np from typing import Callable, Sequence, get_origin -logger = logging.getLogger(__name__) - -from ...utils import trace as trace_utils +import numpy as np from ... import ir # pyright: ignore[reportMissingImports, reportAttributeAccessIssue] - +from ...dialects._aie_ops_gen import ( # pyright: ignore[reportMissingImports] + RuntimeSequenceOp, +) from ...dialects.aiex import ( - npu_load_pdi, # pyright: ignore[reportAttributeAccessIssue] - sync_scratchpad_parameters_from_host, # pyright: ignore[reportAttributeAccessIssue] dma_await_task, dma_free_task, + npu_load_pdi, # pyright: ignore[reportAttributeAccessIssue] + sync_scratchpad_parameters_from_host, # pyright: ignore[reportAttributeAccessIssue] ) -from ...dialects._aie_ops_gen import ( # pyright: ignore[reportMissingImports] - RuntimeSequenceOp, -) +from ...extras.dialects.arith import constant # pyright: ignore[reportMissingImports] from ...helpers.util import ( - try_convert_np_type_to_mlir_type, - np_dtype_to_mlir_type, flatten_fn_args, + np_dtype_to_mlir_type, + try_convert_np_type_to_mlir_type, ) -from ...extras.dialects.arith import constant # pyright: ignore[reportMissingImports] +from ...utils import trace as trace_utils from ..dataflow import ObjectFifoHandle from ..resolvable import Resolvable from ..scratchpad_parameter import ScratchpadParameter -from .dmatask import DMATask +from ._context import active_sequence, active_sequence_scope from .data import RuntimeData +from .dmatask import DMATask from .endpoint import RuntimeEndpoint from .taskgroup import TaskGroup -from ._context import active_sequence, active_sequence_scope + +logger = logging.getLogger(__name__) class IronRuntimeError(Exception): diff --git a/python/iron/runtime/task.py b/python/iron/runtime/task.py index 77519db3bf0..4633dab2da5 100644 --- a/python/iron/runtime/task.py +++ b/python/iron/runtime/task.py @@ -4,7 +4,6 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # from ... import ir # pyright: ignore[reportMissingImports, reportAttributeAccessIssue] - from ..resolvable import Resolvable from .taskgroup import TaskGroup diff --git a/python/iron/scratchpad_parameter.py b/python/iron/scratchpad_parameter.py index e9bee1db57c..9bbca996d1a 100644 --- a/python/iron/scratchpad_parameter.py +++ b/python/iron/scratchpad_parameter.py @@ -5,11 +5,9 @@ # """ScratchpadParameter: a named runtime value set from the host and read by Workers.""" -import numpy as np - from .. import ir # pyright: ignore[reportMissingImports, reportAttributeAccessIssue] from ..dialects import aiex -from ..helpers.util import np_dtype_to_mlir_type, NpuDType +from ..helpers.util import NpuDType, np_dtype_to_mlir_type from .resolvable import Resolvable diff --git a/python/iron/worker.py b/python/iron/worker.py index f9bf5237219..9f19ce82402 100644 --- a/python/iron/worker.py +++ b/python/iron/worker.py @@ -9,26 +9,26 @@ from typing import Callable from .. import ir # pyright: ignore[reportMissingImports, reportAttributeAccessIssue] +from ..dialects._aie_enum_gen import ( # pyright: ignore[reportMissingImports] + AIETileType, +) from ..dialects.aie import ( core, lock, use_lock, # pyright: ignore[reportAttributeAccessIssue] ) from ..dialects.aiex import ( - set_lock_value, LockAction, # pyright: ignore[reportAttributeAccessIssue] + set_lock_value, ) from ..helpers.dialects.scf import _for as range_ from ..helpers.util import flatten_fn_args -from .device import Tile, AnyComputeTile -from ..dialects._aie_enum_gen import ( # pyright: ignore[reportMissingImports] - AIETileType, -) -from .dataflow.objectfifo import ObjectFifoHandle, ObjectFifo -from .dataflow.endpoint import ObjectFifoEndpoint from .buffer import Buffer -from .scratchpad_parameter import ScratchpadParameter +from .dataflow.endpoint import ObjectFifoEndpoint +from .dataflow.objectfifo import ObjectFifo, ObjectFifoHandle +from .device import AnyComputeTile, Tile from .resolvable import Resolvable +from .scratchpad_parameter import ScratchpadParameter class Worker(ObjectFifoEndpoint): @@ -59,7 +59,8 @@ def __init__( tile (Tile, optional): The compute tile for the Worker. Also accepts None (treated as AnyComputeTile). Defaults to AnyComputeTile. while_true (bool, optional): If true, will wrap the core_fn in a while(true) loop to ensure it runs until reconfiguration. Defaults to True. stack_size (int, optional): The stack_size in bytes to be allocated for the worker. Defaults to 1024 bytes. - allocation_scheme (str, optional): The memory allocation scheme to use for the Worker, either 'basic-sequential' or 'bank-aware'. If None, defaults to bank-aware. + allocation_scheme (str, optional): The memory allocation scheme to use for the + Worker, either 'basic-sequential' or 'bank-aware'. If None, defaults to bank-aware. Will override any allocation scheme set on the tile. trace (int, optional): If >0, enable tracing for this worker. trace_events (list | None, optional): Custom list of trace events for this worker. Defaults to None. @@ -248,8 +249,8 @@ def resolve( # Create the necessary locks for the core operation to synchronize with the runtime sequence # and register them in the corresponding barriers. for barrier in self._barriers: - l = lock(my_tile) - barrier._add_worker_lock(l) + barrier_lock = lock(my_tile) + barrier._add_worker_lock(barrier_lock) @core( my_tile, @@ -309,8 +310,8 @@ def _add_worker_lock(self, lock): def _set_barrier_value(self, value: int): """Set the value of the barrier.""" - for lock in self.worker_locks: - set_lock_value(lock, value) + for worker_lock in self.worker_locks: + set_lock_value(worker_lock, value) def release_with_value(self, value: int): """ diff --git a/python/requirements_dev.lock b/python/requirements_dev.lock index 40ba156af2b..16525d110dd 100644 --- a/python/requirements_dev.lock +++ b/python/requirements_dev.lock @@ -1208,6 +1208,26 @@ reuse==6.2.0 \ --hash=sha256:12b68549bb9d5f4957f06d726a83a9780628810008fb732bb0d0f21607f8c6d6 \ --hash=sha256:4feae057a2334c9a513e6933cdb9be819d8b822f3b5b435a36138bd218897d23 # via -r requirements_dev.txt +ruff==0.16.0 \ + --hash=sha256:0ff4a79ce3ec0172f3241943835de1c4cb4e2dcd07f0f8c2d02603dbbbee4b17 \ + --hash=sha256:14296fedcd2705c77ab8235439278bbb38f285cf7da5528b00b3e330c3d4872d \ + --hash=sha256:28ea2b7df8ebf7f9da6b7d47b230ab48f387c0a29be3b474c4d0740e197bb9af \ + --hash=sha256:33a3dfac8c35f81498dea9181bccc2f4c4bc8f1521a1dd9406e77643e0f0fb09 \ + --hash=sha256:3c954b1d580bfa035b41654f7858cc7e71d5fc3ac5b723dd62bd9133830ed522 \ + --hash=sha256:429c117f022bf481fabd9d551e7a3952b24c65e6ef44337ea09d90bebef14472 \ + --hash=sha256:48044c678e9cb8698246c99b14aaccfa6601dea7379eb48a6f8f73f7a6d86cd0 \ + --hash=sha256:4f11a8d11010301d0a398a2fdef67691feca7294da6aef55e2150e8fa2cd520b \ + --hash=sha256:6e364e5ed22ed8dc05082fd78e35308618260907ac2d3c1d637b2e682415b6c9 \ + --hash=sha256:7aa0959bad8eb8bef50340154fc9b58678dae31fa4293afa38b44b6e552c0213 \ + --hash=sha256:7fab76fa065c873f41ff744347c6e77bcc3dfec4bcc754dc26b63d23c0f7f5fb \ + --hash=sha256:a5237a0bda500d30d81b8e07a6973a5cbc772864cbf746ae2f4e8a2e01c9f4ed \ + --hash=sha256:a9b50c55e263103586b3dcf5f73d479eb8cb5fdb6098fec59a62891dab653717 \ + --hash=sha256:d327b8fc113a1d4421a04f3839d3752057c8dd1ee320223a6f3f52d04ada462a \ + --hash=sha256:e01c21d10eb1b29f47b7454e1f4056db9a3f0260c646aa88457c610291db9f81 \ + --hash=sha256:e460aafd5495ec89efaa6ced2e4a9a581116451e1c88b9d37ef497e0f8e93982 \ + --hash=sha256:e5115729eb08c585e5121978ba5d5b60caeae394ce21b9fb5e6cd33a1c6c9b1e \ + --hash=sha256:e95c448fca1fb2a18372a9440926c5a6ee789639bb975c72e7ae6d0b04218ab4 + # via -r requirements_dev.txt setuptools==83.0.0 \ --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \ --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3 diff --git a/python/requirements_dev.txt b/python/requirements_dev.txt index c6b7097d2d4..56f7c62857e 100644 --- a/python/requirements_dev.txt +++ b/python/requirements_dev.txt @@ -14,6 +14,7 @@ patchelf; platform_system == "Linux" pre-commit>=4.6.0,<5.0 black[jupyter] pyright +ruff # nanobind ABI-couples to the prebuilt mlir distro wheel (built against the # same version); a mismatch breaks Value-subclass recognition at runtime. Keep # this constraint identical to utils/mlir_wheels/requirements.in and diff --git a/python/utils/__init__.py b/python/utils/__init__.py index 17426d5e443..2b61da1ae04 100644 --- a/python/utils/__init__.py +++ b/python/utils/__init__.py @@ -5,285 +5,27 @@ # """Tensor factories, device helpers, and re-exports for the IRON runtime.""" -import logging -import os from typing import Any -import numpy as np - -# Prevent "No handlers could be found" warnings when aie is used as a library. -logging.getLogger("aie").addHandler(logging.NullHandler()) - -# Honour AIE_LOG_LEVEL env var (e.g. DEBUG, INFO, WARNING, ERROR, CRITICAL). -# This must be done before any aie submodule emits a log record. -_log_level_str = os.environ.get("AIE_LOG_LEVEL", "").upper() -if _log_level_str: - _log_level = getattr(logging, _log_level_str, None) - if _log_level is not None: - logging.getLogger("aie").setLevel(_log_level) - -_logger = logging.getLogger(__name__) - -from .hostruntime.tensor_class import Tensor - -# Capability probes for the two NPU host backends. Both are memoized and lazy: -# importing ``aie.utils`` no longer eagerly probes either backend. Runtime -# selection below probes only the backend it actually needs (so NPU_RUNTIME=hrx -# never imports pyxrt and NPU_RUNTIME=xrt never runs HRX discovery), and the -# public ``aie.utils.has_xrt`` / ``aie.utils.has_hrx`` attributes are served -# on-demand via module ``__getattr__`` so a bare capability query still works in -# any mode (including the default ``auto``) and pays for at most one probe. -_has_xrt: bool | None = None # tri-state cache; None => not probed yet -_has_hrx: bool | None = None - - -def _probe_xrt() -> bool: - """Whether ``pyxrt`` (the XRT userspace) imports on this host. - - Heavyweight -- importing pyxrt pulls in the XRT stack -- so it runs at most - once and only when XRT is actually needed or explicitly queried. - """ - global _has_xrt - if _has_xrt is None: - try: - import pyxrt # noqa: F401 # pyright: ignore[reportMissingImports] - - _has_xrt = True - except ImportError as e: - _logger.warning( - "Failed to import PyXRT: %s, proceeding without runtime libraries.", - e, - ) - _has_xrt = False - return _has_xrt - - -def _probe_hrx() -> bool: - """Whether ``libhrx.so`` can be located on this host. - - Filesystem-only (no dlopen, no device init), but still memoized so repeated - queries do no extra work. - """ - global _has_hrx - if _has_hrx is None: - try: - from .hostruntime.hrxruntime.discovery import hrx_available - - _has_hrx = hrx_available() - except Exception as e: # discovery must never break importing aie.utils - _logger.debug("HRX discovery probe failed: %s", e) - _has_hrx = False - return _has_hrx - - -# Host-runtime backend selection. ``NPU_RUNTIME`` chooses between the XRT and -# HRX host stacks; both consume the identical aiecc artifacts (final.xclbin + -# insts.bin) and only the dispatch path differs. Accepted values: -# xrt - force the XRT backend (falls back to CPU tensors if pyxrt is missing). -# hrx - force the HRX backend (error here if libhrx is not found). -# auto - (default) prefer XRT when present, else fall back to CPU. -# -# HRX is strictly opt-in: it is selected *only* when NPU_RUNTIME=hrx is set -# explicitly. ``auto`` never selects HRX -- the product contract is "XRT remains -# the default, HRX is opt-in", so an XRT-less host degrades to CPU rather than -# silently switching to HRX. -# -# NPU_RUNTIME is read *before* any capability probe so a forced backend only -# probes itself: 'hrx' never imports pyxrt, and 'xrt'/'auto' never run HRX -# discovery. Each backend's tensor/runtime module is likewise imported lazily -# (it dlopen()s / imports its own runtime on first use). -_NPU_RUNTIME = os.environ.get("NPU_RUNTIME", "auto").lower() - -# Strict product contract: an unset NPU_RUNTIME defaults to 'auto', but an -# explicitly *invalid* value is a hard error rather than a silent fallback -- -# a typo'd backend name must not quietly resolve to something else. -if _NPU_RUNTIME not in ("xrt", "hrx", "auto"): - raise ImportError( - f"Invalid NPU_RUNTIME={_NPU_RUNTIME!r}; expected one of xrt|hrx|auto " - f"(unset defaults to 'auto')." - ) - -if _NPU_RUNTIME == "hrx" and not _probe_hrx(): - raise ImportError( - "NPU_RUNTIME=hrx was requested but libhrx.so could not be located. " - "Install HRX to a standard location, or set HRX_DIR/LIBHRX_DIR. " - "Use NPU_RUNTIME=auto to fall back to XRT/CPU when HRX is absent." - ) - -# Resolve 'auto' to a concrete backend with graceful degradation. HRX is never -# auto-selected (opt-in only via NPU_RUNTIME=hrx), so 'auto' is XRT or CPU. -if _NPU_RUNTIME == "auto": - _NPU_RUNTIME = "xrt" if _probe_xrt() else "cpu" - -if _NPU_RUNTIME == "hrx": - from .hostruntime.hrxruntime.tensor import HRXTensor - - DEFAULT_TENSOR_CLASS = HRXTensor -elif _NPU_RUNTIME == "xrt" and _probe_xrt(): - from .hostruntime.xrtruntime.tensor import XRTTensor - - DEFAULT_TENSOR_CLASS = XRTTensor -else: - from .hostruntime.tensor_class import CPUOnlyTensor - - DEFAULT_TENSOR_CLASS = CPUOnlyTensor - - -def ceildiv(a, b): - """Ceiling division: smallest integer >= a/b.""" - return -(a // -b) - - -def tensor(*args, **kwargs): - """ - Create a tensor using the default tensor class. - - Passing a typed ``ndarray`` together with a mismatched ``dtype=`` - kwarg raises :class:`TypeError`. Matching kwargs are passed through - unchanged (the underlying tensor backend uses ``dtype`` for buffer - allocation, so silently stripping it would surprise callers). - - Args: - *args: Arguments passed to the tensor constructor. ``args[0]`` is - either a shape ``tuple`` or an array-like. - **kwargs: Keyword arguments passed to the tensor constructor. - - Returns: - Tensor: The created tensor. - """ - if args and isinstance(args[0], np.ndarray) and "dtype" in kwargs: - arr_dt = args[0].dtype - kw_dt = np.dtype(kwargs["dtype"]) - if arr_dt != kw_dt: - raise TypeError( - f"iron.tensor: ndarray dtype {arr_dt!r} does not match " - f"dtype= kwarg {kw_dt!r}. Cast the array beforehand " - f"(e.g. arr.astype({kw_dt!r})) or drop the dtype= kwarg." - ) - return DEFAULT_TENSOR_CLASS(*args, **kwargs) - - -def ones(*args, **kwargs): - """ - Create a tensor filled with ones using the default tensor class. - - Args: - *args: Arguments passed to the ones method. - **kwargs: Keyword arguments passed to the ones method. - - Returns: - Tensor: The created tensor. - """ - return DEFAULT_TENSOR_CLASS.ones(*args, **kwargs) - - -def zeros(*args, **kwargs): - """ - Create a tensor filled with zeros using the default tensor class. - - Args: - *args: Arguments passed to the zeros method. - **kwargs: Keyword arguments passed to the zeros method. - - Returns: - Tensor: The created tensor. - """ - return DEFAULT_TENSOR_CLASS.zeros(*args, **kwargs) - - -def full(*args, **kwargs): - """ - Create a tensor filled with a scalar value using the default tensor class. - - Args: - *args: Arguments passed to the full method (size, fill_value). - **kwargs: Keyword arguments passed to the full method. - - Returns: - Tensor: The created tensor. - """ - return DEFAULT_TENSOR_CLASS.full(*args, **kwargs) - - -def randint(*args, **kwargs): - """ - Create a tensor filled with random integers using the default tensor class. - - Args: - *args: Arguments passed to the randint method. - **kwargs: Keyword arguments passed to the randint method. - - Returns: - Tensor: The created tensor. - """ - return DEFAULT_TENSOR_CLASS.randint(*args, **kwargs) - - -def rand(*args, **kwargs): - """ - Create a tensor filled with random values using the default tensor class. - - Args: - *args: Arguments passed to the rand method. - **kwargs: Keyword arguments passed to the rand method. - - Returns: - Tensor: The created tensor. - """ - return DEFAULT_TENSOR_CLASS.rand(*args, **kwargs) - - -def arange(*args, **kwargs): - """ - Create a tensor with a range of values using the default tensor class. - - Args: - *args: Arguments passed to the arange method. - **kwargs: Keyword arguments passed to the arange method. - - Returns: - Tensor: The created tensor. - """ - return DEFAULT_TENSOR_CLASS.arange(*args, **kwargs) - - -def zeros_like(*args, **kwargs): - """ - Create a tensor filled with zeros with the same shape as another tensor using the default tensor class. - - Args: - *args: Arguments passed to the zeros_like method. - **kwargs: Keyword arguments passed to the zeros_like method. - - Returns: - Tensor: The created tensor. - """ - return DEFAULT_TENSOR_CLASS.zeros_like(*args, **kwargs) - - -def set_tensor_class(cls): - """ - Set the default tensor class. - - Args: - cls: The new default tensor class. Must inherit from Tensor. - - Raises: - ValueError: If cls does not inherit from Tensor. - """ - if not issubclass(cls, Tensor): - raise ValueError( - f"Tensors must inherit from the Tensor class but {cls} does not." - ) - global DEFAULT_TENSOR_CLASS - DEFAULT_TENSOR_CLASS = cls - - +from . import ( + _log_setup, # noqa: F401 # side effect: configure "aie" logging first + hostruntime, +) from .hostruntime import set_current_device -from . import hostruntime -from .hostruntime.hostruntime import HostRuntime -from .trace import TraceConfig -from .npukernel import NPUKernel +from .hostruntime.hostruntime import HostRuntime as HostRuntime +from .npukernel import NPUKernel as NPUKernel +from .tensor_factory import _NPU_RUNTIME, _probe_hrx, _probe_xrt +from .tensor_factory import arange as arange +from .tensor_factory import ceildiv as ceildiv +from .tensor_factory import full as full +from .tensor_factory import ones as ones +from .tensor_factory import rand as rand +from .tensor_factory import randint as randint +from .tensor_factory import set_tensor_class as set_tensor_class +from .tensor_factory import tensor as tensor +from .tensor_factory import zeros as zeros +from .tensor_factory import zeros_like as zeros_like +from .trace import TraceConfig as TraceConfig _DefaultNPURuntime = None diff --git a/python/utils/_log_setup.py b/python/utils/_log_setup.py new file mode 100644 index 00000000000..5d5aded3fc0 --- /dev/null +++ b/python/utils/_log_setup.py @@ -0,0 +1,21 @@ +# _log_setup.py -*- Python -*- +# +# Copyright (C) 2025-2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +# +"""Logging setup for the ``aie`` namespace, imported first (for its side +effect) by :mod:`aie.utils` so the level is configured before any +``aie.utils`` submodule import can itself emit a log record.""" + +import logging +import os + +# Prevent "No handlers could be found" warnings when aie is used as a library. +logging.getLogger("aie").addHandler(logging.NullHandler()) + +# Honour AIE_LOG_LEVEL env var (e.g. DEBUG, INFO, WARNING, ERROR, CRITICAL). +_log_level_str = os.environ.get("AIE_LOG_LEVEL", "").upper() +if _log_level_str: + _log_level = getattr(logging, _log_level_str, None) + if _log_level is not None: + logging.getLogger("aie").setLevel(_log_level) diff --git a/python/utils/callabledesign.py b/python/utils/callabledesign.py index 99e71040445..326b82ede05 100644 --- a/python/utils/callabledesign.py +++ b/python/utils/callabledesign.py @@ -32,7 +32,7 @@ import inspect as _inspect import logging from pathlib import Path -from typing import Any, Callable, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Callable from aie.utils.compile.cache.utils import _create_function_cache_key from aie.utils.compile.jit.compilabledesign import CompilableDesign diff --git a/python/utils/compile/__init__.py b/python/utils/compile/__init__.py index 340bad6ab6c..a48748c4a76 100644 --- a/python/utils/compile/__init__.py +++ b/python/utils/compile/__init__.py @@ -10,8 +10,8 @@ from .utils import ( compile_cxx_core_function, - compile_mlir_module, compile_external_kernel, + compile_mlir_module, resolve_target_arch, ) @@ -19,3 +19,11 @@ NPU_CACHE_HOME = Path( os.environ.get("NPU_CACHE_HOME", Path.home() / ".npu" / "cache") ).resolve() + +__all__ = [ + "compile_cxx_core_function", + "compile_mlir_module", + "compile_external_kernel", + "resolve_target_arch", + "NPU_CACHE_HOME", +] diff --git a/python/utils/compile/jit/__init__.py b/python/utils/compile/jit/__init__.py index 04c8ed52d94..0f79e33b6dc 100644 --- a/python/utils/compile/jit/__init__.py +++ b/python/utils/compile/jit/__init__.py @@ -2,10 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception """JIT compilation layer: CompilableDesign, compileconfig, markers, and context.""" -from .context import compile_context, get_compile_arg -from .markers import CompileTime, In, InOut, Out from .compilabledesign import CompilableDesign from .compileconfig import compileconfig +from .context import compile_context, get_compile_arg +from .markers import CompileTime, In, InOut, Out __all__ = [ "CompilableDesign", diff --git a/python/utils/compile/jit/_dma_size_parser.py b/python/utils/compile/jit/_dma_size_parser.py index c03f13c11ab..faba8b2d357 100644 --- a/python/utils/compile/jit/_dma_size_parser.py +++ b/python/utils/compile/jit/_dma_size_parser.py @@ -67,14 +67,14 @@ def parse_dma_sizes(kernel_dir: Path) -> list[int] | None: return None try: # Trigger AIE/aiex dialect registration before constructing the context. - from aie.dialects import aie as _aie # noqa: F401 - from aie.dialects import aiex as _aiex # noqa: F401 from aie import ( # pyright: ignore[reportMissingImports] ir, # pyright: ignore[reportAttributeAccessIssue] ) from aie._mlir_libs import ( # pyright: ignore[reportMissingImports] get_dialect_registry, # pyright: ignore[reportAttributeAccessIssue] ) + from aie.dialects import aie as _aie # noqa: F401 + from aie.dialects import aiex as _aiex # noqa: F401 ir_context = ir.Context # pyright: ignore[reportAttributeAccessIssue] ir_location = ir.Location # pyright: ignore[reportAttributeAccessIssue] diff --git a/python/utils/compile/jit/compilabledesign.py b/python/utils/compile/jit/compilabledesign.py index 5b405d54533..94b4b2d8ea3 100644 --- a/python/utils/compile/jit/compilabledesign.py +++ b/python/utils/compile/jit/compilabledesign.py @@ -37,6 +37,13 @@ from types import MappingProxyType from typing import Any, Callable, Mapping +from aie.extras.context import mlir_mod_ctx # pyright: ignore[reportMissingImports] +from aie.ir import ( # pyright: ignore[reportMissingImports] + Module as _Module, # pyright: ignore[reportAttributeAccessIssue] +) +from aie.ir import ( # pyright: ignore[reportMissingImports] + StringAttr, # pyright: ignore[reportAttributeAccessIssue] +) from aie.utils.compile import ( NPU_CACHE_HOME, compile_external_kernel, @@ -44,11 +51,6 @@ ) from aie.utils.compile.cache.utils import file_lock from aie.utils.compile.utils import _cleanup_failed_compilation -from aie.extras.context import mlir_mod_ctx # pyright: ignore[reportMissingImports] -from aie.ir import ( # pyright: ignore[reportMissingImports] - Module as _Module, # pyright: ignore[reportAttributeAccessIssue] - StringAttr, # pyright: ignore[reportAttributeAccessIssue] -) from ._dma_size_parser import parse_dma_sizes from ._hash import ( @@ -61,11 +63,9 @@ _introspect_generator, _is_compile_param, _is_tensor_param, - split_params, ) -from ._serialization import _TensorPlaceholder, _decode_kwarg, _encode_kwarg +from ._serialization import _decode_kwarg, _encode_kwarg, _TensorPlaceholder from .context import compile_context -from .markers import CompileTime, In, InOut, Out # A waiter on this lock is waiting out someone else's *compile*, not just an # acquisition, so the bound has to exceed a full build rather than a handshake. @@ -603,7 +603,7 @@ def split_runtime_args( Returns: ``(tensor_args, scalar_kwargs)`` """ - from aie.iron.kernel import ExternalFunction, Kernel + from aie.iron.kernel import Kernel if not callable(self.mlir_generator): # Static .mlir file: pass everything through as tensors, diff --git a/python/utils/compile/utils.py b/python/utils/compile/utils.py index d622ad55b54..3e4a5ac003d 100644 --- a/python/utils/compile/utils.py +++ b/python/utils/compile/utils.py @@ -13,6 +13,7 @@ import tempfile from pathlib import Path from typing import TYPE_CHECKING + import aie.utils.config as config if TYPE_CHECKING: @@ -500,15 +501,14 @@ def compile_mlir_module( # the loop in compilabledesign.py but for callers (e.g. low-level # designs using rt.inline_ops) that didn't go through @iron.jit. if work_dir and device is not None: - try: - from aie.iron.kernel import ExternalFunction - except ImportError: - ExternalFunction = None - if ExternalFunction is not None: - target_arch = resolve_target_arch(device) - for func in list(ExternalFunction._instances): - if not func._compiled and getattr(func, "_source_file", None): - compile_external_kernel(func, str(work_dir), target_arch) + # Deferred: aie.iron's __init__ imports back into aie.utils.compile.jit, + # so a module-level import here deadlocks on a cold aie.utils.compile entry. + from aie.iron.kernel import ExternalFunction + + target_arch = resolve_target_arch(device) + for func in list(ExternalFunction._instances): + if not func._compiled and getattr(func, "_source_file", None): + compile_external_kernel(func, str(work_dir), target_arch) # When work_dir is provided, invoke the aiecc binary as a subprocess so # that it resolves relative link_with paths (e.g. "add_one.o") against the diff --git a/python/utils/config.py b/python/utils/config.py index c561332030b..1f4b2197633 100644 --- a/python/utils/config.py +++ b/python/utils/config.py @@ -6,6 +6,7 @@ import os import shutil + import aie.utils.configure as config # pyright: ignore[reportMissingImports] diff --git a/python/utils/hostruntime/__init__.py b/python/utils/hostruntime/__init__.py index f2c96765eab..ce3055f003a 100644 --- a/python/utils/hostruntime/__init__.py +++ b/python/utils/hostruntime/__init__.py @@ -6,8 +6,10 @@ """Host runtime utilities: device selection, tensor allocation, and numerical helpers.""" from typing import TYPE_CHECKING -from ml_dtypes import bfloat16 + import numpy as np +from ml_dtypes import bfloat16 + from .tensor_class import Tensor if TYPE_CHECKING: diff --git a/python/utils/hostruntime/hostruntime.py b/python/utils/hostruntime/hostruntime.py index 3f8a7f98a0a..25b12e5cf9e 100644 --- a/python/utils/hostruntime/hostruntime.py +++ b/python/utils/hostruntime/hostruntime.py @@ -1,24 +1,25 @@ # Copyright (C) 2025-2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: Apache-2.0 -from abc import ABC, abstractmethod import logging -import numpy as np import sys +from abc import ABC, abstractmethod from pathlib import Path from typing import TYPE_CHECKING -logger = logging.getLogger(__name__) +import numpy as np -from .. import tensor +from ..tensor_factory import tensor if TYPE_CHECKING: from aie.iron.device import Device -from .tensor_class import Tensor +from ..npukernel import NPUKernel from ..trace import TraceConfig from ..trace.utils import create_ctrl_pkt, extract_tile -from ..npukernel import NPUKernel from . import bfloat16_safe_allclose +from .tensor_class import Tensor + +logger = logging.getLogger(__name__) class HostRuntimeError(Exception): @@ -82,7 +83,7 @@ def has_trace(self) -> bool: Returns: bool: True if trace configuration is present, False otherwise. """ - return not (self._trace_config is None) + return self._trace_config is not None @abstractmethod def is_success(self) -> bool: diff --git a/python/utils/hostruntime/hrxruntime/__init__.py b/python/utils/hostruntime/hrxruntime/__init__.py index 6f41441645e..4cbee08b85e 100644 --- a/python/utils/hostruntime/hrxruntime/__init__.py +++ b/python/utils/hostruntime/hrxruntime/__init__.py @@ -25,4 +25,5 @@ "HRXContext", "HRXError", "control_code_from_elf", + "_hrx_sync_timeout_s", ] diff --git a/python/utils/hostruntime/hrxruntime/context.py b/python/utils/hostruntime/hrxruntime/context.py index 0717c8cb4cb..e3da145b95e 100644 --- a/python/utils/hostruntime/hrxruntime/context.py +++ b/python/utils/hostruntime/hrxruntime/context.py @@ -32,8 +32,8 @@ HrxBufferRef, HrxConstByteSpan, HrxDispatchConfig, - HrxStringView, HRXError, + HrxStringView, _check, _handle, _hrx_sync_timeout_s, diff --git a/python/utils/hostruntime/hrxruntime/hostruntime.py b/python/utils/hostruntime/hrxruntime/hostruntime.py index e57ed1de04a..56dc8677f8e 100644 --- a/python/utils/hostruntime/hrxruntime/hostruntime.py +++ b/python/utils/hostruntime/hrxruntime/hostruntime.py @@ -24,8 +24,8 @@ from typing import TYPE_CHECKING from ..hostruntime import HostRuntime, HostRuntimeError, KernelHandle, KernelResult -from .tensor import HRXTensor from . import HRXContext, HRXError, control_code_from_elf +from .tensor import HRXTensor if TYPE_CHECKING: from aie.iron.device import Device diff --git a/python/utils/hostruntime/hrxruntime/tensor.py b/python/utils/hostruntime/hrxruntime/tensor.py index d13e6b70506..5404a3f6d37 100644 --- a/python/utils/hostruntime/hrxruntime/tensor.py +++ b/python/utils/hostruntime/hrxruntime/tensor.py @@ -14,10 +14,11 @@ """ import ctypes + import numpy as np +from aie.helpers.util import np_ndarray_type_get_shape from ..tensor_class import Tensor -from aie.helpers.util import np_ndarray_type_get_shape from . import HRXContext diff --git a/python/utils/hostruntime/tensor_class.py b/python/utils/hostruntime/tensor_class.py index 3cb06c80c20..613d521c0b8 100644 --- a/python/utils/hostruntime/tensor_class.py +++ b/python/utils/hostruntime/tensor_class.py @@ -5,6 +5,7 @@ # from abc import ABC, abstractmethod from functools import cached_property + import numpy as np import numpy.typing as npt @@ -18,8 +19,8 @@ def _ml_dtype_to_torch_map(): global _ML_DTYPE_TO_TORCH if _ML_DTYPE_TO_TORCH is None: - import torch # pyright: ignore[reportMissingImports] import ml_dtypes + import torch # pyright: ignore[reportMissingImports] _candidates = { ml_dtypes.bfloat16: torch.bfloat16, diff --git a/python/utils/hostruntime/xrtruntime/__init__.py b/python/utils/hostruntime/xrtruntime/__init__.py index a5336fe858e..2bc30a9edcd 100644 --- a/python/utils/hostruntime/xrtruntime/__init__.py +++ b/python/utils/hostruntime/xrtruntime/__init__.py @@ -5,6 +5,6 @@ # try: - import pyxrt as xrt # pyright: ignore[reportMissingImports] + import pyxrt # noqa: F401 # pyright: ignore[reportMissingImports] except Exception as e: raise ImportError(f"Cannot import pyxrt (err={e})... is XRT installed?") diff --git a/python/utils/hostruntime/xrtruntime/hostruntime.py b/python/utils/hostruntime/xrtruntime/hostruntime.py index 85e48c80aed..a64acaa0904 100644 --- a/python/utils/hostruntime/xrtruntime/hostruntime.py +++ b/python/utils/hostruntime/xrtruntime/hostruntime.py @@ -6,23 +6,24 @@ """ import atexit +import gc import logging -from collections import OrderedDict import os import shutil +import subprocess import time import weakref -import gc -import subprocess +from collections import OrderedDict from pathlib import Path from typing import TYPE_CHECKING -import numpy as np + import pyxrt # pyright: ignore[reportMissingImports] from ..hostruntime import HostRuntime, HostRuntimeError, KernelHandle, KernelResult if TYPE_CHECKING: from aie.iron.device import Device + from ...trace import TraceConfig from .tensor import XRTTensor diff --git a/python/utils/hostruntime/xrtruntime/parameter_scratchpad.py b/python/utils/hostruntime/xrtruntime/parameter_scratchpad.py index 6705fe9264f..87171d823a9 100644 --- a/python/utils/hostruntime/xrtruntime/parameter_scratchpad.py +++ b/python/utils/hostruntime/xrtruntime/parameter_scratchpad.py @@ -26,7 +26,6 @@ class (exposed via pybind11). from pathlib import Path import pyxrt # pyright: ignore[reportMissingImports] - from aie._mlir_libs._parameter_scratchpad import ( # pyright: ignore[reportMissingImports] ParameterScratchpad as _ParameterScratchpadImpl, ) diff --git a/python/utils/hostruntime/xrtruntime/tensor.py b/python/utils/hostruntime/xrtruntime/tensor.py index dcb73799f38..bc2dc6f34d1 100644 --- a/python/utils/hostruntime/xrtruntime/tensor.py +++ b/python/utils/hostruntime/xrtruntime/tensor.py @@ -6,9 +6,9 @@ import numpy as np import pyxrt as xrt # pyright: ignore[reportMissingImports] +from aie.helpers.util import np_ndarray_type_get_shape from ..tensor_class import Tensor -from aie.helpers.util import np_ndarray_type_get_shape class XRTTensor(Tensor): diff --git a/python/utils/jit.py b/python/utils/jit.py index 7f3835fc7c2..1abaa261db7 100644 --- a/python/utils/jit.py +++ b/python/utils/jit.py @@ -86,7 +86,7 @@ def jit(mlir_generator: Callable | None = None, **kwargs): # --- Validate CompileTime[T] params when generator is callable --- if callable(mlir_generator): - from aie.utils.compile.jit.compilabledesign import split_params + from aie.utils.compile.jit._introspect import split_params compile_params, _, scalar_params = split_params(mlir_generator) @@ -162,7 +162,7 @@ def jit(mlir_generator: Callable | None = None, **kwargs): f" def {mlir_generator.__name__}(a: In, b: Out, *, " + ", ".join(f"{n}: CompileTime[...]" for n in non_kw_compile_params) + "):\n" - f" ..." + " ..." ) return _CallableDesign( diff --git a/python/utils/ml.py b/python/utils/ml.py index c5abbdc03fa..100221cd664 100644 --- a/python/utils/ml.py +++ b/python/utils/ml.py @@ -18,7 +18,6 @@ import json import math import os -import sys import numpy as np import torch # pyright: ignore[reportMissingImports] @@ -84,10 +83,8 @@ def unpickle(file): def fuse_single_conv_bn_pair(bn_mean, bn_var, bn_wts, bn_bias, conv_wts): # https://github.com/ChoiDM/Pytorch_BN_Fold/blob/master/bn_fold.py eps = 1e-05 - mu = bn_mean var = bn_var gamma = bn_wts - beta = bn_bias W = conv_wts denom = torch.sqrt(var + eps) @@ -413,7 +410,7 @@ def run_conv_torch_test( Default ``None`` (no dumps). """ import aie.iron as iron - from aie.utils import HostRuntime, NPUKernel, DefaultNPURuntime, TraceConfig + from aie.utils import DefaultNPURuntime, HostRuntime, NPUKernel, TraceConfig dtype_in = np.dtype(dtype_in) dtype_wts = np.dtype(dtype_wts) diff --git a/python/utils/npukernel.py b/python/utils/npukernel.py index 05a25abb603..e31482f2a1f 100644 --- a/python/utils/npukernel.py +++ b/python/utils/npukernel.py @@ -3,7 +3,6 @@ # Copyright (C) 2025-2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # -from pathlib import Path from .trace import TraceConfig diff --git a/python/utils/regdb.py b/python/utils/regdb.py index 2fc35a094e1..bddec491954 100644 --- a/python/utils/regdb.py +++ b/python/utils/regdb.py @@ -646,9 +646,9 @@ def annotate_module(self, module) -> int: Number of operations annotated """ # Import here to avoid circular imports and allow module to load without MLIR - from aie.extras.util import find_ops # pyright: ignore[reportMissingImports] import aie.dialects.aie as aiedialect import aie.dialects.aiex as aiexdialect + from aie.extras.util import find_ops # pyright: ignore[reportMissingImports] # These op classes / enums come through compiled dialect bindings that # pyright can't see; fetch them dynamically so the static checker is happy. @@ -747,13 +747,13 @@ def annotate_file(self, input_path: str, output_path: Optional[str] = None) -> i Number of operations annotated """ # Import here to avoid circular imports and allow module to load without MLIR + from aie._mlir_libs import ( # pyright: ignore[reportMissingImports] + get_dialect_registry, # pyright: ignore[reportAttributeAccessIssue] + ) from aie.ir import ( # pyright: ignore[reportMissingImports] Context, # pyright: ignore[reportAttributeAccessIssue] - Module, # pyright: ignore[reportAttributeAccessIssue] Location, # pyright: ignore[reportAttributeAccessIssue] - ) - from aie._mlir_libs import ( # pyright: ignore[reportMissingImports] - get_dialect_registry, # pyright: ignore[reportAttributeAccessIssue] + Module, # pyright: ignore[reportAttributeAccessIssue] ) # Read input file diff --git a/python/utils/tensor_factory.py b/python/utils/tensor_factory.py new file mode 100644 index 00000000000..0eb25cab207 --- /dev/null +++ b/python/utils/tensor_factory.py @@ -0,0 +1,274 @@ +# tensor_factory.py -*- Python -*- +# +# Copyright (C) 2025-2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +# +"""Tensor factories and NPU host-backend selection. + +Split out from :mod:`aie.utils` so that :mod:`aie.utils.hostruntime.hostruntime` +(which needs the ``tensor()`` factory) can import it without importing back +through ``aie.utils.__init__`` -- that reverse edge is what used to force +``aie.utils.__init__``'s own imports of ``HostRuntime`` etc. to be deferred +past this module's definitions. +""" + +import logging +import os + +import numpy as np + +from .hostruntime.tensor_class import Tensor + +_logger = logging.getLogger(__name__) + +# Capability probes for the two NPU host backends. Both are memoized and lazy: +# importing ``aie.utils`` no longer eagerly probes either backend. Runtime +# selection below probes only the backend it actually needs (so NPU_RUNTIME=hrx +# never imports pyxrt and NPU_RUNTIME=xrt never runs HRX discovery), and the +# public ``aie.utils.has_xrt`` / ``aie.utils.has_hrx`` attributes are served +# on-demand via module ``__getattr__`` so a bare capability query still works in +# any mode (including the default ``auto``) and pays for at most one probe. +_has_xrt: bool | None = None # tri-state cache; None => not probed yet +_has_hrx: bool | None = None + + +def _probe_xrt() -> bool: + """Whether ``pyxrt`` (the XRT userspace) imports on this host. + + Heavyweight -- importing pyxrt pulls in the XRT stack -- so it runs at most + once and only when XRT is actually needed or explicitly queried. + """ + global _has_xrt + if _has_xrt is None: + try: + import pyxrt # noqa: F401 # pyright: ignore[reportMissingImports] + + _has_xrt = True + except ImportError as e: + _logger.warning( + "Failed to import PyXRT: %s, proceeding without runtime libraries.", + e, + ) + _has_xrt = False + return _has_xrt + + +def _probe_hrx() -> bool: + """Whether ``libhrx.so`` can be located on this host. + + Filesystem-only (no dlopen, no device init), but still memoized so repeated + queries do no extra work. + """ + global _has_hrx + if _has_hrx is None: + try: + from .hostruntime.hrxruntime.discovery import hrx_available + + _has_hrx = hrx_available() + except Exception as e: # discovery must never break importing aie.utils + _logger.debug("HRX discovery probe failed: %s", e) + _has_hrx = False + return _has_hrx + + +# Host-runtime backend selection. ``NPU_RUNTIME`` chooses between the XRT and +# HRX host stacks; both consume the identical aiecc artifacts (final.xclbin + +# insts.bin) and only the dispatch path differs. Accepted values: +# xrt - force the XRT backend (falls back to CPU tensors if pyxrt is missing). +# hrx - force the HRX backend (error here if libhrx is not found). +# auto - (default) prefer XRT when present, else fall back to CPU. +# +# HRX is strictly opt-in: it is selected *only* when NPU_RUNTIME=hrx is set +# explicitly. ``auto`` never selects HRX -- the product contract is "XRT remains +# the default, HRX is opt-in", so an XRT-less host degrades to CPU rather than +# silently switching to HRX. +# +# NPU_RUNTIME is read *before* any capability probe so a forced backend only +# probes itself: 'hrx' never imports pyxrt, and 'xrt'/'auto' never run HRX +# discovery. Each backend's tensor/runtime module is likewise imported lazily +# (it dlopen()s / imports its own runtime on first use). +_NPU_RUNTIME = os.environ.get("NPU_RUNTIME", "auto").lower() + +# Strict product contract: an unset NPU_RUNTIME defaults to 'auto', but an +# explicitly *invalid* value is a hard error rather than a silent fallback -- +# a typo'd backend name must not quietly resolve to something else. +if _NPU_RUNTIME not in ("xrt", "hrx", "auto"): + raise ImportError( + f"Invalid NPU_RUNTIME={_NPU_RUNTIME!r}; expected one of xrt|hrx|auto " + f"(unset defaults to 'auto')." + ) + +if _NPU_RUNTIME == "hrx" and not _probe_hrx(): + raise ImportError( + "NPU_RUNTIME=hrx was requested but libhrx.so could not be located. " + "Install HRX to a standard location, or set HRX_DIR/LIBHRX_DIR. " + "Use NPU_RUNTIME=auto to fall back to XRT/CPU when HRX is absent." + ) + +# Resolve 'auto' to a concrete backend with graceful degradation. HRX is never +# auto-selected (opt-in only via NPU_RUNTIME=hrx), so 'auto' is XRT or CPU. +if _NPU_RUNTIME == "auto": + _NPU_RUNTIME = "xrt" if _probe_xrt() else "cpu" + +if _NPU_RUNTIME == "hrx": + from .hostruntime.hrxruntime.tensor import HRXTensor + + DEFAULT_TENSOR_CLASS = HRXTensor +elif _NPU_RUNTIME == "xrt" and _probe_xrt(): + from .hostruntime.xrtruntime.tensor import XRTTensor + + DEFAULT_TENSOR_CLASS = XRTTensor +else: + from .hostruntime.tensor_class import CPUOnlyTensor + + DEFAULT_TENSOR_CLASS = CPUOnlyTensor + + +def ceildiv(a, b): + """Ceiling division: smallest integer >= a/b.""" + return -(a // -b) + + +def tensor(*args, **kwargs): + """ + Create a tensor using the default tensor class. + + Passing a typed ``ndarray`` together with a mismatched ``dtype=`` + kwarg raises :class:`TypeError`. Matching kwargs are passed through + unchanged (the underlying tensor backend uses ``dtype`` for buffer + allocation, so silently stripping it would surprise callers). + + Args: + *args: Arguments passed to the tensor constructor. ``args[0]`` is + either a shape ``tuple`` or an array-like. + **kwargs: Keyword arguments passed to the tensor constructor. + + Returns: + Tensor: The created tensor. + """ + if args and isinstance(args[0], np.ndarray) and "dtype" in kwargs: + arr_dt = args[0].dtype + kw_dt = np.dtype(kwargs["dtype"]) + if arr_dt != kw_dt: + raise TypeError( + f"iron.tensor: ndarray dtype {arr_dt!r} does not match " + f"dtype= kwarg {kw_dt!r}. Cast the array beforehand " + f"(e.g. arr.astype({kw_dt!r})) or drop the dtype= kwarg." + ) + return DEFAULT_TENSOR_CLASS(*args, **kwargs) + + +def ones(*args, **kwargs): + """ + Create a tensor filled with ones using the default tensor class. + + Args: + *args: Arguments passed to the ones method. + **kwargs: Keyword arguments passed to the ones method. + + Returns: + Tensor: The created tensor. + """ + return DEFAULT_TENSOR_CLASS.ones(*args, **kwargs) + + +def zeros(*args, **kwargs): + """ + Create a tensor filled with zeros using the default tensor class. + + Args: + *args: Arguments passed to the zeros method. + **kwargs: Keyword arguments passed to the zeros method. + + Returns: + Tensor: The created tensor. + """ + return DEFAULT_TENSOR_CLASS.zeros(*args, **kwargs) + + +def full(*args, **kwargs): + """ + Create a tensor filled with a scalar value using the default tensor class. + + Args: + *args: Arguments passed to the full method (size, fill_value). + **kwargs: Keyword arguments passed to the full method. + + Returns: + Tensor: The created tensor. + """ + return DEFAULT_TENSOR_CLASS.full(*args, **kwargs) + + +def randint(*args, **kwargs): + """ + Create a tensor filled with random integers using the default tensor class. + + Args: + *args: Arguments passed to the randint method. + **kwargs: Keyword arguments passed to the randint method. + + Returns: + Tensor: The created tensor. + """ + return DEFAULT_TENSOR_CLASS.randint(*args, **kwargs) + + +def rand(*args, **kwargs): + """ + Create a tensor filled with random values using the default tensor class. + + Args: + *args: Arguments passed to the rand method. + **kwargs: Keyword arguments passed to the rand method. + + Returns: + Tensor: The created tensor. + """ + return DEFAULT_TENSOR_CLASS.rand(*args, **kwargs) + + +def arange(*args, **kwargs): + """ + Create a tensor with a range of values using the default tensor class. + + Args: + *args: Arguments passed to the arange method. + **kwargs: Keyword arguments passed to the arange method. + + Returns: + Tensor: The created tensor. + """ + return DEFAULT_TENSOR_CLASS.arange(*args, **kwargs) + + +def zeros_like(*args, **kwargs): + """ + Create a tensor filled with zeros with the same shape as another tensor using the default tensor class. + + Args: + *args: Arguments passed to the zeros_like method. + **kwargs: Keyword arguments passed to the zeros_like method. + + Returns: + Tensor: The created tensor. + """ + return DEFAULT_TENSOR_CLASS.zeros_like(*args, **kwargs) + + +def set_tensor_class(cls): + """ + Set the default tensor class. + + Args: + cls: The new default tensor class. Must inherit from Tensor. + + Raises: + ValueError: If cls does not inherit from Tensor. + """ + if not issubclass(cls, Tensor): + raise ValueError( + f"Tensors must inherit from the Tensor class but {cls} does not." + ) + global DEFAULT_TENSOR_CLASS + DEFAULT_TENSOR_CLASS = cls diff --git a/python/utils/trace/__init__.py b/python/utils/trace/__init__.py index a860132ed61..30633e002e6 100644 --- a/python/utils/trace/__init__.py +++ b/python/utils/trace/__init__.py @@ -434,21 +434,41 @@ from .config import TraceConfig from .parse import parse_trace from .setup import ( - configure_packet_ctrl_flow, config_ctrl_pkts_aie, + configure_packet_ctrl_flow, configure_trace, start_trace, ) from .utils import ( - parity, - extract_tile, - pack4bytes, create_ctrl_pkt, - get_kernel_code, extract_buffers, + extract_tile, get_cycles, get_cycles_summary, - print_cycles_summary, + get_kernel_code, get_vector_time, + pack4bytes, + parity, + print_cycles_summary, split_trace_segments, ) + +__all__ = [ + "TraceConfig", + "parse_trace", + "configure_packet_ctrl_flow", + "config_ctrl_pkts_aie", + "configure_trace", + "start_trace", + "parity", + "extract_tile", + "pack4bytes", + "create_ctrl_pkt", + "get_kernel_code", + "extract_buffers", + "get_cycles", + "get_cycles_summary", + "print_cycles_summary", + "get_vector_time", + "split_trace_segments", +] diff --git a/python/utils/trace/config.py b/python/utils/trace/config.py index 6744f4a60a1..fe24593eac9 100644 --- a/python/utils/trace/config.py +++ b/python/utils/trace/config.py @@ -1,10 +1,11 @@ # Copyright (C) 2025-2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: Apache-2.0 -import numpy as np import json + +import numpy as np + from .parse import parse_trace -from .utils import parity, extract_tile class TraceConfig: diff --git a/python/utils/trace/event_ir.py b/python/utils/trace/event_ir.py index e257f71efdc..422ee0c7196 100755 --- a/python/utils/trace/event_ir.py +++ b/python/utils/trace/event_ir.py @@ -1,37 +1,29 @@ #!/usr/bin/env python3 # Copyright (C) 2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -import json import argparse -import sys +import json +import logging +import os import re -import subprocess import shutil -import os -import logging - -logger = logging.getLogger(__name__) +import subprocess +import sys -from .utils import ( - parse_pkt_hdr_in_stream, - trace_pkts_de_interleave, - convert_to_byte_stream, - convert_to_commands, -) from .events import ( NUM_TRACE_TYPES, CoreEvent, MemEvent, - ShimTileEvent, - MemTileEvent, ) +logger = logging.getLogger(__name__) + NUM_EVENTS = 8 # number of events we can view per trace rowoffset = 1 # TODO: temporary workaround to determine row offset for AIE2 tiles -eventIRFile = "eventIR.txt" -tmpTraceDirName = "tmpTrace" +event_ir_file = "eventIR.txt" +tmp_trace_dir_name = "tmpTrace" def parse_args(): @@ -87,7 +79,7 @@ def deactivate( multiples, active_events, timer, cycles, pid, trace_type, loc, pid_events ): for k in active_events.keys(): # an active event - if cycles > 0 or (cycles == 0 and not k in multiples): + if cycles > 0 or (cycles == 0 and k not in multiples): # if not k in multiples: # active event it not in multiples list if active_events[k] > 0: # trace_event = {'name':"Event"+str(k)} @@ -235,7 +227,7 @@ def convert_commands_to_json(trace_events, commands, pid_events): timer = timer + cycles for k in c.keys(): - if not "event" in k: + if "event" not in k: continue # If its already started, don't start it again ... try: @@ -318,9 +310,9 @@ def parse_mlir_trace_events(lines): # TODO Need to check if this line is commented out, check for // ? (harder to check of /* */) # TODO Need to support value in hex with 0x or decimal - # pattern = r"AIEX.npu.write32\s*\{\s*(\w+)\s*=\s*(\d+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(\d+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(\w+)\s*:\s*\w+\s*\}" - # pattern = r"AIEX.npu.write32\s*\{\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*\}" - pattern = r"aiex.npu.write32\s*\{\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*\}" + # pattern = r"AIEX.npu.write32\s*\{\s*(\w+)\s*=\s*(\d+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(\d+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(\w+)\s*:\s*\w+\s*\}" # noqa: E501 + # pattern = r"AIEX.npu.write32\s*\{\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*\}" # noqa: E501 + pattern = r"aiex.npu.write32\s*\{\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*,\s*(\w+)\s*=\s*(0x)?(\w+)\s*:\s*\w+\s*\}" # noqa: E501 pid_events = list() for t in range(NUM_TRACE_TYPES): @@ -370,7 +362,7 @@ def parse_mlir_trace_events(lines): # core event 0 if address == 0x340E0: # 213216, match ignoring case - if pid_events[0].get(key) == None: + if pid_events[0].get(key) is None: pid_events[0][key] = [ 0, 0, @@ -387,7 +379,7 @@ def parse_mlir_trace_events(lines): pid_events[0][key][3] = (value >> 24) & 0xFF # core event 1 elif address == 0x340E4: # 213220, match ignoring case - if pid_events[0].get(key) == None: + if pid_events[0].get(key) is None: pid_events[0][key] = [ 0, 0, @@ -404,7 +396,7 @@ def parse_mlir_trace_events(lines): pid_events[0][key][7] = (value >> 24) & 0xFF # mem event 0 elif address == 0x140E0: # 82144 - if pid_events[1].get(key) == None: + if pid_events[1].get(key) is None: pid_events[1][key] = [ 0, 0, @@ -421,7 +413,7 @@ def parse_mlir_trace_events(lines): pid_events[1][key][3] = (value >> 24) & 0xFF # mem event 1 elif address == 0x140E4: # 82148 - if pid_events[1].get(key) == None: + if pid_events[1].get(key) is None: pid_events[1][key] = [ 0, 0, @@ -559,7 +551,7 @@ def setup_trace_metadata(trace_events, pid_events): pid = pid + 1 -def convert_eventIR_to_json(trace_events, lines, pid_events): +def convert_event_ir_to_json(trace_events, lines, pid_events): check_time = True time_pattern = r"#(\d+)" signal_pattern = r"(\d)\s+(\d)_(\d)\s+cm\.et\.(\d+)" @@ -625,11 +617,11 @@ def create_target(): def print_config_json(pid_events): loc = None - eventArray = None + event_array = None for key, value in pid_events[0].items(): loc = key - eventArray = value - if loc is None or eventArray is None: + event_array = value + if loc is None or event_array is None: return try: with open("config.json", "wt") as f: @@ -652,14 +644,14 @@ def print_config_json(pid_events): f.write(' "start_event": 1,\n') f.write(' "stop_event": 0,\n') f.write(' "traced_events": [\n') - f.write(" " + str(eventArray[0]) + ",\n") - f.write(" " + str(eventArray[1]) + ",\n") - f.write(" " + str(eventArray[2]) + ",\n") - f.write(" " + str(eventArray[3]) + ",\n") - f.write(" " + str(eventArray[4]) + ",\n") - f.write(" " + str(eventArray[5]) + ",\n") - f.write(" " + str(eventArray[6]) + ",\n") - f.write(" " + str(eventArray[7]) + "\n") + f.write(" " + str(event_array[0]) + ",\n") + f.write(" " + str(event_array[1]) + ",\n") + f.write(" " + str(event_array[2]) + ",\n") + f.write(" " + str(event_array[3]) + ",\n") + f.write(" " + str(event_array[4]) + ",\n") + f.write(" " + str(event_array[5]) + ",\n") + f.write(" " + str(event_array[6]) + ",\n") + f.write(" " + str(event_array[7]) + "\n") f.write(" ],\n") f.write(' "group_event_config": {\n') f.write(' "2": 0,\n') @@ -827,28 +819,28 @@ def print_config_json(pid_events): # Right now, we're just checking if trace file has 0x before it (needed for hwfrontend) # If not, we prepend it -def fix_raw_trace_data(rawTraceFile, srcTraceFile): - with open(rawTraceFile, "rt") as inFile: - first_line = inFile.readline() +def fix_raw_trace_data(raw_trace_file, src_trace_file): + with open(raw_trace_file, "rt") as in_file: + first_line = in_file.readline() if first_line[:2] != "0x": - sed_cmd = "sed 's/^/0x/g' " + rawTraceFile + " > " + srcTraceFile + sed_cmd = "sed 's/^/0x/g' " + raw_trace_file + " > " + src_trace_file subprocess.call([sed_cmd], shell=True) else: - shutil.copy(rawTraceFile, srcTraceFile) + shutil.copy(raw_trace_file, src_trace_file) -def run_hwfrontend(fileInName, fileOutName): +def run_hwfrontend(file_in_name, file_out_name): result = subprocess.run( [ "hwfrontend", "--trace", - fileInName, + file_in_name, "--trace_config", "config.json", "--pkg-dir", ".", "--outfile", - fileOutName, + file_out_name, ], capture_output=True, text=True, @@ -881,19 +873,19 @@ def run_hwfrontend(fileInName, fileOutName): colshift = int(opts.colshift) if opts.colshift else 0 try: - os.mkdir(tmpTraceDirName) + os.mkdir(tmp_trace_dir_name) except FileExistsError: pass - logger.info("created temporary directory %s", tmpTraceDirName) - tmpTraceDir = os.path.abspath(tmpTraceDirName) + logger.info("created temporary directory %s", tmp_trace_dir_name) + tmp_trace_dir = os.path.abspath(tmp_trace_dir_name) - mlirFile = os.path.abspath(opts.mlir) - rawTraceFile = os.path.abspath(opts.filename) - srcTraceFileName = "prep." + str(opts.filename) - srcTraceFile = os.path.join(tmpTraceDir, srcTraceFileName) + mlir_file = os.path.abspath(opts.mlir) + raw_trace_file = os.path.abspath(opts.filename) + src_trace_file_name = "prep." + str(opts.filename) + src_trace_file = os.path.join(tmp_trace_dir, src_trace_file_name) # Check source file and prepend 0x - fix_raw_trace_data(rawTraceFile, srcTraceFile) + fix_raw_trace_data(raw_trace_file, src_trace_file) if opts.mlir: try: @@ -904,20 +896,20 @@ def run_hwfrontend(fileInName, fileOutName): logger.error("%s", e) sys.exit(1) - os.chdir(tmpTraceDirName) + os.chdir(tmp_trace_dir_name) create_target() print_config_json(pid_events) - run_hwfrontend(srcTraceFile, eventIRFile) + run_hwfrontend(src_trace_file, event_ir_file) # with open(opts.filename, "r") as f: try: - with open(eventIRFile, "rt") as f: + with open(event_ir_file, "rt") as f: lines = f.read().split("\n") ignore = [""] - lines = [l for l in lines if not l in ignore] + lines = [line for line in lines if line not in ignore] except Exception as e: logger.error("%s", e) sys.exit(1) @@ -927,6 +919,6 @@ def run_hwfrontend(fileInName, fileOutName): setup_trace_metadata(trace_events, pid_events) logger.debug("pid events: %s", pid_events) - convert_eventIR_to_json(trace_events, lines, pid_events) + convert_event_ir_to_json(trace_events, lines, pid_events) print(json.dumps(trace_events)) diff --git a/python/utils/trace/events/__init__.py b/python/utils/trace/events/__init__.py index d0a62cc38e6..4ed92516f81 100644 --- a/python/utils/trace/events/__init__.py +++ b/python/utils/trace/events/__init__.py @@ -12,27 +12,26 @@ Use get_events_for_device() to select the correct architecture. """ +import typing from enum import IntEnum from types import SimpleNamespace -import typing from aie.dialects._aie_enum_gen import ( # pyright: ignore[reportMissingImports] CoreEventAIE, - MemEventAIE, - ShimTileEventAIE, CoreEventAIE2, - MemEventAIE2, - ShimTileEventAIE2, - MemTileEventAIE2, CoreEventAIE2P, + MemEventAIE, + MemEventAIE2, MemEventAIE2P, - ShimTileEventAIE2P, + MemTileEventAIE2, MemTileEventAIE2P, + ShimTileEventAIE, + ShimTileEventAIE2, + ShimTileEventAIE2P, ) - from aie.dialects.aie import ( - WireBundle, # pyright: ignore[reportAttributeAccessIssue] DMAChannelDir, # pyright: ignore[reportAttributeAccessIssue] + WireBundle, # pyright: ignore[reportAttributeAccessIssue] ) # Default to AIE2 for backwards compatibility diff --git a/python/utils/trace/get_trace_summary.py b/python/utils/trace/get_trace_summary.py index 70434616aa4..2bea4fea41d 100755 --- a/python/utils/trace/get_trace_summary.py +++ b/python/utils/trace/get_trace_summary.py @@ -4,6 +4,7 @@ import argparse import logging import sys + from aie.utils.trace.utils import print_cycles_summary logger = logging.getLogger(__name__) diff --git a/python/utils/trace/parse.py b/python/utils/trace/parse.py index 943f77dc3da..b166d6d443c 100755 --- a/python/utils/trace/parse.py +++ b/python/utils/trace/parse.py @@ -1,45 +1,36 @@ #!/usr/bin/env python3 # Copyright (C) 2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -import json import argparse +import json import logging import sys -import re - -logger = logging.getLogger(__name__) +import aie.dialects.aie as aiedialect +import aie.dialects.aiex as aiexdialect from aie.extras.util import find_ops # pyright: ignore[reportMissingImports] from aie.helpers.util import ( # pyright: ignore[reportMissingImports] fold_constant_operand, ) from aie.ir import ( # pyright: ignore[reportMissingImports] Context, # pyright: ignore[reportAttributeAccessIssue] - Module, # pyright: ignore[reportAttributeAccessIssue] Location, # pyright: ignore[reportAttributeAccessIssue] -) -from aie.utils.trace.utils import ( - parity, - extract_tile, - parse_pkt_hdr_in_stream, - trace_pkts_de_interleave, - convert_to_byte_stream, - convert_to_commands, - trim_trace_pkts, - split_trace_segments, + Module, # pyright: ignore[reportAttributeAccessIssue] ) from aie.utils.trace.events import ( NUM_TRACE_TYPES, PacketType, - CoreEvent, - MemEvent, - ShimTileEvent, - MemTileEvent, get_events_for_device, ) +from aie.utils.trace.utils import ( + convert_to_byte_stream, + convert_to_commands, + split_trace_segments, + trace_pkts_de_interleave, + trim_trace_pkts, +) -import aie.dialects.aie as aiedialect -import aie.dialects.aiex as aiexdialect +logger = logging.getLogger(__name__) NUM_EVENTS = 8 # number of events we can view per trace @@ -138,7 +129,7 @@ def deactivate_events( events_module, ): for k in active_events.keys(): # an active event - if cycles > 0 or (cycles == 0 and not k in multiples): + if cycles > 0 or (cycles == 0 and k not in multiples): # if not k in multiples: # active event it not in multiples list if active_events[k] > 0: trace_event = { @@ -457,7 +448,7 @@ def parse_mlir_trace_events(mlir_module_str, colshift=None): # core event 0 if address == 0x340E0: # 213216, match ignoring case if row == 0: # shim - if pid_events[2].get(key) == None: + if pid_events[2].get(key) is None: pid_events[2][key] = [0] * 8 logger.debug("Trace event 0 configured to be %s", hex(value)) pid_events[2][key][0] = value & 0xFF @@ -465,7 +456,7 @@ def parse_mlir_trace_events(mlir_module_str, colshift=None): pid_events[2][key][2] = (value >> 16) & 0xFF pid_events[2][key][3] = (value >> 24) & 0xFF else: # core - if pid_events[0].get(key) == None: + if pid_events[0].get(key) is None: pid_events[0][key] = [0] * 8 logger.debug("Trace event 0 configured to be %s", hex(value)) pid_events[0][key][0] = value & 0xFF @@ -475,14 +466,14 @@ def parse_mlir_trace_events(mlir_module_str, colshift=None): # core event 1 elif address == 0x340E4: # 213220, match ignoring case if row == 0: # shim - if pid_events[2].get(key) == None: + if pid_events[2].get(key) is None: pid_events[2][key] = [0] * 8 pid_events[2][key][4] = value & 0xFF pid_events[2][key][5] = (value >> 8) & 0xFF pid_events[2][key][6] = (value >> 16) & 0xFF pid_events[2][key][7] = (value >> 24) & 0xFF else: # core - if pid_events[0].get(key) == None: + if pid_events[0].get(key) is None: pid_events[0][key] = [0] * 8 pid_events[0][key][4] = value & 0xFF pid_events[0][key][5] = (value >> 8) & 0xFF @@ -490,7 +481,7 @@ def parse_mlir_trace_events(mlir_module_str, colshift=None): pid_events[0][key][7] = (value >> 24) & 0xFF # mem event 0 elif address == 0x140E0: # 82144 - if pid_events[1].get(key) == None: + if pid_events[1].get(key) is None: pid_events[1][key] = [0] * 8 logger.debug("Trace event 0 configured to be %s", hex(value)) pid_events[1][key][0] = value & 0xFF @@ -499,7 +490,7 @@ def parse_mlir_trace_events(mlir_module_str, colshift=None): pid_events[1][key][3] = (value >> 24) & 0xFF # mem event 1 elif address == 0x140E4: # 82148 - if pid_events[1].get(key) == None: + if pid_events[1].get(key) is None: pid_events[1][key] = [0] * 8 pid_events[1][key][4] = value & 0xFF pid_events[1][key][5] = (value >> 8) & 0xFF @@ -507,7 +498,7 @@ def parse_mlir_trace_events(mlir_module_str, colshift=None): pid_events[1][key][7] = (value >> 24) & 0xFF # memtile event 0 elif address == 0x940E0: # 606432 - if pid_events[3].get(key) == None: + if pid_events[3].get(key) is None: pid_events[3][key] = [0] * 8 logger.debug("Trace event 0 configured to be %s", hex(value)) pid_events[3][key][0] = value & 0xFF @@ -516,7 +507,7 @@ def parse_mlir_trace_events(mlir_module_str, colshift=None): pid_events[3][key][3] = (value >> 24) & 0xFF # memtile event 1 elif address == 0x940E4: # 606436 - if pid_events[3].get(key) == None: + if pid_events[3].get(key) is None: pid_events[3][key] = [0] * 8 pid_events[3][key][4] = value & 0xFF pid_events[3][key][5] = (value >> 8) & 0xFF @@ -681,11 +672,11 @@ def align_column_start_index(events, commands): new_events = [] for t in range(NUM_TRACE_TYPES): updated = {} - for loc, l in events[t].items(): + for loc, event_data in events[t].items(): row, col = map(int, loc.split(",")) new_col = col - colshift new_key = f"{row},{new_col}" - updated[new_key] = l + updated[new_key] = event_data new_events.append(updated) return new_events diff --git a/python/utils/trace/setup.py b/python/utils/trace/setup.py index e423ffcf840..59c9c004363 100644 --- a/python/utils/trace/setup.py +++ b/python/utils/trace/setup.py @@ -3,44 +3,44 @@ import logging -logger = logging.getLogger(__name__) - from aie.dialects.aie import ( - packetflow, + TraceMode, # pyright: ignore[reportAttributeAccessIssue] + TracePacketType, # pyright: ignore[reportAttributeAccessIssue] WireBundle, # pyright: ignore[reportAttributeAccessIssue] + get_target_model, # pyright: ignore[reportAttributeAccessIssue] + packetflow, trace, - trace_mode, trace_event, + trace_host_config, + trace_mode, trace_packet, trace_port, trace_start, - trace_stop, trace_start_config, - trace_host_config, - TraceMode, # pyright: ignore[reportAttributeAccessIssue] - TracePacketType, # pyright: ignore[reportAttributeAccessIssue] - DMAChannelDir, # pyright: ignore[reportAttributeAccessIssue] - get_target_model, # pyright: ignore[reportAttributeAccessIssue] + trace_stop, ) from aie.dialects.aiex import ( - npu_write32, # pyright: ignore[reportAttributeAccessIssue] - npu_writebd, # pyright: ignore[reportAttributeAccessIssue] - npu_maskwrite32, # pyright: ignore[reportAttributeAccessIssue] npu_address_patch, # pyright: ignore[reportAttributeAccessIssue] + npu_maskwrite32, # pyright: ignore[reportAttributeAccessIssue] npu_sync, + npu_write32, # pyright: ignore[reportAttributeAccessIssue] + npu_writebd, # pyright: ignore[reportAttributeAccessIssue] ) + from .events import ( BasePortEvent, - GenericEvent, - PortEvent, CoreEvent, + GenericEvent, MemEvent, - ShimTileEvent, MemTileEvent, MemTilePortEvent, PacketType, + PortEvent, + ShimTileEvent, ) +logger = logging.getLogger(__name__) + # Globally defined constants direction_s2mm = 0 direction_mm2s = 1 diff --git a/python/utils/trace/utils.py b/python/utils/trace/utils.py index df2a767eff1..7e6d5a127b7 100644 --- a/python/utils/trace/utils.py +++ b/python/utils/trace/utils.py @@ -2,13 +2,15 @@ # SPDX-License-Identifier: MIT # from CppHeaderParser import CppHeader -import logging -import numpy as np import json -import re +import logging import os +import re import sys from typing import Optional + +import numpy as np + from .events import NUM_TRACE_TYPES logger = logging.getLogger(__name__) @@ -94,7 +96,7 @@ def extract_buffers(test): output_buffers.append(np.array(array, dtype=dtype)) rtps = [] - if test["test_vectors"].get("rtps") != None: + if test["test_vectors"].get("rtps") is not None: for rtp in test["test_vectors"]["rtps"]: array, dtype = rtp.values() rtps.append(np.array(array, dtype=dtype)) @@ -117,13 +119,12 @@ def get_cycles(trace_path): for x in data: if (x["name"] == "INSTR_EVENT_0") and (x["ph"] == "B"): event0.append(x["ts"]) - tmp = x["ts"] if x["name"] == "INSTR_EVENT_1" and x["ph"] == "B": event1.append(x["ts"]) return event1[0] - event0[0] - except: + except Exception: return np.inf @@ -149,12 +150,12 @@ def get_cycles_summary(trace_path): for x in data: idx = int(x["pid"]) if (x["name"] == "INSTR_EVENT_0") and (x["ph"] == "B"): - if in_kernel[idx] == False: + if not in_kernel[idx]: event0[idx] = x["ts"] in_kernel[idx] = True if x["name"] == "INSTR_EVENT_1" and x["ph"] == "B": - if in_kernel[idx] == True: + if in_kernel[idx]: deltas[idx].append(x["ts"] - event0[idx]) in_kernel[idx] = False @@ -267,7 +268,7 @@ def trace_pkts_de_interleave(word_stream): for tt in range(NUM_TRACE_TYPES): if pkt_hdr["type"] == tt: curr_pkt_type = tt - if trace_pkts_sorted[tt].get(curr_loc) == None: + if trace_pkts_sorted[tt].get(curr_loc) is None: trace_pkts_sorted[tt][curr_loc] = list() valid_type_found = True if not valid_type_found: @@ -281,12 +282,12 @@ def trace_pkts_de_interleave(word_stream): def convert_to_byte_stream(toks_list): byte_stream_list = list() - for l in toks_list: + for toks in toks_list: byte_stream_dict = dict() - for loc, stream in l.items(): + for loc, stream in toks.items(): byte_stream_dict[loc] = list() f = ["", "a5a5a5a5"] - toks = [t for t in stream if not t in f] + toks = [t for t in stream if t not in f] events = [int(t, 16) for t in toks] for event in events: for top in range(4): diff --git a/python/utils/verify.py b/python/utils/verify.py index e5d3de146e7..e75f977f6c9 100644 --- a/python/utils/verify.py +++ b/python/utils/verify.py @@ -20,7 +20,6 @@ import sys import numpy as np - from aie.utils.benchmark import print_benchmark _DEFAULT_RTOL = 0.128 diff --git a/ruff.toml b/ruff.toml new file mode 100644 index 00000000000..aedb851dacf --- /dev/null +++ b/ruff.toml @@ -0,0 +1,60 @@ +# Copyright (C) 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +target-version = "py39" +# Deliberately more generous than black's 88: black already wraps ordinary +# code at 88, so anything ruff still sees over 88 is a long string, comment, +# or docstring line black leaves alone -- most commonly this codebase's +# established one-line-per-arg Google-docstring convention (`name (type, +# optional): description. Defaults to X.`). 160 clears nearly all of those +# without forcing a reformat of that convention across the API surface, or +# hand-wraps of long URLs/qualified names/error-message f-strings for little +# readability benefit. +line-length = 160 + +# Same scope as pyrightconfig.json's "include". +include = [ + "python/iron/**/*.py", + "python/utils/**/*.py", + "python/helpers/**/*.py", + "python/compiler/**/*.py", +] + +# Same generated/vendored dirs excluded from the black pre-commit hook. +extend-exclude = [ + "third_party", + "*build*", + "*install*", + "sandbox*", + "llvm*", + "cmakeModules*", + "ironenv", +] + +[lint] +# E/F: pyflakes + pycodestyle errors (ruff defaults). +# I: isort (no import sorter previously enforced). +# N802/N816: pep8-naming subset that's clean across this codebase (verified: +# zero violations at introduction). N801/N803/N806/N812 are dropped -- each +# was individually confirmed to collide with a real, load-bearing convention +# rather than catching a mistake: +# N801: hardware-format dtype classes (e.g. v8bfp16ebs8) are named to match +# the literal spec string shared with C++ (see helper.h); CapWords would +# break that 1:1 mirroring. +# N812: `types as T` / `_aie as CustomTypes` mirror the upstream MLIR +# python-bindings type-namespace convention (`T` used 30+ times). +# N803/N806: matrix/tensor-dimension math notation (N_div_n, ml.py's W/A/D +# following the referenced Pytorch_BN_Fold algorithm's own variable +# names), true local constants (_RGBA, _ACC_FACTOR), an MLIR attribute +# name mirrored verbatim (via_DMA -- literally `via_DMA` in the IR, see +# aie_ops.py FileCheck tests), a public API parameter used externally +# (hasElse), and dynamically-fetched dialect op/enum classes bound via +# getattr (regdb.py/parse.py) because pyright can't see attributes on the +# compiled bindings module -- renaming those to lowercase would make +# `isinstance(x, device_op)`-style checks read as instance access. +select = ["E", "F", "I", "N802", "N816"] + +[lint.per-file-ignores] +# Markdown-formatted module docstring (headers, tables, code fences); wrapping +# it to fit a line-length limit would break the markdown formatting. +"python/utils/trace/__init__.py" = ["E501"] diff --git a/test/python/test_markers.py b/test/python/test_markers.py index 2ef5a8bf0e6..24e6ac65058 100644 --- a/test/python/test_markers.py +++ b/test/python/test_markers.py @@ -12,12 +12,12 @@ import pytest -from aie.utils.compile.jit.markers import CompileTime, In, InOut, Out -from aie.utils.compile.jit.compilabledesign import ( +from aie.utils.compile.jit._introspect import ( _is_compile_param, _is_tensor_param, split_params, ) +from aie.utils.compile.jit.markers import CompileTime, In, InOut, Out # --------------------------------------------------------------------------- # CompileTime[T] — generic parameterisation