Repository navigation
Develop upstream sync 261104 - #3258
Open
mmakevic-amd wants to merge 3975 commits into
Open
mmakevic-amd wants to merge 3975 commits into
mmakevic-amd wants to merge 3975 commits into
Conversation
…39-linalg-numpy-dtype-hint PiperOrigin-RevId: 1011882847
…31-convert-image-dtype-overflow PiperOrigin-RevId: 1011882896
…iter-zlib-options PiperOrigin-RevId: 1011883016
…37-resize-grad-half-dtypes PiperOrigin-RevId: 1011883070
…r-arithmetic-optimizer-bounds PiperOrigin-RevId: 1011883668
…h-missing-bounds-check PiperOrigin-RevId: 1011883743
…upstream PiperOrigin-RevId: 1011884208
…lue-and-graph-def-bounds PiperOrigin-RevId: 1011886220
Draw each row's indices from indices_shape[-1] + 2 slots, so the test still checks that batch_scatter_update leaves unindexed elements alone. Add @test_util.disable_xla to the two tests that pin ops to the CPU, as batch_gather_op_test.py does.
…file-seek-write-only PiperOrigin-RevId: 1011887074
…rng-eager-parity PiperOrigin-RevId: 1011888274
The previous commit applied offset_ in the DT_INT64 branch of SetValue. SetValue parses both the key and the value, so that shifted int64 vocabulary tokens whenever value_index_offset was non-zero (for example IntegerLookup on an integer vocab file). Restore the plain int64 parse. Guard next_id_ + offset_ in the kLineNumber branch instead. next_id_ is non-negative, so the sum can only overflow for a positive offset_. Include <cstdint>, <memory> and <string> directly in lookup_util.h, and add testOffsetOverflow and testInt64WholeLineIgnoresOffset.
…util-compress-repeated-shape-bounds PiperOrigin-RevId: 1011888303
reduce_variance and reduce_std computed mean((x - mean)**2). Any rounding error in the mean shifts every deviation by the same amount, and squaring that shift adds it to the variance. Under XLA, where a division by a constant becomes a multiplication by its reciprocal, a row of identical float32 values such as [1e10, 1e10, 1e10] got a variance of 1048576 instead of 0. Subtract the squared mean of the deviations, the same correction tf.nn.moments uses, and clamp the result at zero.
- Link XLA into math_ops_test so the jit_compile leg runs, and skip that leg only when the build has no XLA compiler, as MomentsTest does. - Call reduce_variance and reduce_std directly as well as through tf.function, with keepdims both off and on. - Add complex64 and complex128 inputs, which exercise the complex branch of the correction. - Allow 5.0 for complex64. When the final multiply and subtract are fused, the unrounded product leaves 2 per real component, so 4 for complex64 against 2 for float32.
…r-function-and-loop-optimizer-bounds PiperOrigin-RevId: 1011889893
PiperOrigin-RevId: 1011891915
`IsFabricMemory` reads `CU_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES` into a 4-byte `unsigned int`, but the driver writes a 64-bit value. So the call overwrites the 4 bytes after `handle_types` on the stack. In an x86-64 build with `-Os`, these are the low 4 bytes of the caller's saved `rbx` register. This only leads to a crash if the caller keeps a live value in `rbx` across the call. After an unrelated change shifted register allocation, `CudaStream::Memcpy` now keeps its `ActivateContext` pointer in `rbx` and segfaults when it destroys that pointer. So this change makes `handle_types` a `uint64_t`. The CUDA documentation doesn't state the type of this attribute, but NCCL and UCX both read it into a 64-bit integer. I also added tests to `xla/stream_executor/cuda/cuda_driver_test.cc` that only use driver API calls. For every pointer attribute and memory pool attribute that stream executor queries, they measure how many bytes the driver writes and compare that to the size of the variable that stream executor uses. They cover device, pinned host, pageable host, managed and VMM memory, and both `cuPointerGetAttribute` and `cuPointerGetAttributes`. Apart from `CU_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES`, all of these variables already had the right size. PiperOrigin-RevId: 1011934369
`NNFunctionsTest.testDotProductAttentionJitGrad` (added in `jax-ml/jax@282298d408ba`) tests `nn.dot_product_attention(..., implementation='xla')` with `bfloat16`, which uses `DotAlgorithmPreset.BF16_BF16_F32` and fails with `UNIMPLEMENTED: Unsupported algorithm on the current device(s): ALG_DOT_BF16_BF16_F32` on CUDA GPUs with compute capability < 8.0 (including T4 at CC 7.5). Temporarily exclude `NNFunctionsTest.testDotProductAttentionJitGrad` via `JAX_EXCLUDE_TEST_TARGETS` in `BuildType.JAX_LINUX_X86_GPU_T4_GITHUB_ACTIONS` until the upstream JAX test skip lands. PiperOrigin-RevId: 1011977303
Remove unused functions, classes, and overloads from `llvm_util.h` and `llvm_util.cc`: - `EmitLogging` and `LogS64` - `GetSizeInBits` - `AllocateSharedMemoryTile` and `SharedMemoryTile` - `GetOrCreateVariableRngState` (unimplemented declaration) and `GetGlobalMemoryAddressSpace` (moved to an anonymous namespace in `llvm_util.cc` alongside `GetOrCreateVariableForRngState`) - `EmitReturnBlock` (make `return_block` required in `EmitEarlyReturn`) - Variadic `DumpToString` overloads and unused `llvm::TargetOptions` forward declaration PiperOrigin-RevId: 1011996006
`GetBatchRowColumnShape` in `xla/service/gpu/matmul_utils.cc` passed literal
`1` (`int`) as the initial value to `absl::c_accumulate`, deducing a 32-bit
`int` accumulator and narrowing 64-bit GEMM batch, row, and column dimension
products back to `int`. Extents exceeding `INT32_MAX` either wrapped negative
and failed shape validation or wrapped positive (`>= 2^32`) and silently
truncated `MatrixLayout` dimensions and strides. Initialize the accumulator
with `int64_t{1}` to preserve 64-bit extents, and add unit tests covering
`> INT32_MAX` and `> UINT32_MAX` GEMM shapes.
PiperOrigin-RevId: 1011997412
…se-int64-feature-address PiperOrigin-RevId: 1012009799
Pass peer pointer tables for symmetric input and output buffers in CollectiveKernelThunk and cache their weak symmetric memory handles across runs. PiperOrigin-RevId: 1012018035
…cross-sparse-and-splits-validation PiperOrigin-RevId: 1012019381
`hlo_op_profiles` only relied on `xla/types.h` for the `xla::se` namespace alias, so qualify `stream_executor::DeviceDescription` directly and drop the unused `//xla:types` dependency. PiperOrigin-RevId: 1012019474
PiperOrigin-RevId: 1012022062
`IrCompiler::EmitMachineCode` lowers LLVM IR directly to an in-memory object buffer via `TargetMachine::addPassesToEmitMC` (`Target`, `TargetInfo`, `TargetMC`, and `AsmPrinter`), and XLA:CPU never emits CPU inline assembly (`llvm::InlineAsm` or module-level inline asm), which is the only LLVM codegen path that invokes `Target::createMCAsmParser`. Remove the unused `LLVMInitialize*AsmParser()` registrations from `InitializeTargets()` (and `ir_compiler_test`) and replace `//xla/tools:llvm_targets` on `//xla/backends/cpu/codegen:ir_compiler` with the per-architecture `@llvm-project//llvm:*CodeGen` dependencies already used by `:jit_compiler`, avoiding linking ~2.5 MB of unused target MC assembly-parser tables and code. Runtime and memory impact: None (neutral to slightly faster one-time initialization). Target `MCAsmParser` factories were previously registered in `TargetRegistry` during `InitializeTargets()` and never instantiated during compilation or execution; generated machine code and peak memory usage are unchanged. PiperOrigin-RevId: 1012025594
… native packed bf16 emission of a fusion. PiperOrigin-RevId: 1013314202
Include libtcmalloc.so.4 as pre-installed under /usr/lib/<arch>-linux-gnu/libtcmalloc.so.4 and available for LD_PRELOAD in containerized build and test environments. PiperOrigin-RevId: 1013326548
…ations. `MatchCallGraphs` phase 3 groups unmatched callees by caller opcode, op_name and `ComputeProgramShape(include_ids=false).ToString()`. That string includes the parameter names (minus the id suffix), so callees whose parameters are named differently, e.g. `Arg_0` vs `p0`, are never signature-matched, even though parameter names aren't part of the signature. When the callees are left unmapped, the computation graph matcher doesn't anchor their parameters and roots, and the exact subgraph matcher maps identical duplicated computations (e.g. two copies of `log_softmax`) crosswise. On a synthetic mutator pair (46k instructions, 34 mutations + rename) this removes the remaining 52 wrong mappings; mapping precision/recall goes from 0.9989 to 1.0. PiperOrigin-RevId: 1013336140
…a temporary tensor or running two `TensorReverseOp` passes. In `tensorflow/core/kernels/scan_ops.h`, `LogSumExpReducer` and `LogAddExp` are non-associative in floating-point arithmetic, so `Scan<Eigen::ThreadPoolDevice>` cannot use Eigen's parallel prefix scan and falls back to the sequential axis reduction loop (`!Eigen::internal:: ScanLauncher<...>::Implemented`). Previously, that fallback always allocated an `output_shape` temporary tensor (`context->allocate_temp`) and executed two full-tensor `TensorReverseOp` passes over `Eigen::ThreadPoolDevice` even when `reverse_ == false`: 1. `TensorReverseOp(in, [axis])` into `temp` 2. Sequential loop iterating backward from `temp.dimension(1) - 1` down to `0` 3. `TensorReverseOp(temp, [axis])` into `out` When `reverse == false`, iterating forward from `0` to `in.dimension(1) - 1` directly from `in` into `out` computes the exact same recurrence (`out[0] = in[0]` or `Identity`, `accum(out[i], in[i])`) with identical floating-point evaluation order and rounding while eliminating the temporary buffer allocation and both `TensorReverseOp` passes. PiperOrigin-RevId: 1013338889
Imported from GitHub PR openxla/xla#52165 📝 Summary of Changes Update onednn version to 3.14.1 🎯 Justification Provides performance updates, and fixes issues including tensorflow#129498. 🚀 Kind of Contribution Please remove what does not apply: 🐛 Bug Fix, ⚡️ Performance Improvement, Copybara import of the project: -- 1ba72d37f413904a2a5308e55664af0d77d3448d by Gauri Deshpande <gauri1.deshpande@intel.com>: [XLA-CPU][oneDNN] Update oneDNN to v3.14.1 Merging this change closes tensorflow#52165 PiperOrigin-RevId: 1013342834
***Problem:*** `DotStrengthReduction` only rewrites unbatched one-hot `dot`s to `gather`, so batched one-hot `dot`s and 1-D one-hot masked sum reductions (`reduce_sum(select(idx == iota, x, 0))`) still scan all `K` elements to read a single value. ***Solution:*** Share a one-hot lowering that emits a single-slice `gather` guarded by an in-bounds `select` (keeping zero outputs for out-of-range indices) and use it in `DotStrengthReduction` for batched one-hot `dot`s and in `GpuAlgebraicSimplifier` for 1-D one-hot `reduce_sum`s, peeling `broadcast` and up to one widening `convert` so the gather reads the narrow operand directly. PiperOrigin-RevId: 1013355786
XLA's Bazel workspace (`hwloc/workspace.bzl`) was pinned to `hwloc 2.7.1` (from 2022), missing years of upstream topology fixes and performance improvements. In particular, `hwloc 2.7.1` reads Linux `/sys` topology files via stdio `fopen`/`fgets`/`fclose` (allocating and freeing a heap buffer per file during `hwloc_topology_load`), whereas `hwloc 2.13.0` uses stack-buffered `openat`/`read`/`close` (`hwloc_read_path_by_length`), significantly reducing syscall and allocator overhead during NUMA topology initialization. PiperOrigin-RevId: 1013364002
Parameterize `KernelDefinition` with a spec type (defaulting to `KernelSpec`), and update `KernelCompiler::Compile` and `CompileToTargetBinary` to accept `GpuKernelDefinition` and `std::optional<GpuKernelSpec>` respectively. PiperOrigin-RevId: 1013381111
When `PjRtCApiCompiler::Compile` or `PjRtCApiPhaseCompiler::CompileToRelocatable` is called for ahead-of-time compilation without a live `PjRtClient` (`client == nullptr`), program serialization previously could not query `client->plugin_attributes()` and fell back to the default 12-week compatibility StableHLO and Shardy versions (`xla::GetDefaultStablehloVersion()` and `xla::GetDefaultSdyVersion()`). Because `PJRT_Plugin_Attributes` is a plugin-level C API call on `const PJRT_Api*` that does not require a `PJRT_Client*`, fall back to querying `c_api->PJRT_Plugin_Attributes` when `client == nullptr` (or when the client does not report attributes) so ahead-of-time compilation serializes programs using the plugin's actual `stablehlo_current_version` and `sdy_current_version`. PiperOrigin-RevId: 1013382123
When the `Terminate` RPC failed, `SubprocessProfilingSession::Stop()` returned right away. The cleanup cancelled the async `Profile` call, but nothing read its completion from `completion_queue_`, so the call leaked (ASan reports it). `Stop()` now cancels the `Profile` call when `Terminate` fails, waits for it on the completion queue, and then returns the `Terminate` error. PiperOrigin-RevId: 1013391194
…objects. By default, `GetHWLocTopology()` in `numa_hwloc.cc` loads the full machine hardware hierarchy—including CPU cache levels, PCI/OS I/O devices (`linuxio`), and per-core `x86` CPUID probing—even though TSL's NUMA utilities only ever query `HWLOC_OBJ_NUMANODE` and CPU affinity sets. On large multi-socket / high-core-count hosts, this redundant discovery adds unnecessary sysfs traversal and CPUID probing overhead during NUMA initialization. This change configures `HWLOC_TYPE_FILTER_KEEP_NONE` and excludes the `linuxio` and `x86` discovery components before `hwloc_topology_load()` so hwloc only discovers the NUMA node and CPU sets that TSL uses. PiperOrigin-RevId: 1013399877
Static-range quantized `tfl.batch_matmul` ops whose constant weights are produced by `tfl.blockwise_dequantize` previously fell back to floating-point weights because `TFL_BatchMatMulOp` rejected `!tfl.blockwise_quant` RHS operands and `FuseQDQPass` only folded affine `tfl.dequantize` producers. Reshapes that insert size-1 dimensions above `tfl.blockwise_dequantize` (e.g. for single-token decode) also hid the constant weight producer from QDQ fusion. * Allow `!tfl.blockwise_quant` RHS operands on `tfl.batch_matmul`. * `FuseQDQPass`: fold a constant `tfl.blockwise_dequantize` RHS of an int8 SRQ `tfl.batch_matmul` into a `!tfl.blockwise_quant` `tfl.pseudo_qconst`, and only treat `tfl.dequantize` of a blockwise-quantized tensor as a DQ parent for the RHS of an int8 SRQ `tfl.batch_matmul`. * `LargeConstantFoldPass`: push reshapes that only insert or remove size-1 dimensions through constant `tfl.blockwise_dequantize` ops. * `OptimizeBatchMatmulPass`: do not convert `tfl.batch_matmul` to `tfl.fully_connected` when the rank-2 RHS is blockwise quantized, or is a `tfl.blockwise_dequantize` while the LHS comes from `tfl.dequantize`. * `GetQTypeFromDefiningDequantize`: return `std::nullopt` instead of crashing on non-`quant::QuantizedType` dequantize inputs. PiperOrigin-RevId: 1013417495
PiperOrigin-RevId: 1013450356
…MatrixBatchFourVectorMultiplyAccumulate. Guard `per_channel_scale + row` with a null check so `nullptr + row` is not evaluated when `per_channel_scale == nullptr`. This is only to clean up UndefinedBehaviorSanitizer (`nullptr-with-nonzero-offset`) undefined behavior; it is not required to avoid a crash in release builds because the inline assembly already checks `is_channel_scale_nullptr` before dereferencing `channel_scales_ptr`. PiperOrigin-RevId: 1013464233
PiperOrigin-RevId: 1013770697
…pport via TensorIR engine Imported from GitHub PR openxla/xla#47839 ## Summary This PR adds support for lowering non-GEMM HLO pointwise fusions to the cuDNN graph API using the `CUDNN_GENERIC_MEMBOUND_FUSION_TENSOR_IR_ENGINE` (cuDNN's dedicated memory-bound TensorIR engine). The feature is controlled by a new flag `--xla_gpu_cudnn_non_gemm_fusion_level`. ### TensorIR engine selection For non-GEMM (memory-bound) fusions, `PrepareGraph` in `cudnn_fusion_compiler.cc` sets `force_tensor_ir=true` in `EngineOptions`. In `CudnnGraph::Prepare` (`cuda_dnn.cc`), this deselects engines that report `NUMERICAL_NOTE_TENSOR_CORE` (the CUTLASS FALLBACK engines, eng0-eng2), leaving only `CUDNN_GENERIC_MEMBOUND_FUSION_TENSOR_IR_ENGINE` (eng3). GEMM fusions continue to use the full engine set. ### Extended `cudnn_fusion_compiler.cc` - **Non-GEMM parameter setup**: Added a parameter-binding path for fusions with no gemm/conv adapter, using `DimensionsAndStrides` directly from the HLO shape. - **Broadcast lowering**: Fixed same-rank and rank-expanding broadcasts. Non-trivial broadcasts are represented via a logical reshape + pointwise identity so cuDNN sees a full-shape virtual tensor. - **Logical ops**: `kAnd`, `kOr`, `kNot` now map to cuDNN `LOGICAL_{AND,OR,NOT}` for `PRED`-typed operands (bitwise variants remain unsupported). - **Clamp fix**: Corrected operand ordering to match XLA semantics: `min(max(value, lower), upper)`. - **Constant handling**: Extended from scalar-only to any single-element constant (e.g. `s32[1]`, `f32[1,1]`). For non-GEMM/conv fusions, constants are always tagged `COMPILE_TIME_CONST`; GEMM/conv fusions keep the prior behavior of runtime scalar tensors for non-scalar-shaped inputs. - **Compute type for compare**: `kCompare` derives its compute type from the first operand's type (inputs are the arithmetic type; output is `PRED`). - **Reduce disabled for non-GEMM fusions**: `kReduce` (including rank-preserving reduction lowering) is now rejected with `UnimplementedError` when the fusion has no gemm/conv adapter; reduce support is unchanged for GEMM/conv fusions. Variadic (tuple-shaped) reduce remains unsupported everywhere. ### Autotuner integration `IsSupportedCudnnFusion` in `xla/backends/gpu/autotuner/cudnn.cc` is extended to accept non-GEMM fusions that aren't already tagged with kCustom when `xla_gpu_cudnn_non_gemm_fusion_level >= 1`, enabling autotuning over cuDNN/Native emitter and other backends for these fusions. ### New flag ``` --xla_gpu_cudnn_non_gemm_fusion_level=<int> 0: Disabled (default) 1: Compile pointwise fusions with cuDNN ``` ### Tests `xla/backends/gpu/codegen/cudnn_test.cc` has broad coverage including: pointwise elementwise ops, broadcasts, logical ops, clamp, reciprocal (`f32`), concat, slice, reshape/bitcast, and various dtype combinations, with test names/traits reorganized and shapes reduced to keep runtime down. `cudnn_non_gemm_fusion_rewriter_test.cc` covers the rewriter pass itself, including the multi-concatenate skip case. Copybara import of the project: -- 66b84f3bbd8c1b7c1d3cfe01554596484f437282 by ske <ske@nvidia.com>: rebased -- e0f8322f38276a6fab45e1daad90e94638fa6c1d by ske <ske@nvidia.com>: add arch >= sm100 for cudnn non gemm fusion -- 307bef7c651e8e87253188c7b33cba26931617e9 by ske <ske@nvidia.com>: fix clang format -- 62935e1c681eab2d56f246d75429f6cd5155851a by ske <ske@nvidia.com>: add shard count to avoid timeout -- adc374bfcae0ab2299aa68840502b309d5734024 by ske <ske@nvidia.com>: addressed comments & fix format -- 3a4e24586c7e5865c1d9a3bee58f73c0f0cdc389 by ske <ske@nvidia.com>: add do deactivate for tensor ir engine -- af1d07e4565bfe8787a88a4da8f4fe8a0f783086 by ske <ske@nvidia.com>: rebase & fix CI errors -- 504aa01f516c89a22135db6bd754f99394386e0b by ske <ske@nvidia.com>: fix clang format -- 8efeb30deb4bb760c680028b96e1e560f06eb5de by ske <ske@nvidia.com>: address comments -- 45e30c7d7de97d266cb6edd2ba18a11232d705b8 by ske <ske@nvidia.com>: fix clang -- 8cf60c4eb50a017137155f7729c292f69124bdac by ske <ske@nvidia.com>: add deactivate on exit & test for multi concat -- ec3a5b6ec062222eae2a6338d2d5a6dde1bbd279 by ske <ske@nvidia.com>: initialize device ordinal as -1 & fix unit test accordingly Merging this change closes tensorflow#47839 PiperOrigin-RevId: 1013780671
None of our existing GPU benchmarks measures how long it takes to load a serialized executable or how much host memory a loaded executable keeps alive. We want to reduce the host-side cost of loading AOT compiled GPU executables, so we need a baseline first. This change adds `BM_LoadSerializedExecutable` (PjRt `LoadSerializedExecutable`) and `BM_GpuAotFromSerializedAndLoad` (deserializing the AOT compilation result and loading it into an `Executable`, below the PjRt layer) to `//xla/pjrt/gpu:se_gpu_pjrt_client_benchmark_test`. `BM_HloModuleFromProto` measures only the creation of the optimized `HloModule` from its proto plus its fingerprint, which is an upper bound of what we can save by not materializing the module. `BM_FirstExecuteAfterLoad` and `BM_ExecuteAfterDeserialize` cover the first and the steady state executions after loading. The destruction of the loaded objects is not part of the measured time. All of them except `BM_ExecuteAfterDeserialize` also report the heap bytes retained by one result (for `BM_FirstExecuteAfterLoad` after the first execution) as a `heap_bytes_retained` counter. This counter is based on `mallinfo`, so it's only reported with glibc. The benchmarks run on a tiny module, a synthetic chain of 500 dots and two real models from `xla/tools/benchmarks/hlo` (Gemma4 2B and a MaxText train step). The execution benchmarks use zero-initialized arguments and skip the MaxText module since it donates its arguments. The H100 variant now runs on a full GPU since the Gemma4 arguments don't fit into a MIG partition. PiperOrigin-RevId: 1013781571
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 join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Motivation
Technical Details
Weekly sync from upstream tensorflow repo
Test Result
TBA
Submission Checklist