From 078d1e1a1d7e5231e98227187d9f322e688cc645 Mon Sep 17 00:00:00 2001 From: Stella Laurenzo Date: Tue, 12 Sep 2023 19:10:02 -0700 Subject: [PATCH] Remove mlir-hlo (replace with stablehlo). (#2460) We just have to do this: I ran into an issue today where I needed to make a one line patch to stablehlo to work around a compiler issue, and it is completely unapparent how to do so given that the mlir-hlo repo is a read-only export and is at the tail end of a multi-week integration chain from the open-source stablehlo repo. We've discussed this often enough and gotten +1 from everyone that they are ok with taking the e2e testing hit if it becomes necessary: It is necessary as the current situation is unmanageable. Looking at it, I expect it wouldn't actually be very difficult to build a little runner binary out of the stablehlo interpreter and subprocess call that in order to get the testing coverage back. I leave that as an exercise to the users of this part of the stack and recommend following the breadcrumbs from the deleted python/torch_mlir_e2e_test/stablehlo_backends/linalg_on_tensors.py file and the main.py changes. Note that I am pointing us at a stablehlo fork for the moment until it is apparent that we don't need to carry any local patches to it. We can update this in a few days if everything is clear. --- .gitmodules | 6 +-- CMakeLists.txt | 9 ++-- .../python_deploy/build_linux_packages.sh | 3 -- docs/development.md | 7 ++- e2e_testing/main.py | 8 +-- externals/mlir-hlo | 1 - externals/stablehlo | 1 + lib/CMakeLists.txt | 8 --- lib/Conversion/Passes.cpp | 1 - lib/Conversion/TorchToStablehlo/Basic.cpp | 51 ++++++++++++++----- .../TorchToStablehlo/CMakeLists.txt | 3 +- .../TorchConversion/Transforms/CMakeLists.txt | 4 +- lib/InitAll.cpp | 12 ----- .../stablehlo_backends/linalg_on_tensors.py | 50 ------------------ tools/torch-mlir-opt/torch-mlir-opt.cpp | 8 --- 15 files changed, 58 insertions(+), 114 deletions(-) delete mode 160000 externals/mlir-hlo create mode 160000 externals/stablehlo delete mode 100644 python/torch_mlir_e2e_test/stablehlo_backends/linalg_on_tensors.py diff --git a/.gitmodules b/.gitmodules index 81c66a441..8b46098d9 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,6 +1,6 @@ [submodule "externals/llvm-project"] path = externals/llvm-project url = https://github.com/llvm/llvm-project.git -[submodule "externals/mlir-hlo"] - path = externals/mlir-hlo - url = https://github.com/tensorflow/mlir-hlo.git +[submodule "externals/stablehlo"] + path = externals/stablehlo + url = https://github.com/openxla/stablehlo.git diff --git a/CMakeLists.txt b/CMakeLists.txt index a3c636fc6..ac62423db 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -119,13 +119,10 @@ endif() if (TORCH_MLIR_ENABLE_STABLEHLO) set(STABLEHLO_BUILD_EMBEDDED ON) - add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/externals/mlir-hlo - ${CMAKE_CURRENT_BINARY_DIR}/mlir-hlo + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/externals/stablehlo + ${CMAKE_CURRENT_BINARY_DIR}/stablehlo EXCLUDE_FROM_ALL) - include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/mlir-hlo/include) - include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/mlir-hlo) - include_directories(${CMAKE_CURRENT_BINARY_DIR}/mlir-hlo/include) - include_directories(${CMAKE_CURRENT_BINARY_DIR}/mlir-hlo) + include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/stablehlo) endif() set(TORCH_MLIR_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") diff --git a/build_tools/python_deploy/build_linux_packages.sh b/build_tools/python_deploy/build_linux_packages.sh index cf15f48d7..733c79fa3 100755 --- a/build_tools/python_deploy/build_linux_packages.sh +++ b/build_tools/python_deploy/build_linux_packages.sh @@ -321,9 +321,6 @@ function test_in_tree() { echo ":::: Run make_fx + TOSA e2e integration tests" python -m e2e_testing.main --config=make_fx_tosa -v - echo ":::: Run StableHLO e2e integration tests" - python -m e2e_testing.main --config=stablehlo -v - echo ":::: Run TOSA e2e integration tests" python -m e2e_testing.main --config=tosa -v } diff --git a/docs/development.md b/docs/development.md index 048f363c0..323db4d8b 100644 --- a/docs/development.md +++ b/docs/development.md @@ -408,13 +408,18 @@ Torch-MLIR by default builds with the latest nightly PyTorch version. This can b # Updating the LLVM and MLIR-HLO submodules Torch-MLIR depends on `llvm-project` (which contains, among other things, -upstream MLIR) and `mlir-hlo`, both of which are submodules in the `externals/` +upstream MLIR) and `stablehlo`, both of which are submodules in the `externals/` directory. We aim to update these at least weekly to bring in the latest features and spread out over time the effort of updating our code for MLIR API breakages. ## Which LLVM commit should I pick? +NOTE: This section is in flux. Specifically, the `mlir-hlo` dep has been +dropped and the project is running off of a `stablehlo` fork which can be +patched for certain OS combinations. As of 2023-09-12, stellaraccident@ +is massaging this situation. Please reach out for advice updating. + Since downstream projects may want to build Torch-MLIR (and thus LLVM and MLIR-HLO) in various configurations (Release versus Debug builds; on Linux, Windows, or macOS; possibly with Clang, LLD, and LLDB enabled), it is crucial to diff --git a/e2e_testing/main.py b/e2e_testing/main.py index b7b7dc44f..885f34477 100644 --- a/e2e_testing/main.py +++ b/e2e_testing/main.py @@ -24,7 +24,6 @@ from torch_mlir_e2e_test.configs import ( ) from torch_mlir_e2e_test.linalg_on_tensors_backends.refbackend import RefBackendLinalgOnTensorsBackend -from torch_mlir_e2e_test.stablehlo_backends.linalg_on_tensors import LinalgOnTensorsStablehloBackend from torch_mlir_e2e_test.tosa_backends.linalg_on_tensors import LinalgOnTensorsTosaBackend from .xfail_sets import ( @@ -44,7 +43,7 @@ from torch_mlir_e2e_test.test_suite import register_all_tests register_all_tests() def _get_argparse(): - config_choices = ["native_torch", "torchscript", "linalg", "stablehlo", "make_fx_tosa", "tosa", "lazy_tensor_core", "torchdynamo"] + config_choices = ["native_torch", "torchscript", "linalg", "make_fx_tosa", "tosa", "lazy_tensor_core", "torchdynamo"] parser = argparse.ArgumentParser(description="Run torchscript e2e tests.") parser.add_argument("-c", "--config", choices=config_choices, @@ -52,7 +51,6 @@ def _get_argparse(): help=f""" Meaning of options: "linalg": run through torch-mlir"s default Linalg-on-Tensors backend. -"stablehlo": run through torch-mlir"s default StableHLO backend. "tosa": run through torch-mlir"s default TOSA backend. "native_torch": run the torch.nn.Module as-is without compiling (useful for verifying model is deterministic; ALL tests should pass in this configuration). "torchscript": compile the model to a torch.jit.ScriptModule, and then run that as-is (useful for verifying TorchScript is modeling the program correctly). @@ -100,10 +98,6 @@ def main(): config = TosaBackendTestConfig(LinalgOnTensorsTosaBackend(), use_make_fx=True) xfail_set = all_test_unique_names - MAKE_FX_TOSA_PASS_SET crashing_set = set() - elif args.config == "stablehlo": - config = StablehloBackendTestConfig(LinalgOnTensorsStablehloBackend()) - xfail_set = all_test_unique_names - STABLEHLO_PASS_SET - crashing_set = STABLEHLO_CRASHING_SET elif args.config == "native_torch": config = NativeTorchTestConfig() xfail_set = set() diff --git a/externals/mlir-hlo b/externals/mlir-hlo deleted file mode 160000 index 16886a108..000000000 --- a/externals/mlir-hlo +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 16886a108eff5197f816ca0f1950cc5ff1b078d9 diff --git a/externals/stablehlo b/externals/stablehlo new file mode 160000 index 000000000..77a59815a --- /dev/null +++ b/externals/stablehlo @@ -0,0 +1 @@ +Subproject commit 77a59815a82b34f7b08ed2d42a711d9920682d0e diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index fb1da3fb6..03123d2ed 100644 --- a/lib/CMakeLists.txt +++ b/lib/CMakeLists.txt @@ -23,14 +23,6 @@ set(LinkedLibs TorchMLIRRefBackend ) -if(TORCH_MLIR_ENABLE_STABLEHLO) - list(APPEND LinkedLibs - MhloPasses - MhloToLinalg - StablehloToMhlo - ) -endif() - add_mlir_library(TorchMLIRInitAll InitAll.cpp diff --git a/lib/Conversion/Passes.cpp b/lib/Conversion/Passes.cpp index 45714601d..0dae24678 100644 --- a/lib/Conversion/Passes.cpp +++ b/lib/Conversion/Passes.cpp @@ -11,7 +11,6 @@ #ifdef TORCH_MLIR_ENABLE_STABLEHLO #include "torch-mlir/Conversion/TorchToStablehlo/TorchToStablehlo.h" -#include "transforms/passes.h" #endif // TORCH_MLIR_ENABLE_STABLEHLO #include "torch-mlir/Conversion/TorchToLinalg/TorchToLinalg.h" diff --git a/lib/Conversion/TorchToStablehlo/Basic.cpp b/lib/Conversion/TorchToStablehlo/Basic.cpp index 0a2ed02e0..979182ae7 100644 --- a/lib/Conversion/TorchToStablehlo/Basic.cpp +++ b/lib/Conversion/TorchToStablehlo/Basic.cpp @@ -13,6 +13,7 @@ #include "PopulatePatterns.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Complex/IR/Complex.h" #include "mlir/Dialect/Shape/IR/Shape.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "stablehlo/dialect/ChloOps.h" @@ -25,7 +26,6 @@ #include "torch-mlir/Dialect/Torch/Utils/TorchUpstream.h" #include "torch-mlir/Dialect/Torch/Utils/Utils.h" #include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.h" -#include "utils/hlo_utils.h" #include #include @@ -34,6 +34,34 @@ using namespace mlir::torch; using namespace mlir::torch::Torch; using namespace mlir::torch::torch_to_stablehlo; +namespace { + +template +static Value getConstantLike(OpBuilder &b, Location loc, T constant, + Value val) { + Type ty = getElementTypeOrSelf(val.getType()); + auto getAttr = [&]() -> Attribute { + if (ty.isa()) + return b.getIntegerAttr(ty, constant); + if (ty.isa()) + return b.getFloatAttr(ty, constant); + if (auto complexTy = ty.dyn_cast()) + return complex::NumberAttr::get(complexTy, constant, 0); + llvm_unreachable("unhandled element type"); + }; + return b.create(loc, cast(getAttr()), + val); +} + +Value getConstantLike(OpBuilder &b, Location loc, const APFloat &constant, + Value val) { + Type ty = getElementTypeOrSelf(val.getType()); + return b.create(loc, b.getFloatAttr(ty, constant), + val); +} + +} // namespace + LogicalResult broadcastRanks(PatternRewriter &rewriter, Operation *op, mlir::Value &self, mlir::Value &other, size_t dimSizeIndexBits) { @@ -836,7 +864,7 @@ LogicalResult ConvertAtenOp::matchAndRewrite( "for AtenReciprocalOp"); } - Value oneTensor = chlo::getConstantLike(rewriter, op->getLoc(), 1, input); + Value oneTensor = getConstantLike(rewriter, op->getLoc(), 1, input); rewriter.replaceOpWithNewOp(op, outTy, oneTensor, input); return success(); } @@ -945,7 +973,7 @@ LogicalResult ConvertAtenOp::matchAndRewrite( } Value zeroTensor; - zeroTensor = chlo::getConstantLike( + zeroTensor = getConstantLike( rewriter, op->getLoc(), APFloat::getZero(lhsElemTy.cast().getFloatSemantics(), false), @@ -967,9 +995,9 @@ LogicalResult ConvertAtenOp::matchAndRewrite( return op.emitError("only ranked tensor type is supported."); } - Value one = chlo::getConstantLike(rewriter, loc, 1.0, input); - Value two = chlo::getConstantLike(rewriter, loc, 2.0, input); - Value half = chlo::getConstantLike(rewriter, loc, 0.5, input); + Value one = getConstantLike(rewriter, loc, 1.0, input); + Value two = getConstantLike(rewriter, loc, 2.0, input); + Value half = getConstantLike(rewriter, loc, 0.5, input); auto rsqrtTwo = rewriter.create(loc, two); auto erfElement = rewriter.create(loc, input, rsqrtTwo); auto erf = rewriter.create(loc, erfElement); @@ -1485,13 +1513,12 @@ LogicalResult ConvertAtenOp::matchAndRewrite( return rewriter.notifyMatchFailure(op, "Unsupported value of approximate"); } // Create constant value - Value kAlpha = - chlo::getConstantLike(rewriter, loc, 0.70710678118654752440, input); + Value kAlpha = getConstantLike(rewriter, loc, 0.70710678118654752440, input); Value cstAlpha0 = - chlo::getConstantLike(rewriter, loc, 1.12837916709551257390, input); - Value half = chlo::getConstantLike(rewriter, loc, .5, input); - Value one = chlo::getConstantLike(rewriter, loc, 1.0, input); - Value negHalf = chlo::getConstantLike(rewriter, loc, -0.5, input); + getConstantLike(rewriter, loc, 1.12837916709551257390, input); + Value half = getConstantLike(rewriter, loc, .5, input); + Value one = getConstantLike(rewriter, loc, 1.0, input); + Value negHalf = getConstantLike(rewriter, loc, -0.5, input); // Compute Value kBeta0 = diff --git a/lib/Conversion/TorchToStablehlo/CMakeLists.txt b/lib/Conversion/TorchToStablehlo/CMakeLists.txt index 84a560cd7..0f9b8faba 100644 --- a/lib/Conversion/TorchToStablehlo/CMakeLists.txt +++ b/lib/Conversion/TorchToStablehlo/CMakeLists.txt @@ -20,7 +20,8 @@ add_mlir_conversion_library(TorchMLIRTorchToStablehlo LINK_LIBS PUBLIC MLIRIR MLIRPass - MLIRBufferTransforms + MLIRComplexDialect + ChloOps StablehloOps TorchMLIRTorchDialect TorchMLIRConversionUtils diff --git a/lib/Dialect/TorchConversion/Transforms/CMakeLists.txt b/lib/Dialect/TorchConversion/Transforms/CMakeLists.txt index 6495e4682..a286d5bbd 100644 --- a/lib/Dialect/TorchConversion/Transforms/CMakeLists.txt +++ b/lib/Dialect/TorchConversion/Transforms/CMakeLists.txt @@ -18,7 +18,9 @@ set(LinkedLibs ) if(TORCH_MLIR_ENABLE_STABLEHLO) - list(APPEND LinkedLibs ChloPasses) + list(APPEND LinkedLibs + StablehloOps + ) endif() add_mlir_library(TorchMLIRTorchConversionPasses diff --git a/lib/InitAll.cpp b/lib/InitAll.cpp index 88bceb013..1d67bdfe2 100644 --- a/lib/InitAll.cpp +++ b/lib/InitAll.cpp @@ -21,10 +21,6 @@ #include "torch-mlir/Dialect/TorchConversion/Transforms/Passes.h" #include "torch-mlir/RefBackend/Passes.h" -#ifdef TORCH_MLIR_ENABLE_STABLEHLO -#include "mhlo/transforms/passes.h" -#endif - void mlir::torch::registerAllDialects(mlir::DialectRegistry ®istry) { registry.insert(); registry.insert(); @@ -40,12 +36,4 @@ void mlir::torch::registerAllPasses() { mlir::torch::registerConversionPasses(); mlir::torch::RefBackend::registerRefBackendPasses(); mlir::torch::TMTensor::registerPasses(); - -#ifdef TORCH_MLIR_ENABLE_STABLEHLO - mlir::mhlo::registerSymbolicShapeOptimizationPass(); - mlir::mhlo::registerStablehloLegalizeToHloPass(); - mlir::mhlo::registerChloLegalizeToHloPass(); - mlir::mhlo::registerHloLegalizeToLinalgPass(); - mlir::mhlo::registerTestUnfuseBatchNormPass(); -#endif // TORCH_MLIR_ENABLE_STABLEHLO } diff --git a/python/torch_mlir_e2e_test/stablehlo_backends/linalg_on_tensors.py b/python/torch_mlir_e2e_test/stablehlo_backends/linalg_on_tensors.py deleted file mode 100644 index 6a36dd196..000000000 --- a/python/torch_mlir_e2e_test/stablehlo_backends/linalg_on_tensors.py +++ /dev/null @@ -1,50 +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. - -from torch_mlir.ir import * -from torch_mlir.passmanager import * -from torch_mlir.compiler_utils import run_pipeline_with_repro_report - -from torch_mlir_e2e_test.linalg_on_tensors_backends.refbackend import ( - RefBackendLinalgOnTensorsBackend, -) - -from .abc import StablehloBackend - -__all__ = [ - "LinalgOnTensorsStablehloBackend", -] - - -class LinalgOnTensorsStablehloBackend(StablehloBackend): - """Main entry-point for the linalg-on-tensors based StableHLO backend. - - This currently uses the linalg-on-tensors RefBackend for actual execution. - """ - - def __init__(self): - super().__init__() - self.refbackend = RefBackendLinalgOnTensorsBackend() - - def compile(self, imported_module: Module): - """Compiles an imported module that satisfied the StableHLO backend contract. - - Args: - imported_module: The MLIR module consisting of funcs in the StableHLO - dialect. - Returns: - An opaque, backend specific compiled artifact object that can be - passed to `load`. - """ - run_pipeline_with_repro_report( - imported_module, - "builtin.module(func.func(chlo-legalize-to-hlo),stablehlo-legalize-to-hlo,func.func(canonicalize,cse,symbolic-shape-optimization,mhlo-test-unfuse-batch-norm,canonicalize,hlo-legalize-to-linalg,canonicalize))", - "Lowering StableHLO to Linalg-on-Tensors", - ) - return self.refbackend.compile(imported_module) - - def load(self, module): - """Loads a compiled artifact into the runtime.""" - return self.refbackend.load(module) diff --git a/tools/torch-mlir-opt/torch-mlir-opt.cpp b/tools/torch-mlir-opt/torch-mlir-opt.cpp index 0a29edf03..fa6a41a70 100644 --- a/tools/torch-mlir-opt/torch-mlir-opt.cpp +++ b/tools/torch-mlir-opt/torch-mlir-opt.cpp @@ -14,8 +14,6 @@ #include "torch-mlir/InitAll.h" #ifdef TORCH_MLIR_ENABLE_STABLEHLO -#include "mhlo/IR/hlo_ops.h" -#include "mhlo/transforms/passes.h" #include "stablehlo/dialect/Register.h" #endif @@ -32,12 +30,6 @@ int main(int argc, char **argv) { #ifdef TORCH_MLIR_ENABLE_STABLEHLO mlir::stablehlo::registerAllDialects(registry); - registry.insert(); - mlir::mhlo::registerSymbolicShapeOptimizationPass(); - mlir::mhlo::registerStablehloLegalizeToHloPass(); - mlir::mhlo::registerChloLegalizeToHloPass(); - mlir::mhlo::registerHloLegalizeToLinalgPass(); - mlir::mhlo::registerTestUnfuseBatchNormPass(); #endif return mlir::asMainReturnCode(mlir::MlirOptMain( argc, argv, "MLIR modular optimizer driver\n", registry));