Sharpen types, including precise islpy

Bump islpy dependency

Author: inducer

.basedpyright/baseline.json

@@ -61,4 +61,39 @@ Schedule
 .. automodule:: loopy.schedule.tools
 .. automodule:: loopy.schedule.tree
 
+References
+----------
+
+Mostly things that Sphinx (our documentation tool) should resolve but won't.
+
+.. class:: constantdict
+
+    See :class:`constantdict.constantdict`.
+
+.. class:: DTypeLike
+
+    See :data:`numpy.typing.DTypeLike`.
+
+.. currentmodule:: p
+
+.. class:: Call
+
+    See :class:`pymbolic.primitives.Call`.
+
+.. class:: CallWithKwargs
+
+    See :class:`pymbolic.primitives.CallWithKwargs`.
+
+.. currentmodule:: isl
+
+.. class:: Space
+
+    See :class:`islpy.Space`.
+
+.. class:: Aff
+
+    See :class:`islpy.Aff`.
+
+.. class:: PwAff
 
+    See :class:`islpy.PwAff`.

doc/ref_internals.rst

@@ -270,6 +270,12 @@ Instructions
 Assignment objects
 ^^^^^^^^^^^^^^^^^^
 
+.. currentmodule:: loopy.kernel.instruction
+
+.. class:: Assignable
+
+.. currentmodule:: loopy
+
 .. autoclass:: Assignment
 
 .. _assignment-syntax:

doc/ref_kernel.rst

