Repository navigation
[pull] master from tensorflow:master - #8858
Merged
Merged
Conversation
Eigen's scalar_sign_op computes |z|^2 = re^2 + im^2 as a float32 intermediate. When |z| < sqrt(float32_tiny) ~= 1.08e-19, this underflows to 0, causing sign(z) to incorrectly return 0 instead of z/|z|. Add specialized sign functors for complex64 and complex128 that check (re != 0 || im != 0) directly and use std::abs() (which uses hypot internally) to compute the magnitude without intermediate underflow/overflow. Fixes #116945
Add testComplexSignSmallMagnitude to cwise_ops_unary_test.py that verifies tf.math.sign returns correct unit-magnitude results for complex64 and complex128 inputs with magnitudes below sqrt(float32_tiny), where the original |z|^2 computation would underflow. Requested by reviewer in PR #116986.
std::abs(complex) may compute sqrt(re*re + im*im) which underflows to 0 when FTZ is enabled and |re|^2 is subnormal. Using std::hypot directly avoids intermediate underflow/overflow per C++ standard guarantees. Also avoids complex division in favor of scalar re/mag, im/mag for numerical robustness.
…ut to prevent crash on invalid rank
…underflow # Conflicts: # tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
The tiled-transpose kernel used 32-bit indexes, which overflowed for tensors with more than INT32_MAX elements and caused out-of-bounds memory access. Use 64-bit indexes in that kernel and add a test for an index above INT32_MAX.
Host builds compiled the GPU test without a CUDA or ROCm guard, and the test included conv_2d_gpu.h without listing it as a source. TensorFlow's build checks rejected both cases. I added the GPU guard and listed conv_2d_gpu.h in the test sources.
Reverts changelist 989948905 PiperOrigin-RevId: 990219345
PiperOrigin-RevId: 990223584
Previously, libnvjitlink support was gated behind the `enable_libnvjitlink_support` build flag (which defaulted to `false` in OSS) and split across `nvjitlink_impl` and `nvjitlink_stub` targets. Because all supported CUDA versions now include libnvjitlink, the conditional build graph is no longer needed. Compile xla/stream_executor/cuda/nvjitlink.cc unconditionally in //xla/stream_executor/cuda:nvjitlink and remove the `enable_libnvjitlink_support` build flag, `libnvjitlink_support_enabled` config setting, and xla/stream_executor/cuda/nvjitlink_stub.cc. Keep //xla/stream_executor/cuda:nvjitlink_support only to disable libnvjitlink under MemorySanitizer. To support OSS wheel builds that dynamically load CUDA libraries via stubs, add //xla/tsl/cuda:nvjitlink (backed by xla/tsl/cuda/nvjitlink_stub.cc and `GetNvJitLinkDsoHandle()` in xla/tsl/platform/default/dso_loader.cc). In `kAuto` compilation mode, check `GetNvJitLinkVersion()` before `LoadedNvJitLinkHasKnownIssues()` so stub builds without `libnvJitLink.so` at runtime fall back cleanly to `ptxas` and `nvlink`. PiperOrigin-RevId: 990235320
MLIR is migrating dialects to strict property handling in custom assembly formats (`useStrictPropertiesInAssemblyFormat = 1`), where `attr-dict` only parses and prints discardable attributes and inherent attributes must be either explicitly bound in the assembly format or handled via `prop-dict`. Add `prop-dict` before `attr-dict` in the assembly formats of `tf_framework.alloc` and `tf_framework.jit_compile_from_str`, enable `useStrictPropertiesInAssemblyFormat = 1` on `TFFramework_Dialect`, and update tests to use the `<...>` property dictionary syntax. PiperOrigin-RevId: 990243949
…accessors
`LiteralBase::data<T>`, `LiteralBase::untyped_data`, `LiteralBase::size_bytes`, `LiteralBase::total_size_bytes`, and `LiteralBase::element_count` previously took a default argument `const ShapeIndex& shape_index = {}`, causing every 0-arg caller across XLA to construct and destroy a temporary `ShapeIndex` (`absl::InlinedVector<int64_t, 2>`) at the call site and route through `piece(shape_index)` tree traversal.
Similarly, `MutableLiteralBase::PopulateParallel` and `PopulateLinearParallel` checked `data<NativeT>().size() > 32` at every template instantiation site (instantiating `data<NativeT>()` just to obtain the element count) rather than `element_count() > 32` like `PopulateInplaceParallel` and `PopulateLinearInplaceParallel`.
Split the 0-arg overloads of `data<T>()`, `untyped_data()`, `size_bytes()`, `total_size_bytes()`, and `element_count()` to access `root_piece()` / `mutable_root_piece()` directly without constructing a temporary `ShapeIndex`, define `element_count()` and `element_count(const ShapeIndex&)` out-of-line in `literal.cc` delegating to `root_piece().element_count()` / `piece(index).element_count()`, and use `element_count() > 32` in `PopulateParallel` and `PopulateLinearParallel`.
### Runtime & Memory Performance Impact
- **CPU / Latency**: Positive. 0-arg calls to `data<T>()`, `untyped_data()`, `size_bytes()`, `total_size_bytes()`, and `element_count()` bypass temporary `ShapeIndex` (`absl::InlinedVector<int64_t, 2>`) construction/destruction and `piece(shape_index)` loop traversal, accessing `root_piece()` directly in O(1) with zero heap or stack-vector overhead.
- **Memory**: Neutral to positive (eliminates temporary `ShapeIndex` stack frames at call sites).
PiperOrigin-RevId: 990244740
…iton PiperOrigin-RevId: 990263277
…-line to index_util.cc The 2-argument convenience overload `IndexUtil::MultidimensionalIndexToLinearIndex(const Shape&, absl::Span<const int64_t>)` was defined inline in `index_util.h`, causing `LayoutUtil::MinorToMajor(shape)`, `Shape::array_state_maybe_underneath_buffer()` checks, and the stride accumulation loop to be inlined into hundreds of `LiteralBase::Piece::Get<NativeT>` and `Set<NativeT>` template instantiations. Meanwhile, hot-path inner loops in `LiteralBase::Broadcast` and `HloEvaluator` use the 3-argument fast-path overload `IndexUtil::MultidimensionalIndexToLinearIndex(const Shape&, absl::Span<const int64_t> minor_to_major, absl::Span<const int64_t> multi_index)` where `minor_to_major` is pre-extracted outside the loop. Move the 2-argument overload out-of-line to `index_util.cc`, keep the 3-argument fast-path overload inline in `index_util.h`, and hoist `absl::Span<const int64_t> dims = shape.dimensions();` outside the stride accumulation loop in the 3-argument overload so `shape.dimensions(...)` is not re-evaluated on every dimension. ### Runtime & Memory Performance Impact - **CPU / Latency**: Preserves hot-path inlining of the 3-argument overload in `LiteralBase::Broadcast` and `HloEvaluator` (with a slight speedup on higher-rank broadcasts from hoisting `shape.dimensions()` out of the inner loop): ``` Benchmark Base CPU (ns) New CPU (ns) Delta -------------------------------------------------------------------------------- BM_BroadcastVectorToMatrix/16/16 3950 4086 -3.4% BM_BroadcastVectorToMatrix/16/1024 133732 131341 +1.8% BM_BroadcastVectorToMatrix/1024/1024 8646962 8415307 +2.7% BM_BroadcastMatrixToTensor/16/16 19883 18694 +6.0% BM_BroadcastMatrixToTensor/16/1024 1145118 1109215 +3.1% BM_BroadcastMatrixToTensor/1024/1024 73361100 68770331 +6.3% BM_Populate/64/process_time 52482 61484 -17.1% BM_Populate/1024/process_time 13224583 14506707 -9.7% BM_PopulateParallel/1024/process_time 23279761 26892908 -15.5% BM_PopulateLinear/1024/process_time 2123731 2081302 +2.0% BM_PopulateLinearParallel/1024/process_time 3329581 3376813 -1.4% ``` - **Memory**: Neutral (no changes to object layouts or heap allocations). PiperOrigin-RevId: 990268395
Add guidance to `xla/AGENTS.md` covering: - Passing `--config=cuda` when running GPU tests. - Including the OpenXLA Apache 2.0 copyright header with the current calendar year on newly created source and `BUILD` files. PiperOrigin-RevId: 990270007
…nk-validation PiperOrigin-RevId: 990273736
The two trivial replication states own no heap object. Replication trees hold no Shape, propagation merges in place, and post orders are computed once. A computation already marked unique is not marked again, and repeat visits evaluate only instructions whose inputs changed. In a computation marked unique, tuple roots and SPMD infeeds are now unique too; other results are unchanged. PiperOrigin-RevId: 990279838
…flow PiperOrigin-RevId: 990290236
…-index-overflow PiperOrigin-RevId: 990295912
We already had symmetric and kConsumerSymmetric where a single consumer waits for multiple producers on the same gpu. Now we also need producer symmetric for the case where a consumer waits for the same producer block from different GPUs. PiperOrigin-RevId: 990304779
Move some refactoring changes outside to reduce churn in the followup change. PiperOrigin-RevId: 990309442
PiperOrigin-RevId: 990319867
PiperOrigin-RevId: 990335835
Store the order-preserving integer key directly in `Descending::KVT` so `Descending::cmp` performs direct integer comparisons without per-comparison bitwise conversions. Apply `details::ToOrdered` once per element when loading keys from global memory in `PerWarpTopK`, and invert the mapping via `details::FromOrdered` only when thread 0 writes the final k outputs in `MergeTopKs`. This reduces the total-order overhead from 1.35x down to 1.12x in comparison to the original C++ less comparator. PiperOrigin-RevId: 990342348
TryRemoveDeadWhileParams kept a bit set of all tuple indices and a UnionFind entry per instruction in two hash maps and ORed the bit sets along every operand edge. Dense arrays indexed by the instruction's local id replace the maps, the bit sets cover only the removable candidate indices, and one post order pass per computation fills them. Same removed indices, same HLO. PiperOrigin-RevId: 990350581
… integrate with the benchmark driver. Move collectives_benchmark.py and all_gather_benchmark.py from internal benchmarks to xla/benchmarks/collectives/, add CollectiveBenchmarkConfig and CollectiveBenchmark adapters, and register 'collectives' in BENCHMARK_FACTORIES so that run_benchmarks.py --benchmarks=collectives runs the collective microbenchmark suite and outputs CSV metrics. PiperOrigin-RevId: 990370997
…sions in YNNPACK Adds support for kCopy operations in YNNPACK fusion and relaxes layout checks to tolerate out-of-order dimensions when their extents are 1. This addresses fusion breakages and performance regressions identified in openxla/xla#49131: - When reduction pipelines use keepdims=True (e.g. pairwise reductions such as gravitational N-body or RBF kernels), layout assignment can insert reshapes and copies with layouts like f32[M, N, 1]{1,0,2}. - Previously, IsLayoutSupportedByYnn rejected any layout whose minor_to_major was not strictly descending, even if the out-of-order dimensions had extent 1 (which does not affect physical element ordering). - Furthermore, kCopy was unsupported in YNNPACK, causing LibraryRewriter to truncate fusions at copy boundaries and split the pipeline into multiple custom fusions, forcing large O(M*N) intermediate buffers to be materialized in memory. Key changes: 1. IsLayoutSupportedByYnn: Ignores extent-1 dimensions when verifying descending layouts. 2. IsCopyOpSupportedByYnn: Validates copy operations whose non-degenerate dimensions have matching layout and extents. 3. ynn_matcher.h: Registers HloOpcode::kCopy in SupportedOps() and dispatches in IsOpSupported(). 4. ynn_emitter.cc: Implements DefineCopyOp using ynn_define_copy and emits copy operations in YNN subgraphs. 5. ynn_fusion_test.cc: Adds unit tests for copy operations and degenerate layout copies across bf16, f32, and f64. PiperOrigin-RevId: 990372521
MLIR is migrating dialects to strict property handling in custom assembly formats (`useStrictPropertiesInAssemblyFormat = 1`), where `attr-dict` only parses and prints discardable attributes and inherent attributes must be either explicitly bound in the assembly format or handled via `prop-dict`. Add `prop-dict` before `attr-dict` in the assembly formats of `xtile.scan` and `xtile.dot_scaled`, enable `useStrictPropertiesInAssemblyFormat = 1` on `XTileDialect`, and update tests to use the `<...>` property dictionary syntax. PiperOrigin-RevId: 990387232
…the constraint propagator. PiperOrigin-RevId: 990402570
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.
14 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 : )