From 69f3233fc761a2dceb73681cbecd38cfd3835ae8 Mon Sep 17 00:00:00 2001
From: Sandeep Dasgupta <sdasgup@google.com>
Date: Wed, 1 May 2024 14:58:00 -0700
Subject: [PATCH] Integrate StableHLO at openxla/stablehlo@b6406a43

PiperOrigin-RevId: 629847908
---
 third_party/stablehlo/temporary.patch             | 15 ++-------------
 third_party/stablehlo/workspace.bzl               |  4 ++--
 .../tests/Dialect/mhlo/verifier_conv_op.mlir      |  8 ++++----
 3 files changed, 8 insertions(+), 19 deletions(-)

diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
index 0d7ac6d36ed37..6a3e5a80bf397 100755
--- a/third_party/stablehlo/temporary.patch
+++ b/third_party/stablehlo/temporary.patch
@@ -1,7 +1,7 @@
 diff --ruN a/stablehlo/CMakeLists.txt b/stablehlo/CMakeLists.txt
 --- stablehlo/CMakeLists.txt
 +++ stablehlo/CMakeLists.txt
-@@ -13,153 +13,20 @@
+@@ -13,154 +13,20 @@
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
@@ -43,6 +43,7 @@ diff --ruN a/stablehlo/CMakeLists.txt b/stablehlo/CMakeLists.txt
  #-------------------------------------------------------------------------------
 -option(STABLEHLO_BUILD_EMBEDDED "Build StableHLO as part of another project" OFF)
 -option(STABLEHLO_ENABLE_BINDINGS_PYTHON "Enables StableHLO Python bindings" OFF)
+-option(STABLEHLO_ENABLE_PYTHON_TF_TESTS "Enables StableHLO to SavedModel tests requiring TF" OFF)
 -option(STABLEHLO_ENABLE_STRICT_BUILD "Build StableHLO with strict warnings and warnings as errors" OFF)
 -option(STABLEHLO_ENABLE_SANITIZER "Enable a sanitizer [OFF, address]" OFF)
 -option(STABLEHLO_ENABLE_SPLIT_DWARF "Enable split DWARF if the platform supports it" OFF)
@@ -2563,16 +2564,4 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloRefineShapes.c
 +}  // namespace experimental
 +}  // namespace stablehlo
 +}  // namespace mlir
-diff --ruN a/stablehlo/stablehlo/transforms/StablehloAggressiveSimplification.cpp b/stablehlo/stablehlo/transforms/StablehloAggressiveSimplification.cpp
---- stablehlo/stablehlo/transforms/StablehloAggressiveSimplification.cpp
-+++ stablehlo/stablehlo/transforms/StablehloAggressiveSimplification.cpp
-@@ -735,7 +735,7 @@
-       addToWorkList(retOp.getOperand(resultNo));
-       while (!workList.empty()) {
-         auto definition = workList.pop_back_val();
--        if (auto blockArg = definition.dyn_cast<BlockArgument>()) {
-+        if (auto blockArg = dyn_cast<BlockArgument>(definition)) {
-           // using one argument implies using the whole argument pair
-           const auto pairNo = blockArg.getArgNumber() % numOperandPairs;
-           usedArgs.set(pairNo);
 
diff --git a/third_party/stablehlo/workspace.bzl b/third_party/stablehlo/workspace.bzl
index 78d5c83d09367..849cedbe18da0 100644
--- a/third_party/stablehlo/workspace.bzl
+++ b/third_party/stablehlo/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
     # LINT.IfChange
-    STABLEHLO_COMMIT = "5217297204acb9e5a21e40fa825aa0769fb3c33f"
-    STABLEHLO_SHA256 = "9090e7f31420ef2bf5dc17a385f6e828f0072a79fa59d065c01c9a90e46ee730"
+    STABLEHLO_COMMIT = "b6406a43b48b7803f3efdbc235b1fbb5da782449"
+    STABLEHLO_SHA256 = "d2ecc5fe29f4a2fd17723f1f403109d24cafe8479f1e1a95d27d447c84a01976"
     # LINT.ThenChange(Google-internal path)
 
     tf_http_archive(
diff --git a/xla/mlir_hlo/tests/Dialect/mhlo/verifier_conv_op.mlir b/xla/mlir_hlo/tests/Dialect/mhlo/verifier_conv_op.mlir
index 5b22e609fa588..004f295e66001 100644
--- a/xla/mlir_hlo/tests/Dialect/mhlo/verifier_conv_op.mlir
+++ b/xla/mlir_hlo/tests/Dialect/mhlo/verifier_conv_op.mlir
@@ -330,8 +330,8 @@ func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>,
 
 // -----
 
-func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>,
-    %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> {
+func.func @invalid_conv_dimensions(%arg0: tensor<3x8x8x207xf32>,
+    %arg1: tensor<3x3x207x16xf32>) -> tensor<3x8x8x16xf32> {
   // expected-error@+1 {{expects output feature dimension size (16) to be a multiple of batch_group_count. Got batch_group_count = 3.}}
   %0 = mhlo.convolution(%arg0, %arg1)
          dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
@@ -342,8 +342,8 @@ func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>,
            feature_group_count = 1 : i64,
            precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]
          } :
-       (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32>
-  func.return %0 : tensor<1x8x8x16xf32>
+       (tensor<3x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<3x8x8x16xf32>
+  func.return %0 : tensor<3x8x8x16xf32>
 }
 
 // -----