diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index 68c41336cd2d164dde136dec410bbcaee758b70a..fdaa88619d291f5b0d6c16feee77334e2c938a20 100644 --- a/loopy/codegen/__init__.py +++ b/loopy/codegen/__init__.py @@ -312,7 +312,10 @@ class CodeGenerationState: # }}} -code_gen_cache = WriteOncePersistentDict( +code_gen_cache: WriteOncePersistentDict[ + TranslationUnit, + CodeGenerationResult +] = WriteOncePersistentDict( "loopy-code-gen-cache-v3-"+DATA_MODEL_VERSION, key_builder=LoopyKeyBuilder()) diff --git a/loopy/schedule/__init__.py b/loopy/schedule/__init__.py index 5be848e02a74141aa875a0941d7dbb4d034283ea..6a73391490a408ae73d0dfaf544589a3044b3c66 100644 --- a/loopy/schedule/__init__.py +++ b/loopy/schedule/__init__.py @@ -42,6 +42,8 @@ from loopy.version import DATA_MODEL_VERSION if TYPE_CHECKING: from loopy.kernel import LoopKernel + from loopy.translation_unit import CallablesTable, TranslationUnit + logger = logging.getLogger(__name__) @@ -2197,7 +2199,10 @@ def _generate_loop_schedules_inner( # }}} -schedule_cache = WriteOncePersistentDict( +schedule_cache: WriteOncePersistentDict[ + Tuple[LoopKernel, CallablesTable], + LoopKernel +] = WriteOncePersistentDict( "loopy-schedule-cache-v4-"+DATA_MODEL_VERSION, key_builder=LoopyKeyBuilder()) diff --git a/loopy/target/execution.py b/loopy/target/execution.py index 1e49de938dc5340f3ccc7007bd5d4fb50ef555a5..ad49093e256b9bdcc0334cf9caba32d084aa21fd 100644 --- a/loopy/target/execution.py +++ b/loopy/target/execution.py @@ -21,7 +21,7 @@ THE SOFTWARE. """ -from typing import (Callable, Tuple, Union, Set, FrozenSet, List, Dict, +from typing import (Callable, Mapping, Tuple, Union, Set, FrozenSet, List, Dict, Optional, Sequence, Any) from dataclasses import dataclass @@ -721,7 +721,10 @@ class ExecutionWrapperGeneratorBase(ABC): # }}} -typed_and_scheduled_cache = WriteOncePersistentDict( +typed_and_scheduled_cache: WriteOncePersistentDict[ + Tuple[str, TranslationUnit, Optional[Mapping[str, LoopyType]]], + TranslationUnit +] = WriteOncePersistentDict( "loopy-typed-and-scheduled-cache-v1-"+DATA_MODEL_VERSION, key_builder=LoopyKeyBuilder()) @@ -729,7 +732,10 @@ typed_and_scheduled_cache = WriteOncePersistentDict( caches.append(typed_and_scheduled_cache) -invoker_cache = WriteOncePersistentDict( +invoker_cache: WriteOncePersistentDict[ + Tuple[str, TranslationUnit, str], + str +] = WriteOncePersistentDict( "loopy-invoker-cache-v10-"+DATA_MODEL_VERSION, key_builder=LoopyKeyBuilder()) diff --git a/setup.py b/setup.py index d9b8f664366d99fea680fad5f7a51e782ee84f93..17f68968f42bef04725d67bafc09f05a458d1e5a 100644 --- a/setup.py +++ b/setup.py @@ -84,7 +84,7 @@ setup(name="loopy", python_requires="~=3.8", install_requires=[ - "pytools>=2023.1.1", + "pytools>=2024.1.2", "pymbolic>=2022.1", "genpy>=2016.1.2",