Enforce ANN/PGH lints (#315)

One thing we lost in the pyre => pyright switch is enforcing that
everything is typed (I had disabled ANN since the errors were rudundant
with pyre).

Author: jansel

benchmarks/run.py

@@ -263,9 +263,9 @@ def run_kernel(kernel_name: str, tritonbench_args: list[str]) -> None:
 
     # Create the benchmark method
     def helion_method(
-        self: Any,
-        *args: Any,
-    ) -> Callable[..., Any]:
+        self: object,
+        *args: object,
+    ) -> Callable[..., object]:
         """Helion implementation."""
 
         # Reset all Helion kernels before creating the benchmark function

examples/all_gather_matmul.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import os
-from typing import Any
 
 import torch
 import torch.distributed as dist
@@ -118,7 +117,7 @@ def helion_all_gather_matmul(
     b: torch.Tensor,
     a_out: torch.Tensor | None = None,
     progress: torch.Tensor | None = None,
-    **kwargs: Any,
+    **kwargs: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     configs = {
         "SPLITS_PER_RANK": kwargs.get("splits_per_rank", 1),

examples/segment_reduction.py

@@ -48,7 +48,12 @@ def segmented_reduction_helion(
 
 
 @triton.jit
-def combine_fn_triton(left_values, left_indices, right_values, right_indices):
+def combine_fn_triton(
+    left_values: tl.tensor,
+    left_indices: tl.tensor,
+    right_values: tl.tensor,
+    right_indices: tl.tensor,
+) -> tuple[tl.tensor, tl.tensor]:
     same_segment = left_indices == right_indices
     combined_values = tl.where(same_segment, left_values + right_values, right_values)
     combined_indices = right_indices
@@ -67,13 +72,13 @@ def combine_fn_triton(left_values, left_indices, right_values, right_indices):
 )
 @triton.jit
 def _segmented_reduction_triton(
-    index,  # the input index tensor
-    in_ptr,  # the input tensor
-    out_ptr,  # the output value tensor
+    index: tl.tensor,  # the input index tensor
+    in_ptr: tl.tensor,  # the input tensor
+    out_ptr: tl.tensor,  # the output value tensor
     E: tl.constexpr,  # Number of elements in the input tensor (1d)
     C: tl.constexpr,  # Number of features in the input tensor (2d)
     BLOCK_SIZE: tl.constexpr,  # Block size for the scan
-):
+) -> None:
     # Triton version adapted from
     # https://github.com/fishmingyu/GeoT/blob/main/geot/triton/seg_reduction.py
     pid = tl.program_id(axis=0)
@@ -101,20 +106,24 @@ def _segmented_reduction_triton(
     tl.atomic_add(out_ptr + idxs * C + feature_id, result_values, mask & segment_start)
 
 
-def segmented_reduction_triton(indices, input_data, num_nodes):
+def segmented_reduction_triton(
+    indices: torch.Tensor, input_data: torch.Tensor, num_nodes: int
+) -> torch.Tensor:
     E, C = input_data.shape
     output = torch.zeros(
         (num_nodes, C), dtype=input_data.dtype, device=input_data.device
     )
 
-    def grid(META):
+    def grid(META: dict[str, int]) -> tuple[int, ...]:
         return (triton.cdiv(E, META["BLOCK_SIZE"]) * C,)
 
     _segmented_reduction_triton[grid](indices, input_data, output, E, C)
     return output
 
 
-def segmented_reduction_pytorch(indices, input_data, num_nodes):
+def segmented_reduction_pytorch(
+    indices: torch.Tensor, input_data: torch.Tensor, num_nodes: int
+) -> torch.Tensor:
     # Run PyTorch reference (scatter_add equivalent)
     num_features = input_data.size(1)
     pytorch_output = torch.zeros(
@@ -126,7 +135,7 @@ def segmented_reduction_pytorch(indices, input_data, num_nodes):
     return pytorch_output
 
 
-def main():
+def main() -> None:
     num_nodes = 100
     num_edges = 2000
     num_features = 128

helion/_compiler/ast_extension.py

@@ -234,7 +234,7 @@ def visit(self, node: ast.AST) -> ast.AST:
 class _TupleParensRemovedUnparser(
     ast._Unparser  # pyright: ignore[reportAttributeAccessIssue]
 ):
-    def visit_Tuple(self, node) -> None:
+    def visit_Tuple(self, node: ast.Tuple) -> None:
         if _needs_to_remove_tuple_parens and isinstance(
             getattr(node, "ctx", None), ast.Store
         ):

helion/_compiler/ast_read_writes.py

@@ -232,7 +232,7 @@ def visit_Expr(self, node: ast.Expr) -> ast.Expr | None:
 def dead_assignment_elimination(
     body: list[ast.AST],
     dce_vars: list[str],
-    num_iterations=8,
+    num_iterations: int = 8,
     input_rw: ReadWrites | None = None,
 ) -> None:
     """

helion/_compiler/inductor_lowering_extra.py

@@ -46,7 +46,7 @@ def _register_inductor_lowering(
     )
 
     @functools.wraps(decomp_fn)  # pyright: ignore[reportArgumentType]
-    def wrapped(*args: Any, **kwargs: Any) -> object:
+    def wrapped(*args: object, **kwargs: object) -> object:
         args = list(args)  # pyright: ignore[reportAssignmentType]
         kwargs = dict(kwargs)
         unpacked = False

helion/_compiler/traceback_compat.py

@@ -153,7 +153,7 @@ def normalize(off: int) -> int:
     return None  # fallback - no fancy anchors
 
 
-def format_frame_summary(frame_summary):  # type: ignore[override]
+def format_frame_summary(frame_summary: traceback.FrameSummary) -> str:  # type: ignore[override]
     """Backport of Python 3.11's traceback.StackSummary.format_frame_summary()."""
 
     _ensure_original_line(frame_summary)
@@ -170,18 +170,18 @@ def format_frame_summary(frame_summary):  # type: ignore[override]
         stripped_line = frame_summary.line.strip()
         row.append(f"    {stripped_line}\n")
 
-        line = frame_summary._original_line
+        line = frame_summary._original_line  # type: ignore[attr-defined]
         orig_line_len = len(line)
         frame_line_len = len(frame_summary.line.lstrip())
         stripped_characters = orig_line_len - frame_line_len
 
-        if frame_summary.colno is not None and frame_summary.end_colno is not None:
-            start_offset = _byte_offset_to_character_offset(line, frame_summary.colno)
-            end_offset = _byte_offset_to_character_offset(line, frame_summary.end_colno)
+        if frame_summary.colno is not None and frame_summary.end_colno is not None:  # type: ignore[attr-defined]
+            start_offset = _byte_offset_to_character_offset(line, frame_summary.colno)  # type: ignore[attr-defined]
+            end_offset = _byte_offset_to_character_offset(line, frame_summary.end_colno)  # type: ignore[attr-defined]
             code_segment = line[start_offset:end_offset]
 
             anchors = None
-            if frame_summary.lineno == frame_summary.end_lineno:
+            if frame_summary.lineno == frame_summary.end_lineno:  # type: ignore[attr-defined]
                 with suppress(Exception):
                     anchors = _extract_caret_anchors_from_line_segment(code_segment)
             else:

helion/_compiler/type_printer.py

@@ -51,7 +51,7 @@ def append(self, text: str) -> None:
 class ASTPrinter(_TupleParensRemovedUnparser):
     _indent: int
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, *args: object, **kwargs: object) -> None:
         super().__init__(*args, **kwargs)
         assert self._source == []
         self._source = self.output = OutputLines(self)

helion/autotuner/base_search.py

@@ -14,6 +14,10 @@
 import time
 from typing import TYPE_CHECKING
 from typing import NamedTuple
+from typing import NoReturn
+
+if TYPE_CHECKING:
+    from triton.runtime.jit import JITFunction
 
 from torch._inductor.runtime.triton_compat import OutOfResources
 from torch._inductor.runtime.triton_compat import PTXASError
@@ -152,7 +156,7 @@ def extract_launcher(
             grid: tuple[int, ...],
             *args: object,
             **kwargs: object,
-        ):
+        ) -> NoReturn:
             """Custom launcher that extracts arguments instead of executing."""
             raise _ExtractedLaunchArgs(triton_kernel, grid, args, kwargs)
 
@@ -524,7 +528,18 @@ def _mark_complete(self) -> bool:
 class _ExtractedLaunchArgs(Exception):
     """Exception that carries kernel launch arguments for precompiler extraction."""
 
-    def __init__(self, triton_kernel, grid, args, kwargs):
+    kernel: JITFunction[object]
+    grid: object
+    args: tuple[object, ...]
+    kwargs: dict[str, object]
+
+    def __init__(
+        self,
+        triton_kernel: JITFunction[object],
+        grid: object,
+        args: tuple[object, ...],
+        kwargs: dict[str, object],
+    ) -> None:
         super().__init__()
         self.kernel = triton_kernel
         self.grid = grid

helion/runtime/__init__.py

@@ -50,7 +50,7 @@ def default_launcher(
     *args: object,
     num_warps: int,
     num_stages: int,
-):
+) -> object:
     """Default launcher function that executes the kernel immediately."""
     return triton_kernel.run(
         *args, grid=grid, warmup=False, num_warps=num_warps, num_stages=num_stages