From 637a27b412016554052244eeace8d86145d407a3 Mon Sep 17 00:00:00 2001 From: Paul Berg Date: Thu, 2 Jul 2026 12:38:19 +0200 Subject: [PATCH 01/12] Update gb25_commit branch from 'main' to 'gw/up' --- .github/workflows/test-gb-25.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test-gb-25.yml b/.github/workflows/test-gb-25.yml index 9ca71f8933..59082c9c00 100644 --- a/.github/workflows/test-gb-25.yml +++ b/.github/workflows/test-gb-25.yml @@ -48,11 +48,11 @@ jobs: - '' # - 'b25f3cbed2bc88c8ffef85f6a5319e2cf7b0454c' gb25_commit: - - 'main' + - 'gw/up' reactant_commit: - 'main' - - 'multipad' + # - 'multipad' grid: - '--float-type=Float64' - '--float-type=Float32' From 862ba91caeac882e6996b02c9a68c6942cfa4002 Mon Sep 17 00:00:00 2001 From: Paul Berg Date: Thu, 2 Jul 2026 12:57:12 +0200 Subject: [PATCH 02/12] Update gb25_commit reference from 'gw/up' to 'glw/up' --- .github/workflows/test-gb-25.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-gb-25.yml b/.github/workflows/test-gb-25.yml index 59082c9c00..3c9743e9b7 100644 --- a/.github/workflows/test-gb-25.yml +++ b/.github/workflows/test-gb-25.yml @@ -48,7 +48,7 @@ jobs: - '' # - 'b25f3cbed2bc88c8ffef85f6a5319e2cf7b0454c' gb25_commit: - - 'gw/up' + - 'glw/up' reactant_commit: - 'main' From db6ac8c2ef690ce062a753a459da1787fe5201cf Mon Sep 17 00:00:00 2001 From: Paul Berg Date: Fri, 3 Jul 2026 20:42:50 +0200 Subject: [PATCH 03/12] Update reactant_commit branch in workflow file --- .github/workflows/test-gb-25.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-gb-25.yml b/.github/workflows/test-gb-25.yml index 3c9743e9b7..0e51682300 100644 --- a/.github/workflows/test-gb-25.yml +++ b/.github/workflows/test-gb-25.yml @@ -51,7 +51,7 @@ jobs: - 'glw/up' reactant_commit: - - 'main' + - 'pb/aliasing-reads' # - 'multipad' grid: - '--float-type=Float64' From 0bb63665b7bfeec95e52fba3d018b79fa7cc9bc6 Mon Sep 17 00:00:00 2001 From: Paul Berg <9824244+Pangoraw@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:03:50 -0500 Subject: [PATCH 04/12] Patch XLA SPMD partitioner DUS-over-Pad fix for barotropic V=0 bug Applies the fix from openxla/xla#45444 (spmd-enzyme: Be more conservative when walking constant DUS chain) on top of the pinned XLA commit. ProcessUpdatePieceExtractOperand's constant-DUS unwind loop was unconditionally descending into a nested DynamicUpdateSlice's update operand even when that update only partially covered the piece, discarding the real data in the DUS's base operand. This is the root cause traced for the sharded V=0 correctness bug (V = integral v dz coming out zero under Y-sharding of the NormalFlow barotropic velocity field's halo-fill pad). Co-Authored-By: Claude Sonnet 5 --- patches/xla_spmd_dus.patch | 71 +++++++++++++++++++++++++++++++++++ third_party/xla/workspace.bzl | 2 +- 2 files changed, 72 insertions(+), 1 deletion(-) create mode 100644 patches/xla_spmd_dus.patch diff --git a/patches/xla_spmd_dus.patch b/patches/xla_spmd_dus.patch new file mode 100644 index 0000000000..54146e7dc0 --- /dev/null +++ b/patches/xla_spmd_dus.patch @@ -0,0 +1,71 @@ +diff --git a/xla/service/spmd/spmd_partitioner.cc b/xla/service/spmd/spmd_partitioner.cc +index bb88694056..6759f3d3cf 100644 +--- a/xla/service/spmd/spmd_partitioner.cc ++++ b/xla/service/spmd/spmd_partitioner.cc +@@ -4150,7 +4150,6 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( + return b_.AddInstruction(std::move(to_add)); + }; + +- std::vector accumulated_offsets(hlo->shape().dimensions().size(), 0); + const HloInstruction* actual_update = piece_update_tensor; + std::optional reshape_desc; + while (actual_update->user_count() == 1) { +@@ -4172,35 +4171,31 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( + } + } else if (actual_update->opcode() == HloOpcode::kDynamicUpdateSlice && + !reshape_desc.has_value()) { ++ const HloInstruction* update_operand = actual_update->operand(1); ++ bool update_covers_whole_piece = true; + bool all_constant = true; +- for (int i = 2; i < actual_update->operand_count(); ++i) { +- if (actual_update->operand(i)->opcode() != HloOpcode::kConstant) { +- all_constant = false; +- break; +- } +- } +- if (!all_constant) { +- break; +- } + + for (int i = 0; i < actual_update->shape().dimensions().size(); ++i) { +- auto const_op = +- DynCast(actual_update->operand(2 + i)); +- auto val = const_op->literal().GetIntegralAsS64({}); +- +- if (!val.has_value()) { ++ if (actual_update->operand(2 + i)->opcode() != HloOpcode::kConstant) { + all_constant = false; + break; + } +- +- accumulated_offsets[i] += *val; ++ if (update_operand->shape().dimensions(i) != ++ actual_update->shape().dimensions(i)) { ++ update_covers_whole_piece = false; ++ break; ++ } + } + + if (!all_constant) { + break; + } + +- actual_update = actual_update->operand(1); ++ if (!update_covers_whole_piece) { ++ break; ++ } ++ ++ actual_update = update_operand; + } else { + break; + } +@@ -4421,8 +4416,7 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( + for (int64_t i = 0; i < hlo->shape().dimensions().size(); ++i) { + int64_t pre_dim = post_to_pre[i]; + if (pre_dim != -1 && ShardCountAtDim(hlo_sharding, i) > 1) { +- int64_t dus_start = +- piece_dus_starts[i] + accumulated_offsets[i]; ++ int64_t dus_start = piece_dus_starts[i]; + int64_t slice_start = slice->slice_starts(pre_dim); + int64_t slice_size = slice->shape().dimensions(pre_dim); + if (dus_start != slice_start) { diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl index 6f0e7d6b96..823f565f0c 100644 --- a/third_party/xla/workspace.bzl +++ b/third_party/xla/workspace.bzl @@ -17,6 +17,6 @@ def repo(extra_patches = [], override_commit = ""): strip_prefix = "xla-{commit}".format(commit = commit), urls = ["https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = commit)], patch_cmds = XLA_PATCHES + extra_patches, - patches = ["//:patches/xla.patch", "//:patches/xla_win.patch", "//:patches/xla_trainium.patch"], + patches = ["//:patches/xla.patch", "//:patches/xla_win.patch", "//:patches/xla_trainium.patch", "//:patches/xla_spmd_dus.patch"], patch_args = ["-p1"], ) From 60e9663d23ddc00d972f5ab5da195936991d0971 Mon Sep 17 00:00:00 2001 From: Paul Berg Date: Mon, 13 Jul 2026 16:56:54 +0200 Subject: [PATCH 05/12] Update test-gb-25.yml --- .github/workflows/test-gb-25.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-gb-25.yml b/.github/workflows/test-gb-25.yml index 0e51682300..3c9743e9b7 100644 --- a/.github/workflows/test-gb-25.yml +++ b/.github/workflows/test-gb-25.yml @@ -51,7 +51,7 @@ jobs: - 'glw/up' reactant_commit: - - 'pb/aliasing-reads' + - 'main' # - 'multipad' grid: - '--float-type=Float64' From a0bf96fd07efb4ad3831e6a782aa12ad510c5e65 Mon Sep 17 00:00:00 2001 From: Paul Berg Date: Mon, 13 Jul 2026 17:28:02 +0200 Subject: [PATCH 06/12] Update workspace.bzl --- third_party/xla/workspace.bzl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl index 823f565f0c..6f0e7d6b96 100644 --- a/third_party/xla/workspace.bzl +++ b/third_party/xla/workspace.bzl @@ -17,6 +17,6 @@ def repo(extra_patches = [], override_commit = ""): strip_prefix = "xla-{commit}".format(commit = commit), urls = ["https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = commit)], patch_cmds = XLA_PATCHES + extra_patches, - patches = ["//:patches/xla.patch", "//:patches/xla_win.patch", "//:patches/xla_trainium.patch", "//:patches/xla_spmd_dus.patch"], + patches = ["//:patches/xla.patch", "//:patches/xla_win.patch", "//:patches/xla_trainium.patch"], patch_args = ["-p1"], ) From 02118941594f3e74e02222c618b6a9a26bb46ebe Mon Sep 17 00:00:00 2001 From: Paul Berg Date: Mon, 13 Jul 2026 17:28:13 +0200 Subject: [PATCH 07/12] Delete patches/xla_spmd_dus.patch --- patches/xla_spmd_dus.patch | 71 -------------------------------------- 1 file changed, 71 deletions(-) delete mode 100644 patches/xla_spmd_dus.patch diff --git a/patches/xla_spmd_dus.patch b/patches/xla_spmd_dus.patch deleted file mode 100644 index 54146e7dc0..0000000000 --- a/patches/xla_spmd_dus.patch +++ /dev/null @@ -1,71 +0,0 @@ -diff --git a/xla/service/spmd/spmd_partitioner.cc b/xla/service/spmd/spmd_partitioner.cc -index bb88694056..6759f3d3cf 100644 ---- a/xla/service/spmd/spmd_partitioner.cc -+++ b/xla/service/spmd/spmd_partitioner.cc -@@ -4150,7 +4150,6 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( - return b_.AddInstruction(std::move(to_add)); - }; - -- std::vector accumulated_offsets(hlo->shape().dimensions().size(), 0); - const HloInstruction* actual_update = piece_update_tensor; - std::optional reshape_desc; - while (actual_update->user_count() == 1) { -@@ -4172,35 +4171,31 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( - } - } else if (actual_update->opcode() == HloOpcode::kDynamicUpdateSlice && - !reshape_desc.has_value()) { -+ const HloInstruction* update_operand = actual_update->operand(1); -+ bool update_covers_whole_piece = true; - bool all_constant = true; -- for (int i = 2; i < actual_update->operand_count(); ++i) { -- if (actual_update->operand(i)->opcode() != HloOpcode::kConstant) { -- all_constant = false; -- break; -- } -- } -- if (!all_constant) { -- break; -- } - - for (int i = 0; i < actual_update->shape().dimensions().size(); ++i) { -- auto const_op = -- DynCast(actual_update->operand(2 + i)); -- auto val = const_op->literal().GetIntegralAsS64({}); -- -- if (!val.has_value()) { -+ if (actual_update->operand(2 + i)->opcode() != HloOpcode::kConstant) { - all_constant = false; - break; - } -- -- accumulated_offsets[i] += *val; -+ if (update_operand->shape().dimensions(i) != -+ actual_update->shape().dimensions(i)) { -+ update_covers_whole_piece = false; -+ break; -+ } - } - - if (!all_constant) { - break; - } - -- actual_update = actual_update->operand(1); -+ if (!update_covers_whole_piece) { -+ break; -+ } -+ -+ actual_update = update_operand; - } else { - break; - } -@@ -4421,8 +4416,7 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( - for (int64_t i = 0; i < hlo->shape().dimensions().size(); ++i) { - int64_t pre_dim = post_to_pre[i]; - if (pre_dim != -1 && ShardCountAtDim(hlo_sharding, i) > 1) { -- int64_t dus_start = -- piece_dus_starts[i] + accumulated_offsets[i]; -+ int64_t dus_start = piece_dus_starts[i]; - int64_t slice_start = slice->slice_starts(pre_dim); - int64_t slice_size = slice->shape().dimensions(pre_dim); - if (dus_start != slice_start) { From 376864d1a4a413e7f939a622e2a4eab23c094b14 Mon Sep 17 00:00:00 2001 From: Paul Date: Mon, 13 Jul 2026 23:40:11 +0200 Subject: [PATCH 08/12] Revert "Delete patches/xla_spmd_dus.patch" This reverts commit 02118941594f3e74e02222c618b6a9a26bb46ebe. --- patches/xla_spmd_dus.patch | 71 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 patches/xla_spmd_dus.patch diff --git a/patches/xla_spmd_dus.patch b/patches/xla_spmd_dus.patch new file mode 100644 index 0000000000..54146e7dc0 --- /dev/null +++ b/patches/xla_spmd_dus.patch @@ -0,0 +1,71 @@ +diff --git a/xla/service/spmd/spmd_partitioner.cc b/xla/service/spmd/spmd_partitioner.cc +index bb88694056..6759f3d3cf 100644 +--- a/xla/service/spmd/spmd_partitioner.cc ++++ b/xla/service/spmd/spmd_partitioner.cc +@@ -4150,7 +4150,6 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( + return b_.AddInstruction(std::move(to_add)); + }; + +- std::vector accumulated_offsets(hlo->shape().dimensions().size(), 0); + const HloInstruction* actual_update = piece_update_tensor; + std::optional reshape_desc; + while (actual_update->user_count() == 1) { +@@ -4172,35 +4171,31 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( + } + } else if (actual_update->opcode() == HloOpcode::kDynamicUpdateSlice && + !reshape_desc.has_value()) { ++ const HloInstruction* update_operand = actual_update->operand(1); ++ bool update_covers_whole_piece = true; + bool all_constant = true; +- for (int i = 2; i < actual_update->operand_count(); ++i) { +- if (actual_update->operand(i)->opcode() != HloOpcode::kConstant) { +- all_constant = false; +- break; +- } +- } +- if (!all_constant) { +- break; +- } + + for (int i = 0; i < actual_update->shape().dimensions().size(); ++i) { +- auto const_op = +- DynCast(actual_update->operand(2 + i)); +- auto val = const_op->literal().GetIntegralAsS64({}); +- +- if (!val.has_value()) { ++ if (actual_update->operand(2 + i)->opcode() != HloOpcode::kConstant) { + all_constant = false; + break; + } +- +- accumulated_offsets[i] += *val; ++ if (update_operand->shape().dimensions(i) != ++ actual_update->shape().dimensions(i)) { ++ update_covers_whole_piece = false; ++ break; ++ } + } + + if (!all_constant) { + break; + } + +- actual_update = actual_update->operand(1); ++ if (!update_covers_whole_piece) { ++ break; ++ } ++ ++ actual_update = update_operand; + } else { + break; + } +@@ -4421,8 +4416,7 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( + for (int64_t i = 0; i < hlo->shape().dimensions().size(); ++i) { + int64_t pre_dim = post_to_pre[i]; + if (pre_dim != -1 && ShardCountAtDim(hlo_sharding, i) > 1) { +- int64_t dus_start = +- piece_dus_starts[i] + accumulated_offsets[i]; ++ int64_t dus_start = piece_dus_starts[i]; + int64_t slice_start = slice->slice_starts(pre_dim); + int64_t slice_size = slice->shape().dimensions(pre_dim); + if (dus_start != slice_start) { From 804ec5abbcaaa520290d87f59ac3e66ad7a0cb9c Mon Sep 17 00:00:00 2001 From: Paul Date: Mon, 13 Jul 2026 23:40:25 +0200 Subject: [PATCH 09/12] Revert "Update workspace.bzl" This reverts commit a0bf96fd07efb4ad3831e6a782aa12ad510c5e65. --- third_party/xla/workspace.bzl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl index 6f0e7d6b96..823f565f0c 100644 --- a/third_party/xla/workspace.bzl +++ b/third_party/xla/workspace.bzl @@ -17,6 +17,6 @@ def repo(extra_patches = [], override_commit = ""): strip_prefix = "xla-{commit}".format(commit = commit), urls = ["https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = commit)], patch_cmds = XLA_PATCHES + extra_patches, - patches = ["//:patches/xla.patch", "//:patches/xla_win.patch", "//:patches/xla_trainium.patch"], + patches = ["//:patches/xla.patch", "//:patches/xla_win.patch", "//:patches/xla_trainium.patch", "//:patches/xla_spmd_dus.patch"], patch_args = ["-p1"], ) From ff96737784f85118c910dd3068d4eb74ae2f97b3 Mon Sep 17 00:00:00 2001 From: Paul Date: Mon, 13 Jul 2026 23:43:53 +0200 Subject: [PATCH 10/12] Reapply "Update workspace.bzl" This reverts commit 804ec5abbcaaa520290d87f59ac3e66ad7a0cb9c. --- third_party/xla/workspace.bzl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl index 823f565f0c..6f0e7d6b96 100644 --- a/third_party/xla/workspace.bzl +++ b/third_party/xla/workspace.bzl @@ -17,6 +17,6 @@ def repo(extra_patches = [], override_commit = ""): strip_prefix = "xla-{commit}".format(commit = commit), urls = ["https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = commit)], patch_cmds = XLA_PATCHES + extra_patches, - patches = ["//:patches/xla.patch", "//:patches/xla_win.patch", "//:patches/xla_trainium.patch", "//:patches/xla_spmd_dus.patch"], + patches = ["//:patches/xla.patch", "//:patches/xla_win.patch", "//:patches/xla_trainium.patch"], patch_args = ["-p1"], ) From ff91bd1d6b660f2fa08d2f7298f6ae55c54f8ee3 Mon Sep 17 00:00:00 2001 From: Paul Date: Mon, 13 Jul 2026 23:44:07 +0200 Subject: [PATCH 11/12] Repply "Delete patches/xla_spmd_dus.patch" This reverts commit 376864d1a4a413e7f939a622e2a4eab23c094b14. --- patches/xla_spmd_dus.patch | 71 -------------------------------------- 1 file changed, 71 deletions(-) delete mode 100644 patches/xla_spmd_dus.patch diff --git a/patches/xla_spmd_dus.patch b/patches/xla_spmd_dus.patch deleted file mode 100644 index 54146e7dc0..0000000000 --- a/patches/xla_spmd_dus.patch +++ /dev/null @@ -1,71 +0,0 @@ -diff --git a/xla/service/spmd/spmd_partitioner.cc b/xla/service/spmd/spmd_partitioner.cc -index bb88694056..6759f3d3cf 100644 ---- a/xla/service/spmd/spmd_partitioner.cc -+++ b/xla/service/spmd/spmd_partitioner.cc -@@ -4150,7 +4150,6 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( - return b_.AddInstruction(std::move(to_add)); - }; - -- std::vector accumulated_offsets(hlo->shape().dimensions().size(), 0); - const HloInstruction* actual_update = piece_update_tensor; - std::optional reshape_desc; - while (actual_update->user_count() == 1) { -@@ -4172,35 +4171,31 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( - } - } else if (actual_update->opcode() == HloOpcode::kDynamicUpdateSlice && - !reshape_desc.has_value()) { -+ const HloInstruction* update_operand = actual_update->operand(1); -+ bool update_covers_whole_piece = true; - bool all_constant = true; -- for (int i = 2; i < actual_update->operand_count(); ++i) { -- if (actual_update->operand(i)->opcode() != HloOpcode::kConstant) { -- all_constant = false; -- break; -- } -- } -- if (!all_constant) { -- break; -- } - - for (int i = 0; i < actual_update->shape().dimensions().size(); ++i) { -- auto const_op = -- DynCast(actual_update->operand(2 + i)); -- auto val = const_op->literal().GetIntegralAsS64({}); -- -- if (!val.has_value()) { -+ if (actual_update->operand(2 + i)->opcode() != HloOpcode::kConstant) { - all_constant = false; - break; - } -- -- accumulated_offsets[i] += *val; -+ if (update_operand->shape().dimensions(i) != -+ actual_update->shape().dimensions(i)) { -+ update_covers_whole_piece = false; -+ break; -+ } - } - - if (!all_constant) { - break; - } - -- actual_update = actual_update->operand(1); -+ if (!update_covers_whole_piece) { -+ break; -+ } -+ -+ actual_update = update_operand; - } else { - break; - } -@@ -4421,8 +4416,7 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand( - for (int64_t i = 0; i < hlo->shape().dimensions().size(); ++i) { - int64_t pre_dim = post_to_pre[i]; - if (pre_dim != -1 && ShardCountAtDim(hlo_sharding, i) > 1) { -- int64_t dus_start = -- piece_dus_starts[i] + accumulated_offsets[i]; -+ int64_t dus_start = piece_dus_starts[i]; - int64_t slice_start = slice->slice_starts(pre_dim); - int64_t slice_size = slice->shape().dimensions(pre_dim); - if (dus_start != slice_start) { From 251d96e6e2d4e8a3b939155d98cf28dfd2942a54 Mon Sep 17 00:00:00 2001 From: Paul Date: Mon, 13 Jul 2026 23:44:28 +0200 Subject: [PATCH 12/12] update jax --- workspace.bzl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workspace.bzl b/workspace.bzl index d5b0d60b8e..da33c83d83 100644 --- a/workspace.bzl +++ b/workspace.bzl @@ -1,4 +1,4 @@ -JAX_COMMIT = "0db0a8b7e1f88f89fdec3bd2b0ce488e4cd2393f" +JAX_COMMIT = "6b606dd6aeb73adb0c73c3ee2649b1104feb4e1f" JAX_SHA256 = "" ENZYME_COMMIT = "e181366ce7fcd798ac4e04fcc87bf716eb9b5f95"