Remove MHLO support (#14008)

Drop the MHLO input conversion pipeline, which has been deprecated for
over a week. The StableHLO pipeline is the direct replacement. See the
announcement thread for more context:
https://groups.google.com/g/iree-discuss/c/s6dBpDtWhtk.

This still uses the copy of stablehlo from the mlir-hlo repo -- we will
switch to the stablehlo repo in a follow-up PR.

Issue: #12678

CMakeLists.txt

@@ -372,14 +372,14 @@ cmake_dependent_option(IREE_TARGET_BACKEND_WEBGPU "Enables the 'webgpu' compiler
 # Compiler Input Dialects
 #-------------------------------------------------------------------------------
 
-cmake_dependent_option(IREE_INPUT_MHLO "Builds support for compiling MHLO programs" ON ${IREE_BUILD_COMPILER} OFF)
+cmake_dependent_option(IREE_INPUT_STABLEHLO "Builds support for compiling StableHLO programs" ON ${IREE_BUILD_COMPILER} OFF)
 cmake_dependent_option(IREE_INPUT_TORCH "Builds support for compiling Torch MLIR programs" ON ${IREE_BUILD_COMPILER} OFF)
 cmake_dependent_option(IREE_INPUT_TOSA "Builds support for compiling TOSA programs" ON ${IREE_BUILD_COMPILER} OFF)
 
 if(IREE_BUILD_COMPILER)
   message(STATUS "IREE compiler input dialects:")
-  if(IREE_INPUT_MHLO)
-    message(STATUS "  - MHLO")
+  if(IREE_INPUT_STABLEHLO)
+    message(STATUS "  - StableHLO")
   endif()
   if(IREE_INPUT_TORCH)
     message(STATUS "  - Torch MLIR")

build_tools/bazel_to_cmake/bazel_to_cmake_targets.py

@@ -77,47 +77,9 @@ def __init__(self, repo_map: Dict[str, str]):
         "@llvm-project//mlir:MlirOptLib": ["MLIROptLib"],
         "@llvm-project//mlir:VectorOps": ["MLIRVector"],
 
-        # MHLO.
-        # TODO: Rework this upstream so that Bazel and CMake rules match up
-        # better.
-        # All of these have to depend on tensorflow::external_mhlo_includes to
-        # ensure that include directories are inherited.
-        "@mlir-hlo//:chlo_legalize_to_hlo": [
-            "tensorflow::external_mhlo_includes",
-            "ChloPasses",
-        ],
-        "@mlir-hlo//:mlir_hlo": [
-            "tensorflow::external_mhlo_includes",
-            "MhloDialect",
-            "MLIRMhloUtils",
-        ],
-        "@mlir-hlo//:map_chlo_to_hlo_op": [
-            "ChloOps",
-            "MhloDialect",
-        ],
-        "@mlir-hlo//:map_mhlo_to_scalar_op": [
-            "tensorflow::external_mhlo_includes",
-            "MhloDialect",
-        ],
-        "@mlir-hlo//:mhlo_passes": [
-            "tensorflow::external_mhlo_includes",
-            "MhloPasses",
-            "MhloShapeOpsToStandard",
-            "MhloToLinalg",
-            "MhloToStablehlo",
-            "MhloToStandard",
-            "StablehloToMhlo",
-            # Note: We deliberately omit some passes that we do not use in IREE,
-            # e.g.: MhloToArithmeticConversion, MhloToLhloConversion, or
-            # MhloToMemrefConversion.
-        ],
-        "@mlir-hlo//:unfuse_batch_norm": [
-            "tensorflow::external_mhlo_includes",
-            "MhloPasses",
-        ],
+        # StableHLO.
         "@mlir-hlo//stablehlo:chlo_ops": ["ChloOps",],
         "@mlir-hlo//stablehlo:stablehlo_ops": ["StablehloOps",],
-        "@mlir-hlo//:stablehlo_legalize_to_hlo_pass": ["StablehloToMhlo",],
         "@mlir-hlo//stablehlo:broadcast_utils": ["StablehloBroadcastUtils",],
 
         # NCCL

build_tools/cmake/test_riscv.sh

@@ -80,11 +80,10 @@ echo "******** Running tools CTest ********"
 ctest ${tools_ctest_args[@]}
 
 if [[ "${RISCV_PLATFORM}-${RISCV_ARCH}" == "linux-riscv_32" ]]; then
-  # mhlo.power is also disabled because musl math library is not compiled for
+  # stablehlo.power is also disabled because musl math library is not compiled for
   # 32-bit.
   test_exclude_args+=(
     "stablehlo.*llvm-cpu.*pow"
-    "xla.*llvm-cpu.*pow"
   )
 fi
 
@@ -96,7 +95,6 @@ test_exclude_args+=(
   "iree/tests/e2e/tensor_ops/check_llvm-cpu_local-task_pack_dynamic_inner_tiles.mlir"
   # TODO(#13421): Enable the tests
   "iree/tests/e2e/stablehlo_ops/check_llvm-cpu_local-task_dot.mlir"
-  "iree/tests/e2e/xla_ops/check_llvm-cpu_local-task_dot.mlir"
   "iree/tests/e2e/matmul/e2e_matmul_direct_i8_small_llvm-cpu_local-task"
   "iree/tests/e2e/matmul/e2e_matmul_direct_f32_small_llvm-cpu_local-task"
   "iree/tests/e2e/matmul/e2e_matmul_direct_f32_small_no_padding_llvm-cpu_local-task"

compiler/bindings/python/iree/compiler/tools/core.py

@@ -46,8 +46,6 @@ class InputType(Enum):
   STABLEHLO_XLA = "stablehlo_xla"
   TOSA = "tosa"
   TM_TENSOR = "tm_tensor"
-  MHLO_LEGACY = "mhlo_legacy"
-  XLA_LEGACY = "xla_legacy"
 
   @staticmethod
   def parse(spec: Union[str, InputType]) -> InputType:

compiler/bindings/python/iree/compiler/tools/tf.py

@@ -97,7 +97,7 @@ class ImportOptions(CompilerOptions):
   exported_names: Sequence[str] = ()
   import_only: bool = False
   import_type: ImportType = ImportType.OBJECT_GRAPH
-  input_type: Union[InputType, str] = InputType.XLA_LEGACY
+  input_type: Union[InputType, str] = InputType.STABLEHLO_XLA
   saved_model_tags: Set[str] = field(default_factory=set)
   save_temp_iree_input: Optional[str] = None
 

compiler/src/iree/compiler/Dialect/Flow/Transforms/VerifyInputLegality.cpp

@@ -26,8 +26,9 @@ class VerifyInputLegalityPass
     target.addLegalOp<tosa::ApplyScaleOp>();
     // We're already depending on the Tosa Dialect
     target.addIllegalDialect<tosa::TosaDialect>();
-    // Avoid MHLO dependency
-    target.addIllegalDialect("mhlo");
+    // Avoid StableHLO dependency
+    target.addIllegalDialect("chlo");
+    target.addIllegalDialect("stablehlo");
     target.addIllegalOp<UnrealizedConversionCastOp>();
 
     if (failed(iree_compiler::verifyAllOperationsAreLegal(getOperation(),

compiler/src/iree/compiler/Dialect/Flow/Transforms/test/verify_input_ir.mlir

@@ -1,10 +1,12 @@
 // RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-verify-input-legality))" --verify-diagnostics %s -split-input-file
 
 // expected-error@below {{illegal operations still remain}}
-func.func @check_no_mhlo(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
+func.func @check_no_stablehlo(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
   // expected-error@+1 {{illegal op still exists}}
-  %0 = "mhlo.add"(%arg0, %arg1) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
-  return %0 : tensor<?x?xf32>
+  %0 = stablehlo.add %arg0, %arg1 : tensor<?x?xf32>
+  // expected-error@+1 {{illegal op still exists}}
+  %1 = chlo.broadcast_add %0, %arg1 : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
 }
 
 // -----

compiler/src/iree/compiler/InputConversion/CMakeLists.txt

@@ -6,8 +6,7 @@
 
 add_subdirectory(Common)
 
-if(IREE_INPUT_MHLO)
-  add_subdirectory(MHLO)
+if(IREE_INPUT_STABLEHLO)
   add_subdirectory(StableHLO)
 endif()
 if(IREE_INPUT_TORCH)

compiler/src/iree/compiler/InputConversion/Common/AutoInputConversionPipeline.cpp

@@ -14,12 +14,10 @@
 #include "mlir/Pass/PassManager.h"
 
 // Dialect specific
-#ifdef IREE_HAVE_MHLO_INPUT
-#include "iree/compiler/InputConversion/MHLO/Passes.h"
+#ifdef IREE_HAVE_STABLEHLO_INPUT
 #include "iree/compiler/InputConversion/StableHLO/Passes.h"
-#include "mhlo/IR/hlo_ops.h"
 #include "stablehlo/dialect/StablehloOps.h"
-#endif  // IREE_HAVE_MHLO_INPUT
+#endif  // IREE_HAVE_STABLEHLO_INPUT
 #ifdef IREE_HAVE_TOSA_INPUT
 #include "iree/compiler/InputConversion/TOSA/Passes.h"
 #endif  // IREE_HAVE_TOSA_INPUT
@@ -45,7 +43,6 @@ struct AutoInputConversionPipelinePass final
 struct InputFeatures {
   // HLO features.
   bool hasStableHLO = false;
-  bool hasMHLO = false;
   // - XLA import features.
   bool hasTuples = false;
 
@@ -93,7 +90,6 @@ static void populateHloFeatures(Operation* op, InputFeatures& features) {
 }
 
 static void populateFeatures(Operation* op, const Dialect* stablehloDialect,
-                             const Dialect* mhloDialect,
                              const Dialect* tmTensorDialect,
                              const Dialect* tosaDialect,
                              InputFeatures& features) {
@@ -102,10 +98,6 @@ static void populateFeatures(Operation* op, const Dialect* stablehloDialect,
     features.hasStableHLO = true;
     return populateHloFeatures(op, features);
   }
-  if (d == mhloDialect) {
-    features.hasMHLO = true;
-    return populateHloFeatures(op, features);
-  }
   if (d == tosaDialect) {
     features.hasTOSA = true;
     return;
@@ -122,22 +114,20 @@ void AutoInputConversionPipelinePass::runOnOperation() {
 
   InputFeatures features;
   const Dialect* stablehloDialect = ctxt->getLoadedDialect("stablehlo");
-  const Dialect* mhloDialect = ctxt->getLoadedDialect("mhlo");
   const Dialect* tosaDialect = ctxt->getLoadedDialect("tosa");
   const Dialect* tmTensorDialect = ctxt->getLoadedDialect("tm_tensor");
-  if (!stablehloDialect && !mhloDialect && !tosaDialect && !tmTensorDialect) {
+  if (!stablehloDialect && !tosaDialect && !tmTensorDialect) {
     return;
   }
 
   auto res = module.walk([&](Operation* op) {
-    populateFeatures(op, stablehloDialect, mhloDialect, tmTensorDialect,
-                     tosaDialect, features);
-    bool hasAnyHLO = features.hasStableHLO || features.hasMHLO;
-    if (hasAnyHLO && features.hasTOSA) {
+    populateFeatures(op, stablehloDialect, tmTensorDialect, tosaDialect,
+                     features);
+    if (features.hasStableHLO && features.hasTOSA) {
       module.emitError("not yet implemented mixture of *HLO and TOSA");
       return WalkResult::interrupt();
     }
-    if (hasAnyHLO && features.hasTmTensor) {
+    if (features.hasStableHLO && features.hasTmTensor) {
       module.emitError("not yet implemented mixture of *HLO and TM Tensor");
       return WalkResult::interrupt();
     }
@@ -150,15 +140,14 @@ void AutoInputConversionPipelinePass::runOnOperation() {
   if (res.wasInterrupted()) {
     return signalPassFailure();
   }
-  if (!features.hasStableHLO && !features.hasMHLO && !features.hasTOSA &&
-      !features.hasTmTensor) {
+  if (!features.hasStableHLO && !features.hasTOSA && !features.hasTmTensor) {
     return;
   }
 
   OpPassManager pm(ModuleOp::getOperationName(),
                    OpPassManager::Nesting::Explicit);
-#ifdef IREE_HAVE_MHLO_INPUT
-  if (features.hasStableHLO && !features.hasMHLO) {
+#ifdef IREE_HAVE_STABLEHLO_INPUT
+  if (features.hasStableHLO) {
     stablehlo::StableHloOptions options;
     options.demoteI64ToI32 = demoteI64ToI32;
     options.demoteF64ToF32 = demoteF64ToF32;
@@ -169,14 +158,7 @@ void AutoInputConversionPipelinePass::runOnOperation() {
       stablehlo::buildStableHLOInputConversionPassPipeline(pm, options);
     }
   }
-  if (features.hasMHLO) {
-    if (features.hasTuples) {
-      MHLO::buildXLAInputConversionPassPipeline(pm);
-    } else {
-      MHLO::buildMHLOInputConversionPassPipeline(pm);
-    }
-  }
-#endif  // IREE_HAVE_MHLO_INPUT
+#endif  // IREE_HAVE_STABLEHLO_INPUT
 #ifdef IREE_HAVE_TOSA_INPUT
   if (features.hasTOSA) {
     buildTOSAInputConversionPassPipeline(pm);
@@ -209,7 +191,7 @@ void AutoInputConversionPipelinePass::getDependentDialects(
         pm.getDependentDialects(registry);
       };
 
-#ifdef IREE_HAVE_MHLO_INPUT
+#ifdef IREE_HAVE_STABLEHLO_INPUT
   auto appendStablehloPipelineDialects =
       [&registry](function_ref<void(OpPassManager&,
                                     const stablehlo::StableHloOptions& options)>
@@ -224,10 +206,7 @@ void AutoInputConversionPipelinePass::getDependentDialects(
       stablehlo::buildStableHLOInputConversionPassPipeline);
   appendStablehloPipelineDialects(
       stablehlo::buildStableHLOXLAInputConversionPassPipeline);
-
-  appendPipelineDialects(MHLO::buildMHLOInputConversionPassPipeline);
-  appendPipelineDialects(MHLO::buildXLAInputConversionPassPipeline);
-#endif  // IREE_HAVE_MHLO_INPUT
+#endif  // IREE_HAVE_STABLEHLO_INPUT
 
 #ifdef IREE_HAVE_TOSA_INPUT
   appendPipelineDialects(buildTOSAInputConversionPassPipeline);

compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel

@@ -95,7 +95,6 @@ iree_compiler_cc_library(
     deps = [
         ":PassHeaders",
         ":PassesIncGen",
-        "//compiler/src/iree/compiler/InputConversion/MHLO",
         "//compiler/src/iree/compiler/InputConversion/StableHLO",
         "//compiler/src/iree/compiler/InputConversion/TMTensor",
         "//compiler/src/iree/compiler/InputConversion/TOSA",
@@ -106,7 +105,6 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:TosaDialect",
         "@llvm-project//mlir:Transforms",
-        "@mlir-hlo//:mlir_hlo",
         "@mlir-hlo//stablehlo:stablehlo_ops",
         "@torch-mlir-dialects//:TorchMLIRTMTensorDialect",
     ],

compiler/src/iree/compiler/InputConversion/Common/CMakeLists.txt

@@ -6,8 +6,7 @@
 
 # Enable input dialects based on options.
 set(IREE_INPUT_DEPS "")
-if(IREE_INPUT_MHLO)
-  list(APPEND IREE_INPUT_DEPS iree::compiler::InputConversion::MHLO)
+if(IREE_INPUT_STABLEHLO)
   list(APPEND IREE_INPUT_DEPS iree::compiler::InputConversion::StableHLO)
 endif()
 if(IREE_INPUT_TORCH)

compiler/src/iree/compiler/InputConversion/Common/test/auto_input_conversion_pipeline.mlir

@@ -8,12 +8,3 @@ func.func @simple_add_stablehlo(%arg0: tensor<2x2xi32>, %arg1: tensor<2x2xi32>)
   %0 = stablehlo.add %arg0, %arg1 : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
   return %0 : tensor<2x2xi32>
 }
-
-// -----
-
-// CHECK-LABEL: func.func @simple_add_mhlo
-// CHECK:  arith.addi
-func.func @simple_add_mhlo(%arg0: tensor<2x2xi32>, %arg1: tensor<2x2xi32>) -> tensor<2x2xi32> {
-  %0 = "mhlo.add"(%arg0, %arg1) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
-  return %0 : tensor<2x2xi32>
-}