From b90837ee243b086bd487119da4e140b2df397498 Mon Sep 17 00:00:00 2001 From: Bob Adolf Date: Tue, 14 Jun 2022 18:24:40 -0700 Subject: [PATCH] Temporarily revert support for custom op extensions. (#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. --- build_tools/update_shape_lib.sh | 19 ++++--- build_tools/update_torch_ods.sh | 19 ++++--- .../Dialect/Torch/IR/GeneratedTorchOps.td | 23 -------- lib/Conversion/TorchToLinalg/CMakeLists.txt | 3 +- .../TorchToLinalg/CustomOpExample.cpp | 54 ------------------- .../TorchToLinalg/PopulatePatterns.h | 6 +-- .../TorchToLinalg/TorchToLinalg.cpp | 4 +- lib/Dialect/Torch/Transforms/RefineTypes.cpp | 2 +- lib/Dialect/Torch/Transforms/ShapeLibrary.cpp | 4 -- .../jit_ir/build_tools/shape_lib_gen.py | 5 +- .../jit_ir/build_tools/torch_ods_gen.py | 3 +- .../test_suite/__init__.py | 3 +- 12 files changed, 37 insertions(+), 108 deletions(-) delete mode 100644 lib/Conversion/TorchToLinalg/CustomOpExample.cpp diff --git a/build_tools/update_shape_lib.sh b/build_tools/update_shape_lib.sh index 7cff596dd0d4..8912805ed5d2 100755 --- 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 index 330d2095f556..e7620437f9ad 100755 --- 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 index ee156e895e33..91115490a54c 100644 --- 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 index 3fd6df386377..e31a09b8b16d 100644 --- 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 deleted file mode 100644 index b44830fa0440..000000000000 --- a/lib/Conversion/TorchToLinalg/CustomOpExample.cpp +++ /dev/null @@ -1,54 +0,0 @@ -//===----------------------------------------------------------------------===// -// -// Part of the LLVM Project, 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 -// Also available under a BSD-style license. See LICENSE. -// -//===----------------------------------------------------------------------===// - -#include "torch-mlir/Conversion/TorchToLinalg/TorchToLinalg.h" - -#include "../PassDetail.h" -#include "PopulatePatterns.h" -#include "Utils.h" -#include "mlir/IR/Matchers.h" -#include "torch-mlir/Conversion/Utils/Utils.h" -#include "torch-mlir/Dialect/Torch/IR/TorchDialect.h" -#include "torch-mlir/Dialect/Torch/IR/TorchOps.h" -#include "torch-mlir/Dialect/Torch/Utils/TorchUpstream.h" -#include "torch-mlir/Dialect/Torch/Utils/Utils.h" - -using namespace mlir; -using namespace mlir::torch; -using namespace mlir::torch::Torch; - -namespace { -class ConvertCustomOpExample - : public OpConversionPattern { -public: - using OpConversionPattern::OpConversionPattern; - LogicalResult matchAndRewrite(TorchMlirCustomOpExampleIdentityOp op, - OpAdaptor adaptor, - ConversionPatternRewriter &rewriter - ) const override { - // Type checks. - if (failed(verifyLinalgCompatibleTypes(op, rewriter))) - return failure(); - // Since the example op does nothing, we simply replace the uses of the - // return value with its argument, then remove the op. - rewriter.replaceOp(op, op->getOperands()); - - return success(); - } -}; -} // namespace - -void mlir::torch::torch_to_linalg::populateCustomOpExamplePatternsAndLegality( - TypeConverter &typeConverter, - RewritePatternSet &patterns, - ConversionTarget &target) { - MLIRContext *context = patterns.getContext(); - target.addIllegalOp(); - patterns.add(typeConverter, context); -} diff --git a/lib/Conversion/TorchToLinalg/PopulatePatterns.h b/lib/Conversion/TorchToLinalg/PopulatePatterns.h index df0bd9fdef59..384c89d333eb 100644 --- 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 index 3d0e8377cd08..1caa1408f9a0 100644 --- 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 index 3d027ce2cb00..ea0401bbb12c 100644 --- 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 index 01a693e41976..03d485859e07 100644 --- 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, !torch.list, !torch.bool, !torch.any) -> !torch.list return %3 : !torch.list } - func.func @"__torch_mlir_shape_fn._torch_mlir_custom_op_example.identity"(%arg0: !torch.list) -> !torch.list { - %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list) -> !torch.list - return %0 : !torch.list - } } )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 index 7eae57559559..2478eadfff69 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 @@ -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 index 5c1f0511c308..42a266b821b4 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 @@ -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 index 57160207317b..584d4c0310fd 100644 --- 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