Sharpen types, including precise islpy

Bump islpy dependency

Author: inducer

.basedpyright/baseline.json

@@ -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

@@ -34,6 +34,9 @@
 
 import constantdict
 
+from loopy.types import NumpyType
+from loopy.typing import not_none
+
 
 logger = logging.getLogger(__name__)
 
@@ -53,14 +56,15 @@
 
 
 if TYPE_CHECKING:
+    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 +210,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()
@@ -341,8 +345,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,7 +367,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()):
+    for var in [*kernel.args, *kernel.temporary_variables.values()]:
+        assert isinstance(var.dtype, NumpyType)
         if var.dtype.involves_complex():
             allow_complex = True
 
@@ -376,7 +385,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 +398,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 +427,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 +439,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)
 

loopy/codegen/__init__.py

@@ -25,10 +25,11 @@
 """
 
 from functools import partial
+from typing import TYPE_CHECKING
 
 import islpy as isl
 
-from loopy.codegen.result import merge_codegen_results, wrap_in_if
+from loopy.codegen.result import CodeGenerationResult, merge_codegen_results, wrap_in_if
 from loopy.diagnostic import LoopyError
 from loopy.schedule import (
     Barrier,
@@ -41,8 +42,17 @@
 )
 
 
-def generate_code_for_sched_index(codegen_state, sched_index):
+if TYPE_CHECKING:
+    from loopy.codegen import CodeGenerationState
+
+
+def generate_code_for_sched_index(
+            codegen_state: CodeGenerationState,
+            sched_index: int
+        ) -> CodeGenerationResult:
     kernel = codegen_state.kernel
+    assert kernel.linearization is not None
+
     sched_item = kernel.linearization[sched_index]
 
     if isinstance(sched_item, CallKernel):