@@ -24,7 +24,7 @@
 """
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 from warnings import warn
 
 import numpy as np
@@ -37,6 +37,8 @@
 if TYPE_CHECKING:
     import pyopencl.array as cla
 
+    from loopy.types import NumpyType
+
 
 AUTO_TEST_SKIP_RUN = False
 
@@ -142,7 +144,7 @@ def make_ref_args(kernel, queue, parameters):
                         "testing" % arg.name)
 
             shape = evaluate_shape(arg.shape, parameters)
-            dtype = arg.dtype
+            dtype = cast("NumpyType", arg.dtype).dtype
 
             is_output = arg.is_output
 

loopy/auto_test.py

@@ -25,10 +25,12 @@
 
 import logging
 from collections import defaultdict
+from dataclasses import dataclass
 from functools import reduce
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
+from typing_extensions import override
 
 import islpy as isl
 from islpy import dim_type
@@ -76,6 +78,7 @@
 if TYPE_CHECKING:
     from collections.abc import Mapping, Sequence
 
+    import pymbolic.primitives as p
     from pymbolic.typing import Expression
 
     from loopy.kernel import LoopKernel
@@ -707,10 +710,12 @@ def subst_func(x):
 # }}}
 
 
-class _AccessCheckMapper(WalkMapper):
-    def __init__(self, kernel, callables_table):
-        self.kernel = kernel
-        self.callables_table = callables_table
+@dataclass
+class _AccessCheckMapper(WalkMapper[[isl.Set, str]]):
+    kernel: LoopKernel
+    callables_table: CallablesTable
+
+    def __post_init__(self) -> None:
         super().__init__()
 
     @memoize_method
@@ -719,14 +724,20 @@ def _make_slab(self, space, iname, start, stop):
         return make_slab(space, iname, start, stop)
 
     @memoize_method
-    def _get_access_range(self, domain, subscript):
-        from loopy.symbolic import UnableToDetermineAccessRangeError, get_access_map
+    def _get_access_range(
+                self,
+                domain: isl.Set,
+                subscript: tuple[Expression, ...]
+            ):
+        from loopy.diagnostic import UnableToDetermineAccessRangeError
+        from loopy.symbolic import get_access_map
         try:
             return get_access_map(domain, subscript).range()
         except UnableToDetermineAccessRangeError:
             return None
 
-    def map_subscript(self, expr, domain, insn_id):
+    @override
+    def map_subscript(self, expr: p.Subscript, domain: isl.Set, insn_id: str):
         WalkMapper.map_subscript(self, expr, domain, insn_id)
 
         from pymbolic.primitives import Variable
@@ -742,6 +753,7 @@ def map_subscript(self, expr, domain, insn_id):
             shape = tv.shape
 
         if shape is not None:
+            assert isinstance(shape, tuple)
             subscript = expr.index
 
             if not isinstance(subscript, tuple):
@@ -787,7 +799,8 @@ def map_subscript(self, expr, domain, insn_id):
                         " establish '%s' is a subset of '%s')."
                         % (expr, insn_id, access_range, shape_domain))
 
-    def map_if(self, expr, domain, insn_id):
+    @override
+    def map_if(self, expr: p. If, domain: isl.Set, insn_id: str):
         from loopy.symbolic import condition_to_set
         then_set = condition_to_set(domain.space, expr.condition)
         if then_set is None:
@@ -800,7 +813,8 @@ def map_if(self, expr, domain, insn_id):
         self.rec(expr.then, domain & then_set, insn_id)
         self.rec(expr.else_, domain & else_set, insn_id)
 
-    def map_call(self, expr, domain, insn_id):
+    @override
+    def map_call(self, expr: p.Call, domain: isl.Set, insn_id: str):
         # perform access checks on the call arguments
         super().map_call(expr, domain, insn_id)
 
@@ -817,7 +831,9 @@ def map_call(self, expr, domain, insn_id):
             and isinstance(self.callables_table[expr.function.name],
                            CallableKernel)):
 
-            subkernel = self.callables_table[expr.function.name].subkernel
+            subkernel = cast(
+                             "CallableKernel",
+                             self.callables_table[expr.function.name]).subkernel
 
             # The plan here is to add the constraints coming from the values
             # args passed at a call-site as assumptions to the callee. To avoid
@@ -835,8 +851,8 @@ def map_call(self, expr, domain, insn_id):
 
             kw_space = isl.Space.create_from_names(
                 subkernel.isl_context, set=[],
-                params=(get_dependencies(tuple(kwargs.values()))
-                        | set(kwargs.keys())))
+                params=[*get_dependencies(tuple(kwargs.values())),
+                        *kwargs.keys()])
 
             extra_assumptions = isl.BasicSet.universe(kw_space).params()
 
@@ -894,8 +910,8 @@ def _check_bounds_inner(kernel: LoopKernel, callables_table: CallablesTable) ->
             domain, assumptions = isl.align_two(domain, kernel.assumptions)
             domain_with_assumptions = domain & assumptions
 
-        def run_acm(expr):
-            acm(expr, domain_with_assumptions, insn.id)  # noqa: B023
+        def run_acm(expr: Expression):
+            acm(expr, domain_with_assumptions, not_none(insn.id))  # noqa: B023
             return expr
 
         insn.with_transformed_expressions(run_acm)
@@ -1659,7 +1675,10 @@ def _get_sub_array_ref_swept_range(
     from loopy.symbolic import get_access_map
     domain = kernel.get_inames_domain(frozenset({iname_var.name
                                                  for iname_var in sar.swept_inames}))
-    return get_access_map(domain, sar.swept_inames, kernel.assumptions).range()
+    return get_access_map(
+                          domain.to_set(),
+                          sar.swept_inames,
+                          kernel.assumptions.to_set()).range()
 
 
 def _are_sub_array_refs_equivalent(
@@ -1875,7 +1894,7 @@ def pre_codegen_checks(t_unit: TranslationUnit) -> None:
 
 def check_implemented_domains(
             kernel: LoopKernel,
-            implemented_domains: Mapping[str, isl.Set],
+            implemented_domains: Mapping[str, Sequence[isl.Set]],
             code: str | None = None,
         ) -> bool:
     from islpy import align_two, dim_type
@@ -1942,7 +1961,7 @@ def check_implemented_domains(
             d_minus_i = desired_domain - insn_impl_domain
 
             parameter_inames = {
-                    insn_domain.get_dim_name(dim_type.param, i)
+                    not_none(insn_domain.get_dim_name(dim_type.param, i))
                     for i in range(insn_impl_domain.dim(dim_type.param))}
 
             lines = []

loopy/check.py

@@ -28,16 +28,15 @@
 from typing import (
     TYPE_CHECKING,
     Any,
-    Mapping,
-    Sequence,
 )
 
 import constantdict
 
+from loopy.typing import not_none
+
 
 logger = logging.getLogger(__name__)
 
-from functools import reduce
 
 import islpy  # to help out Sphinx
 import islpy as isl
@@ -47,20 +46,22 @@
 
 from loopy.diagnostic import LoopyError, warn
 from loopy.kernel.function_interface import CallableKernel
-from loopy.symbolic import CombineMapper
 from loopy.tools import LoopyKeyBuilder, caches
 from loopy.version import DATA_MODEL_VERSION
 
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+
+    from pymbolic import Expression
+
     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__ = """
