Re-enable custom op support

Author: nithinsubbiah

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} \

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}" \

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

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

lib/Conversion/TorchToLinalg/TorchToLinalg.cpp

@@ -62,8 +62,6 @@ class ConvertTorchToLinalg
 
     RewritePatternSet patterns(context);
 
-    //torch_to_linalg::populateCustomOpExamplePatternsAndLegality(
-    //    typeConverter, patterns, target);
     torch_to_linalg::populateTensorScalarInteropPatternsAndLegality(
         typeConverter, patterns, target);
     torch_to_linalg::populateLinearPatternsAndLegality(typeConverter, patterns,

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)

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)

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().
 # ==============================================================================

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -633,16 +633,6 @@ def emit_with_mutating_variants(key, **kwargs):
         "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()):

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