[Issue]: HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION with concatenate fusions.

[owainkenwayucl]

Problem Description

While working on AlphaFold 3 (see #404) I stumbled across another bug with JAX 0.9.1 on ROCm 7.2.2, specifically:

GPU coredump: execvp failed: No such file or directory
Failed to write segment data to pipe: Broken pipe
GPU coredump: handler exited with error (status: 1)
GPU core dump failed
Kernel Name: input_concatenate_fusion
VGPU=0x808f720 SWq=0x7f394c006000, HWq=0x7f38d8400000, id=2
        Dispatch Header =0xb02 (type=2, barrier=1, acquire=1, release=1), setup=0
        grid=[15, 1, 1], workgroup=[15, 1, 1]
        private_seg_size=0, group_seg_size=0
        kernel_obj=0x7f392c00c600, kernarg_address=0x0x7f3828200200
        completion_signal=0x0, correlation_id=0
        rptr=17, wptr=20
 Kernel Name: input_concatenate_fusion
VGPU=0x80b5fe0 SWq=0x7f394c006000, HWq=0x7f38d8400000, id=2
        Dispatch Header =0xb02 (type=2, barrier=1, acquire=1, release=1), setup=0
        grid=[15, 1, 1], workgroup=[15, 1, 1]
        private_seg_size=0, group_seg_size=0
        kernel_obj=0x7f392c00c600, kernarg_address=0x0x7f3828200200
        completion_signal=0x0, correlation_id=0
        rptr=17, wptr=20
 Kernel Name: input_concatenate_fusion
VGPU=0x80d34c0 SWq=0x7f394c006000, HWq=0x7f38d8400000, id=2
        Dispatch Header =0xb02 (type=2, barrier=1, acquire=1, release=1), setup=0
        grid=[15, 1, 1], workgroup=[15, 1, 1]
        private_seg_size=0, group_seg_size=0
        kernel_obj=0x7f392c00c600, kernarg_address=0x0x7f3828200200
        completion_signal=0x0, correlation_id=0
        rptr=17, wptr=20
 Kernel Name: input_concatenate_fusion
VGPU=0x80f0b40 SWq=0x7f394c006000, HWq=0x7f38d8400000, id=2
        Dispatch Header =0xb02 (type=2, barrier=1, acquire=1, release=1), setup=0
        grid=[15, 1, 1], workgroup=[15, 1, 1]
        private_seg_size=0, group_seg_size=0
        kernel_obj=0x7f392c00c600, kernarg_address=0x0x7f3828200200
        completion_signal=0x0, correlation_id=0
        rptr=17, wptr=20
 Kernel Name: input_concatenate_fusion
VGPU=0x810e200 SWq=0x7f394c006000, HWq=0x7f38d8400000, id=2
        Dispatch Header =0xb02 (type=2, barrier=1, acquire=1, release=1), setup=0
        grid=[15, 1, 1], workgroup=[15, 1, 1]
        private_seg_size=0, group_seg_size=0
        kernel_obj=0x7f392c00c600, kernarg_address=0x0x7f3828200200
        completion_signal=0x0, correlation_id=0
        rptr=17, wptr=20
:0:rocdevice.cpp            :3586: 87036087293 us:  Callback: Queue 0x7f38d8400000 aborting with error : HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION: The agent attempted to access memory beyond the largest legal address. code: 0x29
Aborted (core dumped)

Operating System

NAME="Ubuntu" VERSION="24.04.4 LTS (Noble Numbat)" inside container.

CPU

AMD EPYC 9654 96-Core Processor * 2

GPU

AMD Instinct MI300X

ROCm Version

ROCm 7.2.2

ROCm Component

No response

Steps to Reproduce

I've managed to generate a minimum JAX program that triggers this problem:

import jax, jax.numpy as jnp

def frame(a, b):
    return jnp.concatenate([a[None], b[None], jnp.cross(a, b)[None]], axis=0)

a = jnp.ones((5, 3)); b = jnp.ones((5, 3))
out = jax.jit(jax.vmap(frame))(a, b)  # may need shape tweaks to match [1,5,3]/[5,3]
out.block_until_ready()

To run the code:

root@4fa981d0ce64:~/af_output/repro# export XLA_FLAGS='--xla_dump_hlo_pass_re=.*fusion.* --xla_dump_to=/root/af_output/dump_repro2 --xla_dump_hlo_as_text' 
root@4fa981d0ce64:~/af_output/repro# python3 repro.py 
GPU coredump: execvp failed: No such file or directory
Failed to write segment data to pipe: Broken pipe
GPU coredump: handler exited with error (status: 1)
GPU core dump failed
Kernel Name: input_concatenate_fusion
VGPU=0x92bf7c0 SWq=0x7f6050006000, HWq=0x7f5fb8400000, id=2
        Dispatch Header =0xb02 (type=2, barrier=1, acquire=1, release=1), setup=0
        grid=[15, 1, 1], workgroup=[15, 1, 1]
        private_seg_size=0, group_seg_size=0
        kernel_obj=0x7f6040614600, kernarg_address=0x0x7f5f80200200
        completion_signal=0x0, correlation_id=0
        rptr=18, wptr=21
 Kernel Name: input_concatenate_fusion
VGPU=0x92e6080 SWq=0x7f6050006000, HWq=0x7f5fb8400000, id=2
        Dispatch Header =0xb02 (type=2, barrier=1, acquire=1, release=1), setup=0
        grid=[15, 1, 1], workgroup=[15, 1, 1]
        private_seg_size=0, group_seg_size=0
        kernel_obj=0x7f6040614600, kernarg_address=0x0x7f5f80200200
        completion_signal=0x0, correlation_id=0
        rptr=18, wptr=21
 Kernel Name: input_concatenate_fusion
VGPU=0x9303560 SWq=0x7f6050006000, HWq=0x7f5fb8400000, id=2
        Dispatch Header =0xb02 (type=2, barrier=1, acquire=1, release=1), setup=0
        grid=[15, 1, 1], workgroup=[15, 1, 1]
        private_seg_size=0, group_seg_size=0
        kernel_obj=0x7f6040614600, kernarg_address=0x0x7f5f80200200
        completion_signal=0x0, correlation_id=0
        rptr=18, wptr=21
 Kernel Name: input_concatenate_fusion
VGPU=0x9320be0 SWq=0x7f6050006000, HWq=0x7f5fb8400000, id=2
        Dispatch Header =0xb02 (type=2, barrier=1, acquire=1, release=1), setup=0
        grid=[15, 1, 1], workgroup=[15, 1, 1]
        private_seg_size=0, group_seg_size=0
        kernel_obj=0x7f6040614600, kernarg_address=0x0x7f5f80200200
        completion_signal=0x0, correlation_id=0
        rptr=18, wptr=21
 Kernel Name: input_concatenate_fusion
VGPU=0x933e2a0 SWq=0x7f6050006000, HWq=0x7f5fb8400000, id=2
        Dispatch Header =0xb02 (type=2, barrier=1, acquire=1, release=1), setup=0
        grid=[15, 1, 1], workgroup=[15, 1, 1]
        private_seg_size=0, group_seg_size=0
        kernel_obj=0x7f6040614600, kernarg_address=0x0x7f5f80200200
        completion_signal=0x0, correlation_id=0
        rptr=18, wptr=21
:0:rocdevice.cpp            :3586: 87118002428 us:  Callback: Queue 0x7f5fb8400000 aborting with error : HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION: The agent attempted to access memory beyond the largest legal address. code: 0x29
Aborted (core dumped)

Happy to provide dumps etc as before to help track down the bug.

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

[owainkenwayucl]

Just in addition, I've done a test with docker.io/rocm/jax:rocm7.2.2-jax0.8.2-py3.12 and the issue occurs with the version of ROCm/JAX/XLA in that container as well. It doesn't seem to happen with 0.10.0 though so this might be a "wontfix"?

[owainkenwayucl]

The good news is that I believe I've discovered the exact line in the Alphafold 3 code which triggers the XLA bug: https://github.com/google-deepmind/alphafold3/blob/b712a3d0a136e9209b630dbaa1a3775263fe99a3/src/alphafold3/model/network/diffusion_head.py#L40 which means I can work around it with a jax.lax.optimization_barrier.

[lucbruni-amd]

Hey @owainkenwayucl! Thanks for helping investigate these issues!

I'm also glad to hear you have a workaround with the current JAX version. I believe I've found out why this works with JAX 0.10.0 as well and it appears a fix in XLA (openxla/xla@f127034dfc98) has made the cut for that JAX release.

Here's a summary of where the fix is and where it is not:

  ┌───────────────────────┬─────────────────────┬────────────────────────┐
  │        Version        │       XLA Pin       │     Fix included?      │
  ├───────────────────────┼─────────────────────┼────────────────────────┤
  │ JAX 0.9.1             │ 3cc8846c (Feb 26)   │ No                     │
  ├───────────────────────┼─────────────────────┼────────────────────────┤
  │ JAX 0.9.2 (ROCm)      │ d8b2a5f5 (ROCm/xla) │ No                     │
  ├───────────────────────┼─────────────────────┼────────────────────────┤
  │ JAX commit 3e8675e9b3 │ e570e84f (Mar 31)   │ Yes, first inclusion   │
  ├───────────────────────┼─────────────────────┼────────────────────────┤
  │ JAX 0.10.0            │ b6f37ab7 (Apr 15)   │ Yes                    │
  └───────────────────────┴─────────────────────┴────────────────────────┘

As for the issue in jnp.concatenate([a[None], b[None], jnp.cross(a, b)[None]], axis=0), the [None] creates size-1 dimensions along axis 0 = the concatenation axis. The tile analysis simplified these away, producing scf.if branches with mismatched tile sizes. The resulting fused kernel input_concatenate_fusion then computed incorrect memory offsets, causing the HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION.

As for the fix itself, IsConcatenateWithConcatDimPointDimension() was added which prevents point-dimension simplification when any operand has a size-1 dimension in the concatenate axis. The test case in the commit directly matches the bug pattern: concatenate(f32[10,2], f32[10,1]) along dim 1.

It doesn't seem to happen with 0.10.0 though so this might be a "wontfix"?

I'm going to agree with you here that it will be a "wontfix" in 0.9.x but can be considered as an "already-fixed" as of 0.10.0. If this is alright with you, we can keep this issue open until we officially support 0.10.0 and verify the violation no longer occurs through supported environments (Docker, TheRock builds, etc.).

[owainkenwayucl]

Works for me :)