Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Rocm jaxlib v0.4.31 qa cleanup #40

Merged
merged 9 commits into from
Nov 7, 2024
Merged

Rocm jaxlib v0.4.31 qa cleanup #40

merged 9 commits into from
Nov 7, 2024

Commits on Sep 9, 2024

  1. Limit the number of additional ExecutionStreamIds to 4. 

    PiperOrigin-RevId: 663524714
    @tensorflower-gardener @hsharsha
    tensorflower-gardener authored and hsharsha committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    be07654 View commit details
    Browse the repository at this point in the history

  2. PR openxla#13597: Activation offloading dependency fix: insert a wait 

    Imported from GitHub PR openxla#13597

To ensure proper synchronization for the asynchronous copy, this CL makes the other stream wait for the completion of the operation (res_3) in the main stream.
  %param_1 = f32[1024]{0} parameter(1)
  %param_0 = f32[1024]{0} parameter(0)
  %res_3 = f32[1024]{0} fusion(%param_1, %param_0), kind=kInput, calls=mul
  %copy-start = (f32[1024]{0:S(5)}, f32[1024]{0}, u32[]) copy-start(f32[1024]{0} %res_3)
Copybara import of the project:

--
5e41cc4 by Jane Liu <[email protected]>:

Activation offloading dependency fix: insert a wait

--
c6c8b3e by Jane Liu <[email protected]>:

Assign streams to copy-start instructions and add a Waitfor thunk

--
9ffa12a by Jane Liu <[email protected]>:

copy-start always uses a new stream; add Waitfor for both D2H and H2D

Merging this change closes openxla#13597

COPYBARA_INTEGRATE_REVIEW=openxla#13597 from zhenying-liu:dependency-fix 9ffa12a
PiperOrigin-RevId: 666842831
    @zhenying-liu @hsharsha
    zhenying-liu authored and hsharsha committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    e717b65 View commit details
    Browse the repository at this point in the history

  3. disable workspace

    @pemeliya @hsharsha
    pemeliya authored and hsharsha committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    98994ec View commit details
    Browse the repository at this point in the history

  4. PR openxla#16190: [ROCm] Fix pjrt_c_api_gpu_test 

    Imported from GitHub PR openxla#16190

After these changes openxla@d86502a in cuda async allocator were introduced, `PjrtCApiGpuAllocatorTest.ValidOptionsParsing` started failing on ROCm.
Copybara import of the project:

--
ddc6207 by mmakevic <[email protected]>:

Skip cuda_async option for rocm in ValidOptionsParsing test

--
6d491be by mmakevic <[email protected]>:

Fix formatting issues

Merging this change closes openxla#16190

COPYBARA_INTEGRATE_REVIEW=openxla#16190 from ROCm:ci_pjrt_c_api_test_fix 6d491be
PiperOrigin-RevId: 667883110
    @mmakevic-amd @hsharsha
    mmakevic-amd authored and hsharsha committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    cc4ee79 View commit details
    Browse the repository at this point in the history

  5. Disable FP8 rewrite pattern test on ROCm

    @hsharsha
    hsharsha committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    422213a View commit details
    Browse the repository at this point in the history

  6. Use deterministic ops flag in determinism test

    @hsharsha
    hsharsha committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    5406c76 View commit details
    Browse the repository at this point in the history

  7. PR openxla#16037: [ROCm] Get platform name from PlatformUtil helper 

    Imported from GitHub PR openxla#16037

Copybara import of the project:

--
6303981 by Harsha HS <[email protected]>:

[ROCm] Distinguish between NVIDIA and AMD gps in `gpu_device_info_test`

--
246e410 by Harsha HS <[email protected]>:

Simplify build by removing conditionals as per review comment

Merging this change closes openxla#16037

COPYBARA_INTEGRATE_REVIEW=openxla#16037 from ROCm:ci_fix_gpu_device_info_test_20240813 246e410
PiperOrigin-RevId: 663630606
    @hsharsha
    hsharsha committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    a633a28 View commit details
    Browse the repository at this point in the history

  8. PR openxla#16213: [ROCm] Bump ROCm version to 6.2 and fix gemm_rewrit… 

    …er_test

Imported from GitHub PR openxla#16213

Copybara import of the project:

--
0dba5ad by Harsha HS <[email protected]>:

[ROCm] Bump ROCm version to 6.2 and fix gemm_rewriter_test

--
81bea23 by Harsha HS <[email protected]>:

Fix clang formatting

Merging this change closes openxla#16213

COPYBARA_INTEGRATE_REVIEW=openxla#16213 from ROCm:ci_fix_gemm_rewriter_test_20240819 81bea23
PiperOrigin-RevId: 666859220
    @hsharsha
    hsharsha committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    0e8ee68 View commit details
    Browse the repository at this point in the history

  9. [ROCM] Add basic scaffolding and enable MLIR fusion

    @draganmladjenovic @hsharsha
    draganmladjenovic authored and hsharsha committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    23df4ae View commit details
    Browse the repository at this point in the history