[pull] master from tensorflow:master - #8800
Merged
Merged
Conversation
When axis=-1 and range_given=true, QuantizeAndDequantizeV3Op called .scalar<T>() on input_min/input_max without first verifying they are scalars. This triggered a CHECK failure (process crash) instead of raising an InvalidArgumentError. The V2 op already had this validation. Fixes #99458
The MaxPoolBackward GPU kernel at maxpooling_op_gpu.cu.cc:213 writes to `bottom_diff + offset + mask[index]` where `mask` comes from the user-supplied argmax tensor with no bounds checking. This allows an attacker-controlled OOB write via GpuAtomicAdd. The same file's MaxPoolGradBackward kernel (lines 332-350) already validates `read_index >= 0 && read_index < input_size` for the equivalent operation, confirming this was a missed check. On CPU, the equivalent code at maxpooling_op.cc:1090 uses CHECK() which kills the process with SIGABRT instead of returning an error. Fix: - GPU: Add `input_size` parameter to MaxPoolBackward kernel and validate `write_index >= 0 && write_index < input_size` before the GpuAtomicAdd, matching the existing GradGrad kernel pattern. - CPU: Replace CHECK() with a bounds check that skips invalid indices, preventing process crash on malformed argmax values.
Verifies that out-of-bounds argmax indices no longer crash the process on CPU (previously a CHECK failure / SIGABRT) or trigger silent OOB writes on GPU. Follow-up to the fix in 7003dde addressing dmiltr3's review request.
ScatterNdOp::Compute and TensorScatterOp::Compute index updates.shape().dim_size(i) for every outer dimension of indices without first checking that updates has that many dimensions. When the rank of updates is smaller than the number of outer dimensions of indices (e.g. indices=[4,1,1], updates=[4]), TensorShape::dim_size hits CHECK(d < dims()) and aborts the process instead of returning an error. Add a rank guard before the loop so both ops return InvalidArgumentError, matching the validation already done in ValidateScatterNdUpdateShape on the deeper kernel path, and add regression tests. Fixes #93680
UnbatchResource::Compute read batch_index_t.shape().dim_size(1) without checking that batch_index is a rank-2 matrix. A rank-1 batch_index aborted the process with a fatal CHECK failure (d < dims()) instead of raising a catchable InvalidArgumentError. A scalar batch_index took a different broken path: dim_size(0) returned an arbitrary garbage value, producing a nonsensical error message such as "Expected 0th dimension size to be no greater than 1; Got: 24040". Validate that batch_index is a matrix before any dim_size access. The existing testUnbatchInvalidIdArg passed a rank-3 batch_index that only incidentally survived the old checks; it now uses a valid rank-2 index so it still exercises the id validation it was written for. Fixes #104846
The internal presubmits run these regression tests under XLA compilation,
where tf2xla lowers the op itself and reports the shape mismatch with a
different message ("Must have updates.shape = ...") than the CPU/GPU
kernel ("rank at least ...") or graph-mode shape inference ("must
match"). Accept all three so the assertion holds on every execution path.
Also run the tf.scatter_nd regression test in graph mode as well as
eager, matching the tensor_scatter_update one, so graph-mode shape
inference is covered too.
ReverseV2 returned the input unchanged for scalar and empty tensors before looking at the axes, so out-of-range axes went unreported in eager execution while shape inference rejected the same call in graph mode. A scalar has no valid axis at all, yet tf.reverse(scalar, axis=[1,2,3]) returned the scalar; an empty tensor with an out-of-range axis behaved the same way. Validate the axes before taking the shortcut for inputs that have nothing to reverse, then return early only for the data movement. The checks and their messages are unchanged, so eager execution now reports what graph mode already did. An empty axis list stays valid for every input, and in-range axes on empty tensors keep working. Most of the diff is re-indentation from removing the enclosing else block. Fixes #110038
…ource. PiperOrigin-RevId: 974589473
… error. PiperOrigin-RevId: 974622549
…) and batch group convolutions (`batch_group_count > 1`) in Mosaic TPU and JAX Pallas. This is Part 3 of a 3-part changelist chain: - Part 1 (cl/964822765): MLIR `TPU_ConvOp` dialect definition, verification, Python bindings, and basic N-D convolution lowering. - Part 2 (cl/969761722): Dilations (input/LHS and kernel/RHS), negative padding (cropping), and strided slicing emulation. - Part 3 (this CL): Grouped convolutions (`feature_group_count > 1`, `batch_group_count > 1`). PiperOrigin-RevId: 974631914
…ctivesFacade Imported from GitHub PR openxla/xla#47176 ## Summary This PR routes every MORI GPU collective in `MoriCommunicator` through a single, dedicated `CollectivesFacade` entry point instead of the previous ad-hoc per-op plumbing. It also introduces a thin shared header and an inert stub facade so the XLA side compiles, links, and can be committed **without** depending on the external `@roc_mori` library. The facade exposes a **non-templated, enum-based** public API: `Run*` methods take `mori::collective::DataType` and `ReduceOpKind` enums rather than `<T, Op>` template arguments. Host translation units (the communicator) see decl-only method declarations, while a single device translation unit (`mori_kernels.cu.cc`, compiled as HIP) defines `MORI_KERNELS_IMPL` before including the facade and thereby compiles the method definitions exactly once. This replaces the earlier `template <class = void>` + explicit-instantiation machinery. ## Motivation - Consolidate all collective launches (all-reduce, reduce-scatter, all-gather, all-to-all, barrier, send/recv, collective-permute, quiet, fence) behind one facade owned by the communicator, so staging buffers and group counters have a single, clearly-scoped owner and lifetime. - Present a clean compiled boundary: a small non-templated public API instead of template dispatch leaking through the header, no explicit-instantiation lists to maintain, and the type/op dispatch localized inside the facade. - Decouple the XLA build from the MORI source tree so this wiring can land independently, with a compile-time switch to bring in the real device facade later. ## What changed ### `mori_communicator.cc` / `mori_communicator.h` - Each communicator owns a `std::unique_ptr<mori::collective::CollectivesFacade>`, created in `MoriCommunicator::Create` (records rank identity and allocates the symmetric-heap staging buffer). The `unique_ptr` frees the staging/counters via the facade destructor before `ShmemFinalize`. - All `Launch*` paths now call `facade_->Run*` and convert the returned `hipError_t` to `absl::Status` via `se::gpu::ToStatus`. - Reduction dispatch is now enum translation, not macro/switch-key generation: `ToMoriDataType(PrimitiveType)` and `ToMoriReduceOp(ReductionKind)` map XLA enums to the facade enums (returning `Unimplemented` for unsupported dtypes), then `LaunchAllReduce` / `LaunchReduceScatter` call the non-templated `facade_->RunAllReduce(...)` / `RunReduceScatter(...)` directly. The old `MoriRedKey` key-packing and `MORI_FOR_EACH_DTYPE` x `MORI_FOR_EACH_OP` case tables are gone. - Supported dtypes include the OCP fp8 types `F8E5M2` and `F8E4M3FN` alongside F16/BF16/S8/U8/S32/U32/S64/U64/F32/F64. - Removed the old `P2P`/`P2PType` indirection; `Send`/`Recv` now go straight through `LaunchSend`/`LaunchRecv`. `Send` dropped its unused `recv_buffer` parameter and `Recv` its unused `send_buffer` parameter. - Stream handle helper switched from `AsRocmStream` (intptr) to `AsHipStream` (`hipStream_t`); `ToStream` moved into an anonymous namespace. - `MoriCommunicator::Create` now validates `num_ranks > 0`. ### `mori_kernels.h` (new) - Thin seam included by both the host communicator (decl-only) and the device TU. It selects the facade implementation: by default includes the inert stub (`mori_stub.h`); defining `XLA_GPU_USE_REAL_MORI` pulls in the real `mori/collective/collectives_facade.hpp` instead. - No longer defines any dtype/op expansion macros; the type/op dispatch now lives inside the facade. ### `mori_kernels.cu.cc` (new) - Single HIP device TU. It `#define`s `MORI_KERNELS_IMPL` and includes `mori_kernels.h`, which compiles the facade's device path (kernels + non-templated `Run*` definitions) exactly once and emits the symbols the host references. No explicit template instantiations. In the default (stub) build it is a harmless near-empty TU. ### `mori_stub.h` - Inert, header-only `mori::collective::CollectivesFacade` mirroring the real facade's non-templated enum API. Defines the `DataType` (incl. `F8E5M2`/`F8E4M3FN`) and `ReduceOpKind` enums and the `AddressVector` alias. `Create` returns a valid empty facade; every `Run*` is a no-op returning `hipSuccess`. Copybara import of the project: -- a498a2ead52c3aba6721f5edd85d48b6faa6dce3 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: wired all MORI collectives through dedicated collectives facade added symm mem back fix after rebase update collective facade usage fixing absl macros large changes in mori communicator refactoring mori communicator cosmetics fixing build switched to enum-based interface updated collective permute interface fixed clang tidy Revert "fixed clang tidy" This reverts commit f9750eafa37dd81f55c1fddb19de2b0765942217. Revert "updated collective permute interface" This reverts commit 34ca3ef080dd970e82e15bbf65358c0e8da15d8a. adapted interface -- 47a25ee3e2d04ea54c4f41b4deb8ec3b393bb6f7 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: update after rebase Merging this change closes #47176 PiperOrigin-RevId: 974632888
These generated includes cause build failures when path mapping (bazelbuild/bazel#22658); other strings in this file may also contain bazel-out and configuration mnemonics, but those should only lead to cache misses. PiperOrigin-RevId: 974634576
A view use extends the allocation end time through the view's transitive
readers so the base buffer stays reserved while they read through the view.
The prefetch deadline for such a use, however, must stay at the view
instruction's own schedule position: the prefetched copy is materialized
right before the view instruction itself. Using the extended reader time as
the deadline lets the prefetch interval picker place the copy interval
entirely after unrelated buffers have freed the heap, while the copy
instructions actually run earlier, inside those buffers' live ranges. The
result is two allocations booked on the same offsets and a verifier chunk
overlap ("Value ... overlaps with another chunk").
Bound latest_prefetch_time by the use instruction's schedule time instead of
the alias extended use_time. Adds a regression test in which the extended
deadline places an illegal copy interval and the clamped deadline correctly
keeps the base in default memory.
PiperOrigin-RevId: 974642462
…d no mappings Adds a new constructor to RemapPlan that accepts input_specs, output_specs, and input_devices_for_output_map without mappings. Updates Validate() to support RemapPlan instances that have only input_devices_for_output_map and no mappings. Splits the old constructor into one with just mappings and one with both mappings and input_devices_for_output_map, deprecating the latter in favor of the new constructor. PiperOrigin-RevId: 974643048
PiperOrigin-RevId: 974645403
PiperOrigin-RevId: 974645808
Return the full instruction constraint map from ConstraintPropagator::Run and unpack tuple entry parameters in MakeDataflowConstrainedArguments so that GetTupleElement user constraints are propagated to individual tuple elements during test argument generation. PiperOrigin-RevId: 974649257
…on tile selection Imported from GitHub PR openxla/xla#47769 📝 Summary of Changes Add a shared-memory constraint to the Triton tile-selection search for transpose ops so that tiles requiring more shared memory than the device provides are rejected before compilation, letting the search fall back to a smaller tile instead of failing later with a RESOURCE_EXHAUSTED error. A transpose stages its (padded) operand tile in shared memory to perform the layout conversion. The estimate is `product(power-of-2-padded operand tile sizes) * element_byte_size`, compared against `shared_memory_per_block_optin()`. 🎯 Justification The Triton tile selection process did not take into account the amount of shared memory required when choosing the tile size. As a result, some merged kernels failed with a RESOURCE_EXHAUSTED error. Example of kernel failing on gfx1201: ``` HloModule transpose_shared_memory_repro, entry_computation_layout={(bf16[1024,2080]{0,1})->bf16[1024,16,128]{2,1,0}} transpose_fusion { param_0 = bf16[1024,2080]{0,1} parameter(0) bitcast_0 = bf16[2080,1024]{1,0} bitcast(param_0) slice = bf16[2048,1024]{1,0} slice(bitcast_0), slice={[32:2080], [0:1024]} transpose = bf16[1024,2048]{1,0} transpose(slice), dimensions={1,0} ROOT bitcast_1 = bf16[1024,16,128]{2,1,0} bitcast(transpose) } ENTRY main { param_0 = bf16[1024,2080]{0,1} parameter(0) ROOT fusion = bf16[1024,16,128]{2,1,0} fusion(param_0), kind=kCustom, calls=transpose_fusion, backend_config={"fusion_backend_config":{"kind":"__triton","block_level_fusion_config":{"output_tiles":[{"sizes":["32","16","128"]}],"num_warps":"8","num_ctas":1,"num_stages":1,"is_tma_allowed":false,"is_warp_specialization_allowed":false,"waves_per_eu":0}}} } ``` ``` RESOURCE_EXHAUSTED: Shared memory size limit exceeded: requested 131072, available: 65536, context: [Fusion: fusion = bf16[1024,16,128]{2,1,0} fusion(param_0.1), kind=kCustom, calls=transpose_fusion, backend_config={"fusion_backend_config":{"kind":"__triton","block_level_fusion_config":{"output_tiles":[{"sizes":["32","16","128"]}],"num_warps":"8","num_ctas":1,"num_stages":1,"is_tma_allowed":false,"is_warp_specialization_allowed":false,"waves_per_eu":0}}}Computation: transpose_fusion { param_0 = bf16[1024,2080]{0,1} parameter(0) bitcast_0 = bf16[2080,1024]{1,0} bitcast(param_0) slice = bf16[2048,1024]{1,0} slice(bitcast_0), slice={[32:2080], [0:1024]} transpose = bf16[1024,2048]{1,0} transpose(slice), dimensions={1,0} ROOT bitcast_1 = bf16[1024,16,128]{2,1,0} bitcast(transpose) ``` the tile selection must therefore take into consideration the use of shared memory when selecting the tile sizes to reject incompatible tile sizes cleanly instead of failing with a RESOURCE_EXHAUSTED error. 🚀 Kind of Contribution Please remove what does not apply: 🐛 Bug Fix, 🧪 Unit Tests: New tests included in this PR Copybara import of the project: -- 284e4ae61b7b9b79b209747d597ce5c39163bf1b by Maxime France-Pillois <mfrancep@amd.com>: Reject transpose tiles exceeding shared memory during Triton tile selection Add a shared-memory constraint to the Triton tile-selection search for transpose ops so that tiles requiring more shared memory than the device provides are rejected before compilation, letting the search fall back to a smaller tile instead of failing later with a RESOURCE_EXHAUSTED error. A transpose stages its (padded) operand tile in shared memory to perform the layout conversion. The estimate is product(power-of-2-padded operand tile sizes) * element_byte_size, compared against shared_memory_per_block_optin(). Because the tile search uses different constraint filters per path, the check is added to both: - TritonEmitterConstraints::ParametersSatisfyConstraints (legacy/symbolic tiling path). - experimental::VerifyTritonConstraints (experimental tiling path). Add a shared GetPaddedTileSizeInBytes helper and tests covering both paths. -- e066b323035ce1da9210e252b29e3345bbaf8e0a by Maxime France-Pillois <mfrancep@amd.com>: Remove unnecessacy tile size check Merging this change closes #47769 PiperOrigin-RevId: 974653326
PiperOrigin-RevId: 974670040
…ckend config deduplication Expose `xla_deduplicate_backend_configs_min_size` in `DebugOptions`, defaulting to `MAX_INT` (disabled). Setting this flag to any valid non-negative threshold enables backend config deduplication into payloads at serialization time in `HloModule::ToProto()` across TPU, GPU, CPU, etc., without requiring `HloProtoOptions` to be threaded through downstream compiler APIs. Callers that manually enable deduplication (such as `gpu_executable.cc`) are preserved. PiperOrigin-RevId: 974672733
PiperOrigin-RevId: 974697845
…t_c_api_client_test.cc`. `TF_ASSERT_OK_AND_ASSIGN` has been marked deprecated. PiperOrigin-RevId: 974711143
Updates LLVM usage to match [cbc5a226cbf8](llvm/llvm-project@cbc5a226cbf8) PiperOrigin-RevId: 974715562
PiperOrigin-RevId: 974720189
PiperOrigin-RevId: 974720484
…n-93680 PiperOrigin-RevId: 974720791
…dation PiperOrigin-RevId: 974721021
PiperOrigin-RevId: 974721580
PiperOrigin-RevId: 974721944
PiperOrigin-RevId: 974721953
…input-validation PiperOrigin-RevId: 974722183
…ank-validation PiperOrigin-RevId: 974722244
…onOp across calls PiperOrigin-RevId: 974728954
…ob-write PiperOrigin-RevId: 974739436
…tPropagator. PiperOrigin-RevId: 974744064
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.
20 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 : )