[Feature]: Support for gfx1151

[flowtyone]

Suggestion Description

I couldn't find anywhere if iGPU is supported by this implementation, but it would be nice if we could use it with Strix Halo AI iGPUs. Currently when trying the nightly builds, I'm getting this error from XLA:

Traceback (most recent call last):
  File "/home/u/Desktop/jax-test/.venv/lib/python3.11/site-packages/jax/_src/xla_bridge.py", line 812, in backends
    backend = _init_backend(platform)
              ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/u/Desktop/jax-test/.venv/lib/python3.11/site-packages/jax/_src/xla_bridge.py", line 896, in _init_backend
    backend = registration.factory()
              ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/u/Desktop/jax-test/.venv/lib/python3.11/site-packages/jax/_src/xla_bridge.py", line 549, in make_pjrt_c_api_client
    return xla_client.make_c_api_client(plugin_name, updated_options, None)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/u/Desktop/jax-test/.venv/lib/python3.11/site-packages/jaxlib/xla_client.py", line 156, in make_c_api_client
    return _xla.get_c_api_client(
           ^^^^^^^^^^^^^^^^^^^^^^
jaxlib._jax.XlaRuntimeError: INTERNAL: no supported devices found for platform ROCM

Operating System

Ubuntu 24.04

GPU

AMD Ryzen™ Al Max+ 395 8060s

ROCm Component

No response

[jrhip]

CPU works, basic operations work (+ - * /) but any linear algebra operations cause a segmentation fault.
Which I think means gfx1151 is not yet supported. Fedora 43

jax 0.7.1
jax-rocm7-pjrt 0.7.1
jax-rocm7-plugin 0.7.1
jaxlib 0.7.1

Fatal Python error: Segmentation fault


Main thread:
Current thread 0x00007fa6673f6780 (most recent call first):
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/compiler.py", line 378 in backend_compile_and_load
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/profiler.py", line 364 in wrapper
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/compiler.py", line 762 in _compile_and_write_cache
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/compiler.py", line 494 in compile_or_get_cached
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/interpreters/pxla.py", line 2840 in _cached_compilation
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/interpreters/pxla.py", line 3059 in from_hlo
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/interpreters/pxla.py", line 2513 in compile
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/pjit.py", line 1782 in _pjit_call_impl_python
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/pjit.py", line 149 in _python_pjit_helper
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/pjit.py", line 270 in cache_miss
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/jax/_src/traceback_util.py", line 180 in reraise_with_filtered_traceback
  File "/home/user/Documents/Python Bits/jax_test.py", line 21 in <module>
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/spyder_kernels/customize/utils.py", line 209 in exec_encapsulate_locals
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/spyder_kernels/customize/code_runner.py", line 640 in _exec_code
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/spyder_kernels/customize/code_runner.py", line 470 in _exec_file
  File "/home/user/anaconda3/envs/jax/lib/python3.13/site-packages/spyder_kernels/customize/code_runner.py", line 177 in runfile