Swap to using pyright (#259)

Author: oulgen

.github/workflows/lint.yml

@@ -47,9 +47,6 @@ jobs:
       - name: Install lint dependencies
         run: ./lint.sh install
 
-      - name: Run lint checks
-        run: ./lint.sh check
-
       - name: Install pre-commit
         run: |
           python -m pip install --upgrade pip

.pre-commit-config.yaml

@@ -31,12 +31,13 @@ repos:
 -   repo: local
     hooks:
       - id: check-main-in-examples
-        name: Ensure all example scripts have main()
+        name: ensure all example scripts have main()
         entry: python scripts/lint_examples_main.py
         language: system
         files: ^examples/.*\.py$
 
-- repo: https://github.com/facebook/pyre-check
-  rev: 24abb8681e9b5a93700cfd2a404de2506d0f9689
-  hooks:
-   - id: pyre-check-no-python
+-   repo: https://github.com/RobertCraigie/pyright-python
+    rev: v1.1.403
+    hooks:
+      - id: pyright
+        language: system

.pyre_configuration

@@ -63,7 +63,7 @@ def check_and_setup_tritonbench() -> None:
     """Check if tritonbench is installed and install it from GitHub if not."""
     # Check if tritonbench is already installed
     try:
-        import tritonbench
+        import tritonbench  # pyright: ignore[reportMissingImports]
 
         return  # Already installed
     except ImportError:
@@ -126,7 +126,7 @@ def check_and_setup_tritonbench() -> None:
 
         # Verify installation worked
         try:
-            import tritonbench  # noqa: F401
+            import tritonbench  # noqa: F401  # pyright: ignore[reportMissingImports]
 
             print(
                 f"Tritonbench installed successfully with {install_flag}.",
@@ -196,7 +196,9 @@ def main() -> None:
 
     # Import tritonbench components
     try:
-        from tritonbench.utils.parser import get_parser  # pyre-ignore[21]
+        from tritonbench.utils.parser import (  # pyright: ignore[reportMissingImports]
+            get_parser,
+        )
     except ImportError:
         print(
             "Error: Could not import tritonbench. Make sure it's in the path.",
@@ -223,16 +225,16 @@ def main() -> None:
     tb_args = tb_parser.parse_args(tritonbench_args)
 
     # Register the Helion kernel with tritonbench BEFORE importing the operator
-    from tritonbench.utils.triton_op import (  # pyre-ignore[21]
+    from tritonbench.utils.triton_op import (  # pyright: ignore[reportMissingImports]
         register_benchmark_mannually,
     )
 
     # Create the benchmark method
-    def create_helion_method(  # pyre-ignore[3]
-        kernel_func: Callable[..., Any],  # pyre-ignore[2]
+    def create_helion_method(
+        kernel_func: Callable[..., Any],
     ) -> Callable[..., Any]:
-        def helion_method(  # pyre-ignore[3]
-            self: Any,  # pyre-ignore[2]
+        def helion_method(
+            self: Any,
             *args: Any,
         ) -> Callable[..., Any]:
             """Helion implementation."""
@@ -246,7 +248,7 @@ def helion_method(  # pyre-ignore[3]
                 if isinstance(attr, Kernel):
                     attr.reset()
 
-            def _inner() -> Callable[..., Any]:  # pyre-ignore[3]
+            def _inner() -> Callable[..., Any]:
                 return kernel_func(*args)
 
             return _inner

benchmarks/run.py

@@ -181,7 +181,7 @@ def test(M: int, N: int, K: int, world_size: int, device: torch.device) -> None:
     dist_group = dist.group.WORLD
     if dist_group is None:
         raise RuntimeError("No distributed group available")
-    ag_golden, mm_golden = torch.ops.symm_mem.fused_all_gather_matmul(
+    ag_golden, mm_golden = torch.ops.symm_mem.fused_all_gather_matmul(  # pyright: ignore[reportCallIssue]
         golden_a, [b], gather_dim=0, group_name=dist_group.group_name
     )
     torch.testing.assert_close(c, mm_golden[0], rtol=1e-1, atol=1e-1)

examples/all_gather_matmul.py

@@ -66,10 +66,9 @@ def attention(
     return out.view(q_in.size())
 
 
-attention_dynamic: object = helion.kernel(
+attention_dynamic: object = helion.kernel(  # pyright: ignore[reportCallIssue]
     attention.fn,
-    # pyre-fixme[6]
-    configs=attention.configs,
+    configs=attention.configs,  # pyright: ignore[reportArgumentType]
     static_shapes=False,
 )
 

examples/attention.py

@@ -39,7 +39,6 @@ class ExtendedAST:
     subclassing each AST node class and mixing in this one.
     """
 
-    # pyre-ignore[13]
     _fields: tuple[str, ...]
 
     def __init__(
@@ -132,8 +131,7 @@ class Wrapper(ExtendedAST, cls):
 
 
 def create(cls: type[_T], **fields: object) -> _T:
-    # pyre-ignore[28]
-    result = get_wrapper_cls(cls)(**fields, _location=current_location())
+    result = get_wrapper_cls(cls)(**fields, _location=current_location())  # pyright: ignore[reportCallIssue]
     assert isinstance(result, ExtendedAST)
     result._location.to_ast(result)
     return typing.cast("_T", result)
@@ -165,16 +163,16 @@ def statement_from_string(template: str, **placeholders: ast.AST) -> ast.stmt:
 
     def _replace(node: _R) -> _R:
         if isinstance(node, list):
-            return [_replace(item) for item in node]
+            return [_replace(item) for item in node]  # pyright: ignore[reportReturnType]
         if not isinstance(node, ast.AST):
             return node
         if isinstance(node, ast.Name) and node.id in placeholders:
-            return placeholders[node.id]
+            return placeholders[node.id]  # pyright: ignore[reportReturnType]
         cls = get_wrapper_cls(type(node))
-        return location.to_ast(
+        return location.to_ast(  # pyright: ignore[reportReturnType]
             cls(
                 **{field: _replace(getattr(node, field)) for field in node._fields},
-                _location=location,
+                _location=location,  # pyright: ignore[reportCallIssue]
             )
         )
 
@@ -196,11 +194,10 @@ def convert(node: ast.AST) -> ast.AST:
             # some nodes like arguments lack location information
             location = current_location()
         with location:
-            # pyre-ignore[28]
             return cls(
                 **{field: convert(getattr(node, field)) for field in node._fields},
                 **{attr: getattr(node, attr) for attr in node._attributes},
-                _location=location,
+                _location=location,  # pyright: ignore[reportCallIssue]
             )
     elif isinstance(node, list):
         return [convert(item) for item in node]
@@ -234,23 +231,23 @@ def visit(self, node: ast.AST) -> ast.AST:
 )
 
 
-class _TupleParensRemovedUnparser(ast._Unparser):  # pyre-ignore[11]
-    def visit_Tuple(self, node) -> None:  # pyre-ignore[2]
+class _TupleParensRemovedUnparser(
+    ast._Unparser  # pyright: ignore[reportAttributeAccessIssue]
+):
+    def visit_Tuple(self, node) -> None:
         if _needs_to_remove_tuple_parens and isinstance(
             getattr(node, "ctx", None), ast.Store
         ):
             if len(node.elts) == 1:  # single-element tuple
-                self.traverse(node.elts[0])  # pyre-ignore[16]
-                self.write(",")  # pyre-ignore[16]
+                self.traverse(node.elts[0])
+                self.write(",")
             else:  # multi-element tuple
-                self.interleave(  # pyre-ignore[16]
-                    lambda: self.write(", "), self.traverse, node.elts
-                )
+                self.interleave(lambda: self.write(", "), self.traverse, node.elts)
             return
         # For everything else fall back to default behavior
-        super().visit_Tuple(node)  # pyre-ignore[16]
+        super().visit_Tuple(node)
 
 
 def unparse(ast_obj: ast.AST) -> str:
     unparser = _TupleParensRemovedUnparser()
-    return unparser.visit(ast_obj)  # pyre-ignore[16]
+    return unparser.visit(ast_obj)

helion/_compiler/ast_extension.py

@@ -43,7 +43,7 @@ class ReadWrites(typing.NamedTuple):
     reads: dict[str, int]
     writes: dict[str, int]
 
-    def __iter__(self) -> typing.Iterator[str]:
+    def __iter__(self) -> typing.Iterator[str]:  # pyright: ignore[reportIncompatibleMethodOverride]
         return iter({**self.reads, **self.writes})
 
     @staticmethod

helion/_compiler/ast_read_writes.py

@@ -161,7 +161,7 @@ def create_block_var(self, debug_name: str, hint: int = 64) -> torch.SymInt:
             # )
             # TODO(jansel): I was hoping the above would work, seems like some decomps require concrete values
             #               to determine zeroness.  Figure out a better way to do this.
-            # pyre-ignore[29]
+
             self.shape_env.var_to_val[sym._sympy_()] = sympy.Integer(hint)
         assert isinstance(sym._sympy_(), sympy.Symbol)
         self.debug_shape_renames[sym._sympy_()] = sympy.Symbol(debug_name, integer=True)
@@ -191,8 +191,8 @@ def cached_create_unbacked_symint(
         Returns:
             A consistent unbacked symint for the given key
         """
-        # pyre-ignore[16]
-        key = tuple([x._sympy_() if hasattr(x, "_sympy_") else x for x in key])
+
+        key = tuple([x._sympy_() if hasattr(x, "_sympy_") else x for x in key])  # pyright: ignore[reportAttributeAccessIssue]
         result = self._symint_cache.get(key)
         if result is None:
             result = self.create_unbacked_symint(hint)
@@ -237,9 +237,9 @@ def to_fake(self, obj: object, origin: Origin) -> object:
             return [self.to_fake(e, origin) for e in obj]
         if isinstance(obj, tuple) and hasattr(obj, "_fields"):
             return type(obj)(
-                **{  # pyre-ignore[6]
+                **{
                     k: self.to_fake(e, origin)
-                    for k, e in obj._asdict().items()  # pyre-ignore[16]
+                    for k, e in obj._asdict().items()  # pyright: ignore[reportAttributeAccessIssue]
                 }
             )
         if isinstance(obj, tuple):
@@ -248,10 +248,10 @@ def to_fake(self, obj: object, origin: Origin) -> object:
             return {k: self.to_fake(e, origin) for k, e in obj.items()}
         if dataclasses.is_dataclass(obj):
             return dataclasses.replace(
-                obj,
+                obj,  # pyright: ignore[reportArgumentType]
                 **{
                     k: self.to_fake(getattr(obj, k), origin)
-                    for k in obj.__dataclass_fields__  # pyre-ignore[16]
+                    for k in obj.__dataclass_fields__
                 },
             )
 
@@ -289,8 +289,8 @@ def size_hint(self, n: int | torch.SymInt) -> int:
                 # If the size is a symbolic expression with unbacked symbols, then the shape environment
                 # hint will be wrong since we assign a default value to unbacked symbols.  Return a default hint.
                 return 8192
-            # pyre-ignore[6]
-            return int(self.shape_env.size_hint(n._sympy_()))
+
+            return int(self.shape_env.size_hint(n._sympy_()))  # pyright: ignore[reportArgumentType]
         assert isinstance(n, int)
         return n
 
@@ -514,5 +514,4 @@ def _to_sympy(x: int | torch.SymInt) -> sympy.Expr:
 
 
 def _has_unbacked(expr: sympy.Expr) -> bool:
-    # pyre-ignore[16]
-    return any(n.name.startswith("u") for n in expr.free_symbols)
+    return any(n.name.startswith("u") for n in expr.free_symbols)  # pyright: ignore[reportAttributeAccessIssue]

helion/_compiler/compile_environment.py

@@ -221,7 +221,7 @@ def sympy_expr(self, expr: sympy.Expr) -> str:
         if expr in expr_to_origin:
             return self._lift_sympy_arg(expr)
         replacements = {}
-        for sym in sorted(expr.free_symbols, key=lambda x: x.name):
+        for sym in sorted(expr.free_symbols, key=lambda x: x.name):  # pyright: ignore[reportAttributeAccessIssue]
             assert isinstance(sym, sympy.Symbol)
             if sym in self.expr_to_var_info:
                 replacements[sym] = sympy.Symbol(
@@ -254,7 +254,7 @@ def _lift_sympy_arg(self, expr: sympy.Expr) -> str:
     def user_sympy_expr(self, expr: sympy.Expr) -> str:
         """A sympy expression that flows into user computations."""
         replacements = {}
-        for sym in sorted(expr.free_symbols, key=lambda s: s.name):
+        for sym in sorted(expr.free_symbols, key=lambda s: s.name):  # pyright: ignore[reportAttributeAccessIssue]
             assert isinstance(sym, sympy.Symbol)
             block_idx = CompileEnvironment.current().get_block_id(sym)
             if block_idx is not None:
@@ -474,7 +474,9 @@ def dead_code_elimination(self) -> None:
 
         # drop any unused args
         args_to_remove = {
-            arg.name for arg in self.arguments if arg.name not in rw.reads
+            arg.name
+            for arg in self.arguments
+            if arg.name not in rw.reads  # pyright: ignore[reportPossiblyUnboundVariable]
         }
         if args_to_remove:
             self.arguments = [