Tarry Singh
9488672d4b
Also: - Bump XLA deps : `com_github_grpc_grpc` and `com_google_protobuf` - Inject `rules_ml_toolchain` - Fix `zig_proto_library` rule
125 lines
5.5 KiB
Diff
125 lines
5.5 KiB
Diff
From 6078da86a46b6f0d983dccb9ae4f36fc90640247 Mon Sep 17 00:00:00 2001
|
|
From: Hugo Mano <hugo@zml.ai>
|
|
Date: Fri, 11 Jul 2025 14:05:16 +0200
|
|
Subject: [PATCH] zml patch
|
|
|
|
---
|
|
third_party/stablehlo/workspace.bzl | 1 +
|
|
third_party/stablehlo/zml.patch | 93 +++++++++++++++++++++++++++++
|
|
2 files changed, 94 insertions(+)
|
|
create mode 100644 third_party/stablehlo/zml.patch
|
|
|
|
diff --git a/third_party/stablehlo/workspace.bzl b/third_party/stablehlo/workspace.bzl
|
|
index d9d5063744..44980948d0 100644
|
|
--- a/third_party/stablehlo/workspace.bzl
|
|
+++ b/third_party/stablehlo/workspace.bzl
|
|
@@ -15,5 +15,6 @@ def repo():
|
|
urls = tf_mirror_urls("https://github.com/openxla/stablehlo/archive/{commit}.zip".format(commit = STABLEHLO_COMMIT)),
|
|
patch_file = [
|
|
"//third_party/stablehlo:temporary.patch", # Autogenerated, don't remove.
|
|
+ "//third_party/stablehlo:zml.patch", # Autogenerated, don't remove.
|
|
],
|
|
)
|
|
diff --git a/third_party/stablehlo/zml.patch b/third_party/stablehlo/zml.patch
|
|
new file mode 100644
|
|
index 0000000000..2a09384582
|
|
--- /dev/null
|
|
+++ b/third_party/stablehlo/zml.patch
|
|
@@ -0,0 +1,93 @@
|
|
+From e38ab68376dd8a17ebf4469d2c8350f521310182 Mon Sep 17 00:00:00 2001
|
|
+From: Hugo Mano <hugo@zml.ai>
|
|
+Date: Fri, 11 Jul 2025 12:08:35 +0200
|
|
+Subject: [PATCH] zml patch
|
|
+
|
|
+---
|
|
+ stablehlo/dialect/Serialization.cpp | 5 ++---
|
|
+ stablehlo/dialect/Serialization.h | 3 +--
|
|
+ stablehlo/integrations/c/StablehloDialectApi.cpp | 3 +--
|
|
+ stablehlo/integrations/c/StablehloDialectApi.h | 2 +-
|
|
+ stablehlo/tools/StablehloTranslateMain.cpp | 2 +-
|
|
+ 5 files changed, 6 insertions(+), 9 deletions(-)
|
|
+
|
|
+diff --git a/stablehlo/dialect/Serialization.cpp b/stablehlo/dialect/Serialization.cpp
|
|
+index cb89d673..4370d588 100644
|
|
+--- a/stablehlo/dialect/Serialization.cpp
|
|
++++ b/stablehlo/dialect/Serialization.cpp
|
|
+@@ -39,8 +39,7 @@ namespace stablehlo {
|
|
+
|
|
+ LogicalResult serializePortableArtifact(ModuleOp module,
|
|
+ StringRef targetVersion,
|
|
+- raw_ostream& os,
|
|
+- bool allowOtherDialects) {
|
|
++ raw_ostream& os) {
|
|
+ MLIRContext* context = module.getContext();
|
|
+
|
|
+ // Convert StableHLO --> VHLO.
|
|
+@@ -49,7 +48,7 @@ LogicalResult serializePortableArtifact(ModuleOp module,
|
|
+ {
|
|
+ PassManager pm(context);
|
|
+ StablehloLegalizeToVhloPassOptions options;
|
|
+- options.allowOtherDialects = allowOtherDialects;
|
|
++ options.allowOtherDialects = false;
|
|
+ pm.addPass(stablehlo::createStablehloLegalizeToVhloPass(options));
|
|
+ if (!succeeded(pm.run(module))) {
|
|
+ return failure();
|
|
+diff --git a/stablehlo/dialect/Serialization.h b/stablehlo/dialect/Serialization.h
|
|
+index 811ca97b..abe95e63 100644
|
|
+--- a/stablehlo/dialect/Serialization.h
|
|
++++ b/stablehlo/dialect/Serialization.h
|
|
+@@ -34,8 +34,7 @@ namespace stablehlo {
|
|
+ // unsupported dialects.
|
|
+ LogicalResult serializePortableArtifact(ModuleOp module,
|
|
+ StringRef targetVersion,
|
|
+- raw_ostream& os,
|
|
+- bool allowOtherDialects = false);
|
|
++ raw_ostream& os);
|
|
+
|
|
+ // Read StableHLO portable artifact
|
|
+ //
|
|
+diff --git a/stablehlo/integrations/c/StablehloDialectApi.cpp b/stablehlo/integrations/c/StablehloDialectApi.cpp
|
|
+index 343f8d0b..8f52e4d5 100644
|
|
+--- a/stablehlo/integrations/c/StablehloDialectApi.cpp
|
|
++++ b/stablehlo/integrations/c/StablehloDialectApi.cpp
|
|
+@@ -81,8 +81,7 @@ MlirLogicalResult stablehloSerializePortableArtifactFromModule(
|
|
+ MlirStringCallback callback, void *userData, bool allowOtherDialects) {
|
|
+ mlir::detail::CallbackOstream stream(callback, userData);
|
|
+ if (failed(mlir::stablehlo::serializePortableArtifact(
|
|
+- unwrap(moduleStr), unwrap(targetVersion), stream,
|
|
+- allowOtherDialects)))
|
|
++ unwrap(moduleStr), unwrap(targetVersion), stream)))
|
|
+ return mlirLogicalResultFailure();
|
|
+ return mlirLogicalResultSuccess();
|
|
+ }
|
|
+diff --git a/stablehlo/integrations/c/StablehloDialectApi.h b/stablehlo/integrations/c/StablehloDialectApi.h
|
|
+index 385156bf..24d11c1d 100644
|
|
+--- a/stablehlo/integrations/c/StablehloDialectApi.h
|
|
++++ b/stablehlo/integrations/c/StablehloDialectApi.h
|
|
+@@ -93,7 +93,7 @@ stablehloSerializePortableArtifactFromModule(MlirModule moduleStr,
|
|
+ MlirStringRef targetVersion,
|
|
+ MlirStringCallback callback,
|
|
+ void* userData,
|
|
+- bool allowOtherDialects = false);
|
|
++ bool allowOtherDialects);
|
|
+
|
|
+ // Read a StableHLO program from a portable artifact, returning the module as
|
|
+ // MLIR bytecode. Note, this bytecode returned is not a portable artifact,
|
|
+diff --git a/stablehlo/tools/StablehloTranslateMain.cpp b/stablehlo/tools/StablehloTranslateMain.cpp
|
|
+index fdf0d6a9..8d5c8752 100644
|
|
+--- a/stablehlo/tools/StablehloTranslateMain.cpp
|
|
++++ b/stablehlo/tools/StablehloTranslateMain.cpp
|
|
+@@ -323,7 +323,7 @@ TranslateFromMLIRRegistration serializeRegistration(
|
|
+ }
|
|
+
|
|
+ return stablehlo::serializePortableArtifact(
|
|
+- module, targetVersion, os, allowOtherDialectsOption.getValue());
|
|
++ module, targetVersion, os);
|
|
+ },
|
|
+ [](DialectRegistry ®istry) {
|
|
+ mlir::registerAllDialects(registry);
|
|
+--
|
|
+2.39.5 (Apple Git-154)
|
|
+
|
|
--
|
|
2.39.5 (Apple Git-154)
|
|
|