[pull] master from tensorflow:master - #8785
Merged
Merged
Conversation
… across MLIR and TFLite - Corrected 'overwritting' to 'overwriting' in legalize_tf_collective error message and tests. - Fixed spelling of 'retrieved', 'separate', 'arguments', and 'implementation' in comments, test docstrings, and BUILD files.
Converting default_collective_perf_table.txtpb to binary proto at build time improves build performance. PiperOrigin-RevId: 971890342
Removes direct dependencies on computation_placer from test targets where it is not directly referenced or used by their sources. Also removes unused includes of computation_placer.h. PiperOrigin-RevId: 971892387
PiperOrigin-RevId: 971892553
Fixes openxla/xla#47778 as suggested in openxla/xla#47778 (comment) PiperOrigin-RevId: 971892743
Previously, 2-operand sort where the comparator could not be inlined fell
back to generic multi-operand sorting with `SortIterator` and `Inputs<2>`.
This introduced significant overhead due to offset multiplications on every
dereference, runtime switch statements in `Swap`/`Memcpy`, and oversized
pivot copies (`Value<2>`).
This change optimizes 2-operand sort with custom (non-inlined) comparators:
1. Reuses `ZipIterator` for 2-operand sort with arbitrary `LessThan*`
comparators, adapting `ZipRef` and `std::pair` to the 4-pointer
comparator ABI (`[lhs_k, rhs_k, lhs_v, rhs_v]`).
2. Dispatches based on operand byte widths ({1, 2, 4, 8} bytes), keeping
template instantiations minimal without requiring semantic types.
3. Unifies slice iteration and strided gather/scatter across inlined and
non-inlined 2-operand sort paths into a shared `Sort2DSlices` helper.
4. Moves internal helpers and iterators into the anonymous namespace without
logic changes: `ZipRef`, `ZipIterator`, `SortKeyValueSlice`,
`DSort1DInplace`, and the internal `Sort1DInplace` overloads.
Microbenchmark results (`BM_SortKeyValue2D<F32, S32>`, 1024x512, fallback):
- Single-threaded (`num_threads=0`): 135.8 ms -> 85.2 ms (1.59x speedup)
- Multi-threaded (`num_threads=16`): 17.5 ms -> 7.35 ms (2.38x speedup)
PiperOrigin-RevId: 971892843
…ation. This ensures that the constants created by ZeroSizedHloElimination are folded into their users immediately. Also remove unnecessary call to `ZeroSizedHloElimination` and only keep the one from `RunPreSPMDPartitionerPasses`. PiperOrigin-RevId: 971893469
…-agnostic API Imported from GitHub PR openxla/xla#46911 📝 Summary of Changes Implements the collectives FFI extension for XLA:GPU: GpuCollectivesState + MakeCollectivesExtension map request/get onto GPU clique requests/acquisition; custom_call_thunk attaches the extension to each invocation. 🎯 Justification Wires the extension to XLA:GPU so handlers get XLA's own ncclComm_t and reuse its clique management instead of duplicating it. Enables host-side NCCL collectives from custom calls, and a path to symmetric-memory registration. 🚀 Kind of Contribution ✨ New Feature, 🧪 Tests 📊 Benchmark (for Performance Improvements) Please measure and include speedups for one of the public HLOs in `compiler/xla/tools/benchmarks/hlo/`. 🧪 Unit Tests: //xla/backends/gpu/tests:collective_ops_ffi_test → CollectiveOpsTestFFI.PublicApiAllReduce 🧪 Execution Tests: N/A Copybara import of the project: -- 2950d246ed003183710807a01f7cc1c3b04d979e by Sevin F. Varoglu <svaroglu@nvidia.com>: [XLA:GPU] Add GPU collectives FFI extension on the backend-agnostic API -- 0e84c00a7435cbd294fd660099e758593b849655 by Sevin F. Varoglu <svaroglu@nvidia.com>: Add review feedback -- 0ef61a7af727097872b336853dc998ac068de41e by Sevin F. Varoglu <svaroglu@nvidia.com>: Fix clang tidy errors -- fbc96e977273ea70fad352c285c1ad2c403b31b8 by Sevin F. Varoglu <svaroglu@nvidia.com>: Add review feedback -- dba2dd30196acc6bd0e418b21d8683970d9549e7 by Sevin F. Varoglu <svaroglu@nvidia.com>: Clang format -- 5450156fb240730e69adb70ca3c2808e63984ad8 by Sevin F. Varoglu <svaroglu@nvidia.com>: Add review feedback Merging this change closes #46911 PiperOrigin-RevId: 971906179
…mma for ope… Imported from GitHub PR openxla/xla#47755 Fixes #122047. The shape equality check in XLA operations `Igamma`, `IgammaGradA`, `RandomGammaGrad`, `Igammac`, `Zeta`, and `Polygamma` was using `operator!=` (`Shape::Equal()`). This strict equality check inappropriately compares dynamic-dimension bits and layouts in addition to logical dimensions and element types. When one operand originates from a path that propagates a dynamic-dimension bit while the other operand does not, the two `xla::Shape` objects compare as unequal. This triggers a failure even though their printed representation and logical semantics are identical. The check would confusingly fire with: `"must have equal shapes and types; got f64[6,32,32] and f64[6,32,32]"`. **The Fix:** Replaced `operator!=` with `!ShapeUtil::Compatible()`, which is defined as `Shape::Equal().IgnoreDynamicDimension().IgnoreLayout()`. This safely accepts any pair of shapes that share the same rank, bounded dimensions, and element type, correctly validating the op constraints without falsely failing on dynamic dimension metadata. *Note: This PR ports the C++ side fixes originally authored in #125398 to OpenXLA so that they can be synced via Copybara.* Copybara import of the project: -- a94b626d9fb513a78d7843eacb863b8b15559c37 by arpittkhandelwal <arpitkhandelwal810@gmail.com>: Fix false shape-mismatch in XLA Igamma/Igammac/Zeta/Polygamma for operands with identical logical shape Merging this change closes #47755 PiperOrigin-RevId: 971907708
Imported from GitHub PR openxla/xla#47757 This PR is a follow-up to #47515 and enables the SoL latency estimator for gfx942 (MI300). Key changes: 1. Interpolator data (collective & matmul) ``` xla/service/gpu/model/default_collective_perf_table.txtpb xla/service/gpu/model/default_matmul_perf_table.txtpb xla/service/gpu/model/collective_interpolator_test.cc xla/service/gpu/model/matmul_interpolator_test.cc ``` 2. Enablement for the SoL estimator (gfx942) ``` xla/service/gpu/model/sol_latency_estimator.h xla/service/gpu/model/sol_latency_estimator.cc xla/service/gpu/model/sol_latency_estimator_test.cc xla/backends/gpu/transforms/collectives/collective_ops_utils.cc ``` Unlike gfx950, gfx942 uses native FNUZ (NANOO) FP8 types (`f8e4m3fnuz` and `f8e5m2fnuz`) instead of OCP FP8. Copybara import of the project: -- 8b455136a55c02eb7b9df2c8e1952573a0e25ffa by scxfjiang <xuefei.jiang@amd.com>: enable sol estimator for mi300 -- 052de1cc2f2fb7dbd84ddbdf0034687dd33c5545 by scxfjiang <xuefei.jiang@amd.com>: format Merging this change closes #47757 PiperOrigin-RevId: 971920190
Swapping the convolution operands can lead to flipping the filter operand. A flipped filter operand is not supported in epilogue fusions. PiperOrigin-RevId: 971934304
None of the testOnly libraries actually utilize ComputationPlacer. The dependencies and includes were legacy holdovers from before DeviceAssignment was extracted into its own header and target. Migrate targets and headers that only need DeviceAssignment to device_assignment and remove unused computation_placer dependencies. PiperOrigin-RevId: 971936135
Backport upstream commit 570be4dc70ba8f94af692f0a822a681cdb5bdbe6
("[Bazel] Enable LLVM threading on Windows (#218183)").
PiperOrigin-RevId: 971936364
…kernel Copies now always lower to CopyThunk. CopyThunk handles equal-shape sub-byte copies as a flat memcpy and layout-changing byte-width copies via TransposePlan. Layout-changing sub-byte copies route through the fusion emitters (wrapper change) rather than erroring, since CopyThunk cannot take them. PiperOrigin-RevId: 971943516
… with large input offsets. Fixes #118701 PiperOrigin-RevId: 971948653
…daries PiperOrigin-RevId: 971953596
…ofiler timeline (PR1) Imported from GitHub PR openxla/xla#46933 First of two stacked PRs adding ROCTX support to the ROCm profiler. Design discussion: openxla/xla#46782 This PR is the **listener**: ROCTX ranges emitted by an application — `roctxRangePushA`, `roctxRangePop`, `roctxMarkA` — now appear as named bands in the XPlane host-thread timeline, the ROCm counterpart to CUPTI's NVTX rows. XLA itself emits nothing here; that is PR2, which stacks on this one. ### What it does `RocmTracer::InitProfiling` registers `ROCPROFILER_CALLBACK_TRACING_MARKER_CORE_API` alongside the existing HIP-API, kernel-dispatch and memory-copy buffer services. `MarkerCallback` pairs push/pop into a `Generic` `RocmTracerEvent` and emits `roctxMarkA` as a zero-duration event. Range state is a `thread_local` stack rather than a mutex-guarded per-thread map. roctx push/pop is thread-local by definition and the marker callback runs synchronously on the calling thread, so nothing shared is required — which matters because the HIP-API callback reads the current label on **every** HIP call and must not touch a process-wide lock. Session isolation is by generation, not by clearing, since `Enable()` cannot reach another thread's stack: each frame carries the generation it was pushed in, `Enable()` bumps an atomic counter, and a pop whose frame predates it is dropped rather than emitted into the new session with the previous session's timestamp. `Generic` events bypass `ApiActivityInfoExchange` — they are host-side and have no GPU activity record to merge with — and are capped by `max_callback_api_events` with drops reported through `OnEventsDropped`. This also activates `kNVTXRange` on kernel events. The field and its stat emission already existed upstream but nothing ever wrote it; `GetCurrentRoctxLabel` and `AnnotationMap::LookUpRoctxRange` supply the writers. ### No XPlane schema change Markers land on the existing `kCuptiActivityNvtxPlaneName` plane rather than a new constant. **This PR touches no files under `xla/tsl/`.** The plane is a transient routing token: `PostProcessSingleHostXSpace` merges it into `/host:CPU` and `RemovePlanes`'s it before serialization, `MergePlanes` never reads the source plane name, and nothing downstream branches on it. Reuse is also safer across the PJRT plugin boundary, where the collector's XSpace is produced in the plugin binary and merged in the client binary — a plane name the client does not recognise would leak into the viewer unmerged. Lines are named `Host Threads/<tid>/ROCTX`, mirroring CUPTI's `/NVTX`. Lines are sorted by name after the merge, so this places each marker track directly beneath the thread that produced it. ### Also fixed Two pre-existing collector defects found while working here, both small and separable if you would rather they landed on their own: - a meaningless `kDeviceId=4294967295` stat emitted for events with no device - a race where `annotation_map_.Clear()` ran after `rocprofiler_start_context` and could wipe callbacks arriving in between ### Testing `rocm_tracer_test` — 13 marker tests added (push/pop, instantaneous mark, unmatched pop, null-label drop, null-label stack balance, export routing, the three `GetCurrentRoctxLabel` lifetime cases, the three `AnnotationMap` roctx cases). `rocm_collector_test` — 3 added, driving `RocmTraceCollectorImpl` directly with no rocprofiler context so they run on any host: cap enforcement, the `num_gpus_ == 0` drop path, and marker-vs-API line routing on the same thread. The cap test was verified to fail without the cap (0 drops instead of 17, 25 events retained instead of 8). On ROCm hardware, `RealRoctxCallsProduceNvtxRangeInXSpace` `dlopen`s `librocprofiler-sdk-roctx.so` and drives the real intercept path end to end. Runtime `dlopen`, so it adds no build-time dependency. Note `libroctx64.so` is **not** a substitute — rocprofiler-sdk does not intercept it. Measured: `rocm_tracer_test` 19 passed / 3 failed, `rocm_collector_test` 5 passed. The 3 failures are pre-existing tests requiring GPU device nodes (`hipErrorNoDevice`); `upstream/main` unmodified fails the same 3 on the same host. ### Known gap `roctxRangeStartA` / `roctxRangeStop` are not captured — the documented idiom for ranges that begin and end on different threads or that overlap, which a thread-local LIFO cannot express. `MarkerCallback` warns once rather than dropping them silently. CUPTI drops the NVTX equivalent too, so this is a shared limitation rather than a ROCm-only gap. Copybara import of the project: -- 7a3d96a455bb715127194fe23df1a9603ff199b2 by cj401-amd <chunyjin@amd.com>: [ROCm] Capture application-emitted ROCTX markers in the profiler timeline Adds the listener half of ROCTX support: ranges emitted by roctxRangePushA / roctxRangePop / roctxMarkA now appear as named bands in the XPlane host-thread timeline, the ROCm counterpart to CUPTI's NVTX rows. Listener. RocmTracer::InitProfiling registers ROCPROFILER_CALLBACK_TRACING_MARKER_CORE_API on the session context alongside the existing HIP-API, kernel-dispatch and memory-copy buffer services. MarkerCallback handles roctxRangePushA, roctxRangePop and roctxMarkA, emitting a Generic RocmTracerEvent for each completed range and each instantaneous mark. Range state is a thread_local stack, not a mutex-guarded per-thread map. roctx push/pop is thread-local by definition and the marker callback runs synchronously on the calling thread, so no shared structure is needed -- which matters because the HIP-API callback reads the current label on every HIP call and must not touch a process-wide lock. It also means no per-thread bookkeeping outlives its thread, and there is no lock to order against collector_mutex_. Session isolation is by generation rather than by clearing, since Enable() cannot reach another thread's stack. Each RoctxFrame carries the generation it was pushed in; Enable() bumps an atomic counter and a pop whose frame predates it is dropped rather than emitted into the new session with the previous session's start timestamp. Routing. Generic events bypass ApiActivityInfoExchange via standalone_events_ -- they are host-side and have no GPU activity record to merge with -- and are capped by max_callback_api_events with drops reported through OnEventsDropped. Marker volume is application-driven and the buffer only drains at Flush(), so without the cap a long capture retains every marker for the whole session. They land on the existing kCuptiActivityNvtxPlaneName plane rather than a new schema constant. The plane is a transient routing token: PostProcessSingleHostXSpace merges it into /host:CPU and RemovePlanes's it before serialization, MergePlanes never reads the source plane name, and nothing downstream branches on it. Reuse is also safer across the PJRT plugin boundary, where the collector's XSpace is serialized in the plugin binary and merged in the client binary -- a plane name the client does not recognise would leak into the viewer unmerged. This PR therefore adds no XPlane schema constant and does not touch xla/tsl/. Lines are named "Host Threads/<tid>/ROCTX", mirroring CUPTI's "/NVTX". Lines are sorted by name after the merge, so this form places each marker track directly beneath the thread that produced it; the earlier "ROCTX Threads/<tid>" sorted into a separate alphabetical block, divorcing every marker track from its thread. kNVTXRange on kernel events. RocmTracerEvent::roctx_range and the stat emission in CreateXEvent already existed upstream but nothing ever wrote the field. This supplies the writers via GetCurrentRoctxLabel and AnnotationMap::LookUpRoctxRange, so a kernel dispatched inside an application ROCTX range carries that label. GetCurrentRoctxLabel returns absl::string_view. The HIP runtime API ENTER callback calls it on every HIP API call, so returning by value would allocate and copy a label on the hot path only for AnnotationMap::Add to intern a copy moments later, and would pay that cost even when both the label and the annotation are empty and nothing is stored. The view aliases the thread_local frame and is valid until this thread's next roctx call, which is the whole window the caller needs: the marker callback that pops runs synchronously on the same thread that is currently inside the HIP API callback, so no pop can interleave. RoctxFrame keeps owning its std::string -- the push path runs once per range rather than once per HIP call, so the copy is cheap there, and interning at push time would tie frame lifetime to a pool that Enable() clears and that max_annotation_strings can refuse to grow. AnnotationMap's correlation_map and roctx_range_map store std::reference_wrapper<const std::string> rather than absl::string_view. Both point into map_.annotations, a node_hash_set with pointer and reference stability on rehash, so either is safe; the reference_wrapper makes the backing store explicit. The public LookUp / LookUpRoctxRange still return absl::string_view. XLA itself emits no roctx here. On ROCm nvtx_utils_impl builds nvtx_utils_stub.cc, whose DefaultProfilerDomain() returns null, so ScopedAnnotation and jax.profiler.TraceAnnotation take the AnnotationStack path unchanged. Wiring XLA's own annotations to ROCTX needs a real emitter and brings a link dependency, so it is a separate change. Also fixes two pre-existing collector defects found while working here: a meaningless kDeviceId=4294967295 stat emitted for events with no device, and a race where annotation_map_.Clear() ran after rocprofiler_start_context and could wipe callbacks that arrived in between. Tests: rocm_tracer_test MarkerCallbackPushPopEmitsRoctxRange, MarkerCallbackMarkEmitsInstantaneousEvent, MarkerCallbackUnmatchedPopIsIgnored, MarkerCallbackNullLabelRangeIsDroppedNotEmitted, MarkerCallbackNullLabelRangeKeepsStackBalanced, MarkerEventAppearsInExportedXSpace, GetCurrentRoctxLabel{ReturnsTopOfStack,EmptyAfterPop, ViewIsValidUntilNextRoctxCall}, AnnotationMap{StoresRoctxRange,RoctxRangeEmptyWhenNotProvided, StoresRoctxRangeWhenAnnotationEmpty} rocm_collector_test MarkerEventsRespectMaxCallbackApiEvents, MarkerEventsDroppedWhenNoGpusReported, MarkerAndApiEventsOnSameThreadGetSeparateLines The cap test was verified to fail without the cap (0 drops instead of 17, 25 events retained instead of 8). GetCurrentRoctxLabelViewIsValidUntilNextRoctxCall walks the sequence the real callback performs -- read the view while the frame is live, intern it the way Add does, pop, then confirm the interned copy outlived the frame and the stack reads empty. It does not read the view after the pop. On ROCm hardware: RealRoctxCallsProduceNvtxRangeInXSpace dlopens librocprofiler-sdk-roctx.so and drives the real intercept path end to end. Runtime dlopen, so it adds no build-time dependency. Note libroctx64.so is not a substitute -- rocprofiler-sdk does not intercept it. Known gap: roctxRangeStartA / roctxRangeStop are not captured. They are the documented idiom for ranges that begin and end on different threads or that overlap, and a thread_local LIFO cannot express them. MarkerCallback warns once rather than dropping them silently. CUPTI drops the NVTX equivalent too, so this is a shared limitation rather than a ROCm-only gap. Incorporates review feedback from draganmladjenovic on PR #46933: the reference_wrapper change, the string_view return, and a corrected comment on RocmTracerEvent::roctx_range that referred to a roctx_strings_ member removed when the roctx stack became thread_local. Verified with --config=rocm: rocm_tracer builds clean under the repo's -Wall -Werror, and rocm_tracer_test passes 22/22 on a host with GPU device nodes available. clang-format 17.0.6 against main's .clang-format and buildifier 6.4.0 --lint=warn are both clean. Merging this change closes #46933 PiperOrigin-RevId: 971956931
This simple heuristic estimates warp distribution similar to what Triton does. This method will be used to estimate operand registers. PiperOrigin-RevId: 971959354
… NaNs. PiperOrigin-RevId: 971959429
…revent underflow and overflow. PiperOrigin-RevId: 971962733
Pads odd input and output channel dimensions (e.g. C=1 or C=3) to the nearest multiple of 2 for 16-bit float (BF16 and F16) convolutions. cuDNN requires 32-bit memory alignment for epilogue fusions. PiperOrigin-RevId: 971963939
…d Mac. This change explicitly checks if the .pb file exists before calling CreateFingerprintDefPb. If the file is missing, it directly returns a reduced fingerprint definition, avoiding unnecessary error handling from a failed file read. PiperOrigin-RevId: 971983971
… tolerate errors. - This matches the read behavior from legacy directory based autotune cache. - For cache writes, legacy cache did throw an error, but we extended it to digest the error. - It matches jax's philosophy of warning for cache errors, jax-ml/jax#12582. - I believe we won't have any write issues but we can confirm for a few days, if someone sees the warning and complain. PiperOrigin-RevId: 971996660
We were targeting an old default API version, preventing the use of NumPy 2.x features. PiperOrigin-RevId: 972010985
PiperOrigin-RevId: 972012488
… keyrings 1. Modernize setup.sources.sh: - Use HTTPS key downloads and scoped dearmored keys in /etc/apt/keyrings/ with [signed-by=...] in custom.list (replaces legacy apt-key keyserver lookup which timed out on port 11371). - Add set -euo pipefail for error safety. 2. Add CUDA 12.1 + cuDNN 9.10 container target: - Add cuda12.1_cudnn9.10.packages.txt pinning libcudnn9-*-cuda-12=9.10.2.21-1. - Add cuda12.1cudnn9.10 target and build step in build.sh. PiperOrigin-RevId: 972026025
…mlir-tflite PiperOrigin-RevId: 972052163
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.
19 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 : )