[StableHLO] Add stablehlo pipeline to auto input conversion (#13868)
Prefer it over the legacy MHLO pipeline.
Issue: https://github.com/openxla/iree/issues/12678
diff --git a/compiler/src/iree/compiler/InputConversion/Common/AutoInputConversionPipeline.cpp b/compiler/src/iree/compiler/InputConversion/Common/AutoInputConversionPipeline.cpp
index c7233ad..68a08b5 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/AutoInputConversionPipeline.cpp
+++ b/compiler/src/iree/compiler/InputConversion/Common/AutoInputConversionPipeline.cpp
@@ -16,7 +16,9 @@
// Dialect specific
#ifdef IREE_HAVE_MHLO_INPUT
#include "iree/compiler/InputConversion/MHLO/Passes.h"
+#include "iree/compiler/InputConversion/StableHLO/Passes.h"
#include "mhlo/IR/hlo_ops.h"
+#include "stablehlo/dialect/StablehloOps.h"
#endif // IREE_HAVE_MHLO_INPUT
#ifdef IREE_HAVE_TOSA_INPUT
#include "iree/compiler/InputConversion/TOSA/Passes.h"
@@ -28,17 +30,17 @@
namespace mlir::iree_compiler {
namespace {
-struct AutoInputConversionPipelinePass
- : public AutoInputConversionPipelineBase<AutoInputConversionPipelinePass> {
+struct AutoInputConversionPipelinePass final
+ : AutoInputConversionPipelineBase<AutoInputConversionPipelinePass> {
void runOnOperation() override;
void getDependentDialects(DialectRegistry& registry) const override;
};
// All the features seen that should be handled during input conversion.
struct InputFeatures {
- // MHLO features.
- bool hasMHLO = false;
+ // HLO features.
bool hasStableHLO = false;
+ bool hasMHLO = false;
// - XLA import features.
bool hasTuples = false;
@@ -50,7 +52,6 @@
};
static void populateHloFeatures(Operation* op, InputFeatures& features) {
- features.hasMHLO = true;
if (features.hasTuples) {
return;
}
@@ -86,12 +87,18 @@
}
}
-static void populateFeatures(Operation* op, const Dialect* mhloDialect,
+static void populateFeatures(Operation* op, const Dialect* stablehloDialect,
+ const Dialect* mhloDialect,
const Dialect* tmTensorDialect,
const Dialect* tosaDialect,
InputFeatures& features) {
Dialect* d = op->getDialect();
+ if (d == stablehloDialect) {
+ features.hasStableHLO = true;
+ return populateHloFeatures(op, features);
+ }
if (d == mhloDialect) {
+ features.hasMHLO = true;
return populateHloFeatures(op, features);
}
if (d == tosaDialect) {
@@ -109,20 +116,23 @@
MLIRContext* ctxt = &getContext();
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 (!mhloDialect && !tosaDialect && !tmTensorDialect) {
+ if (!stablehloDialect && !mhloDialect && !tosaDialect && !tmTensorDialect) {
return;
}
auto res = module.walk([&](Operation* op) {
- populateFeatures(op, mhloDialect, tmTensorDialect, tosaDialect, features);
- if (features.hasMHLO && features.hasTOSA) {
+ populateFeatures(op, stablehloDialect, mhloDialect, tmTensorDialect,
+ tosaDialect, features);
+ bool hasAnyHLO = features.hasStableHLO || features.hasMHLO;
+ if (hasAnyHLO && features.hasTOSA) {
module.emitError("not yet implemented mixture of *HLO and TOSA");
return WalkResult::interrupt();
}
- if (features.hasMHLO && features.hasTmTensor) {
+ if (hasAnyHLO && features.hasTmTensor) {
module.emitError("not yet implemented mixture of *HLO and TM Tensor");
return WalkResult::interrupt();
}
@@ -135,13 +145,21 @@
if (res.wasInterrupted()) {
return signalPassFailure();
}
- if (!features.hasMHLO && !features.hasTOSA && !features.hasTmTensor) {
+ if (!features.hasStableHLO && !features.hasMHLO && !features.hasTOSA &&
+ !features.hasTmTensor) {
return;
}
OpPassManager pm(ModuleOp::getOperationName(),
OpPassManager::Nesting::Explicit);
#ifdef IREE_HAVE_MHLO_INPUT
+ if (features.hasStableHLO && !features.hasMHLO) {
+ if (features.hasTuples) {
+ stablehlo::buildStableHLOXLAInputConversionPassPipeline(pm);
+ } else {
+ stablehlo::buildStableHLOInputConversionPassPipeline(pm);
+ }
+ }
if (features.hasMHLO) {
if (features.hasTuples) {
MHLO::buildXLAInputConversionPassPipeline(pm);
@@ -183,6 +201,10 @@
};
#ifdef IREE_HAVE_MHLO_INPUT
+ appendPipelineDialects(stablehlo::buildStableHLOInputConversionPassPipeline);
+ appendPipelineDialects(
+ stablehlo::buildStableHLOXLAInputConversionPassPipeline);
+
appendPipelineDialects(MHLO::buildMHLOInputConversionPassPipeline);
appendPipelineDialects(MHLO::buildXLAInputConversionPassPipeline);
#endif // IREE_HAVE_MHLO_INPUT
diff --git a/compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel b/compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel
index ae10820..8f61d1d 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel
+++ b/compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel
@@ -107,6 +107,7 @@
"@llvm-project//mlir:TosaDialect",
"@llvm-project//mlir:Transforms",
"@mlir-hlo//:mlir_hlo",
+ "@mlir-hlo//stablehlo:stablehlo_ops",
"@torch-mlir-dialects//:TorchMLIRTMTensorDialect",
],
)
diff --git a/compiler/src/iree/compiler/InputConversion/Common/test/auto_input_conversion_pipeline.mlir b/compiler/src/iree/compiler/InputConversion/Common/test/auto_input_conversion_pipeline.mlir
index 1e8e90d..11ac344 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/test/auto_input_conversion_pipeline.mlir
+++ b/compiler/src/iree/compiler/InputConversion/Common/test/auto_input_conversion_pipeline.mlir
@@ -1,10 +1,19 @@
-// RUN: iree-opt --iree-auto-input-conversion %s | FileCheck %s
+// RUN: iree-opt --iree-auto-input-conversion --split-input-file %s | FileCheck %s
// Check that the input conversion pipeline handles a simple input and does not crash.
-// CHECK-LABEL: func.func @simple_add
+// CHECK-LABEL: func.func @simple_add_stablehlo
// CHECK: arith.addi
-func.func @simple_add(%arg0: tensor<2x2xi32>, %arg1: tensor<2x2xi32>) -> tensor<2x2xi32> {
+func.func @simple_add_stablehlo(%arg0: tensor<2x2xi32>, %arg1: tensor<2x2xi32>) -> 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>
}
