From 52b0f4e1bc067c2e0f71aa62118fab2c55502f20 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Wed, 28 Oct 2026 15:54:56 -0500 Subject: [PATCH 1/2] Improve typing around ISPC codegen --- loopy/target/c/codegen/expression.py | 4 +-- loopy/target/ispc.py | 44 +++++++++++++++++++++++----- 2 files changed, 38 insertions(+), 10 deletions(-) diff --git a/loopy/target/c/codegen/expression.py b/loopy/target/c/codegen/expression.py index 1d7d5d2fd..b2ea5c4ba 100644 --- a/loopy/target/c/codegen/expression.py +++ b/loopy/target/c/codegen/expression.py @@ -632,13 +632,13 @@ def map_power(self, expr: p.Power, type_context: TypeContext): @override def map_group_hw_index(self, expr: GroupHardwareAxisIndex, - type_context: TypeContext): + type_context: TypeContext) -> Expression: raise LoopyError("plain C does not have group hw axes") @override def map_local_hw_index(self, expr: LocalHardwareAxisIndex, - type_context: TypeContext): + type_context: TypeContext) -> Expression: raise LoopyError("plain C does not have local hw axes") def map_nan(self, expr: p.NaN, type_context: TypeContext): diff --git a/loopy/target/ispc.py b/loopy/target/ispc.py index 821fe464b..c1ac1da59 100644 --- a/loopy/target/ispc.py +++ b/loopy/target/ispc.py @@ -61,22 +61,27 @@ from loopy.codegen import CodeGenerationState from loopy.codegen.result import CodeGenerationResult + from loopy.expression import TypeContext from loopy.kernel import LoopKernel from loopy.kernel.instruction import Assignment from loopy.schedule import CallKernel + from loopy.target.c.compyte.dtypes import DTypeRegistry from loopy.types import LoopyType class IsVaryingMapper(CombineMapper[bool, []]): # FIXME: Update this if/when ispc reduction support is added. + kernel: LoopKernel def __init__(self, kernel: LoopKernel) -> None: self.kernel = kernel super().__init__() + @override def combine(self, values: Iterable[bool]) -> bool: return reduce(operator.or_, values, False) + @override def map_constant(self, expr): return False @@ -113,18 +118,25 @@ def _get_index_ctype(self): else: raise ValueError("unexpected index_type") - def map_group_hw_index(self, expr, type_context): + @override + def map_group_hw_index(self, + expr: GroupHardwareAxisIndex, + type_context: TypeContext): return var( "((uniform %s) taskIndex%d)" % (self._get_index_ctype(), expr.axis)) - def map_local_hw_index(self, expr, type_context): + @override + def map_local_hw_index(self, + expr: LocalHardwareAxisIndex, + type_context: TypeContext): if expr.axis == 0: return var("(varying %s) programIndex" % self._get_index_ctype()) else: raise LoopyError("ISPC only supports one local axis") - def map_constant(self, expr, type_context): + @override + def map_constant(self, expr: object, type_context: TypeContext): if isinstance(expr, (complex, np.complexfloating)): raise NotImplementedError("complex numbers in ispc") else: @@ -143,7 +155,8 @@ def map_constant(self, expr, type_context): raise RuntimeError("don't know how to generate code " "for constant '%s'" % expr) - def map_variable(self, expr, type_context): + @override + def map_variable(self, expr: p.Variable, type_context: TypeContext): tv = self.kernel.temporary_variables.get(expr.name) if tv is not None and tv.address_space == AddressSpace.PRIVATE: @@ -160,7 +173,8 @@ def map_variable(self, expr, type_context): else: return super().map_variable(expr, type_context) - def map_subscript(self, expr, type_context): + @override + def map_subscript(self, expr: p.Subscript, type_context: TypeContext): from loopy.kernel.data import TemporaryVariable ary = self.find_array(expr) @@ -199,7 +213,11 @@ def wrap_in_typecast(self, actual_type: LoopyType, needed_type: LoopyType, s): raise NotImplementedError("wrap_in_typecast needs uniform-ness information " "for ispc") - def rec(self, expr, type_context=None, needed_type: LoopyType | None = None): + @override + def rec(self, + expr: Expression, + type_context: TypeContext = None, + needed_type: LoopyType | None = None): result = super().rec(expr, type_context) if needed_type is None: @@ -223,7 +241,10 @@ def rec(self, expr, type_context=None, needed_type: LoopyType | None = None): # {{{ type registry -def fill_registry_with_ispc_types(reg, respect_windows, include_bool=True): +def fill_registry_with_ispc_types( + reg: DTypeRegistry, + respect_windows: bool, + include_bool: bool = True): reg.get_or_register_dtype("bool", bool) reg.get_or_register_dtype(["int8", "signed char", "char"], np.int8) @@ -330,6 +351,7 @@ def get_function_declaration( from loopy.target.c import FunctionDeclarationWrapper return [], FunctionDeclarationWrapper(result) + @override def get_kernel_call(self, codegen_state: CodeGenerationState, subkernel_name: str, gsize: tuple[Expression, ...], @@ -369,12 +391,14 @@ def get_kernel_call(self, codegen_state: CodeGenerationState, # {{{ code generation guts - def get_expression_to_c_expression_mapper(self, codegen_state): + @override + def get_expression_to_c_expression_mapper(self, codegen_state: CodeGenerationState): return ExprToISPCExprMapper(codegen_state) def add_vector_access(self, access_expr, index): return access_expr[index] + @override def emit_barrier(self, synchronization_kind, mem_kind, comment): from cgen import Comment, Statement @@ -393,12 +417,14 @@ def emit_barrier(self, synchronization_kind, mem_kind, comment): # {{{ declarators + @override def get_value_arg_declarator( self, name: str, dtype: LoopyType, is_written: bool) -> Declarator: from cgen.ispc import ISPCUniform return ISPCUniform(super().get_value_arg_declarator( name, dtype, is_written)) + @override def get_array_arg_declarator( self, arg: ArrayArg, is_written: bool) -> Declarator: # FIXME restrict? @@ -411,6 +437,7 @@ def get_array_arg_declarator( return decl + @override def get_temporary_var_declarator(self, codegen_state: CodeGenerationState, temp_var: TemporaryVariable) -> Declarator: @@ -442,6 +469,7 @@ def get_temporary_var_declarator(self, # {{{ emit_... + @override def emit_assignment( self, codegen_state: CodeGenerationState, From 860882102f76ddc120afa2e6a2e5d65cbc35238a Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Wed, 28 Oct 2026 15:59:35 -0500 Subject: [PATCH 2/2] Update baseline --- .basedpyright/baseline.json | 456 ------------------------------------ 1 file changed, 456 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 309719deb..24b052692 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -38619,14 +38619,6 @@ "lineCount": 1 } }, - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 26, - "lineCount": 1 - } - }, { "code": "reportImplicitOverride", "range": { @@ -38651,14 +38643,6 @@ "lineCount": 1 } }, - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 26, - "lineCount": 1 - } - }, { "code": "reportImplicitOverride", "range": { @@ -40743,22 +40727,6 @@ } ], "./loopy/target/ispc.py": [ - { - "code": "reportUnannotatedClassAttribute", - "range": { - "startColumn": 13, - "endColumn": 19, - "lineCount": 1 - } - }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 15, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -40775,14 +40743,6 @@ "lineCount": 1 } }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 20, - "lineCount": 1 - } - }, { "code": "reportUnknownParameterType", "range": { @@ -40823,94 +40783,6 @@ "lineCount": 1 } }, - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 26, - "lineCount": 1 - } - }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 26, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 33, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 39, - "endColumn": 51, - "lineCount": 1 - } - }, - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 26, - "lineCount": 1 - } - }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 26, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 33, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 39, - "endColumn": 51, - "lineCount": 1 - } - }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 20, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 27, - "endColumn": 31, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 33, - "endColumn": 45, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -40927,30 +40799,6 @@ "lineCount": 1 } }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 20, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 27, - "endColumn": 31, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 33, - "endColumn": 45, - "lineCount": 1 - } - }, { "code": "reportUnknownParameterType", "range": { @@ -40959,30 +40807,6 @@ "lineCount": 1 } }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 21, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 28, - "endColumn": 32, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 34, - "endColumn": 46, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -40991,14 +40815,6 @@ "lineCount": 1 } }, - { - "code": "reportCallIssue", - "range": { - "startColumn": 28, - "endColumn": 77, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -41007,14 +40823,6 @@ "lineCount": 1 } }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 61, - "endColumn": 76, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -41055,174 +40863,6 @@ "lineCount": 1 } }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 11, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 18, - "endColumn": 22, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 24, - "endColumn": 36, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 34, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 34, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 39, - "endColumn": 54, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 39, - "endColumn": 54, - "lineCount": 1 - } - }, - { - "code": "reportUnusedParameter", - "range": { - "startColumn": 39, - "endColumn": 54, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 56, - "endColumn": 68, - "lineCount": 1 - } - }, - { - "code": "reportUnusedParameter", - "range": { - "startColumn": 56, - "endColumn": 68, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 4, - "endColumn": 29, - "lineCount": 1 - } - }, { "code": "reportUnannotatedClassAttribute", "range": { @@ -41303,14 +40943,6 @@ "lineCount": 1 } }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 23, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -41343,38 +40975,6 @@ "lineCount": 1 } }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportUnknownParameterType", - "range": { - "startColumn": 52, - "endColumn": 65, - "lineCount": 1 - } - }, - { - "code": "reportMissingParameterType", - "range": { - "startColumn": 52, - "endColumn": 65, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 36, - "endColumn": 49, - "lineCount": 1 - } - }, { "code": "reportUnknownParameterType", "range": { @@ -41423,14 +41023,6 @@ "lineCount": 1 } }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 20, - "lineCount": 1 - } - }, { "code": "reportUnknownParameterType", "range": { @@ -41495,30 +41087,6 @@ "lineCount": 1 } }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 32, - "lineCount": 1 - } - }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 32, - "lineCount": 1 - } - }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 36, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -41527,14 +41095,6 @@ "lineCount": 2 } }, - { - "code": "reportImplicitOverride", - "range": { - "startColumn": 8, - "endColumn": 23, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -42113,14 +41673,6 @@ "lineCount": 1 } }, - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 26, - "lineCount": 1 - } - }, { "code": "reportImplicitOverride", "range": { @@ -42137,14 +41689,6 @@ "lineCount": 1 } }, - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 26, - "lineCount": 1 - } - }, { "code": "reportImplicitOverride", "range": {