Refactor FxImporter to new location.

Author: stellaraccident

python/shark_turbine/aot/builtins/jittable.py

@@ -22,15 +22,15 @@
 )
 from torch.fx.passes.shape_prop import TensorMetadata
 
-from ...dynamo.importer import (
-    GraphNodeImporter,
-    FxImporter,
-)
-
 from ...dynamo.passes import (
     DEFAULT_DECOMPOSITIONS,
 )
 
+from ...importers.fx_importer import (
+    GraphNodeImporter,
+    FxImporter,
+)
+
 from ..passes import (
     functorch_functionalize,
 )

python/shark_turbine/aot/support/ir_utils.py

@@ -13,7 +13,7 @@
 import numpy as np
 import torch
 
-from ...dynamo.importer import (
+from ...importers.fx_importer import (
     ContextCache,
     TORCH_DTYPE_TO_MLIR_TYPE_ASM,
 )

python/shark_turbine/importers/README.md

@@ -0,0 +1,7 @@
+# Importers from various systems
+
+This directory is self-contained and intended to be shared with other
+projects with its source-of-truth in torch-mlir.
+
+All MLIR API dependencies must route through the relative `ir.py`, which
+it is expected that sub-projects will customize accordingly.

python/shark_turbine/dynamo/importer.py renamed to python/shark_turbine/importers/fx_importer.py

@@ -3,44 +3,13 @@
 # Licensed under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-import builtins
 import logging
 import operator
 import re
 from types import NoneType, BuiltinMethodType, BuiltinFunctionType
 from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 import numpy as np
 
-from iree.compiler.ir import (
-    Attribute as MlirAttribute,
-    Block,
-    Context,
-    FloatAttr,
-    BF16Type,
-    ComplexType,
-    F16Type,
-    F32Type,
-    F64Type,
-    FunctionType,
-    InsertionPoint,
-    IntegerAttr,
-    IntegerType,
-    RankedTensorType,
-    Location,
-    Module,
-    Operation,
-    StringAttr,
-    Type as MlirType,
-    Value,
-    DenseResourceElementsAttr,
-)
-
-import iree.compiler.dialects.func as func_dialect
-from iree.compiler.ir import SymbolTable
-
-# import iree.compiler.dialects.torch as torch_dialect
-
-
 import torch
 import torch.fx as torch_fx
 from torch.fx.passes.shape_prop import TensorMetadata
@@ -67,6 +36,36 @@
     Argument as NodeArgument,
 )
 
+from .ir import (
+    Attribute,
+    Block,
+    Context,
+    DenseResourceElementsAttr,
+    FloatAttr,
+    BF16Type,
+    ComplexType,
+    F16Type,
+    F32Type,
+    F64Type,
+    FunctionType,
+    InsertionPoint,
+    IntegerAttr,
+    IntegerType,
+    RankedTensorType,
+    Location,
+    Module,
+    Operation,
+    StringAttr,
+    SymbolTable,
+    IrType,
+    Value,
+    func_dialect,
+)
+
+from .utils import (
+    TypeSubclassMap,
+)
+
 __all__ = [
     "FxImporter",
 ]
@@ -100,7 +99,7 @@
     torch.complex128: "complex<f64>",
 }
 
