Repository navigation
[pull] master from tensorflow:master - #8850
Merged
Merged
Conversation
Keeps both this branch's BlockLSTMGrad shape validation tests and the seq_len_max bounds tests added on master in the same region of rnn_cell_test.py.
…n deserialize Imported from GitHub PR openxla/xla#46773 ShardingParam::MinorToMajor::ToDeviceList computed a cumulative axis size product using 32-bit int with no overflow protection, and ShardingParam::verify() computed a dimension-shard product the same unguarded way before using it as a modulo divisor -- if the product wraps to exactly 0, that is a division by zero.\n\nConfirmed via a faithful standalone extraction of both arithmetic patterns, compiled with UBSan+ASan: both overflow with individually plausible input values, and the verify() case produces a real SIGFPE crash when the product wraps to 0.\n\nSeparately, ShardingParam::FromProto never called verify() at all, so a deserialized ShardingParam was never validated before its axis_sizes/dim_shards were used downstream.\n\nThis adds an overflow-safe int64_t accumulation with an explicit range check in verify(), calls verify() from FromProto, and hardens ToDeviceList to fail closed rather than silently propagate a wrapped value. Copybara import of the project: -- c015dcb95f89e46df503008a3de0333867b60dd7 by prasanna8585 <65734642+prasanna8585@users.noreply.github.com>: Fix unguarded int overflow in ShardingParam and validate on deserialize ShardingParam::MinorToMajor::ToDeviceList computed a cumulative axis size product using 32-bit int with no overflow protection (cum_size *= size), and ShardingParam::verify() computed a dimension- shard product the same unguarded way, then used it as a modulo divisor (NumDevices() % total_dim_shards_size) -- if the product wraps to exactly 0, this is a division by zero. Confirmed via a faithful standalone extraction of both arithmetic patterns, compiled with UBSan+ASan: both overflow with individually plausible input values, and the verify() case produces a real SIGFPE crash when the product wraps to 0. Separately, ShardingParam::FromProto never called verify() at all, so a deserialized ShardingParam (e.g. via xla/python/ifrt_proxy/server/ifrt_backend.cc's handling of MakeArraysFromHostBufferShards/Assemble requests) was never validated before its axis_sizes/dim_shards were used downstream. This adds an overflow-safe int64_t accumulation with an explicit range check in verify(), calls verify() from FromProto so deserialized input is actually validated, and hardens ToDeviceList to fail closed (CHECK) rather than silently propagate a wrapped value into device- index computation, on the assumption that verify() should have already caught it upstream. -- 961720e9258f1e54e674d12ecc821119f459cbf3 by prasanna8585 <65734642+prasanna8585@users.noreply.github.com>: Address review: move overflow check to verify(), use __builtin_mul_overflow, reject dim_shard<=0 - MinorToMajor::verify() now validates the axis_sizes cumulative product before it reaches ToDeviceList(), per @hyeontaek: validation should happen close to where unvalidated data enters, not repeated downstream. - Both overflow checks now use __builtin_mul_overflow instead of a post-hoc numeric_limits<int>::max() comparison, matching the pattern in xla/python/ifrt/shape.cc, so correctness doesn't depend on sizeof(int) < sizeof(int64_t). - ShardingParam::verify() now rejects dim_shard <= 0 up front, so the separate zero-divisor check before the modulo is no longer needed. -- 75a4bb962c7687082c65af78b290f4103f39bfb3 by prasanna8585 <65734642+prasanna8585@users.noreply.github.com>: Add unit test for ShardingParam::verify() overflow rejection Covers axis_sizes product overflow, dim_shards product overflow, non-positive dim_shard rejection, and the valid-input path, exercising ShardingParam::verify() directly rather than going through FromProto/deserialization. -- f038531ededf56b60638e9526a04117929169486 by prasanna8585 <prasanna8585@users.noreply.github.com>: Address review: use int directly, trim comments, drop redundant check hyeontaek's review, all five points: 1-3. cum_size and total_dim_shards_size are now plain int rather than int64_t. __builtin_mul_overflow checks overflow against its output argument's type directly, so once the output is int there is no need to widen the inputs to int64_t and separately compare against std::numeric_limits<int>::max() -- the builtin already does exactly that check when the output type is int. Verified directly: an isolated int-only build of this exact pattern correctly flags 2 * 1,500,000,000 (the real regression test's case) and the precise 46340/46341 boundary around INT_MAX, and correctly does not flag values within range. Removed the now-unused <limits> include from this file (the test file's own <limits> include is untouched, out of scope here). 2. Trimmed the three comments this PR added to state the same thing with less text. 4. sharding_param_test.cc: EXPECT_TRUE(sharding.verify().ok()) -> EXPECT_OK(sharding.verify()), matching the existing absl/status/status_matchers.h include already used by the rest of this file's assertions. 5. Removed the redundant overflow check in ToDeviceList(). Checked its only two real callers (sharding_spec.cc, sharding.cc) before removing it: both operate on an already-constructed ShardingParam member, consistent with the original commit's own stated design ("validate on deserialize") and with verify() being the sole enforcement point. cum_size is now plain int throughout, updated in the same place it's pushed, with no separate overflow check -- trusting the invariant verify() already established, per the review's own reasoning for why the duplicate check adds nothing. -- ef7d14e8662c4d4d8207890023cbeb66d67f205a by prasanna8585 <prasanna8585@users.noreply.github.com>: Address review: remove two comments explaining assumptions established elsewhere dimitar-asenov's feedback: 1. "Already rejected by verify() at deserialization time." above the cum_size *= size multiplication in ToDeviceList() was confusing -- the reviewer's own question ("What is rejected?") is the problem: the comment references an invariant enforced by a verify() call this function doesn't itself make, doesn't say which of the file's two verify() functions, and gives no way to confirm from this function's own scope that the invariant actually holds for every caller. Removed rather than clarified, since the underlying claim isn't something ToDeviceList() can locally establish or enforce. 2. "`total_dim_shards_size` is positive here, so no zero-divisor check is needed below." above the NumDevices() % total_dim_shards_size check in verify() -- explicit ask: delete it. As the reviewer put it, divisors not being zero is an assumption we don't generally need to explain, and if the code doesn't actually assure that on its own, that's a separate issue from this comment. No logic changed in either case -- both are comment-only deletions. -- f9ab8c868ba4a02a9d4f9d223dc9844f0a3c00b8 by prasanna8585 <prasanna8585@users.noreply.github.com>: Fix build: add missing tsl/platform/errors.h include thcmbs reported a integration failure: sharding_param_test.cc used tsl::error::INVALID_ARGUMENT (via the StatusIs matcher) without including the header that declares the tsl::error namespace, so it only resolved by chance in the environment this test was originally written against. hyeontaek's fix: add #include "xla/tsl/platform/errors.h". Also added the corresponding //xla/tsl/platform:errors Bazel dependency to the sharding_param_test target in BUILD -- the include alone isn't sufficient without it. Confirmed both the include ordering (alphabetical within the absl/llvm/xla group, matching this file's existing style) and the BUILD dep ordering (// targets before @ targets, matching every other target in this same file) against the existing conventions already used elsewhere in this same BUILD file and by sibling files (layout_test.cc, serdes.cc, serdes_test.cc) that already depend on this same header. Merging this change closes #46773 PiperOrigin-RevId: 986704240
…s FFI extension Imported from GitHub PR openxla/xla#49012 📝 Summary of Changes Adds a collective-memory window API to the backend-agnostic collectives FFI extension: `register_window` (Prepare) registers a batch of already-allocated buffers with the clique, and `get_window` (Init/Execute) returns an opaque `XLA_FFI_Window*` per buffer along with the byte offset of the queried address within that window. Provides a C++ wrapper (`ffi::Communicator::RegisterWindow` / `GetWindow` returning `WindowLookup{window, offset}`). Bumps the extension minor version to 2. 🎯 Justification Backend-agnostic API surface for external custom-call handlers to reuse XLA-managed collective memory windows instead of allocating and registering their own. The offset is what handlers pass to the backend's window accessors (e.g. `ncclGetPeerDevicePointer(win, offset, peer, &out)`). Split out from #48557 per review request; GPU wiring + tests remain in #48557 and will merge after this. 🚀 Kind of Contribution ✨ New Feature 📊 Benchmark (for Performance Improvements) N/A 🧪 Unit Tests: N/A in this PR — window API is exercised by `CollectiveOpsTestFFI.PublicApiWindow` in the GPU-wiring follow-up (#48557). 🧪 Execution Tests: N/A Copybara import of the project: -- f4a4dd69232de42f2049e0f30a1b9fa24f732ba2 by Sevin F. Varoglu <svaroglu@nvidia.com>: [XLA:FFI] Add RegisterWindow + GetWindow to the collectives FFI extension -- c35f0ded55c1e24354cca4aa7261f04b5e2c1958 by Sevin F. Varoglu <svaroglu@nvidia.com>: Rename WindowRegister and add unit tests -- 73bffbd02c09682fde3b1293773f79a67580a88a by Sevin F. Varoglu <svaroglu@nvidia.com>: Remove SupportsWindow Merging this change closes #49012 PiperOrigin-RevId: 986709932
…undaries Imported from GitHub PR openxla/xla#49155 ## 📝 Summary of Changes Command buffer conversion can split an async region and, in doing so, lose the ordering between operations queued on the same async stream. An async start/done pair runs its body on an async stream and joins it back to the main stream at the done. When a region cannot be captured (for example a collective that is disabled for command buffers), the pass still captured any convertible region nested inside it. That moves the nested body into a graph launched on the main stream, where nothing waits for the outer operation that is still running on the async stream. Notation: `CB{...}` is the set of thunks captured into one command buffer; minimum graph size is 1. **Example 1, nested pair on the same stream.** `A` and `B` both run on async stream `S`, so `A` must finish before `B` starts. `A` cannot be captured; `B` and the ordinary main-stream command `C` can. ```text Input: Start(A); Start(B); C; Done(B); Done(A) Before: Start(A); CB{Start(B); C; Done(B)}; Done(A) After: Start(A); Start(B); CB{C}; Done(B); Done(A) ``` Before, `B`'s body moves into the graph on the main stream, which has no wait for `A` to complete. After, `B` stays on stream `S` and `A → B` is preserved, while `C` still benefits from capture. **Example 2, nested pair on a different stream (unchanged).** Same shape, but `B` runs on stream `T`. `B`'s body only ever waited on the main stream, never on `S`, so capturing it loses no ordering. The fix is scoped to the stream of the outstanding operation, so this case keeps the behaviour it has today. The same holds after an unmatched start such as a pipelined recv: only its own stream is affected. ```text Input: Start(A on S); Start(B on T); C; Done(B); Done(A) Before: Start(A); CB{Start(B); C; Done(B)}; Done(A) After: Start(A); CB{Start(B); C; Done(B)}; Done(A) ``` **Example 3, loop body joining an outer operation.** `A` cannot be captured; its done sits inside a loop body together with `C`, which reads `A`'s result. With `WHILE` enabled, the body alone looks convertible. ```text Input: Start(A); While{ Done(A); C } Before: Start(A); CB{ While{ Done(A); C } } After: Start(A); While{ Done(A); CB{C} } ``` Inside a graph the done is a no-op, so before this change `C` could run before `A` finishes. After, the done stays a thunk and only `C` is captured. This affects partially convertible workloads, such as compute overlapped with collectives that are disabled for command buffers. Part 1 of 3: #49158 moves LHS dependency construction into the emitter; #48900 then preserves async fork/join ordering inside captured command buffers. ## 🚀 Kind of Contribution 🐛 Bug Fix ## 🧪 Unit Tests: `command_buffer_async_conversion_test` Copybara import of the project: -- ce5cde2dc8c619e58d6dcfd0220c8a523029c478 by Shawn Wang <shawnw@nvidia.com>: [XLA:GPU] Preserve async ordering across command buffer boundaries -- deb8a3494242f599c95b91385e6ef4b4c56a03d1 by Shawn Wang <shawnw@nvidia.com>: [XLA:GPU] Track open async streams during command buffer conversion Replace the positional open async region with a per-stream view. When a start stays a thunk, only its execution stream is open, and only until the done that joins it. Async regions on other streams can still be captured: their bodies only ever waited on the main stream, so a command buffer loses no ordering. A region is refused when it contains a nested start on an open stream, and loop or branch bodies are converted with the streams open at that point. Reject loop and branch bodies that contain an async done whose start is not in the same body. Capturing such a body whole would drop the join with the outstanding operation, while the top-level loop already keeps these thunks in place. -- 9986e62199ce6fb7662381bede9212bcbab09949 by Shawn Wang <shawnw@nvidia.com>: [XLA:GPU] Fix include order in command buffer conversion pass clang-format places tsl/ headers after xla/ headers; the presubmit flagged the block touched by the previous commits. -- 53d4322cc4233c1e4c181bb0e6a1b29cbf03637b by Shawn Wang <shawnw@nvidia.com>: [XLA:GPU] Tidy the command buffer async region handling Stop the nested-start walk at the first match instead of visiting the whole tree, inline the single-use region span helper, and let async start and done thunks fall through to the shared flush-and-append tail. Drop the unused status matcher include and dependency from the test. No behavior change. -- d9e5145b90e8f7d666ec4317359a57a978e1349e by Shawn Wang <shawnw@nvidia.com>: [XLA:GPU] Check for duplicate outstanding async starts in conversion AsyncExecution::Start fails at run time when an execution is started again before its done, so a thunk sequence with two outstanding starts on one execution can never run. Make that explicit in AsyncRegionSize with a DCHECK instead of folding the case silently into the "no closed region" result; optimized builds keep leaving such a region uncaptured. Document the two remaining reasons the function returns zero. Reconcile the AsyncExecution::Start comment with its counter check: pipelined send/recv alternates Start and Done on a shared execution; it does not start it twice before a done. Merging this change closes #49155 PiperOrigin-RevId: 986711311
The thunk-folding API now has enough moving parts that reading the headers is no longer a reasonable way to learn it: a handler author has to know when to prefer a native handler over XLA FFI, which of the three index spaces around a kernel launch a given number lives in, why buffers must be resolved through the emitter context rather than through the buffer assignment, and how to test the result without a GPU. Add a README covering all of that, with the `write_value_thunk_folded` demonstrator as a worked example, and state the two things the API deliberately does not do: validate operand element types, and let a handler decline an instruction and fall back to the FFI path. Also drop the pointer to `g3doc/native_custom_call_handlers.md` from `native_custom_call_handler_utils.h`. That path never existed. PiperOrigin-RevId: 986712400
…rsion Imported from GitHub PR openxla/xla#49199 This PR fixes `StreamExecutorGpuClientTest.PlatformVersionIsDerivedAtRuntime` for the oneAPI backend by adding `oneapi` to the list of accepted platform-version prefixes. It also fixes a crash when running all tests in `StreamExecutorGpuClientTest`, caused by an uncaught SYCL exception during `SyclStream` teardown. Copybara import of the project: -- 231dd5a0e5a3a6b10d2a39def262a1ae0c018aec by Bhavani Subramanian <bhavani1.subramanian@intel.com>: Fix unit test failure due to platform version on oneAPI backend Merging this change closes #49199 PiperOrigin-RevId: 986714208
PiperOrigin-RevId: 986716923
Imported from GitHub PR openxla/xla#48355 📝 Summary of Changes Replace gfx1250 with mi450 and add mi350 test backend 🚀 Kind of Contribution 📚 Documentation, 🧪 Tests Copybara import of the project: -- 7b2c0050c55bded52b02312ddc8fe0ce8e8e25e0 by Harsha Havanur Shamsundara <Harsha.HavanurShamsundara@amd.com>: Add mi350 backend and replace gfx1250 with mi450 -- 5d0781bfc582a2d31213686c0b7c0c4f7396ccba by Harsha Havanur Shamsundara <Harsha.HavanurShamsundara@amd.com>: Fix errors -- acb1275081c47d5c4883d2d0a8d5cb9ee84bf51f by Harsha Havanur Shamsundara <Harsha.HavanurShamsundara@amd.com>: Address review comments -- 8ed01ad35415530bad0d150644936de5a3e811ab by Harsha HS <Harsha.HavanurShamsundara@amd.com>: Fix lit tests to include mi300 and mi450 gpu configs -- 09e10b8fc26ad401a7ee0965895dc15c2b65e3d7 by Harsha HS <Harsha.HavanurShamsundara@amd.com>: Add mi300 config spec -- 3cb49cc2d1322bf33b7ca1143e8629444c17625d by Harsha HS <Harsha.HavanurShamsundara@amd.com>: Fix HBM3 and device memory size Merging this change closes #48355 PiperOrigin-RevId: 986736971
…utation evals from HloConstantFolding. Specialisation cache on hlo evaluator is thread-safe by mutex-locking. PiperOrigin-RevId: 986737880
Fixes jax-ml/jax#40849 PiperOrigin-RevId: 986738186
Migrate remaining callers of Comparison::GetType() and HloCompareInstruction::type() (including XlaBuilder::Compare, hlo_to_mhlo, mhlo_to_hlo, and TPU emitters) to use ComparisonOrder and PrimitiveType directly. Also stop setting the legacy comparison_type field in HloCompareInstruction::ToProto. PiperOrigin-RevId: 986754916
…o host compiler Imported from GitHub PR openxla/xla#49177 📝 Summary of Changes rocm_path and offload_archs compiler arguments were passed to host compiler which fails to compile host side code specifically external dependencies like boringssl or net_zstd while building TF. 🎯 Justification Remove these compiler flags that were passed to host compiler 🚀 Kind of Contribution: 🐛 Bug Fix Copybara import of the project: -- a65d3df05790142b5cb9cbffa6cbf500d4d9cd1e by Harsha Havanur Shamsundara <Harsha.HavanurShamsundara@amd.com>: Remove extrenious arguments like rocm_path passed to host compiler Merging this change closes #49177 PiperOrigin-RevId: 986776069
PiperOrigin-RevId: 986778048
…ss the entire async duration Imported from GitHub PR openxla/xla#48643 📝 Summary of Changes The live range of buffers in async ops need to reflect the duration of async-start and async-done. There are some corner cases where xla's buffer assignment would re-use buffers wrapped in the async ops for other parallel ops within the async-start/done interval which is incorrect. This could lead to silent data corruption. 🎯 Justification Incorrect live range leads to data corruption. 🚀 Kind of Contribution 🐛 Bug Fix 🧪 Unit Tests: added unit test. Copybara import of the project: -- 8b6b61fcb4dbd99a1ad3cbb4a935f58176d96f67 by TJ Xu <tjx@nvidia.com>: make sure the live range of buffers in async ops span across the entire async duration -- d85fd8e108c38dec2b820584f50ab907ff97cb18 by TJ Xu <tjx@nvidia.com>: fix cpu test failure -- 46d8fc531ff617309312f01ff4108209be9335d8 by TJ Xu <tjx@nvidia.com>: address pr comments -- 5a433749da77b3652dab1a307daa9cffcdc60166 by TJ Xu <tjx@nvidia.com>: address comment and fix a crash when computation has multiple users Merging this change closes #48643 PiperOrigin-RevId: 986790742
PiperOrigin-RevId: 986791525
PiperOrigin-RevId: 986801854
…tomCallTopK PiperOrigin-RevId: 986808802
PiperOrigin-RevId: 986848921
PiperOrigin-RevId: 986858417
Add `tpu.expand_load_vreg`, the physical (single vreg) expanded load, as the mirror of `tpu.compress_store_vreg`. The mask is mandatory and one dimensional over the lanes. Packed types keep a 2-D `(lane_count, packing)` vreg, so only dimension 0 has to match the mask. PiperOrigin-RevId: 986863897
MLIR's serialization APIs can't check for I/O errors themselves because raw_ostream doesn't have an API for it. So we have to do it ourselves. PiperOrigin-RevId: 986867518
they have different compute capabilities. Previously that was not an issue but new tiling is able to emit module with triton and autotuner tries to compile the model and fails with "Compute capability mismatch; expected {12.0}, but was {10.0a}"
PiperOrigin-RevId: 986876297
PiperOrigin-RevId: 986881086
PiperOrigin-RevId: 986883135
After GSPMD generates device local, it invokes passes such as DCE to clean up the code and then clears sharding. Previously, it first cleans up the code then clears sharding. This CL changes GSPMD to first clear sharding then clean up the code, so that the clean up passes will not see sharding as the code is already in device local representation. PiperOrigin-RevId: 986883751
PiperOrigin-RevId: 986914259
PiperOrigin-RevId: 986916502
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
19 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )