Erase shape_assertion ops (iree-org#18167)

Just dropping shape_assertion custom call ops.
nod-ai · Aug 9, 2024 · df3d588 · df3d588
1 parent 6f88125
commit df3d588
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 1 deletion.
diff --git a/compiler/plugins/input/StableHLO/Conversion/StableHLOCustomCalls.cpp b/compiler/plugins/input/StableHLO/Conversion/StableHLOCustomCalls.cpp
@@ -221,6 +221,21 @@ struct HouseholderReflectorRewriter final
   }
 };
 
+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 @@ struct LegalizeStableHLOCustomCalls final
     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
@@ -65,6 +65,7 @@ ALL_SRCS = enforce_glob(
         "scatter.mlir",
         "scatter_dynamic.mlir",
         "select.mlir",
+        "shape_assertion.mlir",
         "sine.mlir",
         "slice.mlir",
         "sort.mlir",
@@ -169,6 +170,7 @@ iree_check_single_backend_test_suite(
             "scatter.mlir",
             "scatter_dynamic.mlir",
             "select.mlir",
+            "shape_assertion.mlir",
             "sine.mlir",
             "slice.mlir",
             "sort.mlir",
@@ -247,6 +249,7 @@ iree_check_single_backend_test_suite(
             "scatter.mlir",
             "scatter_dynamic.mlir",
             "select.mlir",
+            "shape_assertion.mlir",
             "sine.mlir",
             "slice.mlir",
             "sort.mlir",
@@ -381,6 +384,7 @@ iree_check_single_backend_test_suite(
             "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
@@ -65,6 +65,7 @@ iree_check_single_backend_test_suite(
     "scatter.mlir"
     "scatter_dynamic.mlir"
     "select.mlir"
+    "shape_assertion.mlir"
     "sine.mlir"
     "slice.mlir"
     "sort.mlir"
@@ -142,6 +143,7 @@ iree_check_single_backend_test_suite(
     "scatter.mlir"
     "scatter_dynamic.mlir"
     "select.mlir"
+    "shape_assertion.mlir"
     "sine.mlir"
     "slice.mlir"
     "sort.mlir"
@@ -219,6 +221,7 @@ iree_check_single_backend_test_suite(
     "scatter.mlir"
     "scatter_dynamic.mlir"
     "select.mlir"
+    "shape_assertion.mlir"
     "sine.mlir"
     "slice.mlir"
     "sort.mlir"
@@ -291,6 +294,7 @@ iree_check_single_backend_test_suite(
     "scatter.mlir"
     "scatter_dynamic.mlir"
     "select.mlir"
+    "shape_assertion.mlir"
     "sine.mlir"
     "slice.mlir"
     "sort.mlir"
@@ -366,6 +370,7 @@ iree_check_single_backend_test_suite(
     "scatter.mlir"
     "scatter_dynamic.mlir"
     "select.mlir"
+    "shape_assertion.mlir"
     "sine.mlir"
     "slice.mlir"
     "sort.mlir"
@@ -450,6 +455,7 @@ iree_check_single_backend_test_suite(
     "scatter.mlir"
     "scatter_dynamic.mlir"
     "select.mlir"
+    "shape_assertion.mlir"
     "sine.mlir"
     "slice.mlir"
     "sort.mlir"
@@ -533,6 +539,7 @@ iree_check_single_backend_test_suite(
     "scatter.mlir"
     "scatter_dynamic.mlir"
     "select.mlir"
+    "shape_assertion.mlir"
     "sine.mlir"
     "slice.mlir"
     "sort.mlir"
@@ -607,6 +614,7 @@ iree_check_single_backend_test_suite(
     "scatter.mlir"
     "scatter_dynamic.mlir"
     "select.mlir"
+    "shape_assertion.mlir"
     "sine.mlir"
     "slice.mlir"
     "sort.mlir"
@@ -680,6 +688,7 @@ iree_check_single_backend_test_suite(
     "scatter.mlir"
     "scatter_dynamic.mlir"
     "select.mlir"
+    "shape_assertion.mlir"
     "sine.mlir"
     "slice.mlir"
     "sort.mlir"
@@ -754,6 +763,7 @@ iree_check_single_backend_test_suite(
     "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
@@ -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
+}