[Wave] Migrate iree.turbine.aot -> wave.aot (#45)

Also run ruff on the code to make sure it
passes pre-commit.

Signed-off-by: Harsh Menon <[email protected]>

Author: harsh-nod

docs/core/aot.rst

@@ -1,38 +1,38 @@
-`iree.turbine.aot`
+`wave.aot`
 =====================
 
 aot
 --------------
 
-.. automodule:: iree.turbine.aot
+.. automodule:: wave.aot
   :imported-members:
   :members:
   :undoc-members:
 
 passes
 --------------
 
-.. automodule:: iree.turbine.aot.passes
+.. automodule:: wave.aot.passes
   :imported-members:
   :members:
   :undoc-members:
 
 support
 --------------
 
-.. automodule:: iree.turbine.aot.support.procedural
+.. automodule:: wave.aot.support.procedural
   :imported-members:
   :members:
   :undoc-members:
 
-.. automodule:: iree.turbine.aot.support.procedural.exported_program
+.. automodule:: wave.aot.support.procedural.exported_program
   :members:
   :undoc-members:
 
 build actions
 --------------
 
-.. py:module:: iree.turbine.aot.build_actions
+.. py:module:: wave.aot.build_actions
 
 .. autofunction:: turbine_generate
 .. autoclass:: RemoteGenerator

iree/turbine/kernel/wave/codegen/emitter.py

@@ -14,7 +14,7 @@
 import torch.fx as fx
 
 from iree.turbine.kernel.lang.global_symbols import *
-from iree.turbine.aot.support.ir_utils import (
+from wave.aot.support.ir_utils import (
     _is_float_type,
     _is_integer_like_type,
 )

iree/turbine/kernel/wave/codegen/handlers.py

@@ -37,7 +37,7 @@
     scf_d,
     vector_d,
 )
-from iree.turbine.aot.support.ir_utils import (
+from wave.aot.support.ir_utils import (
     _is_float_type,
     _is_integer_like_type,
     _is_signed_or_signless_type,

lit_tests/kernel/wave/sharktank_integration.py

@@ -1,13 +1,13 @@
 # RUN: python %s | FileCheck %s
 
 import textwrap
+import wave.aot as aot
 from typing import Optional
 from wave.transforms.merger import Merger
 
 import torch
 from jinja2 import BaseLoader, Environment
 
-import iree.turbine.aot as aot
 import iree.turbine.kernel.wave as tkw
 from iree.compiler.ir import (
     Context,

iree/turbine/aot/__init__.py renamed to wave/aot/__init__.py

@@ -1,6 +1,4 @@
-"""
-Toolkit for ahead-of-time (AOT) compilation and export of PyTorch programs.
-"""
+"""Toolkit for ahead-of-time (AOT) compilation and export of PyTorch programs."""
 
 # Copyright 2023 Nod Labs, Inc
 #
@@ -13,5 +11,10 @@
 from .decompositions import *
 from .exporter import *
 from .fx_programs import FxPrograms, FxProgramsBuilder
-from .tensor_traits import *
 from .params import *
+from .tensor_traits import *
+
+__all__ = [
+    "FxPrograms",
+    "FxProgramsBuilder",
+]

iree/turbine/aot/builtins/__init__.py renamed to wave/aot/builtins/__init__.py

@@ -4,8 +4,6 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from .globals import *
-from .jittable import jittable
 from ..support.procedural import (
     AbstractBool,
     AbstractF32,
@@ -19,23 +17,25 @@
 
 # Export the instantiated IREEEmitter as "IREE"
 from ..support.procedural.iree_emitter import IREEEmitter as _IREEEmitter
+from .globals import *
+from .jittable import jittable
 
 IREE = _IREEEmitter()
 del _IREEEmitter
 
 __all__ = [
+    "IREE",
     "AbstractBool",
     "AbstractF32",
     "AbstractF64",
     "AbstractI32",
     "AbstractI64",
     "AbstractIndex",
     "AbstractTensor",
-    "IREE",
     "abstractify",
+    "export_buffers",
     "export_global",
     "export_global_tree",
     "export_parameters",
-    "export_buffers",
     "jittable",
 ]

iree/turbine/aot/builtins/globals.py renamed to wave/aot/builtins/globals.py

@@ -5,42 +5,39 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from typing import Any, Callable, Optional
+from typing import Any, Optional
 
-import torch.nn as nn
+from torch import nn
+from torch.utils._pytree import (
+    TreeSpec,
+    tree_flatten,
+    tree_map,
+)
 
+from ..support.ir_utils import (
+    GlobalAttributes,
+    NameMapCallback,
+)
 from ..support.procedural import (
-    AbstractTypedef,
     Abstractifiable,
+    AbstractTypedef,
     GlobalsDef,
     TreeAbstractifiable,
     abstractify_single_value,
 )
 
-from ..support.ir_utils import (
-    NameMapCallback,
-    GlobalAttributes,
-)
-
-from torch.utils._pytree import (
-    TreeSpec,
-    tree_flatten,
-    tree_map,
-)
-
-
 __all__ = [
+    "export_buffers",
     "export_global",
     "export_global_tree",
     "export_parameters",
-    "export_buffers",
 ]
 
 
 class export_global(GlobalsDef, Abstractifiable):
     """Exports a single global into a CompiledModule."""
 
-    __slots__ = ["_name", "_value", "_schema"]
+    __slots__ = ["_name", "_schema", "_value"]
 
     def __init__(
         self,
@@ -104,7 +101,7 @@ def __init__(
         self._items, self._schema = tree_flatten(tree)
         self._names, _ = tree_flatten(_transform_tree_to_names("", tree))
         assert len(self._items) == len(
-            self._names
+            self._names,
         ), f"Name and value tree are different sizes: {len(self._items)} != {len(self._names)}"
 
     def items(self):
@@ -235,10 +232,9 @@ def _transform_tree_to_names(prefix: str, tree):
         return tree.__class__(
             (k, _transform_tree_to_names(join(k), v)) for k, v in tree.items()
         )
-    elif isinstance(tree, (list, tuple)):
+    if isinstance(tree, (list, tuple)):
         return tree.__class__(
             _transform_tree_to_names(join(str(index)), v)
             for index, v in enumerate(tree)
         )
-    else:
-        return prefix
+    return prefix

iree/turbine/aot/builtins/jittable.py renamed to wave/aot/builtins/jittable.py

@@ -7,13 +7,12 @@
 
 """Tracing builtins."""
 
-from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
-
-import warnings
+from collections.abc import Sequence
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import torch
-from torch._decomp import get_decompositions
 import torch._dynamo as dynamo
+from torch._decomp import get_decompositions
 from torch.fx import (
     GraphModule,
 )
@@ -23,13 +22,12 @@
 )
 
 from iree.compiler.extras.fx_importer import (
-    GraphNodeImporter,
     FxImporter,
     FxImporterHooks,
+    GraphNodeImporter,
     InputInfo,
 )
-
-from ...support.ir_imports import (
+from iree.turbine.support.ir_imports import (
     FlatSymbolRefAttr,
     FunctionType,
     Operation,
@@ -40,18 +38,15 @@
     func_d,
     util_d,
 )
-
-from ...support.logging import aot_logger as logger
+from iree.turbine.support.logging import aot_logger as logger
 
 from ..decompositions import current_aot_decompositions
 from ..passes import (
     functorch_functionalize,
 )
-
 from ..support.ir_utils import (
     ModuleBuilder,
 )
-
 from ..support.procedural import (
     CallableIntrinsic,
     IrImmediateTensor,
@@ -79,7 +74,10 @@ def __init__(self, module_builder: ModuleBuilder):
         self.cloned_global_symbols: set[str] = set()
 
     def resolve_literal(
-        self, gni: GraphNodeImporter, literal: Any, info: Optional[InputInfo] = None
+        self,
+        gni: GraphNodeImporter,
+        literal: Any,
+        info: Optional[InputInfo] = None,
     ) -> Optional[Value]:
         module_builder = self.module_builder
         cloned_global_symbols = self.cloned_global_symbols
@@ -110,7 +108,8 @@ def resolve_literal(
         # Emit a global load and conversion.
         vtensor_type = gni._cc.tensor_to_vtensor_type(literal)
         loaded_value = util_d.GlobalLoadOp(
-            materialized_global.ir_type, materialized_global.symbol_name
+            materialized_global.ir_type,
+            materialized_global.symbol_name,
         ).result
         converted_value = Operation.create(
             "torch_c.from_builtin_tensor",
@@ -131,11 +130,11 @@ class jittable(CallableIntrinsic):
     """
 
     __slots__ = [
-        "dynamic_shapes",
+        "_passes",
         "decomposition_table",
-        "wrapped_f",
+        "dynamic_shapes",
         "function_name",
-        "_passes",
+        "wrapped_f",
     ]
 
     def __init__(
@@ -266,10 +265,10 @@ def flat_wrapped_f(*args):
         # TODO: Debug upstream why iteration over children isn't creating a typed view.
         # This should just be `target_op.function_type`
         target_ftype = FunctionType(
-            TypeAttr(target_op.attributes["function_type"]).value
+            TypeAttr(target_op.attributes["function_type"]).value,
         )
         target_symbol_ref = FlatSymbolRefAttr.get(
-            StringAttr(target_op.attributes["sym_name"]).value
+            StringAttr(target_op.attributes["sym_name"]).value,
         )
 
         assert len(flat_ir_args) == len(target_ftype.inputs), (
@@ -287,7 +286,9 @@ def flat_wrapped_f(*args):
 
         with proc_trace.ip, proc_trace.loc:
             flat_ir_results = func_d.CallOp(
-                target_ftype.results, target_symbol_ref, flat_ir_args
+                target_ftype.results,
+                target_symbol_ref,
+                flat_ir_args,
             ).results
 
         assert len(flat_ir_results) == len(result_tensor_infos)
@@ -300,7 +301,7 @@ def flat_wrapped_f(*args):
                 flat_py_results.append(IrImmediateTensor(native_ir_result, dtype))
             else:
                 raise TypeError(
-                    f"Unknown PyTorch->IREE value mapping for jittable result: {result_tensor_info}->{native_ir_result}"
+                    f"Unknown PyTorch->IREE value mapping for jittable result: {result_tensor_info}->{native_ir_result}",
                 )
 
         tree_py_results = tree_unflatten(flat_py_results, out_spec)
@@ -317,14 +318,14 @@ def _split_py_arg(self, arg) -> Tuple[Value, Any]:
 class _Merger:
     __slots__ = [
         "context",
-        "to_module_builder",
         "from_module_op",
         "from_symbol_table",
         "import_function_name",
-        "rename_map",
         "nested_symbol_ops",
         "nested_symbol_table_ops",
         "private_attr",
+        "rename_map",
+        "to_module_builder",
     ]
 
     def __init__(
@@ -440,7 +441,7 @@ def _extract_graph_output_metadata(
         out_spec = gm._out_spec
     except AttributeError:
         raise AssertionError(
-            "Expected PyTorch to add an _out_spec attribute to the GraphModule"
+            "Expected PyTorch to add an _out_spec attribute to the GraphModule",
         )
 
     output_nodes = []