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