Updates `build-linux-x86-cpu`, `build-linux-x86-cuda`, `build-linux-x86-cuda13`, and `build-arm64` in `.github/workflows/ci-build.yml` to run on `linux-x86-n4-16` self-hosted runners.
PiperOrigin-RevId: 1003163639
Generating and updating empirical HBM bandwidth scaling tables for the GPU dot
fusion cost model across hardware architectures requires running bandwidth
sweeps directly on target GPU devices outside of unit test runners, with
configurable transfer size ranges and iteration counts.
Add a standalone `gpu_bandwidth_benchmark_main` CLI binary that runs
`MeasureBandwidthTable` for configurable GPU device ordinals, power-of-two
transfer size exponents, and iteration counts, and prints the formatted table.
PiperOrigin-RevId: 1003160442
Estimating HLO dot fusion memory transfer times requires empirical bandwidth
scaling curves across transfer sizes from 8 KiB to 8 GiB. Previously,
`MeasureD2dBandwidthBytesPerSec` timed repeated `Stream::Memcpy` calls without
flushing L2, measuring warm L2 hits at small sizes and 2x HBM traffic above L2.
Replace `MeasureD2dBandwidthBytesPerSec` in
`xla/tools/cost_model/gpu_bandwidth_benchmark.h` with `MeasureBandwidthTable`,
which pre-allocates the target and L2 flush buffers, zeroes a `max(2 *
l2_cache_size, 64 MiB)` buffer before each timed `Stream::MemZero` (capturing
sub-L2 launch/write ramp-up and cold-L2 HBM write saturation above L2), and
returns mean bandwidth fractions in `[0.0, 1.0]` across power-of-two sizes.
PiperOrigin-RevId: 1003121761
This change moves the definition of TfLiteInterpreterCreate(), which is the
only function that uses CreateOpResolver(), into a separate "c_api_create.cc" file.
It adds a separate build target ":c_api_nocreate"
in `tensorflow/lite/core/c` and a corresponding forwarding target in
`tensorflow/lite/c`.
The purpose of this is to allow library targets (such as custom ops and delegates)
that use the C API without calling TfLiteInterpreterCreate() to
depend on ":c_api_nocreate" to avoid taking a
dependency (defined or undefined) on CreateOpResolver(),
and thus avoid the need to decide whether to depend on the builtin CreateOpResolver or to require clients to define their own CreateOpResolver(),
each of which have their own drawbacks. This is important for migration of
custom ops from the unstable custom ops API to the stable custom ops API,
since that migration introduces new dependencies on the C API,
which can then run into additional difficulties due to that implying a dependency
on CreateOpResolver().
The opaque C API never calls TfLiteInterpreterCreate(), so ":c_api_opaque" now
depends on ":c_api_nocreate". ":c_api_opaque_without_op_resolver" is now
redundant and is kept only as a deprecated alias of ":c_api_opaque".
PiperOrigin-RevId: 1003104081
Previously, `Resource::kCollectiveCommunicator` mapped to `NodeEdge::Kind::kScheduling`, which only waited for an earlier thunk's `Execute()` call to return before starting the next collective. When a `WhileThunk` or `ConditionalThunk` ran asynchronously on worker threads, its `Execute()` returned while collectives inside the loop or branch were still running or not yet started. Downstream collectives then ran at the same time as the inner collectives, causing message swaps over Gloo (which reuses tag 0) and rendezvous crashes in `InProcessCommunicator`.
Map `Resource::kCollectiveCommunicator` to `NodeEdge::Kind::kExecution` so each collective and each container thunk holding a collective completes before the next collective starts. Also bump `JAX_IFRT_VERSION_NUMBER` to `79` in `version.h` and add single-process (`shard_map_test.py`) and multi-process (`tests/multiprocess/all_reduce_test.py` and `tests/multiprocess/ppermute_test.py`) JAX regression tests guarded on `ifrt_version >= 79`.
Fixes https://github.com/openxla/xla/issues/50882
PiperOrigin-RevId: 1003082004
The test intermittently aborts with HSA_STATUS_ERROR_ILLEGAL_INSTRUCTION on the MI350X CI runners (crashed in about 1 in 3 executed runs since Oct 9, across unrelated PRs). The `flaky` tag on the amdgpu_any backend excludes it from ROCm CI via --config=no_flaky; other backends are unaffected. Also documents the `flaky` tag in build_tools/lint/tags.py, which the bazel-tags lint requires.
PiperOrigin-RevId: 1003056581
Triton can now codegen collectives such as all-to-all, so GemmFusionV2 started
pulling them into `__triton_gemm` fusions. These fusions need collective
parameters and a device assignment at runtime, which the autotuner does not
provide when profiling a fusion in isolation, so every config failed and
compilation aborted.
Exclude all collective communication ops from the GEMM fusion search space.
Fixes https://github.com/openxla/xla/issues/51073.
PiperOrigin-RevId: 1003054327
When `AlgebraicSimplifierVisitor::HandleTranspose` deletes an identity
transpose, record a recovery entry in `OriginalValueRecoveryTable` before
calling `ReplaceInstruction` so original value correspondence is preserved.
PiperOrigin-RevId: 1003053645
Handle negative remainders and CHECK that divisors are positive in
getExplicitRemainder, CHECK that divisors are positive and guard multiplication
overflow in isStructurallyDivisible, and extend isStructurallyDivisible to
handle unsigned min and signed and unsigned max operations.
PiperOrigin-RevId: 1003049190
Update tf_gpu_kernel_library and tf_custom_op_library in tensorflow.bzl to
dispatch to rocm_library from @config_rocm_hipcc//rocm:build_defs.bzl when ROCm
is configured instead of using cuda_library with rocm_copts. Bump
rules_ml_toolchain in TensorFlow's workspace3.bzl and MODULE.bazel so
rocm_library is available in @config_rocm_hipcc.
PiperOrigin-RevId: 1003047308
Enables direct lowering of `chlo.asin` and `chlo.acos` to `mhlo.asin` / `mhlo.acos` and `HloOpcode::kAsin` / `HloOpcode::kAcos` for CPU and TPU backends, rather than decomposing them early in `ChloLegalizeToStablehloPass` (with `enable_f64_asin` and `enable_f64_acos` disabled on CPU and TPU so `f64` lowers via CHLO expansion).
* Adds `getCpuChloToHighLevelMhloOptions()` and `getTpuChloToHighLevelMhloOptions()` in `chlo_legalize_to_hlo_pass.cc` enabling `acos` and `asin`.
* Updates `xla::Asin` and `xla::Acos` in `xla/hlo/builder/lib/math.cc` to emit `HloOpcode::kAsin` and `HloOpcode::kAcos` directly, and moves the `F64`, `C64`, and `C128` HLO expansions into `ElementaryFunctionExpander` (with `F64` `acos` using `select(x^2 <= 0.5, 1 - x^2, (1 - x) * (1 + x))` to avoid cancellation near `|x| = 1`).
* Implements degree-4 minimax polynomial approximations for `asin` (`F32`) and `acos` (`F32`) accurate to <= 1 ULP across `[-1, 1]`:
- CPU: Defines C++ vector intrinsics in `xla/codegen/intrinsic/cpp/trig.{h,cc}` using `HornerPoly`, `Fmuladd`, and 13-bit mantissa splitting (1 implicit + 12 fraction bits) supporting scalar and vector `F32`, registers them in `intrinsic_lib`, adds `LowerIntrinsicPattern<ci::Asin, mm::AsinOp>` and `LowerIntrinsicPattern<ci::Acos, mm::AcosOp>` in `LowerXlaIntrinsicLibPass`, and updates `CpuElementalIrEmitter::EmitAsin` and `CpuElementalIrEmitter::EmitAcos` to call the `F32` intrinsics for `F32`/`F16`/`BF16` and return `Unimplemented` for other types.
* Updates `lax.asin_p` in JAX `lax.py` to lower directly to `chlo.asin` when `jaxlib_extension_version >= 512` and not in forward-compatibility mode, also gating TPU lowering on `libtpu >= 0.0.51` now that StableHLO PR #3024 is integrated.
* Updates `AsinOp` and `AcosOp` reference evaluation in `exhaustive_unary_test_ops.inc` to evaluate `std::asin` and `std::acos` in `double` precision, and tightens exhaustive ULP bounds for both `asin` and `acos` in XLA CPU exhaustive tests (`bf16`/`f16`: 0 ULP, `f32`: 1 ULP, `f64`: 2 ULP) and JAX `trig_test.py` and `accuracy.md` (`bf16`: 0.5 ULP, `f16`: 0.5 ULP for `asin` and 1.0 ULP for `acos`, `f32`: 1.0 ULP, `f64`: 1.5 ULP).
PiperOrigin-RevId: 1003045300
When `FindElementwiseSubgraphSurroundedByReshapesAndBroadcasts` and `ReplaceElementwiseGroupSurroundedByReshapesAndBroadcasts` encounter a non-elementwise operand for a `kReshape` rewrite with a known target shape, allow the rewrite if the operand already has a `kReshape` user with matching dimensions, and reuse that existing reshape instruction:
```
Before: After:
p1 [2] p0 [6] p1 [2] p0 [6]
| / \ | |
b [2,3] / r_other [2,3] b [2,3] r_other [2,3]
| / | \ / \
r_b [6] / ... d [2,3] ...
\ /
d [6]
|
r_d [2,3]
```
PiperOrigin-RevId: 1003013712
`GpuExecutableProto.buffer_assignment` and
`GpuExecutableProto.buffer_allocations_debug_summary` are populated for
debugging, profiling, and OOM logging only; `GpuExecutable::FromProto`
reconstructs runtime allocations exclusively from `buffer_allocations` (which
is already ignored by `CompareGPUExecutables`), while thunk buffer slices,
output allocations, and constant allocations are verified through `thunks`,
`output_info_map`, and `constants`.
Ignore `buffer_assignment` and `buffer_allocations_debug_summary` in
`AOTInterceptionPjrtClient::CompareGPUExecutables` so compiler-internal changes
to `BufferAssignmentProto` or verbose debug formatting do not cause spurious
AOT golden verification failures.
PiperOrigin-RevId: 1002992849
Imported from GitHub PR https://github.com/openxla/xla/pull/49679
This PR registers oneAPI topology deserializer with the global pjrt topology description registry. It also updates the `FromProto` deserializers to accept the oneAPI platform so oneAPI topologies are deserialized the same way as other GPU platforms. It fixes errors like:
`NOT_FOUND: No PjRtTopologyDescriptionDeserializer registered for platform_name 'oneapi' or platform_id '13790883982447909111'.`
and
`The platform is not a GPU platform. platform_id: 13790883982447909111, platform_name: 'oneapi'`
Copybara import of the project:
--
6f4fe9bbd22de63a5f8f62adb2495e570249cec9 by Akhil Goel <akhil.goel@intel.com>:
Register oneAPI PJRT topology deserializer
Merging this change closes#49679
PiperOrigin-RevId: 1002986974
Imported from GitHub PR https://github.com/openxla/xla/pull/49794📝 Summary of Changes
Use rocm_library macro provided from withing rules_ml_toolchain
🎯 Justification
We need a full control over this macro withing rules_ml_toolchain
as it might be transformed to a custom rule for compilation if
we decide to remove the wrapper script
🚀 Kind of Contribution
Please remove what does not apply: ♻️ Cleanup
📊 Benchmark (for Performance Improvements)
Not relevant
🧪 Unit Tests:
CI
🧪 Execution Tests:
CI
Copybara import of the project:
--
430c4d26fe11b9ffbd864136bb81077872824cb6 by Alexandros Theodoridis <atheodor@amd.com>:
Migrate to use rocm_library_from_rules_ml_toolchain
--
abea2ee46b94a19d82ae824e04a29b67d7898fe6 by Alexandros Theodoridis <atheodor@amd.com>:
Update reference to rules_ml_toolchain
--
11789c8b9a4c8cbd9b9cf0cb26c98b820c287162 by Alexandros Theodoridis <atheodor@amd.com>:
Fix
--
9c5b76c83f8c61ecd7233fcec0ee6e4ef84c152e by Alexandros Theodoridis <atheodor@amd.com>:
Fix ident
--
13b8cd1f2e94f13b133b7b90af2950e3b30b25ab by Alexandros Theodoridis <atheodor@amd.com>:
Address review comments
--
b4fedfc82684719e7a7c321acc5e8788cfc5a5f7 by Alexandros Theodoridis <atheodor@amd.com>:
Mirror rocm_library impl in rules_ml_toolchain
--
ceb0f052f57f1c3614ec21b0d60584794f334389 by Alexandros Theodoridis <atheodor@amd.com>:
Update reference to rules_ml_toolchain
--
41a1848c2aaa0fe6247b1a9c01ddc29257b63fdd by Alexandros Theodoridis <atheodor@amd.com>:
Update reference to rules ml toolchain
Merging this change closes#49794
PiperOrigin-RevId: 1002986678
In `ArCrsCombiner::RunImpl`, `CallGraph::Build(module)` ran unconditionally
even though `call_graph_` is only read by `WhileFromBodyParameter` and
`ConditionalFromBodyParameter` in the `!spmd_partition_` (MPMD) path. Guard
`CallGraph::Build(module)` behind `if (!spmd_partition_)`.
In `ArCrsCombiner::KeepProvablyEqualInstructionGroupsSPMD`, return early when
`all_reduce_map_` is empty before running `HloReplicationAnalysis::Run` across
the full module.
PiperOrigin-RevId: 1002974253
`ExecutionGraph` checks `ResourceUse` conflicts by `Resource` pointer identity, and container thunks (`CallThunk`, `ConditionalThunk`, `WhileThunk`) report the `resource_uses()` of their nested thunk sequences. Previously, deserializing a nested thunk sequence in `ThunkSequenceSerDesProtobuf` always allocated new `Resource` instances instead of reusing the enclosing container thunk's resources, dropping resource dependency edges between container thunks and outer thunks after a proto round-trip.
Pass the enclosing container thunk's resolved resources as `parent_resources` when deserializing nested sequences in `CallThunk`, `ConditionalThunk`, and `WhileThunk`, and add a round-trip unit test in `thunk_sequence_serdes_test.cc`.
PiperOrigin-RevId: 1002974153
Imported from GitHub PR https://github.com/openxla/xla/pull/50617📝 Summary of Changes
Layout assignment currently handles async-start ops as follows:
- The callee computation determines its computation layout first.
- `AddAsyncStartConstraints` can add mandatory buffer constraints on the outputs of the async-op based on the callee computation's interface.
However, async-start ops' first output is a tuple that aliases all its operands. If there is a mismatch in the layout between the input to the async-start and the interface required by the callee computation, this means that `AddAsyncStartConstraints` will improperly constrain the operand with a mandatory buffer constraint. If the operand was constrained for some other reason (eg due to being a parameter of an entry computation, or due to being produced by a custom call), this can lead to layout assignment incorrectly constraining that value.
This PR excludes the operand-aliasing subtree at output index {0} from buffer layout constraints, including nested tuple leaves. Operand constraints still enforce the callee's required input layouts, allowing copies to resolve mismatches.
🎯 Justification
The newly added unit tests (`AsyncStartPreservesMandatoryCallerParameterLayout`, `AsyncStartPreservesMandatoryCallerCustomCallLayout`) demonstrate failures without this fix. Without this fix, the HLO inside `AsyncStartPreservesMandatoryCallerParameterLayout` crashes due to a mismatch between the required entry computation layout and the layout incorrectly set by LayoutAssignment due to adding a buffer constraint on the async-start's outputs. `AsyncStartPreservesMandatoryCallerCustomCallLayout` fails because a custom-call's layout is improperly overridden by LayoutAssignment.
🚀 Kind of Contribution
🐛 Bug Fix
🧪 Unit Tests:
- `AsyncStartPreservesMandatoryCallerParameterLayout`
- `AsyncStartPreservesMandatoryCallerCustomCallLayout`
Copybara import of the project:
--
b8c2cfbe7d2915bd89a8392208638a699f864dcb by Ashish Rao <asrao@nvidia.com>:
Do not add buffer constraints to aliased async-start outputs in AddAsyncStartConstraints
--
9ff52d727c66b9365b3fec386968548c9a54d76c by Ashish Rao <asrao@nvidia.com>:
Ensure layouts are forwarded correctly for AsyncStart
Merging this change closes#50617
PiperOrigin-RevId: 1002970548
Imported from GitHub PR https://github.com/openxla/xla/pull/50180
### 📝 Summary of Changes
Enable GPU BFC memory pools to grow beyond their initial allocation when explicitly requested. This connects the memory policies added in #50179 to actual allocator growth.
Set `--xla_gpu_memory_fraction_policy` through `XLA_FLAGS`:
- `0.75-0.85`: start at 75% of device memory and grow up to 85%.
- `0.75+`: start at 75% and grow up to 100%.
- `0.75`: keep a fixed 75% cap, matching the default behavior.
Growth applies to preallocated, spatially partitioned BFC pools. On CUDA, additional memory is mapped without moving existing allocations. Collective allocations stay within the initial region. Automatic registration of default-space NCCL buffers keeps the pool fixed.
### 🎯 Justification
Let workloads use available GPU memory during memory peaks while preserving stable collective-buffer addresses and existing defaults.
### 🚀 Kind of Contribution
✨ New Feature, 🐛 Bug Fix, 🧪 Tests
### 🧪 Unit Tests
Cover policy resolution, flag restoration, cap rounding, equal endpoints, growth limits, and registration pinning.
### 🧪 Execution Tests
One-GPU tests cover pool growth and error reporting for cross-host transfers spanning physical mappings. A two-GPU test checks NCCL windows remain valid when one GPU's pool grows.
Validation: all 40 selected tests passed on Prenyx with at most two GPUs, including the two-GPU NCCL and cross-mapping receive tests. No tests were skipped. Formatting and diff checks passed.
Copybara import of the project:
--
297e4e5f9db69d03cae2ad617928f5661e522e66 by Shaogang Wang <shawnw@nvidia.com>:
[XLA:GPU] PJRT: enable BFC growth for growable memory fractions
CreateBFCAllocator reserves the growth cap in device VA for FlexMemFraction
policies, maps the start fraction (or gpu_system_memory_size) up front and
lets BFC extend default memory on demand, rounding the cap down and the
initial allocation up to the mapping granularity and rejecting an initial
allocation that does not fit under the cap. The default policy is
Flex{0.75, 1.0}, so a plain XLA_PYTHON_CLIENT_MEM_FRACTION=0.75 grows;
"0.75-0.85" caps growth and "0.75-0.75" restores the fixed pool.
--xla_gpu_memory_fraction_policy overrides the client's policy for the BFC
allocator so XLA_FLAGS users can cap or pin the pool without a JAX change.
ResolveBfcMemFraction combines config, flag and the registration opt-in:
--xla_gpu_enable_nccl_user_buffers_in_default_space pins a growable policy
to a fixed pool with a warning, since automatic registration needs a fixed,
fully registered arena. Cross-host fabric transfers reject buffers that
span physical mappings. docs/gpu_bfc_growth.md documents the design, the
grammar and the run-ahead/headroom trade-off; a two-GPU test checks that
NCCL symmetric windows survive asymmetric growth.
Behavior change: preallocated spatial pools now grow by default. JAX's
gpu_memory_flags_test asserts pool_bytes == bytes_limit and needs updating
alongside this change. Last of five stacked changes (see #48985).
--
3074788acaa08bc94ded791de35ddf01027988a5 by Shaogang Wang <shawnw@nvidia.com>:
[XLA:GPU] Notify descriptor consumers when cross-host receive fails
--
6f8a709a5cda1b079894050c253bbfb833f292f6 by Shaogang Wang <shawnw@nvidia.com>:
[XLA:GPU] Align BFC growth tests and docs with explicit opt-in
--
842fc04f2cea16bffb58abd3b3b1f5cea13d1e40 by Shaogang Wang <shawnw@nvidia.com>:
[XLA:GPU] Fix BFC growth policy consistency and test isolation
Apply the resolved memory-fraction override to both default and separate
collective BFC pools while preserving explicit collective size overrides.
Restore the original XLA_FLAGS after scoped test overrides.
Keep strict cap rounding and test narrow ranges, alignment padding, and
invalid sizes. Treat valid equal-endpoint C++ policies like parsed fixed
policies, and avoid floating-point undershoot in byte-derived test caps.
Use xla_test for the NCCL growth test and split its single-GPU case into a
separate target.
--
61f99bc374e8cf99c1cd79dd64333422d518fa81 by Shaogang Wang <shawnw@nvidia.com>:
[XLA:GPU] Fix BFC growth include ordering
--
9f4d7e7894081bd4e809f7e7d09297b04872f3af by Shaogang Wang <shawnw@nvidia.com>:
[XLA:GPU] Use TSL threads in BFC multi-GPU test
--
b7b3c4289d94167f0179e0c71f540afb8a49b995 by Shaogang Wang <shawnw@nvidia.com>:
[XLA:GPU] Fix BFC growth receive formatting
Merging this change closes#50180
PiperOrigin-RevId: 1002957525
Imported from GitHub PR https://github.com/openxla/xla/pull/51011📝 Summary of Changes
Fix a deadlock when printing the communicator in the destructor, locks are acquired multiple times there and it's easy to get a deadlock.
Copybara import of the project:
--
528416f3407a16367e37f82bc17c0444ad00d36e by TJ Xu <tjx@nvidia.com>:
fix deadlock when logging is enabled for nccl communicator
--
a1ddd5ed026e523d19b5c72e6bd8f7485143a840 by TJ Xu <tjx@nvidia.com>:
added a quick unit test
--
4bd0b673dc2f10325b56df147c18558f788fd616 by TJ Xu <tjx@nvidia.com>:
run formatting
Merging this change closes#51011
PiperOrigin-RevId: 1002954859
In StableHLO, integer and boolean signedness is encoded directly in the operand
element types, making the SIGNED and UNSIGNED ComparisonType values redundant.
- Mark SIGNED and UNSIGNED as deprecated in StablehloEnums.td and ChloEnums.td,
and update the chlo.broadcast_compare and sign documentation.
- Update CompareOpCanon in StablehloAggressiveSimplification to simplify
integer comparisons when compare_type is omitted and actively strip redundant
SIGNED, UNSIGNED, and NOTYPE compare_type attributes.
- Update StableHLO tests to omit redundant SIGNED and UNSIGNED attributes.
PiperOrigin-RevId: 1002909951
Imported from GitHub PR https://github.com/openxla/xla/pull/50598📝 Summary of Changes
`RunGpuAsync` built the argument shapes from the executable's own entry computation layout and passed them to `LocalExecutable::RunHelper`, which compared them against that same layout. The check could never fail but cost a vector allocation and a `MatchesLayoutInShape` call per parameter on every launch. Input buffers are already validated by PJRT in `CommonPjRtLoadedExecutable::ExecutePrepare` before the handler runs.
Split execution-option validation and stream setup out of `RunHelper` into `LocalExecutable::PrepareRunOptions` and call it from `RunGpuAsync`. `LocalExecutable::Run/RunAsync` still validate arguments through `RunHelper`.
🎯 Justification
Unnecessary time spent on the host before launching the GPU work.
🚀 Kind of Contribution
Please remove what does not apply: ⚡️ Performance Improvement
📊 Benchmark (for Performance Improvements)
I don't have exact measurements for this commit, but layout matching alone on 1000+ arguments costs tens of microseconds.
Copybara import of the project:
--
0086af545e6af0a5cadbc4c83ed718880e9ad615 by Benjamin Rabier <benjamin@zml.ai>:
[XLA:GPU] Remove no-op argument shape check from RunGpuAsync
RunGpuAsync built the argument shapes from the executable's own entry
computation layout and passed them to LocalExecutable::RunHelper, which
compared them against that same layout. The check could never fail but
cost a vector allocation and a MatchesLayoutInShape call per parameter on
every launch. Input buffers are already validated by PJRT in
CommonPjRtLoadedExecutable::ExecutePrepare before the handler runs.
Split execution-option validation and stream setup out of RunHelper into
LocalExecutable::PrepareRunOptions and call it from RunGpuAsync.
LocalExecutable::Run/RunAsync still validate arguments through RunHelper.
Merging this change closes#50598
PiperOrigin-RevId: 1002907872
Imported from GitHub PR https://github.com/openxla/xla/pull/49601📝 Summary of Changes
Move buffer comparator mismatch counter to device memory.
Allocate the mismatch counter in device (VRAM) memory instead, so the atomics are serviced by the L2 cache. The single 8-byte result is still copied back to the host via the existing Memcpy, and the buffer is freed after BlockHostUntilDone.
Parallelize the host-side fallback comparison across a thread pool for large buffers, but only when a full scan is actually required (verbose logging or a non-null error_report). When only a boolean result is needed, the serial path is kept so it can early-exit on the first mismatch.
🎯 Justification
The GPU buffer comparison kernel performs one `atomicAdd(mismatch_count, 1)` per mismatching element. The counter was allocated in pinned host memory (HostMemoryAllocate), so every atomic became a PCIe read-modify-write to system RAM. When two buffers differ across (nearly) all elements - e.g. an 8.59e9-element / 32 GiB attention output during autotuning on MI350X those billions of serialized PCIe atomics took ~60s per comparison, ballooning autotuning by ~20 minutes.
Furthermore, when a report is required (which is the default behavior in level 4 “autotuner” mode), a second comparison between the reference buffer and the candidate output is performed on the host. This comparison was previously performed in a single thread and took a long time for large kernels. This PR therefore proposes to parallelize the comparison on the host side using a thread pool.
With [this](https://github.com/user-attachments/files/32739612/vae_autotune_segfault_repro.sh) small reproducer, we obtained the following performance numbers on MI350:
| Case | Default | Mismatch counter on device | Mismatch counter on device + parallelized host compare |
|------|--------:|---------------------------:|-------------------------------------------------------:|
| A_b2_at4_lt_shmap | 4599.600s | 2622.550s | 687.298s |
🚀 Kind of Contribution
Please remove what does not apply:⚡️ Performance Improvement
Copybara import of the project:
--
af0b4787e664f832b9f972dfe24ec6ad55a3f2c2 by Maxime France-Pillois <mfrancep@amd.com>:
Move buffer comparator mismatch counter to device memory and parallelize host compare
The GPU buffer comparison kernel performs one atomicAdd(mismatch_count, 1)
per mismatching element. The counter was allocated in pinned host memory
(HostMemoryAllocate), so every atomic became a PCIe read-modify-write to
system RAM. When two buffers differ across (nearly) all elements - e.g. an
8.59e9-element / 32 GiB attention output during autotuning on MI350X with
XNACK disabled - those billions of serialized PCIe atomics took ~60s per
comparison, ballooning autotuning by ~20 minutes.
Allocate the mismatch counter in device (VRAM) memory instead, so the
atomics are serviced by the L2 cache. The single 8-byte result is still
copied back to the host via the existing Memcpy, and the buffer is freed
after BlockHostUntilDone. This makes the mismatch case as fast as the
match case (~60s -> ~13ms per comparison).
Additionally, parallelize the host-side fallback comparison across a thread
pool for large buffers, but only when a full scan is actually required
(verbose logging or a non-null error_report). When only a boolean result is
needed, the serial path is kept so it can early-exit on the first mismatch.
--
50e5149641c7a911ca0186840e0ef9a5c067e799 by Maxime France-Pillois <mfrancep@amd.com>:
Fix clang tidy errors
--
0cb209af5c7cb07c6ea1ecb8033f91b3173419fb by Maxime France-Pillois <mfrancep@amd.com>:
Cache buffer comparator mismatch counter as a per-executor resource
The previous change allocated the device-side mismatch counter via
AllocateScalar/Deallocate on every comparison. On CUDA this goes through the
VMM allocator (cuMemCreate/cuMemMap/cuMemRelease) plus a context-wide
cuCtxSynchronize, which is very expensive at autotuning scale (hundreds of
comparisons per module) and regressed compile time by ~13% geomean (up to 2x
on some models).
Attach the 8-byte counter to the StreamExecutor as a Resource via
GetOrCreateResource, so it is allocated once per executor, reused across all
comparisons, and freed when the executor is destroyed. Reuse is safe because
DeviceCompare MemZero's the counter before each kernel launch. This keeps the
device-side (L2-served) atomic-counter benefit without the per-call allocation
cost.
Add a regression test pinning that the reused counter is correctly re-zeroed
between comparisons.
--
19b8c466dd016fa9e887f2fe7933ce90549a275f by Maxime France-Pillois <mfrancep@amd.com>:
Optimize comparator kernel to perform only 1 global memory access per block
--
0ba8415a815b77e94c34c2a061f9a6cb4b78d5f2 by Maxime France-Pillois <mfrancep@amd.com>:
Add mutex + Add missing dependencies
--
dbfc0aa42c793435dc72b4d018bd0ae3faac1bc2 by Maxime France-Pillois <mfrancep@amd.com>:
Use host memory instead of the device memory + mutex
--
36337019114e8d59459f38c3013243ef992d3284 by Maxime France-Pillois <mfrancep@amd.com>:
Detect block mismatch in kernels (instead of element mismatch) to avoid 32-bit overflow
Merging this change closes#49601
PiperOrigin-RevId: 1002901818
Deviceless cuDNN plans embed device properties built from the target spec, and cuDNN refuses to load them (`CUDNN_STATUS_NOT_SUPPORTED_ARCH_MISMATCH`) if they don't match the live device. The A100 specs don't set `reserved_shared_memory_per_block`, so every deviceless plan compiled for them fails to load.
The new `SpecMatchesLiveDevice` test compares the properties built from the spec against cuDNN's for the attached GPU, and fails on A100 without this fix (the existing test builds both from the live device).
I took the a100_sxm_40 values from the live device properties that the new test prints when it fails on an A100-SXM4-40GB:
Spec | Field | Value | Source
:--- | :--- | :--- | :---
a100_sxm_40 | `reserved_shared_memory_per_block` | 1024 | Test output (`reservedSharedMemoryPerBlock`)
a100_sxm_40 | `mem_clock_ghz` | 1.215 | Test output (`memClockRateKHz`), also in `model_str`
a100_sxm_40 | `max_blocks_per_multiprocessor` | 32 | Test output (`maxCtasPerSM`)
a100_sxm_40 | `device_memory_size` | 42298834944 | Test output (`totalGlobalMem`), also in `model_str`. The old value was larger than the device
a100_sxm_80, a100_pcie_80 | `reserved_shared_memory_per_block` | 1024 | Same sm80 chip, 1 KB reserved per block on sm80+ ([CUDA guide](https://docs.nvidia.com/cuda/archive/12.4.0/cuda-c-programming-guide/index.html#shared-memory-8-x))
PiperOrigin-RevId: 1002899247
Now that RemovePropagationBarriersPass runs in sdy::addExportPipeline before runShardyPartitioner, ReshardOp inputs no longer need to look through sdy.propagation_barrier.
Also add TODO(b/420837831) noting that PropagationBarrierPattern can be removed once ExportShardyForGSPMD is removed.
PiperOrigin-RevId: 1002898209
Import upstream tensor-ir v0.1.6, drop upstreamed patches (`static_grid_getter.patch`, `symbol_op_interface.patch`, and `tiling_analysis.patch`), drop `tiling_hotfixes.patch` (its `Enumerate.cpp` hunk was upstreamed and its `TileCandidateGenerator.cpp` hunk is obsolete since `TileCandidateGenerator` is deprecated and unused when tile sizes are specified), rename `pow_op.patch` to `cuda_tile_test_fixes.patch` after its `OutlinedComputeOps.cpp` hunk was upstreamed, refresh remaining patches, and update BUILD rules and XLA GPU TensorIR codegen callers for the new `CompilerInvocation`, `NVTensorIRCompilerOptions`, `TensorIRArtifact`, and `TilingEvaluator::create` APIs.
PiperOrigin-RevId: 1002891983
Imported from GitHub PR https://github.com/openxla/xla/pull/50873📝 Summary of Changes
Makes (again) JAX/XLA built with older CUDA use newer PTX versions supported by newer local ptxas.
nvptx_backend: Expose existing PTX version calculation as a new method.
nvptx_compiler: Fix target selection using the provider's PTX version.
🎯 Justification
Makes official JAX builds built with CUDA 13.0 support Rubin when running in CUDA 13.4+ environments.
🚀 Kind of Contribution
🐛 Bug Fix
🧪 Unit Tests:
yes
Copybara import of the project:
--
29c46e12dc56fbf61bf752c343b099592d082b8b by Ilia Sergachev <isergachev@nvidia.com>:
[GPU] Align PTX compilation and linking targets with the generated code.
nvptx_backend: Expose existing PTX version calculation as a new method.
nvptx_compiler: Fix target selection using the provider's PTX version.
--
389de2523d9f52c3e7235de254fbf05114b0ec6a by Ilia Sergachev <isergachev@nvidia.com>:
fix clang-tidy warnings
--
c10513cb8015c8372751f7d4fe8fd6123046d74d by Ilia Sergachev <isergachev@nvidia.com>:
Address review feedback.
Merging this change closes#50873
PiperOrigin-RevId: 1002889211
`LinkComputation` now names linked computations through the module's name uniquer. A nested computation already named `linked.N` no longer produces a duplicate name.
PiperOrigin-RevId: 1002870722
Import upstream cuda-tile v13.4.2, drop upstreamed patches (`llvm_cl1000840178.patch`, `llvm_cl_989877215.patch`, and `pow_op_alias.patch`), and grant `//third_party/gpus/tensor_ir:__subpackages__` visibility to `//third_party/gpus/cuda_tile/tools:cuda-tile-translate`.
PiperOrigin-RevId: 1002860839
Quoc Truong CA