From e1d8cc92ded98a8f1d14a3b0451152be8b6d937d Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Thu, 19 Dec 2024 12:14:33 -0600 Subject: [PATCH] Enable, fix TC rules --- doc/conf.py | 14 ------ loopy/check.py | 12 +++-- loopy/codegen/__init__.py | 22 +++------- loopy/codegen/bounds.py | 8 +++- loopy/codegen/result.py | 10 ++--- loopy/codegen/tools.py | 10 +++-- loopy/frontend/fortran/expression.py | 10 +++-- loopy/kernel/__init__.py | 15 ++++--- loopy/kernel/array.py | 12 ++--- loopy/kernel/data.py | 20 ++++++--- loopy/kernel/function_interface.py | 7 +-- loopy/kernel/instruction.py | 8 +++- loopy/kernel/tools.py | 10 +++-- loopy/library/function.py | 7 ++- loopy/library/random123.py | 6 ++- loopy/library/reduction.py | 7 ++- loopy/match.py | 66 +++++++++++++++------------- loopy/options.py | 7 ++- loopy/preprocess.py | 13 +++--- loopy/schedule/__init__.py | 10 ++--- loopy/schedule/tools.py | 23 +++++++--- loopy/statistics.py | 7 ++- loopy/symbolic.py | 14 +++--- loopy/target/__init__.py | 9 ---- loopy/target/c/__init__.py | 21 +++++---- loopy/target/c/c_execution.py | 20 +++++---- loopy/target/c/codegen/expression.py | 7 ++- loopy/target/cuda.py | 9 ++-- loopy/target/execution.py | 10 +++-- loopy/target/ispc.py | 17 ++++--- loopy/target/opencl.py | 12 +++-- loopy/target/pyopencl.py | 19 ++++---- loopy/target/pyopencl_execution.py | 12 ++--- loopy/target/python.py | 9 ++-- loopy/transform/array_buffer_map.py | 7 ++- loopy/transform/callable.py | 7 ++- loopy/transform/data.py | 9 ++-- loopy/transform/iname.py | 9 ++-- loopy/transform/instruction.py | 7 ++- loopy/transform/precompute.py | 14 +++--- loopy/transform/realize_reduction.py | 12 +++-- loopy/translation_unit.py | 2 +- loopy/typing.py | 2 + pyproject.toml | 3 +- 44 files changed, 315 insertions(+), 220 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index 59d9a2c94..cd2fad122 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -35,20 +35,6 @@ "pyrsistent": ("https://pyrsistent.readthedocs.io/en/latest/", None), } -# Some modules need to import things just so that sphinx can resolve symbols in -# type annotations. Often, we do not want these imports (e.g. of PyOpenCL) when -# in normal use (because they would introduce unintended side effects or hard -# dependencies). This flag exists so that these imports only occur during doc -# build. Since sphinx appears to resolve type hints lexically (as it should), -# this needs to be cross-module (since, e.g. an inherited arraycontext -# docstring can be read by sphinx when building meshmode, a dependent package), -# this needs a setting of the same name across all packages involved, that's -# why this name is as global-sounding as it is. -import sys - - -sys._BUILDING_SPHINX_DOCS = True - nitpicky = True nitpick_ignore_regex = [ diff --git a/loopy/check.py b/loopy/check.py index b33ad4ce3..1a63c90bc 100644 --- a/loopy/check.py +++ b/loopy/check.py @@ -25,8 +25,8 @@ import logging from collections import defaultdict -from collections.abc import Mapping, Sequence from functools import reduce +from typing import TYPE_CHECKING import numpy as np @@ -41,7 +41,6 @@ WriteRaceConditionWarning, warn_with_kernel, ) -from loopy.kernel import LoopKernel from loopy.kernel.array import ( ArrayBase, FixedStrideArrayDimTag, @@ -73,6 +72,14 @@ from loopy.typing import not_none +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from pymbolic.typing import Expression + + from loopy.kernel import LoopKernel + + logger = logging.getLogger(__name__) @@ -216,7 +223,6 @@ def check_separated_array_consistency(kernel: LoopKernel) -> None: @check_each_kernel def check_offsets_and_dim_tags(kernel: LoopKernel) -> None: from pymbolic.primitives import ExpressionNode, Variable - from pymbolic.typing import Expression from loopy.symbolic import DependencyMapper diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index ddc5442a9..ca44cb3a1 100644 --- a/loopy/codegen/__init__.py +++ b/loopy/codegen/__init__.py @@ -24,7 +24,6 @@ """ import logging -import sys from dataclasses import dataclass, replace from typing import ( TYPE_CHECKING, @@ -35,10 +34,6 @@ import immutables -from loopy.codegen.result import CodeGenerationResult -from loopy.library.reduction import ReductionOpFunction -from loopy.translation_unit import CallablesTable, TranslationUnit - logger = logging.getLogger(__name__) @@ -51,24 +46,21 @@ from pytools.persistent_dict import WriteOncePersistentDict from loopy.diagnostic import LoopyError, warn -from loopy.kernel import LoopKernel from loopy.kernel.function_interface import CallableKernel from loopy.symbolic import CombineMapper -from loopy.target import TargetBase from loopy.tools import LoopyKeyBuilder, caches -from loopy.types import LoopyType -from loopy.typing import Expression from loopy.version import DATA_MODEL_VERSION if TYPE_CHECKING: - from loopy.codegen.result import GeneratedProgram - from loopy.codegen.tools import CodegenOperationCacheManager - - -if getattr(sys, "_BUILDING_SPHINX_DOCS", False): - from loopy.codegen.result import GeneratedProgram + from loopy.codegen.result import CodeGenerationResult, GeneratedProgram from loopy.codegen.tools import CodegenOperationCacheManager + from loopy.kernel import LoopKernel + from loopy.library.reduction import ReductionOpFunction + from loopy.target import TargetBase + from loopy.translation_unit import CallablesTable, TranslationUnit + from loopy.types import LoopyType + from loopy.typing import Expression __doc__ = """ diff --git a/loopy/codegen/bounds.py b/loopy/codegen/bounds.py index e8708b9ad..0f3bdba41 100644 --- a/loopy/codegen/bounds.py +++ b/loopy/codegen/bounds.py @@ -24,11 +24,15 @@ """ +from typing import TYPE_CHECKING + import islpy as isl from islpy import dim_type -from loopy.codegen.tools import CodegenOperationCacheManager -from loopy.kernel import LoopKernel + +if TYPE_CHECKING: + from loopy.codegen.tools import CodegenOperationCacheManager + from loopy.kernel import LoopKernel # {{{ approximate, convex bounds check generator diff --git a/loopy/codegen/result.py b/loopy/codegen/result.py index faf8664c3..02b5ce27a 100644 --- a/loopy/codegen/result.py +++ b/loopy/codegen/result.py @@ -31,10 +31,10 @@ Sequence, ) -import islpy as isl - if TYPE_CHECKING: + import islpy + from loopy.codegen import CodeGenerationState @@ -58,8 +58,6 @@ def process_preambles(preambles: Sequence[tuple[int, str]]) -> Sequence[str]: __doc__ = """ .. currentmodule:: loopy.codegen.result -.. autoclass:: GeneratedProgram - .. autoclass:: CodeGenerationResult .. autofunction:: merge_codegen_results @@ -121,7 +119,7 @@ class CodeGenerationResult: """ host_program: GeneratedProgram | None device_programs: Sequence[GeneratedProgram] - implemented_domains: Mapping[str, isl.Set] + implemented_domains: Mapping[str, islpy.Set] host_preambles: Sequence[tuple[str, str]] = () device_preambles: Sequence[tuple[str, str]] = () @@ -249,7 +247,7 @@ def merge_codegen_results( new_device_programs = [] new_device_preambles: list[tuple[str, str]] = [] dev_program_names = set() - implemented_domains: dict[str, isl.Set] = {} + implemented_domains: dict[str, islpy.Set] = {} codegen_result = None block_cls = codegen_state.ast_builder.ast_block_class diff --git a/loopy/codegen/tools.py b/loopy/codegen/tools.py index a2076480b..be3d6ade5 100644 --- a/loopy/codegen/tools.py +++ b/loopy/codegen/tools.py @@ -25,12 +25,11 @@ from dataclasses import dataclass from functools import cached_property +from typing import TYPE_CHECKING from pytools import memoize_method from loopy.kernel import LoopKernel -from loopy.kernel.data import Iname -from loopy.kernel.instruction import InstructionBase from loopy.schedule import ( Barrier, BeginBlockItem, @@ -43,6 +42,11 @@ ) +if TYPE_CHECKING: + import loopy.kernel.data + from loopy.kernel.instruction import InstructionBase + + __doc__ = """ .. autoclass:: KernelProxyForCodegenOperationCacheManager @@ -58,7 +62,7 @@ class KernelProxyForCodegenOperationCacheManager: """ instructions: list[InstructionBase] linearization: list[ScheduleItem] - inames: dict[str, Iname] + inames: dict[str, loopy.kernel.data.Iname] @cached_property def id_to_insn(self): diff --git a/loopy/frontend/fortran/expression.py b/loopy/frontend/fortran/expression.py index ec0d36ba7..54bca20bb 100644 --- a/loopy/frontend/fortran/expression.py +++ b/loopy/frontend/fortran/expression.py @@ -24,9 +24,8 @@ """ import re -from collections.abc import Mapping from sys import intern -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar import numpy as np @@ -34,7 +33,12 @@ from pymbolic.parser import Parser as ExpressionParserBase from loopy.frontend.fortran.diagnostic import TranslationError -from loopy.symbolic import LexTable + + +if TYPE_CHECKING: + from collections.abc import Mapping + + from loopy.symbolic import LexTable _less_than = intern("less_than") diff --git a/loopy/kernel/__init__.py b/loopy/kernel/__init__.py index 42e648123..d612b5db3 100644 --- a/loopy/kernel/__init__.py +++ b/loopy/kernel/__init__.py @@ -53,7 +53,6 @@ import islpy # to help out Sphinx import islpy as isl from islpy import dim_type -from pymbolic import ArithmeticExpression from pytools import ( UniqueNameGenerator, generate_unique_names, @@ -62,6 +61,7 @@ ) from pytools.tag import Tag, Taggable +import loopy.codegen import loopy.kernel.data # to help out Sphinx from loopy.diagnostic import CannotBranchDomainTree, LoopyError, StaticValueFindingError from loopy.kernel.data import ( @@ -73,18 +73,19 @@ _ArraySeparationInfo, filter_iname_tags_by_type, ) -from loopy.kernel.instruction import InstructionBase -from loopy.options import Options -from loopy.schedule import ScheduleItem -from loopy.target import TargetBase from loopy.tools import update_persistent_hash from loopy.types import LoopyType, NumpyType -from loopy.typing import Expression, InameStr if TYPE_CHECKING: - import loopy.codegen # to help out Sphinx + from pymbolic import ArithmeticExpression + from loopy.kernel.function_interface import InKernelCallable + from loopy.kernel.instruction import InstructionBase + from loopy.options import Options + from loopy.schedule import ScheduleItem + from loopy.target import TargetBase + from loopy.typing import Expression, InameStr # {{{ loop kernel object diff --git a/loopy/kernel/array.py b/loopy/kernel/array.py index a91f7925b..9895685fb 100644 --- a/loopy/kernel/array.py +++ b/loopy/kernel/array.py @@ -24,7 +24,6 @@ """ import re -import sys from dataclasses import dataclass from typing import ( TYPE_CHECKING, @@ -41,7 +40,6 @@ import numpy as np # noqa from typing_extensions import Self, TypeAlias -from pymbolic import ArithmeticExpression from pymbolic.primitives import is_arithmetic_expression from pytools import ImmutableRecord from pytools.tag import Tag, Taggable @@ -53,15 +51,13 @@ if TYPE_CHECKING: + from pymbolic import ArithmeticExpression + from loopy.codegen import VectorizationInfo from loopy.kernel import LoopKernel from loopy.kernel.data import ArrayArg, TemporaryVariable from loopy.target import TargetBase -if getattr(sys, "_BUILDING_SPHINX_DOCS", False): - from loopy.target import TargetBase - - T = TypeVar("T") @@ -629,7 +625,7 @@ def _parse_shape_or_strides( x_tup: tuple[Expression | str, ...] = x_parsed else: assert x_parsed is not auto - x_tup = (cast(Expression, x_parsed),) + x_tup = (cast("Expression", x_parsed),) def parse_arith(x: Expression | str) -> ArithmeticExpression: if isinstance(x, str): @@ -1296,7 +1292,7 @@ def eval_expr_assert_integer_constant(i, expr) -> int: index = tuple(remaining_index) # only arguments (not temporaries) may be sep-tagged - ary = cast(ArrayArg, + ary = cast("ArrayArg", kernel.arg_dict[ary._separation_info.subarray_names[tuple(sep_index)]]) # }}} diff --git a/loopy/kernel/data.py b/loopy/kernel/data.py index ac4c1120d..d10401e51 100644 --- a/loopy/kernel/data.py +++ b/loopy/kernel/data.py @@ -30,6 +30,7 @@ from enum import IntEnum from sys import intern from typing import ( + TYPE_CHECKING, Any, ClassVar, Sequence, @@ -40,9 +41,7 @@ import numpy # FIXME: imported as numpy to allow sphinx to resolve things import numpy as np -from immutables import Map -from pymbolic import ArithmeticExpression, Variable from pytools import ImmutableRecord from pytools.tag import Tag, Taggable, UniqueTag as UniqueTagBase @@ -61,10 +60,17 @@ VarAtomicity, make_assignment, ) -from loopy.types import LoopyType, ToLoopyTypeConvertible from loopy.typing import Expression, ShapeType, auto +if TYPE_CHECKING: + from immutables import Map + + from pymbolic import ArithmeticExpression, Variable + + from loopy.types import LoopyType, ToLoopyTypeConvertible + + __doc__ = """ .. autofunction:: filter_iname_tags_by_type @@ -110,7 +116,7 @@ def _names_from_expr(expr: Expression | str | None) -> frozenset[str]: if isinstance(expr, str): return frozenset({expr}) elif isinstance(expr, ExpressionNode): - return frozenset(cast(Variable, v).name for v in dep_mapper(expr)) + return frozenset(cast("Variable", v).name for v in dep_mapper(expr)) elif expr is None: return frozenset() elif isinstance(expr, Number): @@ -435,7 +441,7 @@ class _ArraySeparationInfo: class ArrayArg(ArrayBase, KernelArgument): - __doc__ = cast(str, ArrayBase.__doc__) + ( + __doc__ = cast("str", ArrayBase.__doc__) + ( """ .. attribute:: address_space @@ -637,7 +643,7 @@ def get_arg_decl(self, ast_builder): # {{{ temporary variable class TemporaryVariable(ArrayBase): - __doc__ = cast(str, ArrayBase.__doc__) + """ + __doc__ = cast("str", ArrayBase.__doc__) + """ .. autoattribute:: storage_shape .. autoattribute:: base_indices .. autoattribute:: address_space @@ -814,7 +820,7 @@ def nbytes(self) -> Expression: raise ValueError("shape is None") if self.shape is auto: raise ValueError("shape is auto") - shape = cast(Tuple[ArithmeticExpression], self.shape) + shape = cast("Tuple[ArithmeticExpression]", self.shape) if self.dtype is None: raise ValueError("data type is indeterminate") diff --git a/loopy/kernel/function_interface.py b/loopy/kernel/function_interface.py index ded26c9c9..146d40f4f 100644 --- a/loopy/kernel/function_interface.py +++ b/loopy/kernel/function_interface.py @@ -23,7 +23,6 @@ THE SOFTWARE. """ from abc import ABC, abstractmethod -from collections.abc import Mapping, Sequence from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Callable, TypeVar from warnings import warn @@ -35,15 +34,17 @@ from loopy.kernel.array import ArrayBase, ArrayDimImplementationTag from loopy.kernel.data import AddressSpace, ArrayArg, ValueArg from loopy.symbolic import DependencyMapper, WalkMapper -from loopy.types import LoopyType -from loopy.typing import ShapeType if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + from typing_extensions import Self from loopy.kernel import LoopKernel from loopy.translation_unit import CallablesTable, FunctionIdT + from loopy.types import LoopyType + from loopy.typing import ShapeType __doc__ = """ .. currentmodule:: loopy.kernel.function_interface diff --git a/loopy/kernel/instruction.py b/loopy/kernel/instruction.py index 362e3799f..f882c09f3 100644 --- a/loopy/kernel/instruction.py +++ b/loopy/kernel/instruction.py @@ -31,6 +31,7 @@ from functools import cached_property from sys import intern from typing import ( + TYPE_CHECKING, Any, ClassVar, Mapping, @@ -44,8 +45,11 @@ from loopy.diagnostic import LoopyError from loopy.tools import Optional as LoopyOptional -from loopy.types import LoopyType -from loopy.typing import Expression, InameStr + + +if TYPE_CHECKING: + from loopy.types import LoopyType + from loopy.typing import Expression, InameStr # {{{ instruction tags diff --git a/loopy/kernel/tools.py b/loopy/kernel/tools.py index 9d73fe419..c48da4be9 100644 --- a/loopy/kernel/tools.py +++ b/loopy/kernel/tools.py @@ -28,14 +28,13 @@ import sys from functools import reduce from sys import intern -from typing import AbstractSet, Mapping, Sequence +from typing import TYPE_CHECKING, AbstractSet, Mapping, Sequence import numpy as np import islpy as isl from islpy import dim_type from pytools import memoize_on_first_arg, natsorted -from pytools.tag import Tag from loopy.diagnostic import LoopyError, warn_with_kernel from loopy.kernel import LoopKernel @@ -47,7 +46,12 @@ ) from loopy.symbolic import CombineMapper from loopy.translation_unit import TranslationUnit, TUnitOrKernelT, for_each_kernel -from loopy.types import ToLoopyTypeConvertible + + +if TYPE_CHECKING: + from pytools.tag import Tag + + from loopy.types import ToLoopyTypeConvertible logger = logging.getLogger(__name__) diff --git a/loopy/library/function.py b/loopy/library/function.py index c7802444e..8b61ad41a 100644 --- a/loopy/library/function.py +++ b/loopy/library/function.py @@ -23,14 +23,19 @@ THE SOFTWARE. """ +from typing import TYPE_CHECKING + import numpy as np from loopy.diagnostic import LoopyError from loopy.kernel.function_interface import ScalarCallable -from loopy.translation_unit import CallablesTable from loopy.types import NumpyType +if TYPE_CHECKING: + from loopy.translation_unit import CallablesTable + + class MakeTupleCallable(ScalarCallable): def with_types(self, arg_id_to_dtype, callables_table): new_arg_id_to_dtype = arg_id_to_dtype.copy() diff --git a/loopy/library/random123.py b/loopy/library/random123.py index f47652000..f65fa7600 100644 --- a/loopy/library/random123.py +++ b/loopy/library/random123.py @@ -26,6 +26,7 @@ from dataclasses import dataclass, replace +from typing import TYPE_CHECKING import numpy as np from mako.template import Template @@ -33,7 +34,10 @@ from pymbolic.typing import not_none from loopy.kernel.function_interface import ScalarCallable -from loopy.target import TargetBase + + +if TYPE_CHECKING: + from loopy.target import TargetBase # {{{ rng metadata diff --git a/loopy/library/reduction.py b/loopy/library/reduction.py index 2935a6629..6ddc3fb86 100644 --- a/loopy/library/reduction.py +++ b/loopy/library/reduction.py @@ -24,11 +24,12 @@ """ +from typing import TYPE_CHECKING + import numpy as np from pymbolic import var from pymbolic.primitives import expr_dataclass -from pytools.persistent_dict import Hash, KeyBuilder from loopy.diagnostic import LoopyError from loopy.kernel.function_interface import ScalarCallable @@ -37,6 +38,10 @@ from loopy.types import NumpyType +if TYPE_CHECKING: + from pytools.persistent_dict import Hash, KeyBuilder + + __doc__ = """ .. currentmodule:: loopy.library.reduction diff --git a/loopy/match.py b/loopy/match.py index cbd3c4331..c73c1ed1b 100644 --- a/loopy/match.py +++ b/loopy/match.py @@ -1,5 +1,32 @@ -"""Matching functionality for instruction ids and substitution -rule invocations stacks.""" +""" +.. autoclass:: Matchable +.. autoclass:: StackMatchComponent +.. autoclass:: StackMatch + +.. autofunction:: parse_match + +.. autofunction:: parse_stack_match + +.. autodata:: ToStackMatchCovertible + +Match expressions +^^^^^^^^^^^^^^^^^ + +.. autoclass:: MatchExpressionBase +.. autoclass:: All +.. autoclass:: And +.. autoclass:: Or +.. autoclass:: Not +.. autoclass:: Id +.. autoclass:: ObjTagged +.. autoclass:: Tagged +.. autoclass:: Writes +.. autoclass:: Reads +.. autoclass:: InKernel +.. autoclass:: Iname + +""" + from __future__ import annotations @@ -29,43 +56,22 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from sys import intern -from typing import Protocol, Sequence, Union +from typing import TYPE_CHECKING, Protocol, Sequence, Union + +from typing_extensions import TypeAlias -from loopy.kernel import LoopKernel from loopy.kernel.instruction import InstructionBase NoneType = type(None) -import pytools.tag from pytools.lex import RE -__doc__ = """ -.. autoclass:: Matchable -.. autoclass:: StackMatchComponent -.. autoclass:: StackMatch - -.. autofunction:: parse_match - -.. autofunction:: parse_stack_match - -Match expressions -^^^^^^^^^^^^^^^^^ +if TYPE_CHECKING: + import pytools.tag -.. autoclass:: MatchExpressionBase -.. autoclass:: All -.. autoclass:: And -.. autoclass:: Or -.. autoclass:: Not -.. autoclass:: Id -.. autoclass:: ObjTagged -.. autoclass:: Tagged -.. autoclass:: Writes -.. autoclass:: Reads -.. autoclass:: InKernel -.. autoclass:: Iname -""" + from loopy.kernel import LoopKernel def re_from_glob(s: str) -> re.Pattern: @@ -532,7 +538,7 @@ def __call__( # {{{ stack match parsing -ToStackMatchCovertible = Union[StackMatch, str, None] +ToStackMatchCovertible: TypeAlias = Union[StackMatch, str, None] def parse_stack_match(smatch: ToStackMatchCovertible) -> StackMatch: diff --git a/loopy/options.py b/loopy/options.py index ebf0b147c..9996bb119 100644 --- a/loopy/options.py +++ b/loopy/options.py @@ -26,13 +26,16 @@ import os import re -from collections.abc import Mapping -from typing import Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar from warnings import warn from pytools import ImmutableRecord +if TYPE_CHECKING: + from collections.abc import Mapping + + ALLOW_TERMINAL_COLORS = True diff --git a/loopy/preprocess.py b/loopy/preprocess.py index a555ec938..aee4044be 100644 --- a/loopy/preprocess.py +++ b/loopy/preprocess.py @@ -24,7 +24,7 @@ """ import logging -from typing import Iterable, List, TypeVar, cast +from typing import TYPE_CHECKING, Iterable, List, TypeVar, cast logger = logging.getLogger(__name__) @@ -42,8 +42,6 @@ WriteRaceConditionWarning, warn_with_kernel, ) -from loopy.kernel import LoopKernel -from loopy.kernel.array import ArrayDimImplementationTag from loopy.kernel.data import ( ArrayArg, KernelArgument, @@ -71,7 +69,12 @@ # for the benefit of loopy.statistics, for now from loopy.type_inference import infer_unknown_types -from loopy.typing import Expression + + +if TYPE_CHECKING: + from loopy.kernel import LoopKernel + from loopy.kernel.array import ArrayDimImplementationTag + from loopy.typing import Expression # {{{ check for writes to predicates @@ -196,7 +199,7 @@ def make_arrays_for_sep_arrays(kernel: LoopKernel) -> LoopKernel: sep_axis_indices_set=sep_axis_indices_set, subarray_names=Map({ ind: vng(f"{arg.name}_s{'_'.join(str(i) for i in ind)}") - for ind in np.ndindex(*cast(List[int], sep_shape))})) + for ind in np.ndindex(*cast("List[int]", sep_shape))})) new_args.append(arg.copy(_separation_info=sep_info)) diff --git a/loopy/schedule/__init__.py b/loopy/schedule/__init__.py index 73a23a983..55e0a197b 100644 --- a/loopy/schedule/__init__.py +++ b/loopy/schedule/__init__.py @@ -1,7 +1,5 @@ from __future__ import annotations -from loopy.kernel.function_interface import InKernelCallable -from loopy.translation_unit import FunctionIdT from loopy.typing import not_none @@ -29,7 +27,6 @@ import logging import sys -from collections.abc import Hashable, Iterator, Mapping, Sequence, Set from dataclasses import dataclass, replace from typing import ( TYPE_CHECKING, @@ -44,19 +41,22 @@ from pytools.persistent_dict import WriteOncePersistentDict from loopy.diagnostic import LoopyError, ScheduleDebugInputError, warn_with_kernel -from loopy.kernel.instruction import InstructionBase from loopy.tools import LoopyKeyBuilder, caches from loopy.typing import InameStr from loopy.version import DATA_MODEL_VERSION if TYPE_CHECKING: + from collections.abc import Hashable, Iterator, Mapping, Sequence, Set + from loopy.kernel import LoopKernel + from loopy.kernel.function_interface import InKernelCallable + from loopy.kernel.instruction import InstructionBase from loopy.schedule.tools import ( InameStrSet, LoopTree, ) - from loopy.translation_unit import CallablesTable, TranslationUnit + from loopy.translation_unit import CallablesTable, FunctionIdT, TranslationUnit logger = logging.getLogger(__name__) diff --git a/loopy/schedule/tools.py b/loopy/schedule/tools.py index d09c14210..b659ee7b7 100644 --- a/loopy/schedule/tools.py +++ b/loopy/schedule/tools.py @@ -25,13 +25,19 @@ .. autoclass:: AccessMapDescriptor .. autoclass:: WriteRaceChecker -.. autoclass:: InameStrSet .. autoclass:: LoopNestTree .. autoclass:: LoopTree .. autofunction:: separate_loop_nest .. autofunction:: get_partial_loop_nest_tree .. autofunction:: get_loop_tree + +References +^^^^^^^^^^ + +.. class:: InameStrSet + + See :class:`loopy.typing.InameStrSet` """ __license__ = """ @@ -55,10 +61,9 @@ """ import enum -from collections.abc import Callable, Collection, Mapping from dataclasses import dataclass from functools import cached_property, reduce -from typing import AbstractSet, FrozenSet, Sequence +from typing import TYPE_CHECKING, AbstractSet, Sequence from immutables import Map from typing_extensions import TypeAlias @@ -67,11 +72,16 @@ from pytools import memoize_method, memoize_on_first_arg from loopy.diagnostic import LoopyError -from loopy.kernel import LoopKernel from loopy.kernel.data import AddressSpace, ArrayArg, TemporaryVariable -from loopy.schedule import ScheduleItem from loopy.schedule.tree import Tree -from loopy.typing import InameStr, not_none +from loopy.typing import InameStr, InameStrSet, not_none + + +if TYPE_CHECKING: + from collections.abc import Callable, Collection, Mapping + + from loopy.kernel import LoopKernel + from loopy.schedule import ScheduleItem # {{{ block boundary finder @@ -674,7 +684,6 @@ def do_accesses_result_in_races(self, insn1, insn1_dir, insn2, insn2_dir, # }}} -InameStrSet: TypeAlias = FrozenSet[InameStr] LoopNestTree: TypeAlias = Tree[InameStrSet] LoopTree: TypeAlias = Tree[InameStr] diff --git a/loopy/statistics.py b/loopy/statistics.py index 1b0a36489..e0987f9e1 100755 --- a/loopy/statistics.py +++ b/loopy/statistics.py @@ -28,9 +28,8 @@ THE SOFTWARE. """ -from collections.abc import Sequence from functools import cached_property, partial -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar import islpy as isl from islpy import dim_type @@ -45,6 +44,10 @@ from loopy.translation_unit import TranslationUnit +if TYPE_CHECKING: + from collections.abc import Sequence + + __doc__ = """ .. currentmodule:: loopy diff --git a/loopy/symbolic.py b/loopy/symbolic.py index e6f9a0488..7b4b9805f 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -25,7 +25,6 @@ THE SOFTWARE. """ import re -from collections.abc import Callable, Collection, Iterable from dataclasses import dataclass, replace from functools import cached_property, reduce from sys import intern @@ -80,7 +79,6 @@ ) from pymbolic.mapper.unifier import UnidirectionalUnifier as UnidirectionalUnifierBase from pymbolic.parser import Parser as ParserBase -from pymbolic.typing import ArithmeticOrExpressionT from pytools import memoize, memoize_method, memoize_on_first_arg from pytools.tag import Tag, Taggable, ToTagSetConvertible @@ -89,15 +87,19 @@ LoopyError, UnableToDetermineAccessRangeError, ) -from loopy.types import LoopyType, NumpyType, ToLoopyTypeConvertible from loopy.typing import Expression, not_none if TYPE_CHECKING: + from collections.abc import Callable, Collection, Iterable + + from pymbolic.typing import ArithmeticOrExpressionT + from loopy.kernel import LoopKernel from loopy.kernel.data import KernelArgument, SubstitutionRule, TemporaryVariable from loopy.kernel.instruction import InstructionBase from loopy.library.reduction import ReductionOperation, ReductionOpFunction + from loopy.types import LoopyType, NumpyType, ToLoopyTypeConvertible __doc__ = """ @@ -226,7 +228,7 @@ def map_sub_array_ref(self, expr, *args: P.args, **kwargs: P.kwargs): and new_subscript is expr.subscript): return expr - return SubArrayRef(cast(tuple[Variable, ...], new_inames), new_subscript) + return SubArrayRef(cast("tuple[Variable, ...]", new_inames), new_subscript) def map_resolved_function(self, expr, *args: P.args, **kwargs: P.kwargs): # leaf, doesn't change @@ -248,7 +250,7 @@ def is_expr_integer_valued(self, expr: Expression) -> bool: def flatten(expr: ArithmeticOrExpressionT) -> ArithmeticOrExpressionT: - return cast(ArithmeticOrExpressionT, FlattenMapper()(expr)) + return cast("ArithmeticOrExpressionT", FlattenMapper()(expr)) class IdentityMapper(IdentityMapperBase, IdentityMapperMixin, Generic[P]): @@ -1044,7 +1046,7 @@ def _get_dependencies_and_reduction_inames( ) -> tuple[AbstractSet[str], AbstractSet[str]]: dep_mapper: DependencyMapperWithReductionInames[[]] = \ DependencyMapperWithReductionInames(composite_leaves=False) - deps = frozenset(cast(Variable, dep).name for dep in dep_mapper(expr)) + deps = frozenset(cast("Variable", dep).name for dep in dep_mapper(expr)) reduction_inames = dep_mapper.reduction_inames return deps, reduction_inames diff --git a/loopy/target/__init__.py b/loopy/target/__init__.py index 6cfd77498..00f1891bd 100644 --- a/loopy/target/__init__.py +++ b/loopy/target/__init__.py @@ -10,15 +10,6 @@ .. autoclass:: OpenCLTarget .. autoclass:: PyOpenCLTarget .. autoclass:: ISPCTarget - -References to Canonical Names ------------------------------ - -.. currentmodule:: loopy.target - -.. class:: TargetBase - - See :class:`loopy.TargetBase`. """ from __future__ import annotations diff --git a/loopy/target/c/__init__.py b/loopy/target/c/__init__.py index 0bc9617ee..e37194e4a 100644 --- a/loopy/target/c/__init__.py +++ b/loopy/target/c/__init__.py @@ -25,7 +25,7 @@ """ import re -from typing import Any, Sequence, cast +from typing import TYPE_CHECKING, Any, Sequence, cast import numpy as np @@ -44,10 +44,7 @@ from pymbolic.mapper.stringifier import PREC_NONE from pytools import memoize_method -from loopy.codegen import CodeGenerationState -from loopy.codegen.result import CodeGenerationResult from loopy.diagnostic import LoopyError, LoopyTypeError -from loopy.kernel import LoopKernel from loopy.kernel.array import ArrayBase, FixedStrideArrayDimTag from loopy.kernel.data import ( AddressSpace, @@ -58,16 +55,22 @@ ValueArg, ) from loopy.kernel.function_interface import ScalarCallable -from loopy.schedule import CallKernel from loopy.symbolic import IdentityMapper from loopy.target import ASTBuilderBase, DummyHostASTBuilder, TargetBase -from loopy.target.execution import ExecutorBase from loopy.tools import remove_common_indentation -from loopy.translation_unit import FunctionIdT, TranslationUnit from loopy.types import LoopyType, NumpyType, to_loopy_type from loopy.typing import Expression, auto +if TYPE_CHECKING: + from loopy.codegen import CodeGenerationState + from loopy.codegen.result import CodeGenerationResult + from loopy.kernel import LoopKernel + from loopy.schedule import CallKernel + from loopy.target.execution import ExecutorBase + from loopy.translation_unit import FunctionIdT, TranslationUnit + + __doc__ = """ .. currentmodule loopy.target.c @@ -845,7 +848,7 @@ def get_function_declaration( assert codegen_state.kernel.linearization is not None subkernel_name = cast( - CallKernel, + "CallKernel", codegen_state.kernel.linearization[schedule_index] ).kernel_name @@ -1072,7 +1075,7 @@ def get_temporary_arg_decl( arg_decl: Declarator = RestrictPointer( self.wrap_decl_for_address_space( self.get_array_base_declarator(temp_var), - cast(AddressSpace, temp_var.address_space))) + cast("AddressSpace", temp_var.address_space))) if not is_written: arg_decl = Const(arg_decl) diff --git a/loopy/target/c/c_execution.py b/loopy/target/c/c_execution.py index 43bfc56d4..270c3d0dc 100644 --- a/loopy/target/c/c_execution.py +++ b/loopy/target/c/c_execution.py @@ -28,30 +28,34 @@ import os import tempfile from dataclasses import dataclass -from typing import Any, Callable, ClassVar, Sequence +from typing import TYPE_CHECKING, Any, Callable, ClassVar, Sequence import numpy as np from codepy.jit import compile_from_string from codepy.toolchain import GCCToolchain, ToolchainGuessError, guess_toolchain -from immutables import Map from pytools import memoize_method from pytools.codegen import CodeGenerator, Indentation from pytools.prefork import ExecError -from loopy.codegen.result import GeneratedProgram -from loopy.kernel import LoopKernel from loopy.kernel.array import ArrayBase -from loopy.kernel.data import ArrayArg -from loopy.schedule.tools import KernelArgInfo from loopy.target.execution import ( ExecutionWrapperGeneratorBase, ExecutorBase, get_highlighted_code, ) -from loopy.translation_unit import TranslationUnit from loopy.types import LoopyType -from loopy.typing import Expression + + +if TYPE_CHECKING: + from immutables import Map + + from loopy.codegen.result import GeneratedProgram + from loopy.kernel import LoopKernel + from loopy.kernel.data import ArrayArg + from loopy.schedule.tools import KernelArgInfo + from loopy.translation_unit import TranslationUnit + from loopy.typing import Expression logger = logging.getLogger(__name__) diff --git a/loopy/target/c/codegen/expression.py b/loopy/target/c/codegen/expression.py index e8b6d0ffb..b1723e9d6 100644 --- a/loopy/target/c/codegen/expression.py +++ b/loopy/target/c/codegen/expression.py @@ -24,6 +24,8 @@ """ +from typing import TYPE_CHECKING + import numpy as np import islpy as isl @@ -45,13 +47,16 @@ from loopy.diagnostic import LoopyError from loopy.expression import dtype_to_type_context -from loopy.symbolic import TypeCast from loopy.target.c import CExpression from loopy.type_inference import TypeReader from loopy.types import LoopyType from loopy.typing import Expression, is_integer +if TYPE_CHECKING: + from loopy.symbolic import TypeCast + + __doc__ = """ .. currentmodule:: loopy.target.c.codegen.expression diff --git a/loopy/target/cuda.py b/loopy/target/cuda.py index 3003480b1..50d2ac7fe 100644 --- a/loopy/target/cuda.py +++ b/loopy/target/cuda.py @@ -24,7 +24,7 @@ THE SOFTWARE. """ -from typing import Sequence +from typing import TYPE_CHECKING, Sequence import numpy as np @@ -32,8 +32,6 @@ from pymbolic import var from pytools import memoize_method -from loopy.codegen import CodeGenerationState -from loopy.codegen.result import CodeGenerationResult from loopy.diagnostic import LoopyError, LoopyTypeError from loopy.kernel.array import ArrayBase, FixedStrideArrayDimTag, VectorArrayDimTag from loopy.kernel.data import ( @@ -49,6 +47,11 @@ from loopy.types import NumpyType +if TYPE_CHECKING: + from loopy.codegen import CodeGenerationState + from loopy.codegen.result import CodeGenerationResult + + # {{{ vector types class vec: # noqa diff --git a/loopy/target/execution.py b/loopy/target/execution.py index 3a89c06fe..cb737f95a 100644 --- a/loopy/target/execution.py +++ b/loopy/target/execution.py @@ -28,6 +28,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from typing import ( + TYPE_CHECKING, Any, Callable, Mapping, @@ -50,14 +51,17 @@ from loopy.kernel import KernelState, LoopKernel from loopy.kernel.data import ArrayArg, _ArraySeparationInfo, auto -from loopy.schedule.tools import KernelArgInfo from loopy.tools import LoopyKeyBuilder, caches -from loopy.translation_unit import TranslationUnit from loopy.types import LoopyType, NumpyType from loopy.typing import Expression, integer_expr_or_err from loopy.version import DATA_MODEL_VERSION +if TYPE_CHECKING: + from loopy.schedule.tools import KernelArgInfo + from loopy.translation_unit import TranslationUnit + + # {{{ object array argument packing class SeparateArrayPackingController: @@ -257,7 +261,7 @@ def generate_integer_arg_finding_from_array_data( unknown_var, = deps order_to_unknown_to_equations \ .setdefault(eqn.order, {}) \ - .setdefault(cast(Variable, unknown_var).name, []) \ + .setdefault(cast("Variable", unknown_var).name, []) \ .append(eqn) else: # Zero deps: nothing to determine, forget about it. diff --git a/loopy/target/ispc.py b/loopy/target/ispc.py index d8f8df26d..e493ee3e9 100644 --- a/loopy/target/ispc.py +++ b/loopy/target/ispc.py @@ -25,7 +25,7 @@ """ -from typing import Sequence, cast +from typing import TYPE_CHECKING, Sequence, cast import numpy as np @@ -35,16 +35,19 @@ from pymbolic.mapper.stringifier import PREC_NONE from pytools import memoize_method -from loopy.codegen import CodeGenerationState -from loopy.codegen.result import CodeGenerationResult from loopy.diagnostic import LoopyError from loopy.kernel.data import AddressSpace, ArrayArg, TemporaryVariable -from loopy.schedule import CallKernel from loopy.symbolic import Literal from loopy.target.c import CFamilyASTBuilder, CFamilyTarget from loopy.target.c.codegen.expression import ExpressionToCExpressionMapper -from loopy.types import LoopyType -from loopy.typing import Expression + + +if TYPE_CHECKING: + from loopy.codegen import CodeGenerationState + from loopy.codegen.result import CodeGenerationResult + from loopy.schedule import CallKernel + from loopy.types import LoopyType + from loopy.typing import Expression # {{{ expression mapper @@ -215,7 +218,7 @@ def get_function_declaration( assert codegen_state.kernel.linearization is not None subkernel_name = cast( - CallKernel, + "CallKernel", codegen_state.kernel.linearization[schedule_index] ).kernel_name diff --git a/loopy/target/opencl.py b/loopy/target/opencl.py index 6b5a0d313..4218ae9fa 100644 --- a/loopy/target/opencl.py +++ b/loopy/target/opencl.py @@ -24,16 +24,13 @@ THE SOFTWARE. """ -from typing import Sequence +from typing import TYPE_CHECKING, Sequence import numpy as np -from cgen import Declarator, Generable from pymbolic import var from pytools import memoize_method -from loopy.codegen import CodeGenerationState -from loopy.codegen.result import CodeGenerationResult from loopy.diagnostic import LoopyError, LoopyTypeError from loopy.kernel.array import ArrayBase, FixedStrideArrayDimTag, VectorArrayDimTag from loopy.kernel.data import AddressSpace, ConstantArg, ImageArg @@ -43,6 +40,13 @@ from loopy.types import NumpyType +if TYPE_CHECKING: + from cgen import Declarator, Generable + + from loopy.codegen import CodeGenerationState + from loopy.codegen.result import CodeGenerationResult + + # {{{ dtype registry wrappers diff --git a/loopy/target/pyopencl.py b/loopy/target/pyopencl.py index dd12dedb5..d3a2373ac 100644 --- a/loopy/target/pyopencl.py +++ b/loopy/target/pyopencl.py @@ -31,7 +31,6 @@ import numpy as np -import genpy import pymbolic.primitives as p from cgen import ( Block, @@ -46,10 +45,7 @@ ) from cgen.opencl import CLGlobal -from loopy.codegen import CodeGenerationState -from loopy.codegen.result import CodeGenerationResult from loopy.diagnostic import LoopyError, LoopyTypeError -from loopy.kernel import LoopKernel from loopy.kernel.data import ( ArrayArg, ConstantArg, @@ -64,18 +60,23 @@ OpenCLCASTBuilder, OpenCLTarget, ) -from loopy.target.pyopencl_execution import PyOpenCLExecutor from loopy.target.python import PythonASTBuilderBase -from loopy.translation_unit import FunctionIdT, TranslationUnit from loopy.types import NumpyType -from loopy.typing import Expression logger = logging.getLogger(__name__) if TYPE_CHECKING: + import genpy import pyopencl as cl + from loopy.codegen import CodeGenerationState + from loopy.codegen.result import CodeGenerationResult + from loopy.kernel import LoopKernel + from loopy.target.pyopencl_execution import PyOpenCLExecutor + from loopy.translation_unit import FunctionIdT, TranslationUnit + from loopy.typing import Expression + # {{{ pyopencl function scopers @@ -1033,7 +1034,7 @@ def get_function_definition( kernel = codegen_state.kernel assert kernel.linearization is not None - subkernel_name = cast(CallKernel, + subkernel_name = cast("CallKernel", kernel.linearization[schedule_index]).kernel_name result = [] @@ -1104,7 +1105,7 @@ def get_function_declaration( assert codegen_state.kernel.linearization is not None subkernel_name = cast( - CallKernel, + "CallKernel", codegen_state.kernel.linearization[schedule_index] ).kernel_name diff --git a/loopy/target/pyopencl_execution.py b/loopy/target/pyopencl_execution.py index 170122e3b..c9191e1d1 100644 --- a/loopy/target/pyopencl_execution.py +++ b/loopy/target/pyopencl_execution.py @@ -29,17 +29,12 @@ from typing import TYPE_CHECKING, Any, Callable, Sequence import numpy as np -from immutables import Map from pytools import memoize_method from pytools.codegen import CodeGenerator, Indentation -from loopy.codegen.result import CodeGenerationResult -from loopy.kernel import LoopKernel from loopy.kernel.data import ArrayArg -from loopy.schedule.tools import KernelArgInfo from loopy.target.execution import ExecutionWrapperGeneratorBase, ExecutorBase -from loopy.types import LoopyType from loopy.typing import Expression, integer_expr_or_err @@ -47,8 +42,15 @@ if TYPE_CHECKING: + from immutables import Map + import pyopencl as cl + from loopy.codegen.result import CodeGenerationResult + from loopy.kernel import LoopKernel + from loopy.schedule.tools import KernelArgInfo + from loopy.types import LoopyType + # {{{ invoker generation diff --git a/loopy/target/python.py b/loopy/target/python.py index 1881c465b..1b2560402 100644 --- a/loopy/target/python.py +++ b/loopy/target/python.py @@ -24,7 +24,7 @@ THE SOFTWARE. """ -from typing import Sequence +from typing import TYPE_CHECKING, Sequence import numpy as np @@ -32,14 +32,17 @@ from pymbolic.mapper import Mapper from pymbolic.mapper.stringifier import StringifyMapper -from loopy.codegen import CodeGenerationState -from loopy.codegen.result import CodeGenerationResult from loopy.diagnostic import LoopyError from loopy.kernel.data import ValueArg from loopy.target import ASTBuilderBase from loopy.type_inference import TypeReader +if TYPE_CHECKING: + from loopy.codegen import CodeGenerationState + from loopy.codegen.result import CodeGenerationResult + + # {{{ expression to code class ExpressionToPythonMapper(StringifyMapper): diff --git a/loopy/transform/array_buffer_map.py b/loopy/transform/array_buffer_map.py index 89acc7566..fb54dedd4 100644 --- a/loopy/transform/array_buffer_map.py +++ b/loopy/transform/array_buffer_map.py @@ -26,7 +26,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, replace -from typing import Any, Callable, Sequence +from typing import TYPE_CHECKING, Any, Callable, Sequence from typing_extensions import Self @@ -37,7 +37,10 @@ from pytools import memoize_method from loopy.symbolic import SubstitutionMapper, get_dependencies -from loopy.typing import Expression + + +if TYPE_CHECKING: + from loopy.typing import Expression @dataclass(frozen=True) diff --git a/loopy/transform/callable.py b/loopy/transform/callable.py index e5ca4d0a4..8669a4abb 100644 --- a/loopy/transform/callable.py +++ b/loopy/transform/callable.py @@ -23,7 +23,8 @@ THE SOFTWARE. """ -from collections.abc import Sequence + +from typing import TYPE_CHECKING from immutables import Map @@ -52,6 +53,10 @@ from loopy.translation_unit import FunctionIdT, TranslationUnit, for_each_kernel +if TYPE_CHECKING: + from collections.abc import Sequence + + __doc__ = """ .. currentmodule:: loopy diff --git a/loopy/transform/data.py b/loopy/transform/data.py index 9fb7fa700..80a0c4a12 100644 --- a/loopy/transform/data.py +++ b/loopy/transform/data.py @@ -24,7 +24,7 @@ """ from dataclasses import dataclass, replace -from typing import cast +from typing import TYPE_CHECKING, cast from warnings import warn import numpy as np @@ -39,7 +39,10 @@ from loopy.kernel.function_interface import CallableKernel, ScalarCallable from loopy.translation_unit import TranslationUnit, for_each_kernel from loopy.types import LoopyType -from loopy.typing import Expression + + +if TYPE_CHECKING: + from loopy.typing import Expression # {{{ convenience: add_prefetch @@ -1055,7 +1058,7 @@ def allocate_temporaries_for_base_storage(kernel: LoopKernel, approx_array_nbytes = 0 bs_key = (tv.base_storage, - cast(AddressSpace, tv.address_space), tv.dtype) + cast("AddressSpace", tv.address_space), tv.dtype) bsi = name_aspace_dtype_to_bsi.get(bs_key) if bsi is None or ( diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index aa16b197c..7587dd531 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -25,7 +25,7 @@ from collections.abc import Collection, Iterable, Mapping, Sequence -from typing import Any +from typing import TYPE_CHECKING, Any from typing_extensions import TypeAlias @@ -36,8 +36,6 @@ from loopy.diagnostic import LoopyError from loopy.kernel import LoopKernel from loopy.kernel.function_interface import CallableKernel -from loopy.kernel.instruction import InstructionBase -from loopy.match import ToStackMatchCovertible from loopy.symbolic import ( RuleAwareIdentityMapper, RuleAwareSubstitutionMapper, @@ -46,6 +44,11 @@ from loopy.translation_unit import TranslationUnit, for_each_kernel +if TYPE_CHECKING: + from loopy.kernel.instruction import InstructionBase + from loopy.match import ToStackMatchCovertible + + __doc__ = """ .. currentmodule:: loopy diff --git a/loopy/transform/instruction.py b/loopy/transform/instruction.py index 653ab2f9c..ec876ea03 100644 --- a/loopy/transform/instruction.py +++ b/loopy/transform/instruction.py @@ -23,16 +23,19 @@ THE SOFTWARE. """ -from typing import Mapping, Sequence +from typing import TYPE_CHECKING, Mapping, Sequence from loopy.diagnostic import LoopyError from loopy.kernel import LoopKernel from loopy.kernel.function_interface import CallableKernel, ScalarCallable -from loopy.kernel.instruction import InstructionBase from loopy.symbolic import RuleAwareIdentityMapper from loopy.translation_unit import TranslationUnit, for_each_kernel +if TYPE_CHECKING: + from loopy.kernel.instruction import InstructionBase + + # {{{ find_instructions def find_instructions_in_single_kernel(kernel, insn_match): diff --git a/loopy/transform/precompute.py b/loopy/transform/precompute.py index 88634fc48..a3d3113c2 100644 --- a/loopy/transform/precompute.py +++ b/loopy/transform/precompute.py @@ -25,7 +25,7 @@ from dataclasses import dataclass -from typing import Sequence, cast +from typing import TYPE_CHECKING, Sequence, cast import numpy as np from immutables import Map @@ -34,10 +34,8 @@ from pymbolic import ArithmeticExpression, var from pymbolic.mapper.substitutor import make_subst_func from pytools import memoize_on_first_arg -from pytools.tag import Tag from loopy.diagnostic import LoopyError -from loopy.kernel import LoopKernel from loopy.kernel.data import AddressSpace from loopy.kernel.function_interface import CallableKernel, ScalarCallable from loopy.kernel.instruction import InstructionBase, MultiAssignmentBase @@ -45,7 +43,6 @@ find_most_recent_global_barrier, kernel_has_global_barriers, ) -from loopy.match import ToStackMatchCovertible from loopy.symbolic import ( CombineMapper, RuleAwareIdentityMapper, @@ -71,6 +68,13 @@ ) +if TYPE_CHECKING: + from pytools.tag import Tag + + from loopy.kernel import LoopKernel + from loopy.match import ToStackMatchCovertible + + # {{{ contains_subst_rule_invocation class FunctionNameCollector(CombineMapper): @@ -582,7 +586,7 @@ def precompute_for_single_kernel( if isinstance(fpg, Variable): args: tuple[ArithmeticExpression, ...] = () elif isinstance(fpg, Call): - args = cast(tuple[ArithmeticExpression, ...], fpg.parameters) + args = cast("tuple[ArithmeticExpression, ...]", fpg.parameters) else: raise ValueError("footprint generator must " "be substitution rule invocation") diff --git a/loopy/transform/realize_reduction.py b/loopy/transform/realize_reduction.py index 60f1e8044..5f504e722 100644 --- a/loopy/transform/realize_reduction.py +++ b/loopy/transform/realize_reduction.py @@ -29,7 +29,7 @@ import logging from dataclasses import dataclass, replace -from typing import Callable, Sequence +from typing import TYPE_CHECKING, Callable, Sequence logger = logging.getLogger(__name__) @@ -37,12 +37,9 @@ from immutables import Map import islpy as isl -from pymbolic.primitives import ExpressionNode from pytools import memoize_on_first_arg -from pytools.tag import Tag from loopy.diagnostic import LoopyError, ReductionIsNotTriangularError, warn_with_kernel -from loopy.kernel import LoopKernel from loopy.kernel.data import AddressSpace, TemporaryVariable, make_assignment from loopy.kernel.function_interface import CallableKernel from loopy.kernel.instruction import Assignment, InstructionBase, MultiAssignmentBase @@ -51,6 +48,13 @@ from loopy.translation_unit import ConcreteCallablesTable, TranslationUnit +if TYPE_CHECKING: + from pymbolic.primitives import ExpressionNode + from pytools.tag import Tag + + from loopy.kernel import LoopKernel + + # {{{ reduction realization context @dataclass diff --git a/loopy/translation_unit.py b/loopy/translation_unit.py index 9fab7d8ad..37211ebc1 100644 --- a/loopy/translation_unit.py +++ b/loopy/translation_unit.py @@ -54,11 +54,11 @@ RuleAwareIdentityMapper, SubstitutionRuleMappingContext, ) -from loopy.target import TargetBase if TYPE_CHECKING: from loopy.kernel import LoopKernel + from loopy.target import TargetBase from loopy.target.execution import ExecutorBase diff --git a/loopy/typing.py b/loopy/typing.py index 9b11720f4..5316c356b 100644 --- a/loopy/typing.py +++ b/loopy/typing.py @@ -2,6 +2,7 @@ .. autoclass:: Expression .. autoclass:: ShapeType .. autodata:: InameStr +.. autodata:: InameStrSet .. currentmodule:: loopy @@ -50,6 +51,7 @@ StridesType: TypeAlias = ShapeType InameStr: TypeAlias = str +InameStrSet: TypeAlias = frozenset[InameStr] class auto: # noqa diff --git a/pyproject.toml b/pyproject.toml index 06eeea2b6..e299848b2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,10 +97,10 @@ extend-select = [ "Q", # flake8-quotes "W", # pycodestyle - # TODO "UP", # pyupgrade "RUF", # ruff "FA", + "TC", ] extend-ignore = [ "C90", # McCabe complexity @@ -157,7 +157,6 @@ warn_unused_ignores = true # check_untyped_defs = true exclude = [ - "loopy/target/c/compyte/ndarray/.*", "loopy/target/c/compyte/array.py", ]