From fde390c7669e29362b18388448ef2b188713383f Mon Sep 17 00:00:00 2001 From: nithinsubbiah Date: Tue, 2 Aug 2022 20:35:56 +0000 Subject: [PATCH] Re-enable custom op support --- build_tools/update_shape_lib.sh | 20 ++++++++---------- build_tools/update_torch_ods.sh | 21 +++++++++---------- lib/Conversion/TorchToLinalg/CMakeLists.txt | 2 -- .../TorchToLinalg/PopulatePatterns.h | 3 --- .../TorchToLinalg/TorchToLinalg.cpp | 2 -- python/CMakeLists.txt | 6 +----- .../CMakeLists.txt | 2 +- .../jit_ir/build_tools/shape_lib_gen.py | 4 ---- .../jit_ir/build_tools/torch_ods_gen.py | 10 --------- .../test_suite/__init__.py | 2 -- 10 files changed, 21 insertions(+), 51 deletions(-) diff --git a/build_tools/update_shape_lib.sh b/build_tools/update_shape_lib.sh index e0de9f6cb..16526569c 100755 --- a/build_tools/update_shape_lib.sh +++ b/build_tools/update_shape_lib.sh @@ -16,19 +16,17 @@ build_dir="$(realpath "${TORCH_MLIR_BUILD_DIR:-$src_dir/build}")" torch_transforms_cpp_dir="${src_dir}/lib/Dialect/Torch/Transforms" python_packages_dir="${build_dir}/tools/torch-mlir/python_packages" +TORCH_MLIR_EXT_PYTHONPATH="${TORCH_MLIR_EXT_PYTHONPATH:-""}" pypath="${python_packages_dir}/torch_mlir" -# TODO: Re-enable once custom op support is back. -#if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then -# pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}" -#fi -#ext_module="torch_mlir._torch_mlir_custom_op_example" -#if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then -# ext_module="${ext_module},${TORCH_MLIR_EXT_MODULES} " -#fi +if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then + pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}" +fi +TORCH_MLIR_EXT_MODULES="${TORCH_MLIR_EXT_MODULES:-""}" +if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then + ext_module="${TORCH_MLIR_EXT_MODULES} " +fi PYTHONPATH="${pypath}" python \ -m torch_mlir.dialects.torch.importer.jit_ir.build_tools.shape_lib_gen \ + --pytorch_op_extensions=${ext_module} \ --torch_transforms_cpp_dir="${torch_transforms_cpp_dir}" - -# TODO: Add back to shape_lib_gen invocation once custom op support is back. -# --pytorch_op_extensions=${ext_module} \ diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh index 2e22ab12e..2b30ffe66 100755 --- a/build_tools/update_torch_ods.sh +++ b/build_tools/update_torch_ods.sh @@ -16,20 +16,19 @@ build_dir="$(realpath "${TORCH_MLIR_BUILD_DIR:-$src_dir/build}")" torch_ir_include_dir="${src_dir}/include/torch-mlir/Dialect/Torch/IR" python_packages_dir="${build_dir}/tools/torch-mlir/python_packages" +TORCH_MLIR_EXT_PYTHONPATH="${TORCH_MLIR_EXT_PYTHONPATH:-""}" pypath="${python_packages_dir}/torch_mlir" -# TODO: Re-enable once custom op support is back. -#if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then -# pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}" -#fi -#ext_module="torch_mlir._torch_mlir_custom_op_example" -#if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then -# ext_module="${ext_module},${TORCH_MLIR_EXT_MODULES}" -#fi +if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then + pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}" +fi +TORCH_MLIR_EXT_MODULES="${TORCH_MLIR_EXT_MODULES:-""}" +ext_module="${ext_module:-""}" +if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then + ext_module="${TORCH_MLIR_EXT_MODULES}" +fi PYTHONPATH="${pypath}" python \ -m torch_mlir.dialects.torch.importer.jit_ir.build_tools.torch_ods_gen \ --torch_ir_include_dir="${torch_ir_include_dir}" \ + --pytorch_op_extensions="${ext_module}" \ --debug_registry_dump="${torch_ir_include_dir}/JITOperatorRegistryDump.txt" - -# TODO: Add back to torch_ods_gen invocation once custom op support is back. -# --pytorch_op_extensions="${ext_module}" \ diff --git a/lib/Conversion/TorchToLinalg/CMakeLists.txt b/lib/Conversion/TorchToLinalg/CMakeLists.txt index 7a86e43a6..ece929597 100644 --- a/lib/Conversion/TorchToLinalg/CMakeLists.txt +++ b/lib/Conversion/TorchToLinalg/CMakeLists.txt @@ -1,6 +1,4 @@ add_mlir_conversion_library(TorchMLIRTorchToLinalg -# TODO: Re-enable after MacOS support is fixed for the custom op extension. -# CustomOpExample.cpp DataMovement.cpp IndirectDataMovement.cpp Linear.cpp diff --git a/lib/Conversion/TorchToLinalg/PopulatePatterns.h b/lib/Conversion/TorchToLinalg/PopulatePatterns.h index 384c89d33..56691c82c 100644 --- a/lib/Conversion/TorchToLinalg/PopulatePatterns.h +++ b/lib/Conversion/TorchToLinalg/PopulatePatterns.h @@ -63,9 +63,6 @@ void populateIndirectDataMovementPatternsAndLegality( void populateTensorConstructorsPatternsAndLegality(TypeConverter &typeConverter, RewritePatternSet &patterns, ConversionTarget &target); -//void populateCustomOpExamplePatternsAndLegality(TypeConverter &typeConverter, -// RewritePatternSet &patterns, -// ConversionTarget &target); } // namespace torch_to_linalg } // namespace torch diff --git a/lib/Conversion/TorchToLinalg/TorchToLinalg.cpp b/lib/Conversion/TorchToLinalg/TorchToLinalg.cpp index 1caa1408f..f8ebc349f 100644 --- a/lib/Conversion/TorchToLinalg/TorchToLinalg.cpp +++ b/lib/Conversion/TorchToLinalg/TorchToLinalg.cpp @@ -62,8 +62,6 @@ public: RewritePatternSet patterns(context); - //torch_to_linalg::populateCustomOpExamplePatternsAndLegality( - // typeConverter, patterns, target); torch_to_linalg::populateTensorScalarInteropPatternsAndLegality( typeConverter, patterns, target); torch_to_linalg::populateLinearPatternsAndLegality(typeConverter, patterns, diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 6b389d042..2c584ef4a 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -111,8 +111,7 @@ add_subdirectory(torch_mlir/eager_mode) # Required for running the update_torch_ods.sh and update_shape_lib.sh scripts. ################################################################################ -# TODO: renable once it build on macOS Intel / M1 -#add_subdirectory(torch_mlir/_torch_mlir_custom_op_example) +# add_subdirectory(torch_mlir/_torch_mlir_custom_op_example) ################################################################################ # Generate packages and shared library @@ -159,9 +158,6 @@ if(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER) add_dependencies(TorchMLIRPythonModules TorchMLIRE2ETestPythonModules) endif() -# TODO: Add after macOS builds are fixed -#add_dependencies(TorchMLIRPythonModules torch_mlir_custom_op_example) - if(TORCH_MLIR_ENABLE_LTC) # Add Torch-MLIR LTC backend as dependency add_dependencies(TorchMLIRPythonModules torch_mlir_ltc_backend) diff --git a/python/torch_mlir/_torch_mlir_custom_op_example/CMakeLists.txt b/python/torch_mlir/_torch_mlir_custom_op_example/CMakeLists.txt index fdd5997f4..a5a011b55 100644 --- a/python/torch_mlir/_torch_mlir_custom_op_example/CMakeLists.txt +++ b/python/torch_mlir/_torch_mlir_custom_op_example/CMakeLists.txt @@ -1,5 +1,5 @@ # Setup PyTorch -list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/../dialects/torch/importer/jit_ir/cmake/modules") +list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/../cmake/modules") include(TorchMLIRPyTorch) TorchMLIRProbeForPyTorchInstall() find_package(Torch 1.8 REQUIRED) diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py index 8f756cec2..d990404b2 100644 --- a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py +++ b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py @@ -1173,10 +1173,6 @@ def aten〇linalg_vector_norm(self: List[int], ord: float = 2, dim: Optional[Lis dim = list(range(len(self))) return upstream_shape_functions.mean_dim(self, dim, keepdim, dtype) -# TODO: Re-enable after MacOS support is fixed for the extension. -#def _torch_mlir_custom_op_example〇identity(t: List[int]) -> List[int]: -# return upstream_shape_functions.unary(t) - # ============================================================================== # Shape library generator main(). # ============================================================================== diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py index 9c84728f3..a2151fdec 100644 --- a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py +++ b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py @@ -633,16 +633,6 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry): "quantized::linear : (Tensor, __torch__.torch.classes.quantized.LinearPackedParamsBase, float, int) -> (Tensor)", traits=["HasValueSemantics"]) - # ========================================================================== - # `_torch_mlir_custom_op_example::` namespace. - # - # This is a demonstration of supporting an operation defined in a PyTorch - # extension. - # ========================================================================== - - # TODO: Re-enable after MacOS support is fixed for the extension. - #emit("_torch_mlir_custom_op_example::identity : (Tensor) -> (Tensor)") - def dump_registered_ops(outfile: TextIO, registry: Registry): for _, v in sorted(registry.by_unique_key.items()): diff --git a/python/torch_mlir_e2e_test/test_suite/__init__.py b/python/torch_mlir_e2e_test/test_suite/__init__.py index b7454e910..63d30d8fd 100644 --- a/python/torch_mlir_e2e_test/test_suite/__init__.py +++ b/python/torch_mlir_e2e_test/test_suite/__init__.py @@ -51,5 +51,3 @@ def register_all_tests(): from . import return_types from . import control_flow from . import stats - # TODO: Re-enable after MacOS support is fixed for the extension. - #from . import custom_op_example