For batched functions with many large outputs, unbatching dominates the host CPU time and LLC misses of `BatchResourceBase`. `BatchResourceBase::SplitOutputTensors` calls `tensor::Split`, which allocates a new buffer for every request's outputs and copies that request's rows into it, even when the rows could be handed out as aligned views of the batched output.
Return zero-copy `Tensor::Slice` views of the batched output when every task's rows start at an aligned address (`Tensor::IsAligned()`) and the batch has no more padding rows than task rows. Each view holds a reference to the whole batched output, which is freed only after every request has released its view. So a request that holds on to its outputs keeps the entire batch's outputs and padding allocated, not only its own rows, and `Tensor::AllocatedBytes()` on those outputs reports the size of the whole batched buffer. The padding limit keeps that buffer at most twice the size of the task rows. Otherwise, fall back to `tensor::Split`, which keeps the previous copying behavior and dtype coverage (for example, quantized outputs), but copy only the task rows; the padding rows were previously copied and then discarded.
In `bs=2` microbenchmarks on a 64-vCPU AMD Milan host (`intra_op_parallelism_threads = 16`, process CPU time, median of 9 runs), splitting a 12.23 MB/request output tensor drops from 1,397 µs to 38 µs of CPU per batch, about the cost of the batch round trip itself, and end-to-end host batch/unbatch overhead across 20 inputs and 62 outputs drops from 2.79 ms to 0.22 ms of CPU (13x).
PiperOrigin-RevId: 1005968722
Replaces the custom PartialIgnore criteria in tsl::proto_testing with MessageDifferencer::PARTIAL scope, correctly ignoring extra map entries as well as extra elements in repeated fields when combined with IgnoringRepeatedFieldOrdering.
PiperOrigin-RevId: 1005963876
llvm/llvm-project#232791 changed MLIR's ComplexToStandard lowering of
complex.exp, which the XLA CPU and GPU backends use, to return the C99 Annex G
values when the real part is infinite and the imaginary part is not finite:
exp(-inf + I*y) = 0 + 0*I and exp(+inf + I*y) = inf + NaN*I. Before, it
returned NaN + NaN*I, which is what complex_unary_op_samples.h expects.
ExpTest now skips those samples, and a new ExpInfRealNonFiniteImag test checks
these inputs (including y = NaN) and accepts either the old NaN + NaN*I result
or the C99 value. The test then passes with both old and new LLVM.
PiperOrigin-RevId: 1005838599
In `GetMatchingPaths`, each `parent` directory in the queue has already been matched at previous levels, so we only need to match each `child` entry against the current segment's basename (`base_pattern`) rather than joining `parent` and `child` and re-matching the full path prefix. This avoids both re-evaluating earlier path components in `Match` and allocating joined path strings for non-matching entries.
Additionally, skip `GetChildren` when `base_pattern` contains no glob characters, avoid redundant `FileExists` calls on entries returned by `GetChildren`, reuse the `children` vector across iterations, and remove the redundant `FileExists` check before `Stat` in `FileSystem::IsDirectory`.
PiperOrigin-RevId: 1005739160
Upstream LLVM PR https://github.com/llvm/llvm-project/pull/221605 (commit https://github.com/llvm/llvm-project/commit/6cd72b33224f0c3002b41e0e356eb2ab94ab3545) enabled AST documentation comment retention and checking whenever any subgroup of -Wdocumentation (including -Wdocumentation-html) is active.
In third_party/tensorflow/lite/nnapi/NeuralNetworksTypes.h, three doc comments use Javadoc-style {@link ...} inside <ul>...</ul> blocks. Because Clang's comment lexer registers @link as a VerbatimBlockCommand ending at @/link, an unclosed @link consumes the closing </ul> tag until the end of the comment, failing the build with:
error: HTML tag 'ul' requires an end tag [-Werror,-Wdocumentation-html]
Add explicit @/link closing delimiters to {@link ... @/link} inside the three <ul>...</ul> blocks.
PiperOrigin-RevId: 1005563736
For non-integer scale factors (e.g. 688x1520 -> 459x1013), the legacy
ResizeNearest XLA custom call is slow on TPU v6e (~152 ms -> ~0.2 ms per image,
~750x speedup) while separable 1D gathers (`xla::TorchIndexSelect`) preserve
bit-exact numerics against the CPU reference.
The gather lowering applies when H and W are static, positive, within 32-bit
index bounds, not both `align_corners` and `half_pixel_centers`, and at least
one spatial axis has a non-integer scale ratio. Integer ratios, dynamic spatial
dims, 1x1 inputs, identity resizes, and gradients keep their existing paths.
PiperOrigin-RevId: 1005216345
Tests that register `SimpleMultiPassTracer` via `RegisterSimpleMultiPassTracerForTest` previously lacked a way to clear the global TSL multi-pass profiler factory registry during test teardown, causing duplicate tracer registration across multiple test cases in the same test binary.
Add `ClearMultiPassTracerRegisterForTest` in `xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test` to forward to `tsl::profiler::ClearRegisteredMultiPassProfilersForTest`, verified by multi-pass profiler lifecycle unit tests.
PiperOrigin-RevId: 1004995739
Fixes data race in oneDnnThreadPool::parallel_for() for both sync and async cases, when we're running in multi-threaded mode. The issue was caused by the fact, that oneDnnThreadPool::get_in_parallel() was always returning false. Because of that, when some oneDNN kernels started inner parallel_for(), those functions were unable to see that they were running inside of a thread already. This caused them to overwrite done_event_ from outer parallel_for().
The fix implements needed checks to satisfy oneDNN contract, which is stated in `src/common/dnnl_thread.hpp`:
Since oneDNN does not allow nested parallelism, inside a parallel region the number of available threads is 1.
PiperOrigin-RevId: 1004992090
`XlaBuilder::AddSubComputation` copied the embedded computation's
`HloComputationProto`s but dropped its `stack_frame_index`, leaving the
instructions' `stack_frame_id`s dangling. `HloModule::CanonicalizeStackFrameIds`
then reset them to 0, so source locations were lost for any computation built
via `ConvertMlirHloToHlo` and then embedded through `xla::Call`.
`XlaBuilder` now keeps a `StackFrames` table for imported computations:
- `AddSubComputation` re-interns each imported frame (and its parent chain)
into it, so identical frames are stored once, and rewrites the imported
instructions' ids accordingly.
- `BuildSubComputation` re-interns a child builder's table into its parent, only for entries that were imported with a stack frame index.
- `Build()` / `BuildConstantSubGraph()` emit it as the module's
`stack_frame_index`.
- Ids set directly via `SetOpMetadata` (i.e. on the builder's own
instructions, or on imported computations that carry no table) refer to a
caller-owned table and are left untouched.
PiperOrigin-RevId: 1004589040
1. In `GpuSolver::CheckLapackInfoAndDeleteSolverAsync` (`cuda_solvers.cc` and `rocm_solvers.cc`), if `dev_lapack_info.CopyToHost(&success)` fails, drain the GPU stream via `BlockHostUntilDone()` (so any previously enqueued async D2H memcpy into `host_lapack_infos` completes before `host_lapack_infos` is destroyed), destroy `solver`, and invoke `info_checker_callback(status, host_lapack_infos)` instead of returning early via synchronous `OP_REQUIRES` (which skipped `info_checker_callback` and the caller `AsyncOpKernel::DoneCallback`). Also destroy `solver` before invoking `info_checker_callback` on the `dev_lapack_infos.empty()` early return path.
2. In `DeterminantOpGpu::ComputeAsync`, `LogDeterminantOpGpu::ComputeAsync` (`determinant_op.cc`), and `CholeskyOpGpu::ComputeAsync` (`cholesky_op_gpu.cu.cc`), propagate non-`InvalidArgument` errors (such as `InternalError` from `CopyToHost` failure) to `context` via `OP_REQUIRES_OK_ASYNC(context, status, done)` in `info_checker`.
3. Add unit tests in `determinant_op_gpu_test.cc` verifying `CheckLapackInfoAndDeleteSolverAsync` (including partial `CopyToHost` failure across multiple `DeviceLapackInfo` entries), `MatrixDeterminant`, `LogMatrixDeterminant`, and `Cholesky` on GPU when `CopyToHost` fails.
PiperOrigin-RevId: 1004418238
When `ConvertMover` moves a `kConvert` across a commuting operation, preserve the original `kConvert` frontend attributes on the new `kConvert` instructions and keep the commuting operation's attributes on the cloned instruction instead of copying the outer instruction's attributes onto the replacement root.
PiperOrigin-RevId: 1004414597
Add six monitoring metrics under `/tensorflow/data/service/cross_trainer_cache/`:
- `trainer_queries`: Counter of cache hits and misses per `trainer_id` and `source` (`"cache"`, `"cache_waited"`, or `"prefetch"`).
- `blocked_wait_usecs`: Counter of microseconds spent waiting for another thread to extend the cache per `trainer_id`.
- `elements_skipped`: Counter of elements skipped due to eviction per `trainer_id`.
- `cursor_offset`: Distribution of each trainer's element offset from the eviction edge (`cache_start_index_`).
- `size_bytes`: Gauge tracking total cached byte size across active `CrossTrainerCache` instances.
- `max_size_bytes`: Gauge recording the configured maximum cache size in bytes.
PiperOrigin-RevId: 1004358006
HLO diff can report matched instructions that are unchanged themselves but used
differently (e.g. swapped operands) as moved instead of unchanged. The
comparison manager and the original tensor summary comparator only mapped
unchanged and changed pairs, so the tensors of moved instructions would no
longer be compared. Map moved pairs too, and count them as unchanged in the
creation metrics.
PiperOrigin-RevId: 1004193597
tfl.floor_div had no folder, so a FLOOR_DIV on constant operands survived
conversion as a runtime op. GPU delegates that reject FLOOR_DIV then see a
CPU-only op in the middle of an otherwise delegable graph.
The fold matches the TFLite kernel: f32 operands are divided and floored in
double, and integer division rounds toward negative infinity. Division by
zero (for any element type) is left to the runtime, which reports it as an
error.
HF DETR models hit this in the sine position embedding (dim_t // 2 on a
constant arange); it is their only FLOOR_DIV.
Tests: const-fold.mlir adds f32 (with broadcast, and a 1.0 // 0.1 case that
needs the double division) and i32 (all sign combinations) folds, plus
divide-by-zero (i32 and f32) and runtime-operand negatives.
PiperOrigin-RevId: 1004127675
1. Use `gtl::MakeCleanup` in `TPUPartitionedCallOp::ComputeAsync` so `DoneCallback` (and `DequeueFromCoreSelector` if a core was selected) is always invoked even when an early error return occurs before async execution begins.
2. Store `init_status_` as a member variable of `TPUPartitionedCallOp` so initialization errors inside `absl::call_once` (`once_` and `ordinal_selector_once_`) are preserved across subsequent `ComputeAsync` calls.
3. In `InstantiateFunctionsFromSubgraphs`, populate `partition_cache_[cache_hash]` only after all subgraphs have been rewritten and instantiated without error, avoiding a partial/empty cache entry on failure.
PiperOrigin-RevId: 1004036684
`BatchMatMulOp` had no folder, so a matmul of two constants survived into the
flatbuffer and was recomputed on every inference. Add `BatchMatMulOp::fold`
for f32 operands with static shapes, honoring `adj_x` / `adj_y` and
right-aligned batch broadcasting, subject to the usual `ShouldFoldOperation`
size limit.
PiperOrigin-RevId: 1003936085
cuDNN v9 internally captures temporary CUDA graphs during heuristic queries and execution plan creation (e.g. via `cudnn::backend::cudnnBackendGetMetadataFromGraph` inside `cudnnBackendFinalize`), which does not go through `CudaStream::CaptureHandle`.
During parallel autotuning (`ConfigAssigner::GetConfigsForAll`), a worker thread running `GpuProfiler::Profile` can deallocate VMM buffers via `CudaDeviceMemoryAllocation::~CudaDeviceMemoryAllocation` -> `DrainPendingVmmDeallocations` -> `cuCtxSynchronize()` concurrently with the main thread querying cuDNN heuristics in `CudnnSupport::GetConvolveRunners`. Calling `cuCtxSynchronize()` while cuDNN is capturing a CUDA graph on the same CUDA context either fails with `CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED` / `cudaErrorStreamCaptureInvalidated` or segfaults inside `libcuda.so` (`cuCtxSynchronize`).
Call `CudaDeviceAllocator::EnterStreamCapture` and `CudaDeviceAllocator::ExitStreamCapture` in `CudnnHandle` (`xla/stream_executor/cuda/cuda_dnn.cc`) so that VMM deallocations are deferred until the active cuDNN operation completes.
PiperOrigin-RevId: 1003883162
Introduces SubchannelCostModel to model subchannel quantization workloads and extends PlatformInfo to track cross lane unit timing capabilities.
PiperOrigin-RevId: 1003267036
With TF_XLA_FUNCTIONALIZE_CONTROL_FLOW=1, run FunctionalizeControlFlow before
ValidateGraph so graphs with v1 control flow (e.g. Enter) compile instead of
failing validation. Off by default.
PiperOrigin-RevId: 1003176296
