Erase shape_assertion ops (#18167)
Just dropping shape_assertion custom call ops.
diff --git a/compiler/plugins/input/StableHLO/Conversion/StableHLOCustomCalls.cpp b/compiler/plugins/input/StableHLO/Conversion/StableHLOCustomCalls.cpp
index c410f71..a411868 100644
--- a/compiler/plugins/input/StableHLO/Conversion/StableHLOCustomCalls.cpp
+++ b/compiler/plugins/input/StableHLO/Conversion/StableHLOCustomCalls.cpp
@@ -221,6 +221,21 @@
}
};
+struct ShapeAssertionDrop final
+ : OpRewritePattern<mlir::stablehlo::CustomCallOp> {
+ using OpRewritePattern<mlir::stablehlo::CustomCallOp>::OpRewritePattern;
+ using OpAdaptor = mlir::stablehlo::CustomCallOp::Adaptor;
+
+ LogicalResult matchAndRewrite(mlir::stablehlo::CustomCallOp op,
+ PatternRewriter &rewriter) const final {
+ if (op.getCallTargetName() != "shape_assertion") {
+ return rewriter.notifyMatchFailure(op, "not shape_assertion");
+ }
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+
//===----------------------------------------------------------------------===//
// Pass Definition.
//===----------------------------------------------------------------------===//
@@ -237,7 +252,7 @@
MLIRContext *ctx = f.getContext();
RewritePatternSet patterns(ctx);
- patterns.add<HouseholderReflectorRewriter>(ctx);
+ patterns.add<HouseholderReflectorRewriter, ShapeAssertionDrop>(ctx);
if (failed(applyPatternsAndFoldGreedily(f, std::move(patterns)))) {
signalPassFailure();
}
diff --git a/tests/e2e/stablehlo_ops/BUILD.bazel b/tests/e2e/stablehlo_ops/BUILD.bazel
index e10b217..fb9a65d 100644
--- a/tests/e2e/stablehlo_ops/BUILD.bazel
+++ b/tests/e2e/stablehlo_ops/BUILD.bazel
@@ -65,6 +65,7 @@
"scatter.mlir",
"scatter_dynamic.mlir",
"select.mlir",
+ "shape_assertion.mlir",
"sine.mlir",
"slice.mlir",
"sort.mlir",
@@ -169,6 +170,7 @@
"scatter.mlir",
"scatter_dynamic.mlir",
"select.mlir",
+ "shape_assertion.mlir",
"sine.mlir",
"slice.mlir",
"sort.mlir",
@@ -247,6 +249,7 @@
"scatter.mlir",
"scatter_dynamic.mlir",
"select.mlir",
+ "shape_assertion.mlir",
"sine.mlir",
"slice.mlir",
"sort.mlir",
@@ -381,6 +384,7 @@
"scatter.mlir",
"scatter_dynamic.mlir",
"select.mlir",
+ "shape_assertion.mlir",
"sine.mlir",
"slice.mlir",
"sort.mlir",
diff --git a/tests/e2e/stablehlo_ops/CMakeLists.txt b/tests/e2e/stablehlo_ops/CMakeLists.txt
index 2050c52..4ebc7b5 100644
--- a/tests/e2e/stablehlo_ops/CMakeLists.txt
+++ b/tests/e2e/stablehlo_ops/CMakeLists.txt
@@ -65,6 +65,7 @@
"scatter.mlir"
"scatter_dynamic.mlir"
"select.mlir"
+ "shape_assertion.mlir"
"sine.mlir"
"slice.mlir"
"sort.mlir"
@@ -142,6 +143,7 @@
"scatter.mlir"
"scatter_dynamic.mlir"
"select.mlir"
+ "shape_assertion.mlir"
"sine.mlir"
"slice.mlir"
"sort.mlir"
@@ -219,6 +221,7 @@
"scatter.mlir"
"scatter_dynamic.mlir"
"select.mlir"
+ "shape_assertion.mlir"
"sine.mlir"
"slice.mlir"
"sort.mlir"
@@ -291,6 +294,7 @@
"scatter.mlir"
"scatter_dynamic.mlir"
"select.mlir"
+ "shape_assertion.mlir"
"sine.mlir"
"slice.mlir"
"sort.mlir"
@@ -366,6 +370,7 @@
"scatter.mlir"
"scatter_dynamic.mlir"
"select.mlir"
+ "shape_assertion.mlir"
"sine.mlir"
"slice.mlir"
"sort.mlir"
@@ -450,6 +455,7 @@
"scatter.mlir"
"scatter_dynamic.mlir"
"select.mlir"
+ "shape_assertion.mlir"
"sine.mlir"
"slice.mlir"
"sort.mlir"
@@ -533,6 +539,7 @@
"scatter.mlir"
"scatter_dynamic.mlir"
"select.mlir"
+ "shape_assertion.mlir"
"sine.mlir"
"slice.mlir"
"sort.mlir"
@@ -607,6 +614,7 @@
"scatter.mlir"
"scatter_dynamic.mlir"
"select.mlir"
+ "shape_assertion.mlir"
"sine.mlir"
"slice.mlir"
"sort.mlir"
@@ -680,6 +688,7 @@
"scatter.mlir"
"scatter_dynamic.mlir"
"select.mlir"
+ "shape_assertion.mlir"
"sine.mlir"
"slice.mlir"
"sort.mlir"
@@ -754,6 +763,7 @@
"scatter.mlir"
"scatter_dynamic.mlir"
"select.mlir"
+ "shape_assertion.mlir"
"sine.mlir"
"slice.mlir"
"sort.mlir"
diff --git a/tests/e2e/stablehlo_ops/shape_assertion.mlir b/tests/e2e/stablehlo_ops/shape_assertion.mlir
new file mode 100644
index 0000000..0d5de86
--- /dev/null
+++ b/tests/e2e/stablehlo_ops/shape_assertion.mlir
@@ -0,0 +1,9 @@
+func.func @tensor() {
+ %0 = util.unfoldable_constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
+ %1 = util.unfoldable_constant dense<[5.0, 6.0, 7.0, 8.0]> : tensor<4xf32>
+ %4 = stablehlo.compare EQ, %0, %1, NOTYPE : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
+ stablehlo.custom_call @shape_assertion(%4) {error_message = "Shape assertion failed", has_side_effect = true} : (tensor<4xi1>) -> ()
+ %result = "stablehlo.add"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
+ check.expect_almost_eq_const(%result, dense<[6.0, 8.0, 10.0, 12.0]> : tensor<4xf32>) : tensor<4xf32>
+ return
+}