@@ -206,7 +207,7 @@ def intersect(self, other):
         new_impl, new_other = isl.align_two(self.implemented_domain, other)
         return self.copy(implemented_domain=new_impl & new_other)
 
-    def fix(self, iname, aff):
+    def fix(self, iname: str, aff: isl.Aff) -> CodeGenerationState:
         new_impl_domain = self.implemented_domain
 
         impl_space = self.implemented_domain.get_space()
@@ -296,32 +297,6 @@ def ast_builder(self):
 caches.append(code_gen_cache)
 
 
-class InKernelCallablesCollector(CombineMapper):
-    """
-    Returns an instance of :class:`frozenset` containing instances of
-    :class:`loopy.kernel.function_interface.InKernelCallable` in the
-    :attr:``kernel`.
-    """
-    def __init__(self, kernel):
-        self.kernel = kernel
-
-    def combine(self, values):
-        import operator
-        return reduce(operator.or_, values, frozenset())
-
-    def map_resolved_function(self, expr):
-        return frozenset([self.kernel.scoped_functions[
-            expr.name]])
-
-    def map_constant(self, expr):
-        return frozenset()
-
-    map_variable = map_constant
-    map_function_symbol = map_constant
-    map_tagged_variable = map_constant
-    map_type_cast = map_constant
-
-
 @dataclass(frozen=True)
 class PreambleInfo:
     """
@@ -341,8 +316,12 @@ class PreambleInfo:
 
 # {{{ main code generation entrypoint
 
-def generate_code_for_a_single_kernel(kernel, callables_table, target,
-        is_entrypoint):
+def generate_code_for_a_single_kernel(
+            kernel: LoopKernel,
+            callables_table: CallablesTable,
+            target: TargetBase,
+            is_entrypoint: bool,
+        ) -> CodeGenerationResult:
     """
     :returns: a :class:`CodeGenerationResult`
 
@@ -359,8 +338,8 @@ def generate_code_for_a_single_kernel(kernel, callables_table, target,
     # {{{ examine arg list
 
     allow_complex = False
-    for var in kernel.args + list(kernel.temporary_variables.values()):
-        if var.dtype.involves_complex():
+    for var in [*kernel.args, *kernel.temporary_variables.values()]:
+        if not_none(var.dtype).involves_complex():
             allow_complex = True
 
     # }}}
@@ -376,7 +355,7 @@ def generate_code_for_a_single_kernel(kernel, callables_table, target,
     codegen_state = CodeGenerationState(
             kernel=kernel,
             target=target,
-            implemented_domain=initial_implemented_domain,
+            implemented_domain=isl.Set.from_basic_set(initial_implemented_domain),
             implemented_predicates=frozenset(),
             seen_dtypes=seen_dtypes,
             seen_functions=seen_functions,
@@ -389,7 +368,7 @@ def generate_code_for_a_single_kernel(kernel, callables_table, target,
                 target.host_program_name_prefix
                 + kernel.name
                 + kernel.target.host_program_name_suffix),
-            schedule_index_end=len(kernel.linearization),
+            schedule_index_end=len(not_none(kernel.linearization)),
             callables_table=callables_table,
             is_entrypoint=is_entrypoint,
             codegen_cache_manager=CodegenOperationCacheManager.from_kernel(kernel),
@@ -418,7 +397,8 @@ def generate_code_for_a_single_kernel(kernel, callables_table, target,
     if kernel.all_inames():
         seen_dtypes.add(kernel.index_dtype)
 
-    preambles = kernel.preambles + codegen_result.device_preambles
+    preambles = [
+        *kernel.preambles, *codegen_result.device_preambles]
 
     preamble_info = PreambleInfo(
             kernel=kernel,
@@ -429,10 +409,10 @@ def generate_code_for_a_single_kernel(kernel, callables_table, target,
             codegen_state=codegen_state
             )
 
-    preamble_generators = (list(kernel.preamble_generators)
-            + list(target.get_device_ast_builder().preamble_generators()))
-    for prea_gen in preamble_generators:
-        preambles = preambles + tuple(prea_gen(preamble_info))
+    for prea_gen in [
+            *kernel.preamble_generators,
+            *target.get_device_ast_builder().preamble_generators()]:
+        preambles.extend(prea_gen(preamble_info))
 
     codegen_result = codegen_result.copy(device_preambles=preambles)