Add EvalOrPattern to StablehloRefineShapes pass (#1867)

Port of tensorflow/mlir-hlo@c132265#diff-c4f3cba740037085e1b02f7f7de83522697ba55cf0577ca9d3e63a7a0f71e1ecR321
openxla · Dec 1, 2023 · 57e5a4a · 57e5a4a
1 parent b3d018f
commit 57e5a4a
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/stablehlo/tests/stablehlo_refine_shapes.mlir b/stablehlo/tests/stablehlo_refine_shapes.mlir
@@ -344,6 +344,19 @@ func.func @eval_multiply() -> tensor<i64> {
 
 // -----
 
+// CHECK-LABEL: func @eval_or
+func.func @eval_or() -> tensor<i1> {
+  // CHECK-NOT: stablehlo.or
+  // CHECK: [[RESULT:%.*]] = stablehlo.constant dense<true> : tensor<i1>
+  // CHECK: return [[RESULT]]
+  %0 = stablehlo.constant dense<true> : tensor<i1>
+  %1 = stablehlo.constant dense<false> : tensor<i1>
+  %2 = stablehlo.or %0, %1 : tensor<i1>
+  func.return %2 : tensor<i1>
+}
+
+// -----
+
 // CHECK-LABEL: func @eval_remainder
 func.func @eval_remainder() -> tensor<i64> {
   // CHECK-NOT: stablehlo.remainder

diff --git a/stablehlo/transforms/StablehloRefineShapes.cpp b/stablehlo/transforms/StablehloRefineShapes.cpp
@@ -304,6 +304,20 @@ struct EvalMulOpPattern : public OpRewritePattern<MulOp> {
   }
 };
 
+struct EvalOrOpPattern : public OpRewritePattern<OrOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(OrOp op,
+                                PatternRewriter& rewriter) const override {
+    auto resultType = op.getType();
+    if (!resultType.getElementType().isInteger(1))
+      return rewriter.notifyMatchFailure(op, "expected boolean element type");
+
+    return evalElementwise(rewriter, op, [&](APSInt lhsInt, APSInt rhsInt) {
+      return getAPSInt(resultType.getElementType(), lhsInt != 0 || rhsInt != 0);
+    });
+  }
+};
+
 struct EvalRemOpPattern : public OpRewritePattern<RemOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(RemOp op,
@@ -1165,6 +1179,7 @@ struct StablehloRefineShapesPass
     patterns.add<EvalMaxOpPattern>(&getContext());
     patterns.add<EvalMinOpPattern>(&getContext());
     patterns.add<EvalMulOpPattern>(&getContext());
+    patterns.add<EvalOrOpPattern>(&getContext());
     patterns.add<EvalRemOpPattern>(&getContext());
     patterns.add<EvalReshapeOpPattern>(&getContext());
     patterns.add<EvalSelectOpPattern>(&getContext());