Temporarily revert support for custom op extensions. (llvm#944)

The MacOS builders are having linking trouble with the extension library. Until it's fixed, all support for op extensions is disabled. It should be easy to restore once the issue is resolved.
qedawkins · Jun 15, 2022 · b90837e · b90837e
1 parent 8fd0843
commit b90837e
Show file tree

Hide file tree

Showing 12 changed files with 37 additions and 108 deletions.
diff --git a/build_tools/update_shape_lib.sh b/build_tools/update_shape_lib.sh
@@ -17,15 +17,18 @@ torch_transforms_cpp_dir="${src_dir}/lib/Dialect/Torch/Transforms"
 python_packages_dir="${build_dir}/tools/torch-mlir/python_packages"
 
 pypath="${python_packages_dir}/torch_mlir"
-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
+# 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
 
 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
@@ -17,16 +17,19 @@ torch_ir_include_dir="${src_dir}/include/torch-mlir/Dialect/Torch/IR"
 python_packages_dir="${build_dir}/tools/torch-mlir/python_packages"
 
 pypath="${python_packages_dir}/torch_mlir"
-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
+# 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
 
 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/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
@@ -8316,26 +8316,3 @@ def Torch_QuantizedLinearOp : Torch_Op<"quantized.linear", [
   }];
 }
 
-def Torch_TorchMlirCustomOpExampleIdentityOp : Torch_Op<"_torch_mlir_custom_op_example.identity", [
-    AllowsTypeRefinement,
-    HasValueSemantics,
-    ReadOnly
-  ]> {
-  let summary = "Generated op for `_torch_mlir_custom_op_example::identity : (Tensor) -> (Tensor)`";
-  let arguments = (ins
-    AnyTorchTensorType:$t
-  );
-  let results = (outs
-    AnyTorchTensorType:$result
-  );
-  let hasCustomAssemblyFormat = 1;
-  let extraClassDefinition = [{
-    ParseResult TorchMlirCustomOpExampleIdentityOp::parse(OpAsmParser &parser, OperationState &result) {
-      return parseDefaultTorchOp(parser, result, 1, 1);
-    }
-    void TorchMlirCustomOpExampleIdentityOp::print(OpAsmPrinter &printer) {
-      printDefaultTorchOp(printer, *this, 1, 1);
-    }
-  }];
-}
-
diff --git a/lib/Conversion/TorchToLinalg/CMakeLists.txt b/lib/Conversion/TorchToLinalg/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_mlir_conversion_library(TorchMLIRTorchToLinalg
-  CustomOpExample.cpp
+# 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/CustomOpExample.cpp b/lib/Conversion/TorchToLinalg/CustomOpExample.cpp
diff --git a/lib/Conversion/TorchToLinalg/PopulatePatterns.h b/lib/Conversion/TorchToLinalg/PopulatePatterns.h
@@ -63,9 +63,9 @@ void populateIndirectDataMovementPatternsAndLegality(
 void populateTensorConstructorsPatternsAndLegality(TypeConverter &typeConverter,
                                                    RewritePatternSet &patterns,
                                                    ConversionTarget &target);
-void populateCustomOpExamplePatternsAndLegality(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
@@ -62,8 +62,8 @@ class ConvertTorchToLinalg
 
     RewritePatternSet patterns(context);
 
-    torch_to_linalg::populateCustomOpExamplePatternsAndLegality(
-        typeConverter, patterns, target);
+    //torch_to_linalg::populateCustomOpExamplePatternsAndLegality(
+    //    typeConverter, patterns, target);
     torch_to_linalg::populateTensorScalarInteropPatternsAndLegality(
         typeConverter, patterns, target);
     torch_to_linalg::populateLinearPatternsAndLegality(typeConverter, patterns,

diff --git a/lib/Dialect/Torch/Transforms/RefineTypes.cpp b/lib/Dialect/Torch/Transforms/RefineTypes.cpp
@@ -642,7 +642,7 @@ ChangeResult TypeAnalyzer::visitOperation(
           AtenZero_Op, AtenIndexTensorOp, ValsemVariantAtenIndexPutImplOp,
           AtenIndexPutOp, ValsemVariantAtenCopyOp, AtenZeroFunctionalOp,
           AtenIndexPutHackedTwinOp, AtenMaskedFillScalarOp, AtenFlipOp,
-          PrimAbsScalarOp, TorchMlirCustomOpExampleIdentityOp>(op)) {
+          PrimAbsScalarOp>(op)) {
     return incorporateKnowledge(op->getResult(0), operands[0]->getValue());
   }
 

diff --git a/lib/Dialect/Torch/Transforms/ShapeLibrary.cpp b/lib/Dialect/Torch/Transforms/ShapeLibrary.cpp
@@ -6512,10 +6512,6 @@ module {
     %3 = call @__torch__.torch.jit._shape_functions.mean_dim(%arg0, %1, %arg3, %2) : (!torch.list<int>, !torch.list<int>, !torch.bool, !torch.any) -> !torch.list<int>
     return %3 : !torch.list<int>
   }
-  func.func @"__torch_mlir_shape_fn._torch_mlir_custom_op_example.identity"(%arg0: !torch.list<int>) -> !torch.list<int> {
-    %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list<int>) -> !torch.list<int>
-    return %0 : !torch.list<int>
-  }
 }
 )mlir");
 #pragma clang diagnostic pop

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
@@ -1034,8 +1034,9 @@ 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)
 
-def _torch_mlir_custom_op_example〇identity(t: List[int]) -> List[int]:
-    return upstream_shape_functions.unary(t)
+# 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
@@ -597,7 +597,8 @@ def emit_with_mutating_variants(key, **kwargs):
     # extension.
     # ==========================================================================
 
-    emit("_torch_mlir_custom_op_example::identity : (Tensor) -> (Tensor)")
+    # 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):

diff --git a/python/torch_mlir_e2e_test/test_suite/__init__.py b/python/torch_mlir_e2e_test/test_suite/__init__.py
@@ -53,4 +53,5 @@ def register_all_tests():
     from . import return_types
     from . import control_flow
     from . import stats
-    from . import custom_op_example
+    # TODO: Re-enable after MacOS support is fixed for the extension.
+    #from . import custom_op_example