[StableHLO] Guard convolution widen operand fusion (#24424)
The motivation for this PR is to keep StableHLO preprocessing valid for
mixed-precision models that contain bf16/f32 convolutions. IREE
currently folds widening converts into dense ops to reduce explicit
casts, but doing this for stablehlo.convolution can create invalid
operand type combinations that StableHLO type inference rejects. This PR
avoid folding widening converts into stablehlo.convolution when doing so
would make the lhs and rhs element types invalid for StableHLO type
inference.
---------
Signed-off-by: Hao Ren <rhao8608@gmail.com>
Co-authored-by: Hao Ren <rhao8608@gmail.com>
diff --git a/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py b/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
index 922e2f6..9801560 100644
--- a/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
+++ b/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
@@ -90,6 +90,9 @@
"@llvm-project//mlir:TargetSMTLIB": ["MLIRExportSMTLIB"],
"@llvm-project//mlir:VectorOps": ["MLIRVector"],
# StableHLO.
+ "@stablehlo//:base": [
+ "StablehloBase",
+ ],
"@stablehlo//:chlo_ops": [
"ChloOps",
],
diff --git a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/BUILD.bazel b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/BUILD.bazel
index 1b08fb7..a31bb0e 100644
--- a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/BUILD.bazel
+++ b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/BUILD.bazel
@@ -94,6 +94,7 @@
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:TransformUtils",
"@llvm-project//mlir:Transforms",
+ "@stablehlo//:base",
"@stablehlo//:chlo_ops",
"@stablehlo//:stablehlo_ops",
],
diff --git a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/CMakeLists.txt b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/CMakeLists.txt
index bc428e3..ed085b1 100644
--- a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/CMakeLists.txt
+++ b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/CMakeLists.txt
@@ -78,6 +78,7 @@
MLIRTensorDialect
MLIRTransformUtils
MLIRTransforms
+ StablehloBase
StablehloOps
PUBLIC
)
diff --git a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp
index c4b227e..93f5c4b 100644
--- a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp
+++ b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp
@@ -22,6 +22,7 @@
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "stablehlo/dialect/Base.h"
#include "stablehlo/dialect/ChloOps.h"
#include "stablehlo/dialect/StablehloOps.h"
@@ -1437,10 +1438,30 @@
}
};
+// Accepts fused widening operands for ops that have no additional type
+// constraints.
+template <typename Op>
+LogicalResult verifyFusedWidenOperands(Op op, ArrayRef<Value> operands) {
+ return success();
+}
+
+// Rejects convolution fusion when the resulting operand types violate
+// StableHLO.
+LogicalResult verifyFusedWidenOperands(mlir::stablehlo::ConvolutionOp op,
+ ArrayRef<Value> operands) {
+ Type lhsType = getElementTypeOrSelf(operands[0].getType());
+ Type rhsType = getElementTypeOrSelf(operands[1].getType());
+ if (!mlir::hlo::isCompatibleForHloTypeInference(lhsType, rhsType)) {
+ return failure();
+ }
+ return success();
+}
+
// Identifies cases where a dense operation has inputs that come from widening
// operations. For instance, a dot product widening from FP16 to FP32 is better
// to have the casting operation fused into the dot operation. This decreases
-// the loading required during a dense computation.
+// the loading required during a dense computation. Fusion is skipped when it
+// would produce invalid operand types.
template <class Op>
struct FuseWidenOperands final : OpRewritePattern<Op> {
using OpRewritePattern<Op>::OpRewritePattern;
@@ -1476,6 +1497,13 @@
return failure();
}
+ // Recheck type constraints after replacing widened operands with their
+ // inputs.
+ if (failed(verifyFusedWidenOperands(op, operands))) {
+ return rewriter.notifyMatchFailure(
+ op, "fused operands do not satisfy op type constraints");
+ }
+
rewriter.replaceOpWithNewOp<Op>(op, op->getResultTypes(), operands,
op->getAttrs());
return success();
diff --git a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/test/stablehlo_to_stablehlo.mlir b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/test/stablehlo_to_stablehlo.mlir
index 6692be7..97bcedb 100644
--- a/compiler/plugins/input/StableHLO/Conversion/Preprocessing/test/stablehlo_to_stablehlo.mlir
+++ b/compiler/plugins/input/StableHLO/Conversion/Preprocessing/test/stablehlo_to_stablehlo.mlir
@@ -376,6 +376,37 @@
// -----
+// Checks that mixed bf16/f32 convolutions preserve the required lhs widening convert.
+// CHECK-LABEL: @convolution_mixed_precision
+// CHECK-SAME: (%[[ARG0:.+]]: tensor<1x6144x13xbf16>, %[[ARG1:.+]]: tensor<6144x1x4xf32>)
+func.func @convolution_mixed_precision(%arg0: tensor<1x6144x13xbf16>, %arg1: tensor<6144x1x4xf32>) -> tensor<1x6144x16xf32> {
+ // A convolution requires compatible lhs/rhs element types. If only the lhs
+ // has a widening convert, preprocessing must preserve it instead of creating
+ // an invalid bf16/f32 convolution.
+ // CHECK: %[[CAST:.+]] = stablehlo.convert %[[ARG0]]
+ // CHECK-SAME: (tensor<1x6144x13xbf16>) -> tensor<1x6144x13xf32>
+ %cast0 = stablehlo.convert %arg0 : (tensor<1x6144x13xbf16>) -> tensor<1x6144x13xf32>
+ // CHECK: %[[LHS:.+]] = stablehlo.transpose %[[CAST]]
+ // CHECK: %[[RHS:.+]] = stablehlo.transpose %[[ARG1]]
+ // CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[LHS]], %[[RHS]])
+ // CHECK-SAME: : (tensor<1x13x6144xf32>, tensor<4x1x6144xf32>) -> tensor<1x16x6144xf32>
+ %0 = "stablehlo.convolution"(%cast0, %arg1) {
+ batch_group_count = 1 : i64,
+ dimension_numbers = #stablehlo.conv<[b, f, 0]x[o, i, 0]->[b, f, 0]>,
+ feature_group_count = 6144 : i64,
+ lhs_dilation = array<i64: 1>,
+ padding = dense<3> : tensor<1x2xi64>,
+ precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
+ rhs_dilation = array<i64: 1>,
+ window_strides = array<i64: 1>
+ } : (tensor<1x6144x13xf32>, tensor<6144x1x4xf32>) -> tensor<1x6144x16xf32>
+ // CHECK: %[[RESULT:.+]] = stablehlo.transpose %[[CONV]]
+ // CHECK: return %[[RESULT]]
+ func.return %0 : tensor<1x6144x16xf32>
+}
+
+// -----
+
// CHECK-LABEL: @dynamic_dot_general
// This verifies non-crashing, the lowering to linalg happens elsewhere.
func.func @dynamic_dot_general(%arg1: tensor<?x1024x16x64xf32>, %arg2: tensor<?x1024x16x64xf32>) -> tensor<?x16x1024x1024xf32> {
