Type lp.map_instructions.

Author: kaushikcfd

loopy/transform/instruction.py

@@ -39,7 +39,7 @@
 
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping, Sequence
+    from collections.abc import Callable, Mapping, Sequence
 
     from pymbolic import ArithmeticExpression
     from pymbolic.primitives import Subscript
@@ -84,7 +84,8 @@ def find_instructions(
 
 # {{{ map_instructions
 
-def map_instructions(kernel, insn_match, f):
+def map_instructions(kernel: LoopKernel, insn_match: ToMatchConvertible,
+                     f: Callable[[InstructionBase], InstructionBase]) -> LoopKernel:
     from loopy.match import parse_match
     match = parse_match(insn_match)
 
@@ -104,14 +105,16 @@ def map_instructions(kernel, insn_match, f):
 # {{{ set_instruction_priority
 
 @for_each_kernel
-def set_instruction_priority(kernel, insn_match, priority):
+def set_instruction_priority(
+    kernel: LoopKernel, insn_match: ToMatchConvertible, priority: int
+) -> LoopKernel:
     """Set the priority of instructions matching *insn_match* to *priority*.
 
     *insn_match* may be any instruction id match understood by
     :func:`loopy.match.parse_match`.
     """
 
-    def set_prio(insn):
+    def set_prio(insn: InstructionBase):
         return insn.copy(priority=priority)
 
     return map_instructions(kernel, insn_match, set_prio)
@@ -122,7 +125,9 @@ def set_prio(insn):
 # {{{ add_dependency
 
 @for_each_kernel
-def add_dependency(kernel, insn_match, depends_on):
+def add_dependency(
+    kernel: LoopKernel, insn_match: ToMatchConvertible, depends_on: ToMatchConvertible
+) -> LoopKernel:
     """Add the instruction dependency *dependency* to the instructions matched
     by *insn_match*.
 
@@ -148,7 +153,7 @@ def add_dependency(kernel, insn_match, depends_on):
 
     matched = [False]
 
-    def add_dep(insn):
+    def add_dep(insn: InstructionBase):
         new_deps = insn.depends_on
         matched[0] = True
         new_deps = added_deps if new_deps is None else new_deps | added_deps