-TORCH_DTYPE_TO_MLIR_TYPE: Dict[torch.dtype, Callable[[], MlirType]] = {
+TORCH_DTYPE_TO_MLIR_TYPE: Dict[torch.dtype, Callable[[], IrType]] = {
     torch.float16: lambda: F16Type.get(),
     torch.bfloat16: lambda: BF16Type.get(),
     torch.float32: lambda: F32Type.get(),
@@ -313,7 +312,7 @@ def _graph_to_function_meta(self, g: Graph) -> Tuple[FunctionType, Location]:
                 for result_node in node.args[0]:
                     if result_node is None:
                         result_types.append(
-                            MlirType.parse("!torch.none", context=self._c)
+                            IrType.parse("!torch.none", context=self._c)
                         )
                     else:
                         result_types.append(self._cc.node_val_to_type(result_node))
@@ -341,19 +340,19 @@ class ContextCache:
 
     def __init__(self, context: Context):
         self._c = context
-        self._dtype_to_type: Dict[TorchDtype, MlirType] = {}
-        self._tensor_metadata_cache: Dict[Tuple[torch.Size, torch.dtype], MlirType] = {}
+        self._dtype_to_type: Dict[TorchDtype, IrType] = {}
+        self._tensor_metadata_cache: Dict[Tuple[torch.Size, torch.dtype], IrType] = {}
 
         # Common types.
         with context:
-            self.torch_bool_type = MlirType.parse("!torch.bool")
-            self.torch_float_type = MlirType.parse("!torch.float")
-            self.torch_int_type = MlirType.parse("!torch.int")
-            self.torch_none_type = MlirType.parse("!torch.none")
-            self.torch_str_type = MlirType.parse("!torch.str")
-            self.torch_device_type = MlirType.parse("!torch.Device")
-
-    def integer_attr(self, value: int, bits: int) -> MlirAttribute:
+            self.torch_bool_type = IrType.parse("!torch.bool")
+            self.torch_float_type = IrType.parse("!torch.float")
+            self.torch_int_type = IrType.parse("!torch.int")
+            self.torch_none_type = IrType.parse("!torch.none")
+            self.torch_str_type = IrType.parse("!torch.str")
+            self.torch_device_type = IrType.parse("!torch.Device")
+
+    def integer_attr(self, value: int, bits: int) -> Attribute:
         c = self._c
         return IntegerAttr.get(IntegerType.get_signless(bits, c), value)
 
@@ -362,16 +361,16 @@ def integer_attr(self, value: int, bits: int) -> MlirAttribute:
     def format_asm_shape(self, shape: torch.Size) -> str:
         return ",".join("?" if is_symbolic(d) else str(d) for d in list(shape))
 
-    """Return MlirType for !torch.vtensor with the given shape and dtype"""
+    """Return IrType for !torch.vtensor with the given shape and dtype"""
 
     def get_vtensor_type(self, shape: torch.Size, dtype: torch.dtype):
         shape_asm = self.format_asm_shape(shape)
         mlir_dtype = str(self.dtype_to_type(dtype))
-        return MlirType.parse(
+        return IrType.parse(
             f"!torch.vtensor<[{shape_asm}],{str(mlir_dtype)}>", context=self._c
         )
 
-    def node_val_to_type(self, node: torch_fx.Node) -> MlirType:
+    def node_val_to_type(self, node: torch_fx.Node) -> IrType:
         try:
             tensor_meta = node.meta.get("tensor_meta")
             val = node.meta.get("val")
@@ -393,7 +392,7 @@ def node_val_to_type(self, node: torch_fx.Node) -> MlirType:
 
                 t = SCALAR_TYPE_TO_TORCH_MLIR_TYPE.get(type(val))
                 if t is not None:
-                    return MlirType.parse(t, self._c)
+                    return IrType.parse(t, self._c)
 
             raise NotImplementedError(
                 f"FIXME: Unsupported placeholder node (this often indicates that a necessary) "
@@ -404,7 +403,7 @@ def node_val_to_type(self, node: torch_fx.Node) -> MlirType:
                 f"FIXME: Illegal access to torch.fx.Node.meta: {e} ({node.meta.keys()} : {node.meta})"
             )
 
-    def tensor_metadata_to_type(self, tm: TensorMetadata) -> MlirType:
+    def tensor_metadata_to_type(self, tm: TensorMetadata) -> IrType:
         tm_shape = tuple(
             item.node if is_symbolic(item) else item for item in list(tm.shape)
         )
@@ -416,20 +415,20 @@ def tensor_metadata_to_type(self, tm: TensorMetadata) -> MlirType:
             self._tensor_metadata_cache[key] = t
         return t
 
-    def dtype_to_type(self, dtype: TorchDtype) -> MlirType:
+    def dtype_to_type(self, dtype: TorchDtype) -> IrType:
         t = self._dtype_to_type.get(dtype)
         if t is None:
             try:
                 asm = TORCH_DTYPE_TO_MLIR_TYPE_ASM[dtype]
             except IndexError:
                 raise ValueError(f"Unknown conversion from {dtype} to IREE type")
-            t = MlirType.parse(asm, self._c)
+            t = IrType.parse(asm, self._c)
             self._dtype_to_type[dtype] = t
         return t
 
-    def tensor_to_vtensor_type(self, tensor: torch.Tensor) -> MlirType:
+    def tensor_to_vtensor_type(self, tensor: torch.Tensor) -> IrType:
         dtype_asm = str(self.dtype_to_type(tensor.dtype))
-        return MlirType.parse(f"!torch.vtensor<{list(tensor.size())},{dtype_asm}>")
+        return IrType.parse(f"!torch.vtensor<{list(tensor.size())},{dtype_asm}>")
 
     def get_node_location(self, node: torch_fx.Node) -> Optional[Location]:
         stack_trace = node.meta.get("stack_trace")
@@ -844,7 +843,7 @@ def _import_list_argument(
         else:
             list_type = PY_TYPE_TO_TORCH_LIST_TYPE[element_type]
 
-        result_type = MlirType.parse(list_type, context=self._c)
+        result_type = IrType.parse(list_type, context=self._c)
         operation = Operation.create(
             "torch.prim.ListConstruct",
             results=[result_type],
@@ -869,44 +868,8 @@ def _import_default_value(self, loc: Location, arg, expected_jit_type) -> Value:
             return cvt(arg, self, self._cc)
 
 
-class TypeSubclassMap:
-    """Mapping of super-types to values.
-
-    Maintains a cache of actual types seen and uses that instead of a linear
-    scan.
-    """
-
-    __slots__ = [
-        "_cache",
-        "_mapping",
-    ]
-
-    def __init__(self):
-        # The linear list of converters.
-        self._mapping: List[Tuple[type, Any]] = []
-        # When there is a hit on the linear mapping, memoize it here.
-        self._cache: Dict[type, Any] = {}
-
-    def map(self, t: type, value: Any):
-        self._mapping.append((t, value))
-        self._cache[t] = value
-
-    def lookup(self, t: type) -> Any:
-        try:
-            return self._cache[t]
-        except KeyError:
-            pass
-        for t_super, value in self._mapping:
-            if issubclass(t, t_super):
-                self._cache[t] = value
-                return value
-        else:
-            self._cache[t] = None
-            return None
-
-
 def _make_constant_op(
-    op_name: str, value_attr: MlirAttribute, result_type: Optional[MlirType] = None
+    op_name: str, value_attr: Attribute, result_type: Optional[IrType] = None
 ) -> Operation:
     return Operation.create(
         op_name,
@@ -915,7 +878,7 @@ def _make_constant_op(
     )
 
 
-def create_mlir_tensor_type(tensor: torch.Tensor) -> MlirType:
+def create_mlir_tensor_type(tensor: torch.Tensor) -> IrType:
     try:
         dtype = tensor.dtype
         element_type = TORCH_DTYPE_TO_MLIR_TYPE[dtype]()
@@ -925,7 +888,7 @@ def create_mlir_tensor_type(tensor: torch.Tensor) -> MlirType:
         raise TypeError(f"Could not map Torch dtype {dtype} to an IREE type")
 
 
-def _make_vtensor_literal_op(tensor: torch.Tensor, vtensor_type: MlirType) -> Operation:
+def _make_vtensor_literal_op(tensor: torch.Tensor, vtensor_type: IrType) -> Operation:
     npy_dtype = TORCH_DTYPE_TO_NPY_TYPE.get(tensor.dtype)
     assert (
         npy_dtype is not None

python/shark_turbine/importers/ir.py

@@ -0,0 +1,35 @@
+# Copyright 2023 Nod Labs, Inc
+# Portions Copyright 2022 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from iree.compiler.ir import (
+    Attribute as Attribute,
+    Block,
+    Context,
+    DenseResourceElementsAttr,
+    FloatAttr,
+    BF16Type,
+    ComplexType,
+    F16Type,
+    F32Type,
+    F64Type,
+    FunctionType,
+    InsertionPoint,
+    IntegerAttr,
+    IntegerType,
+    RankedTensorType,
+    Location,
+    Module,
+    Operation,
+    StringAttr,
+    SymbolTable,
+    Type as IrType,
+    Value,
+)
+
+from iree.compiler.dialects import (
+    func as func_dialect,
+)

python/shark_turbine/importers/utils.py

@@ -0,0 +1,43 @@
+# Copyright 2023 Nod Labs, Inc
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from typing import Any, Dict, List, Tuple
+
+
+class TypeSubclassMap:
+    """Mapping of super-types to values.
+
+    Maintains a cache of actual types seen and uses that instead of a linear
+    scan.
+    """
+
+    __slots__ = [
+        "_cache",
+        "_mapping",
+    ]
+
+    def __init__(self):
+        # The linear list of converters.
+        self._mapping: List[Tuple[type, Any]] = []
+        # When there is a hit on the linear mapping, memoize it here.
+        self._cache: Dict[type, Any] = {}
+
+    def map(self, t: type, value: Any):
+        self._mapping.append((t, value))
+        self._cache[t] = value
+
+    def lookup(self, t: type) -> Any:
+        try:
+            return self._cache[t]
+        except KeyError:
+            pass
+        for t_super, value in self._mapping:
+            if issubclass(t, t_super):
+                self._cache[t] = value
+                return value
+        else:
+            self._cache[t] = None
+            return None

tests/dynamo/importer_dynamic_test.py

@@ -13,7 +13,7 @@
 from torch._export import dynamic_dim
 
 # from torch._export.constraints import constrain_as_size, constrain_as_value
-from shark_turbine.dynamo.importer import FxImporter
+from shark_turbine.importers.fx_importer import FxImporter
 from shark_turbine.dynamo.passes import turbine_cpu_pass_pipeline
 import torch
 import torch._dynamo as dynamo