Fix approx 150 type errors and enable mypy in core. (#508)

Excludes kernel for the moment as it has quite a number of typing
abnormalities.

(I wasn't keeping detailed count but there were a non trivial number of
bugs -- not just typing adjustments)

Author: stellaraccident

.github/workflows/test.yml

@@ -57,7 +57,12 @@ jobs:
         run: |
           pytest -n 4 serving/
 
-      - name: MyPy Type Checking
+      - name: MyPy Type Checking Core
         if: ${{ !cancelled() }}
         run: |
-          mypy serving/
+          (cd core && mypy)
+
+      - name: MyPy Type Checking Serving
+        if: ${{ !cancelled() }}
+        run: |
+          (cd serving && mypy)

core/mypy.ini

@@ -1,4 +1,21 @@
 [mypy]
 
+explicit_package_bases = True
 mypy_path = $MYPY_CONFIG_FILE_DIR
-packages = shark_turbine.aot,shark_turbine.dynamo,shark_turbine.support
+packages = shark_turbine
+
+# Missing typing stubs for iree.compiler.
+[mypy-iree.compiler.*]
+ignore_missing_imports = True
+
+# Missing typing stubs for iree.runtime.
+[mypy-iree.runtime.*]
+ignore_missing_imports = True
+
+# fx_importer needs to be fixed upstream.
+[mypy-shark_turbine.importers.fx_importer.*]
+ignore_errors = True
+
+# TODO: Fix all typing errors in TK.
+[mypy-shark_turbine.kernel.*]
+ignore_errors = True

core/shark_turbine/aot/builtins/globals.py

@@ -22,7 +22,7 @@
     GlobalAttributes,
 )
 
