This PR intends to remove XLA as a bzlmod and transfer it as a non bzlmod dep. This is because this module will never be upstreamed as is, so keep it private. Also, we fetch llvm-raw and stablehlo from it, which is fine. While there, dummify the various local_config XLA symbols to please the imports, as we don't use those parts in ZML itself. Closes
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)
|
|
|