Integrate StableHLO at openxla/stablehlo@04c5a341

PiperOrigin-RevId: 726196621
openxla · Feb 12, 2025 · 9e845f0 · 9e845f0
1 parent 866ebe7
commit 9e845f0
Show file tree

Hide file tree

Showing 6 changed files with 73 additions and 5 deletions.
diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
@@ -1,3 +1,68 @@
+diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
+--- stablehlo/BUILD.bazel
++++ stablehlo/BUILD.bazel
+@@ -519,8 +519,6 @@
+         "stablehlo/conversions/linalg/transforms/StablehloToLinalgRandom.cpp",
+         "stablehlo/conversions/linalg/transforms/StablehloToLinalgReduce.cpp",
+         "stablehlo/conversions/linalg/transforms/TypeConversion.cpp",
+-        "stablehlo/transforms/conversions/TypeConversion.cpp",
+-        "stablehlo/transforms/conversions/TypeConversion.h",
+     ],
+     hdrs = [
+         "stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h",
+@@ -531,6 +529,7 @@
+     ],
+     strip_include_prefix = ".",
+     deps = [
++        "stablehlo_type_conversions",
+         ":base",
+         ":chlo_ops",
+         ":linalg_pass_inc_gen",
+@@ -943,6 +942,35 @@
+     ],
+ )
+
++# Header-only target, used when using the C API from a separate shared library.
++cc_library(
++    name = "stablehlo_dialect_capi_headers",
++    hdrs = STABLEHLO_DIALECT_CAPI_HEADERS,
++    strip_include_prefix = ".",
++    deps = [
++        "@llvm-project//mlir:CAPIIRHeaders",
++    ],
++)
++
++# Alwayslink target, used when exporting the C API from a shared library.
++cc_library(
++    name = "stablehlo_dialect_capi_objects",
++    srcs = STABLEHLO_DIALECT_CAPI_SOURCES,
++    hdrs = STABLEHLO_DIALECT_CAPI_HEADERS,
++    strip_include_prefix = ".",
++    deps = [
++        ":stablehlo_ops",
++        ":stablehlo_portable_api",
++        ":stablehlo_serialization",
++        ":version",
++        "@llvm-project//llvm:Support",
++        "@llvm-project//mlir:CAPIIRObjects",
++        "@llvm-project//mlir:IR",
++        "@llvm-project//mlir:Support",
++    ],
++    alwayslink = True,
++)
++
+ STABLEHLO_UNIFIED_CAPI_SOURCES = [
+     "stablehlo/integrations/c/StablehloPasses.cpp",
+     "stablehlo/integrations/c/StablehloUnifiedApi.cpp",
+@@ -1010,7 +1038,7 @@
+         ":linalg_passes",
+         ":reference_api",
+         ":reference_configuration",
+-        ":stablehlo_dialect_capi",
++        ":stablehlo_dialect_capi_objects",
+         ":stablehlo_ops",
+         ":stablehlo_passes",
+         ":stablehlo_portable_api",
 diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
 --- stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
 +++ stablehlo/stablehlo/conversions/tosa/tests/binary.mlir

diff --git 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 = "7775e3e2f98fe458f227cdc41aec567535d1c796"
-    STABLEHLO_SHA256 = "a79617c9b234cc8758a6c36a85beed0991bd5d3f927bd34cc0c395e9178ab228"
+    STABLEHLO_COMMIT = "04c5a341abe43f50ea2f8eca97fcc07763230c00"
+    STABLEHLO_SHA256 = "3eb4968cd57544850bc3b2ea2f605501ef43e9f245aee0bf0c141d4748c2d277"
     # LINT.ThenChange(Google-internal path)
 
     tf_http_archive(

diff --git a/xla/mlir_hlo/BUILD b/xla/mlir_hlo/BUILD
@@ -1113,6 +1113,7 @@ cc_library(
         "@stablehlo//:stablehlo_ops",
         "@stablehlo//:stablehlo_ops_inc_gen",
         "@stablehlo//:stablehlo_passes",
+        "@stablehlo//:stablehlo_passes_optimization",
         "@stablehlo//:stablehlo_type_inference",
     ],
 )
diff --git a/xla/mlir_hlo/stablehlo_ext/transforms/CMakeLists.txt b/xla/mlir_hlo/stablehlo_ext/transforms/CMakeLists.txt
@@ -36,6 +36,7 @@ add_mlir_dialect_library(StablehloExtensionPasses
   StablehloExtensionOps
   StablehloBase
   StablehloOps
+  StablehloOptimizationPasses
   StablehloPasses
   StablehloTypeInference
 )
diff --git a/xla/mlir_hlo/stablehlo_ext/transforms/sdy_refine_shapes.cpp b/xla/mlir_hlo/stablehlo_ext/transforms/sdy_refine_shapes.cpp
@@ -18,10 +18,9 @@ limitations under the License.
 #include <cstdint>
 #include <utility>
 
-#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -35,15 +34,16 @@ limitations under the License.
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "shardy/dialect/sdy/ir/dialect.h"
 #include "shardy/dialect/sdy/ir/utils.h"
 #include "stablehlo/dialect/Base.h"
 #include "stablehlo/dialect/ChloOps.h"
 #include "stablehlo/dialect/StablehloOps.h"
-#include "stablehlo/transforms/Passes.h"
 #include "stablehlo/transforms/StablehloRefineShapes.h"
+#include "stablehlo/transforms/optimization/Passes.h"
 #include "stablehlo_ext/transforms/stablehlo_refine_shapes.h"
 
 namespace mlir {

diff --git a/xla/mlir_hlo/stablehlo_ext/transforms/stablehlo_refine_shapes.cpp b/xla/mlir_hlo/stablehlo_ext/transforms/stablehlo_refine_shapes.cpp
@@ -30,6 +30,7 @@ limitations under the License.
 #include "stablehlo/dialect/TypeInference.h"
 #include "stablehlo/transforms/Passes.h"
 #include "stablehlo/transforms/StablehloRefineShapes.h"
+#include "stablehlo/transforms/optimization/Passes.h"
 #include "stablehlo_ext/IR/base.h"
 #include "stablehlo_ext/IR/stablehlo_ops.h"
 #include "stablehlo_ext/transforms/passes.h"  // NOLINT: Used in passes.h.inc