CI: 12/23/24 upstream sync#190
Merged
github-actions[bot] merged 52 commits intorocm-mainfrom Jan 2, 2025
Merged
Conversation
The `build.py` script uses Clang compiler by default, and JAX doesn't support building with GCC officially. However, experimental GCC support is still present. Command examples: ``` python build/build.py build --wheels=jaxlib,jax-cuda-plugin --use_clang=false python build/build.py build --wheels=jaxlib,jax-cuda-plugin --use_clang=false --gcc_path=/use/bin/gcc ``` This change addresses the request in jax-ml#25488. PiperOrigin-RevId: 707930913
PiperOrigin-RevId: 707934273
PiperOrigin-RevId: 707934347
PiperOrigin-RevId: 707935849
http://github.com/openxla/xla/commit/b44f55da3dac449f03466815ac431474f86fd73f. PiperOrigin-RevId: 707970174
PiperOrigin-RevId: 707973428
PiperOrigin-RevId: 707991933
PiperOrigin-RevId: 707992245
…aic GPU. PiperOrigin-RevId: 708010809
infer-vector-layout won't use the full generality anytime soon, but we could reuse this logic for relayouts PiperOrigin-RevId: 708011538
PiperOrigin-RevId: 708016513
…gardless the value of config.enable_x64. PiperOrigin-RevId: 708031525
… together This commit modifies the behavior of the build CLI when building jaxlib and GPU plugin artifacts together (for instance `python build --wheels=jaxlib,jax-cuda-plugin`. Before, CUDA/ROCm build options were only passed when building the CUDA/ROCm artifacts. However, this leads to inefficient use of the build cache as it looks like Bazel tries to rebuild some targets that has already been built in the previous run. This seems to be because the GPU plugin artifacts have a different set of build options compared to `jaxlib` which for some reason causes Bazel to invalidate/ignore certain cache hits. Therefore, this commit makes it so that the build options remain the same when the `jaxlib` and GPU artifacts are being built together so that we can better utilize the build cache. As an example, this means that if `python build --wheels=jaxlib,jax-cuda-plugin` is run, the following build options will apply to both `jaxlib` and `jax-cuda-plugin` builds: ``` /usr/local/bin/bazel run --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --verbose_failures=true --action_env=CLANG_COMPILER_PATH="/usr/lib/llvm-16/bin/clang" \ --repo_env=CC="/usr/lib/llvm-16/bin/clang" \ --repo_env=BAZEL_COMPILER="/usr/lib/llvm-16/bin/clang" \ --config=clang --config=mkl_open_source_only --config=avx_posix \ --config=cuda --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm-16/bin/clang" \ --config=build_cuda_with_nvcc ``` Note, this commit shouldn't affect the content of the wheel it self. It is only meant to give a performance boost when building `jalxib`+plugin aritfacts together. Also, this removes code that was used to build (now deprecated) monolithic `jaxlib` build from `build_wheel.py` PiperOrigin-RevId: 708035062
PiperOrigin-RevId: 708045439
PiperOrigin-RevId: 708045694
Followup to jax-ml#25614. PiperOrigin-RevId: 708077981
…ut-validation PiperOrigin-RevId: 708087898
For (1, 128) tiling 32-bit input, it assigns (1, 128) tiling at output, which can be invalid (e.g. it should be (1, 256) for bf16) PiperOrigin-RevId: 708112341
…for i1 vector relayout We can use relayout-insertion pass to insert necessary ops and their layouts for relayout before unrolling in apply-vector-layout pass. PiperOrigin-RevId: 708143852
Reverts 307c8d3 PiperOrigin-RevId: 708173083
PiperOrigin-RevId: 708196762
PiperOrigin-RevId: 708259645
…th and fix the bug. The bug was that bounds were dropped ctx.avals_in and then they were being extracted. Extract them before dropping them. PiperOrigin-RevId: 708266659
Otherwise, the cycle can be broken by clearing the references of the helper objects, at which points the deallocation of arrays proceeds through regular reference counting (and does not trigger logs!). I have not verified that this is what happens, but the test has been mysteriously failing under a number of configurations and this seems to fix it. I added a note to the garbage collection guard to clarify that it's not guaranteed to report all cycles. PiperOrigin-RevId: 708320953
PiperOrigin-RevId: 708321802
PiperOrigin-RevId: 708340432
… ops. Only untransformed and unsliced loads/stores are supported for now. The rest will be a follow up. PiperOrigin-RevId: 708347442
http://github.com/openxla/xla/commit/750765b1c996fcf15c82eb499ffec3441112235f. PiperOrigin-RevId: 708360218
PiperOrigin-RevId: 708390390
…inus changes in infer-vector-layout We can enable them later but at least this way the support is available to build on (e.g. in the new insert relayouts pass) Reverts 05f3a70 PiperOrigin-RevId: 708397219
warnings.catch_warnings is not thread-safe. However it is always used to avoid complex-to-real conversion warnings, which we can avoid in other ways.
PiperOrigin-RevId: 708404847
PiperOrigin-RevId: 708415848
PiperOrigin-RevId: 708419819
In jax-ml#24370, `ffi_call` was updated to return a callable, and the original calling convention was deprecated. This change is part of the deprecation cycle for this calling convention. PiperOrigin-RevId: 708424223
…space. PiperOrigin-RevId: 708450293
…le colocated Python call PiperOrigin-RevId: 708461989
This was causing an issue when building multiple wheels in editable mode. i.e instead of wheels being stored as: ``` # jax-cuda12-pjrt 0.4.36.dev20241125 ./dist/jax-cuda-pjrt # jax-cuda12-plugin 0.4.36.dev20241125 ./dist/jax-cuda-plugin # jaxlib 0.4.36.dev20241125 ./dist/jaxlib ``` they were being stored as: ``` # jaxlib 0.4.36.dev20241125 ./dist/jaxlib # jax-cuda12-pjrt 0.4.36.dev20241125 ./dist/jaxlib/jax-cuda-pjrt # jax-cuda12-plugin 0.4.36.dev20241125 ./dist/jaxlib/jax-cuda-plugin ``` PiperOrigin-RevId: 708468522
PiperOrigin-RevId: 708468858
http://github.com/openxla/xla/commit/dc7aaf834a0bb5a543f6cf98626284783a4a921c. PiperOrigin-RevId: 708629379
…selection driven by constant prop in mosaic lowering. This CL builds out a simple sketch of constant prop by construction in mosaic - we walk the graph up from cond, collecting the values and either const propping or failing out of const prop. Failure out of const prop is not a bug, but hitting an unimplemented const prop func is for now, in order to drive better coverage. This then allows us to pick a single branch, and ignore branches which do not have a viable mosaic implementation. And, finally, for diag, this means we can replace the initial gather-dependent implementation in lax with a mosaic specific one that avoids gather. PiperOrigin-RevId: 708752566
PiperOrigin-RevId: 708811348
charleshofer
approved these changes
Jan 2, 2025
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
20 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.
Daily sync with upstream