-from ..support.utils import (
+from torch.utils._pytree import (
     TreeSpec,
     tree_flatten,
     tree_map,
@@ -48,7 +48,7 @@ def __init__(
     ):
         if attrs is None:
             attrs = GlobalAttributes(
-                mutable=mutable,
+                mutable=bool(mutable),
                 external=external,
                 external_scope=external_scope,
                 name_mapper=name_mapper,
@@ -85,7 +85,7 @@ def __init__(
     ):
         if attrs is None:
             attrs = GlobalAttributes(
-                mutable=mutable,
+                mutable=bool(mutable),
                 external=external,
                 external_scope=external_scope,
                 name_mapper=name_mapper,
@@ -135,7 +135,7 @@ def __init__(
     ):
         if attrs is None:
             attrs = GlobalAttributes(
-                mutable=mutable,
+                mutable=bool(mutable),
                 external=external,
                 external_scope=external_scope,
                 name_mapper=name_mapper,

core/shark_turbine/aot/builtins/jittable.py

@@ -21,6 +21,10 @@
     GraphModule,
 )
 from torch.fx.passes.shape_prop import TensorMetadata
+from torch.utils._pytree import (
+    tree_flatten,
+    tree_unflatten,
+)
 
 # TODO: Switch to upstream fx_importer vs local fork when ready.
 # from iree.compiler.extras.fx_importer import (
@@ -51,16 +55,12 @@
     util_d,
 )
 
+from ...support.logging import aot_logger as logger
+
 from ..passes import (
     functorch_functionalize,
 )
 
-from ..support.utils import (
-    logger,
-    tree_flatten,
-    tree_unflatten,
-)
-
 from ..support.ir_utils import (
     ModuleBuilder,
 )
@@ -153,10 +153,8 @@ def __init__(
         self,
         wrapped_f,
         *,
-        decompose_ops: Optional[List[torch._ops.OpOverload]] = None,
-        decomposition_table: Optional[
-            Dict[torch._ops.OpOverload, Callable[..., Any]]
-        ] = None,
+        decompose_ops: Optional[List[Any]] = None,
+        decomposition_table: Optional[Dict[Any, Callable[..., Any]]] = None,
         constraints: Optional[List[Constraint]] = None,
         function_name: Optional[str] = None,
         passes: Sequence[str] = DEFAULT_PASSES,
@@ -311,6 +309,7 @@ def flat_wrapped_f(*args):
         assert len(flat_ir_results) == len(result_tensor_infos)
         flat_py_results = []
         for ir_result, result_tensor_info in zip(flat_ir_results, result_tensor_infos):
+            assert result_tensor_info is not None
             (dtype,) = result_tensor_info
             native_ir_result = type_converter.materialize_torch_to_native(ir_result)
             if dtype is not None:
@@ -478,5 +477,5 @@ def _extract_graph_output_metadata(
                 dtype = tensor_meta.dtype
             elif fake_val is not None:
                 dtype = fake_val.dtype
-            output_metadata.append((dtype,))
+            output_metadata.append((dtype,) if dtype is not None else None)
     return out_spec, output_metadata

core/shark_turbine/aot/compiled_module.py

@@ -5,7 +5,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from typing import Callable, Dict, Generator, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
 
 import enum
 import inspect
@@ -26,6 +26,7 @@
     PassManager,
     StringAttr,
 )
+from ..support.logging import aot_logger as logger
 from ..transforms.general.custom_op_expansion import ExpandCustomOpsPass
 
 from .support.procedural import (
@@ -38,9 +39,6 @@
     ModuleBuilder,
 )
 
-from .support.utils import (
-    logger,
-)
 
 __all__ = [
     "CompiledModule",
@@ -155,21 +153,21 @@ def export_procs(self) -> Generator[Tuple[str, ExportProcDef], None, None]:
         return filter(
             lambda kv_tuple: isinstance(kv_tuple[1], ExportProcDef),
             self.all_exports.items(),
-        )
+        )  # type: ignore
 
     @property
     def py_only_defs(self) -> Generator[Tuple[str, PyOnlyDef], None, None]:
         return filter(
             lambda kv_tuple: isinstance(kv_tuple[1], PyOnlyDef),
             self.all_exports.items(),
-        )
+        )  # type: ignore
 
     @property
     def globals_defs(self) -> Generator[Tuple[str, GlobalsDef], None, None]:
         return filter(
             lambda kv_tuple: isinstance(kv_tuple[1], GlobalsDef),
             self.all_exports.items(),
-        )
+        )  # type: ignore
 
     def def_attribute(self, key, value):
         # Some decorators, the only thing we do is convert them to PyOnlyDef.
@@ -209,11 +207,11 @@ def def_export_proc(self, name, f) -> ExportProcDef:
         file_line_loc = None
         try:
             sourcefile = inspect.getsourcefile(f)
-            _, linenums = sourcelines = inspect.getsourcelines(f)
-            if sourcefile and linenums:
-                file_line_loc = [sourcefile, linenums[0]]
-        except TypeError:
-            pass
+            _, linenum = sourcelines = inspect.getsourcelines(f)
+        except OSError:
+            ...
+        else:
+            file_line_loc = (sourcefile or "<unnamed>", linenum)
 
         sig = inspect.signature(f)
         if len(sig.parameters) < 1:
@@ -267,19 +265,18 @@ def __init__(
         self.module_builder = module_builder
         # The shadow dict holds instance attributes. We stash them here and the
         # Program instance itself arbitrates access via getattr/setattr.
-        self.shadow_dict = dict()
+        self.shadow_dict: dict[str, Any] = dict()
         self.current_import_phase = ImportPhase.TORCH_IR
 
 
 ################################################################################
 # Live reference accounting
 ################################################################################
 
-_all_compiled_module_class_infos: Dict[
+_all_compiled_module_class_infos: weakref.WeakKeyDictionary[
     "CompiledModuleMeta", CompiledModuleClassInfo
 ] = weakref.WeakKeyDictionary()
-
-_all_compiled_module_instance_infos: Dict[
+_all_compiled_module_instance_infos: weakref.WeakKeyDictionary[
     "CompiledModule", CompiledModuleInstanceInfo
 ] = weakref.WeakKeyDictionary()
 
@@ -292,7 +289,7 @@ def __init__(
 _metaclass_setup_complete = False
 
 
-@property
+@property  # type: ignore
 def _blackhole_instance_attribute(self):
     # We're not here.
     raise AttributeError

core/shark_turbine/aot/exporter.py

@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from typing import Optional, Sequence, Union
+from typing import Any, Optional, Sequence, Union
 import functools
 import io
 from pathlib import Path
@@ -116,8 +116,8 @@ def compile(
             save_to = Path(save_to)
             output = Output.open_file(str(save_to))
         else:
-            assert isinstance(output, Output)
             output = save_to
+            assert isinstance(output, Output)
 
         target_backends = (
             target_backends
@@ -153,7 +153,7 @@ def compile(
 # Decorator which explicitly exports a function.
 # TODO: Make this a public API on CompiledModule.
 # See https://github.com/nod-ai/SHARK-Turbine/issues/126
-def export_proc(f=None, *, signature: Sequence[AbstractTypedef]) -> ExportProcDef:
+def export_proc(f=None, *, signature: Sequence[AbstractTypedef]) -> Any:
     if f is None:
         return functools.partial(export_proc, signature=signature)
     return ExportProcDef(f.__name__, f, signature=signature)
@@ -177,19 +177,22 @@ def export(mdl: ModuleLike, *example_args: torch.Tensor) -> ExportOutput:
       An ExportOutput object that wraps the compilation and provides
       easy access.
     """
+    TransformedModule: Any
     if isinstance(mdl, torch.nn.Module):
+        nn_module = mdl
         signature = [abstractify(t) for t in example_args]
 
-        class Exported(CompiledModule, export_name=mdl._get_name()):
-            params = export_parameters(mdl)
+        class Exported(CompiledModule, export_name=nn_module._get_name()):
+            params = export_parameters(nn_module)
 
             @export_proc(signature=signature)
             def main(self, *args):
-                return jittable(mdl.forward)(*args)
+                return jittable(nn_module.forward)(*args)
 
+        TransformedModule = Exported
     else:
         assert isinstance(mdl, CompiledModuleMeta)
-        Exported = mdl
+        TransformedModule = mdl
 
     session = Session()
     # There are some bugs with respect to Session/context interop that we
@@ -201,5 +204,5 @@ def main(self, *args):
     else:
         context = Context()
 
-    cm = Exported(context=context, import_to="import")
+    cm = TransformedModule(context=context, import_to="import")
     return ExportOutput(session, cm, importer_uses_session=importer_uses_session)

core/shark_turbine/aot/passes/functorch.py

@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from typing import Callable
+from typing import Any, Callable
 
 import torch
 from torch.fx import (
@@ -40,8 +40,8 @@
 # since that does not result in load bearing information loss. Note that
 # ONNX applies this post export, which suffers from the loss of output
 # destructuring rewrites that torch.export does.
-def functorch_functionalize(gm: GraphModule, *args) -> GraphModule:
-    functionalized_callable = _functionalize_callabale(gm)
+def functorch_functionalize(gm_callable: Any, *args) -> GraphModule:
+    functionalized_callable = _functionalize_callabale(gm_callable)
     # TODO: There is more of a dance needed if the user has entered with a fake_mode.
     with proxy_tensor.maybe_disable_fake_tensor_mode():
         new_gm = proxy_tensor.make_fx(

core/shark_turbine/aot/support/ir_utils.py

@@ -15,10 +15,9 @@
 
 from ...importers.fx_importer import (
     ContextCache,
-)
-
-from ...importers.utils import (
-    RefTracker as FxRefTracker,
+    Empty,
+    EmptyType,
+    RefTracker,
 )
 
 from ...dynamo.type_conversion import (
@@ -58,10 +57,7 @@
     TORCH_DTYPE_TO_IREE_TYPE,
 )
 
-from .utils import (
-    RefTracker,
-    logger,
-)
+from ...support.logging import aot_logger as logger
 
 ###############################################################################
 # Configuration
@@ -150,7 +146,7 @@ def __init__(self, module_op: Operation):
         # Usually the FxImporter makes a new ref tracker for each invocation,
         # but we want to preserve it across individual JIT evaluations so
         # as to better intern tensors to attributes.
-        self.fx_py_attr_tracker = FxRefTracker()
+        self.fx_py_attr_tracker = RefTracker()
         self.native_type_converter = NativeTypeConverter(self.context)
 
     def handle_mlir_error(self, op: Operation, e: MLIRError, message: str):

core/shark_turbine/aot/support/procedural/base.py

@@ -14,8 +14,10 @@
 )
 
 from contextlib import contextmanager
+import threading
 
 import torch
+from torch.utils._pytree import tree_map
 
 from ....support.ir_imports import (
     F32Type,
@@ -34,12 +36,8 @@
     ModuleBuilder,
 )
 
-from ..utils import (
-    thread_state,
-    tree_map,
-)
-
 ShapedTypeDynamicSizeSentinel = ShapedType.get_dynamic_size()
+_thread_state = threading.local()
 
 ###############################################################################
 # Tracing intrinsics
@@ -71,9 +69,9 @@ def handle_assignment(self, scope, target, updated_value):
 
 def _trace_scopes() -> List[IrTrace]:
     try:
-        trace_scopes = thread_state.trace_scopes
+        trace_scopes = _thread_state.trace_scopes
     except AttributeError:
-        trace_scopes = thread_state.trace_scopes = []
+        trace_scopes = _thread_state.trace_scopes = []
     return trace_scopes