Repository navigation
[pull] master from tensorflow:master - #8857
Merged
Merged
Conversation
Extend `triton_xla.block_barrier` with optional `signal_slot` operand and `signal_stride` attributes to support asymmetric block synchronization across GPUs. Before this change Block B on R0 waited for Block B on other GPUs. With this now we can do something like Block B waits for a strided subset of R (world_size) blocks from another GPU. PiperOrigin-RevId: 989715923
…terized runs PiperOrigin-RevId: 989761006
tflite/c/CMakeLists.txt still expected a TensorFlow checkout layout. - Derived TFLITE_SOURCE_DIR from the tflite/c directory and dropped TF_SOURCE_DIR. - Took tflite_types.h from tflite/converter. - Moved the TensorFlow header download pin from v2.21.0-rc0 to v2.21.0. The file hashes are unchanged. PiperOrigin-RevId: 989763450
… regions. This case is different from existing tests that cover duplications in nested regions. This one adds coverage for sibling regions, which shouldn't see each other during emission. PiperOrigin-RevId: 989774195
…memory partitions - Set `NVIDIA_TF32_OVERRIDE=0` for DTensor GPU tests so FP32 tests run with IEEE FP32 precision rather than TF32 on Ampere+ GPUs. - Disable `NCCL_CUMEM_ENABLE`, `NCCL_CUMEM_HOST_ENABLE`, and `NCCL_NVLS_ENABLE` in 2-GPU tests for multi-process NCCL compatibility in containerized test environments. - Reduce `_DEFAULT_GPU_MEMORY_LIMIT` from 1024MB to 512MB per virtual GPU in `test_util.py` so 8 virtual devices plus CUDA/cuSolver runtime overhead fit within smaller GPU memory partitions. PiperOrigin-RevId: 989774536
Also, fix a minor formatting issue in `Api.cpp`. PiperOrigin-RevId: 989776315
…ast reference. Prevents a thread pool starvation deadlock during `EagerContext` destruction by unconditionally offloading the `Unref()` call via `env_->StartDetachedThread` when called from an executor thread pool. PiperOrigin-RevId: 989776895
PiperOrigin-RevId: 989779426
…d all `HloOpcode::kCustomCall` instructions, isolating every typed FFI custom kernel into its own singleton concurrent region with barrier dependencies before and after. Additionally, `CommandBufferConversionPass` did not map `Thunk::kCollectiveKernel` to `DebugOptions::COLLECTIVES`, causing one-shot collective kernels to break and split surrounding CUDA Graphs. Update `CanRunConcurrently` in `gpu_hlo_ordering.cc` to allow `API_VERSION_TYPED_FFI` custom calls whose registered GPU FFI handler specifies `xla::ffi::Traits::kCmdBufferCompatible`, while continuing to isolate legacy custom calls (such as cuBLAS and cuDNN). Map `Thunk::kCollectiveKernel` to `DebugOptions::COLLECTIVES` and `DebugOptions::ALLREDUCE` in `command_buffer_conversion_pass.cc` so one-shot collective kernels are captured inside command buffers. PiperOrigin-RevId: 989779912
…yTest.ExtractsFuncAttributes PiperOrigin-RevId: 989782323
…rom openxla/xla#48616) PiperOrigin-RevId: 989782947
…kTest PiperOrigin-RevId: 989784865
PiperOrigin-RevId: 989793164
PiperOrigin-RevId: 989802032
… devices(). `xla::ifrt::LoadedExecutable::addressable_devices()` is superseded by `xla::ifrt::LoadedExecutable::devices()`, which returns `std::optional<DeviceListRef>` and yields `std::nullopt` for portable executables. Update remaining callers across JAX, IFRT IR, and IFRT Proxy to query `LoadedExecutable::devices()` and derive the addressable device list (`(*devices)->AddressableDeviceList()->devices()`) from the returned `DeviceListRef`. When `devices()` returns `std::nullopt`, callers fall back to an empty addressable device list (or return `UnimplementedError` where portable executables are unsupported). Remove `LoadedExecutable::addressable_devices()` from the base interface, mocks, and all subclass implementations. PiperOrigin-RevId: 989806232
PiperOrigin-RevId: 989812768
…done custom call pairs into async instructions. In legacy mode, the pair becomes an `all-gather-start` / `all-gather-done` pair configured from the all-gather inside the called computation, and in generic mode it becomes an async call wrapping that computation. Only computations built from all-gather and parameter instructions with exactly one single-operand all-gather are supported; anything else returns `UNIMPLEMENTED`. More support will be added after we agree on the general architecture in this CL. Backends can attach backend-specific metadata to the rewritten start op via the new `AsyncComputeOnHelper` extension point. PiperOrigin-RevId: 989821690
…plicit nodes. Currently `CollectiveKernelThunk` can be recorded only with stream capture which takes around 0.25 ms to update at each iteration. With explicit API usage an update time can be reduced to 0.05 ms. This change is also a first step to move invocation counter to the device memory which will make graph update time almost zero. Add e2e tests that run 10 consecutive all-gathers and 10 consecutive all-reduces. PiperOrigin-RevId: 989838065
…e and PjRtMemorySpace PiperOrigin-RevId: 989842055
`xla::ifrt::LoadedExecutable::num_devices()` is superseded by `xla::ifrt::LoadedExecutable::devices()`. Callers have been updated: - In JAX (`pjit.cc`), when `devices()` returns `std::nullopt` (portable executables), `num_global_devices` evaluates to 1. This is guarded with `#if JAX_IFRT_VERSION_NUMBER >= 74`. - In IFRT IR pass and tests, device counts are derived directly from `(*devices)->size()`. - Bump `JAX_IFRT_VERSION_NUMBER` to 74 in `version.h`. - Remove `LoadedExecutable::num_devices()` from the base interface, mocks, and all subclass implementations. PiperOrigin-RevId: 989844342
PiperOrigin-RevId: 989853661
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.
12 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 : )