Extract common code

llvm · Apr 24, 2024 · b4fd52d · b4fd52d
1 parent 4516c1f
commit b4fd52d
Show file tree

Hide file tree

Showing 5 changed files with 4 additions and 4 deletions.
diff --git a/projects/pt1/python/CMakeLists.txt b/projects/pt1/python/CMakeLists.txt
@@ -20,7 +20,6 @@ declare_mlir_python_sources(TorchMLIRPythonSources.TopLevel
   SOURCES
     torchscript.py
     _dynamo_fx_importer.py
-    compiler_utils.py
     dynamo.py
     _version.py
 )

diff --git a/projects/pt1/python/torch_mlir/torchscript.py b/projects/pt1/python/torch_mlir/torchscript.py
@@ -16,7 +16,7 @@
 from torch_mlir.dynamo import _get_decomposition_table
 from torch.fx.experimental.proxy_tensor import make_fx
 
-from .compiler_utils import (
+from torch_mlir.compiler_utils import (
     OutputType,
     run_pipeline_with_repro_report,
     lower_mlir_module,

diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
@@ -44,6 +44,7 @@ declare_mlir_python_sources(TorchMLIRPythonSources.PublicAPI
   ADD_TO_PARENT TorchMLIRPythonSources
   SOURCES
     fx.py
+    compiler_utils.py
     extras/fx_decomp_util.py
 )
 

diff --git a/...s/pt1/python/torch_mlir/compiler_utils.py → python/torch_mlir/compiler_utils.py b/...s/pt1/python/torch_mlir/compiler_utils.py → python/torch_mlir/compiler_utils.py
diff --git a/python/torch_mlir/fx.py b/python/torch_mlir/fx.py
@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 # Also available under a BSD-style license. See LICENSE.
 
-from typing import Optional, Union, Dict, Tuple, Any
+from typing import Optional, Union, Dict, Tuple, Any, Callable
 
 import warnings
 
@@ -55,7 +55,7 @@ def export_and_import(
     dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
     experimental_support_mutation: bool = False,
     hooks: Optional[FxImporterHooks] = None,
-    decomposition_table: Optional[list] = None,
+    decomposition_table: Optional[Dict[torch._ops.OperatorBase, Callable]] = None,
     func_name: str = "main",
     verbose: bool = False,
     **kwargs,
-Original file line number
+Diff line change
@@ Expand Up / @@ -44,6 +44,7 @@ declare_mlir_python_sources(TorchMLIRPythonSources.PublicAPI @@
       ADD_TO_PARENT TorchMLIRPythonSources
       SOURCES
         fx.py
+        compiler_utils.py
         extras/fx_decomp_util.py
     )
@@ Expand Down @@