Silence C23 warning, causing error in upb target for XLA (#230)

* Silence C23 warning, creating error in upb target for XLA

* Copy CI workfow from rocm-main

* Add option to ci_build to run different tests

* Fix tty

---------

Co-authored-by: Charles Hofer <Charles.Hofer@amd.com>

.bazelrc

@@ -158,6 +158,7 @@ build:rocm --config=rocm_base
 build:rocm --action_env=TF_ROCM_CLANG="1"
 build:rocm --action_env=CLANG_COMPILER_PATH="/usr/lib/llvm-18/bin/clang"
 build:rocm --copt=-Wno-gnu-offsetof-extensions
+build:rocm --copt=-Wno-c23-extensions
 build:rocm --copt=-Qunused-arguments
 build:rocm --action_env=TF_HIPCC_CLANG="1"
 
@@ -379,4 +380,4 @@ build:debug --config debug_symbols -c fastbuild
 try-import %workspace%/.jax_configure.bazelrc
 
 # Load rc file with user-specific options.
-try-import %workspace%/.bazelrc.user
+try-import %workspace%/.bazelrc.user

.github/workflows/rocm-ci.yml

@@ -0,0 +1,65 @@
+name: ROCm GPU CI
+
+on:
+  # Trigger the workflow on push or pull request,
+  # but only for the rocm-main branch
+  push:
+    branches:
+      - rocm-main
+      - 'rocm-jaxlib-v*'
+  pull_request:
+    branches:
+      - rocm-main
+      - 'rocm-jaxlib-v*'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  build-jax-in-docker: # strategy and matrix come here
+    runs-on:  mi-250
+    env:
+      BASE_IMAGE: "ubuntu:22.04"
+      TEST_IMAGE: ubuntu-jax-${{ github.run_id }}_${{ github.run_number }}_${{ github.run_attempt }}
+      PYTHON_VERSION: "3.10"
+      ROCM_VERSION: "6.2.4"
+      WORKSPACE_DIR: workdir_${{ github.run_id }}_${{ github.run_number }}_${{ github.run_attempt }}
+    steps:
+      - name: Clean up old runs
+        run: |
+          ls
+          # Make sure that we own all of the files so that we have permissions to delete them
+          docker run -v "./:/jax" ubuntu /bin/bash -c "chown -R $UID /jax/workdir_* || true"
+          # Remove any old work directories from this machine
+          rm -rf workdir_*
+          ls
+      - name: Print system info
+        run: |
+          whoami
+          printenv
+          df -h
+          rocm-smi
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+        with:
+          path: ${{ env.WORKSPACE_DIR }}
+      - name: Build JAX
+        run: |
+          pushd $WORKSPACE_DIR
+          python3 build/rocm/ci_build \
+            --rocm-version $ROCM_VERSION \
+            --base-docker $BASE_IMAGE \
+            --python-versions $PYTHON_VERSION \
+            --compiler=clang \
+            dist_docker \
+            --image-tag $TEST_IMAGE
+      - name: Archive jax wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: rocm_jax_r${{ env.ROCM_VERSION }}_py${{ env.PYTHON_VERSION  }}_id${{ github.run_id }}
+          path: ./dist/*.whl
+      - name: Run tests
+        run: |
+          cd $WORKSPACE_DIR
+          python3 build/rocm/ci_build test $TEST_IMAGE --test-cmd "pytest tests/core_test.py"
+

build/rocm/ci_build

@@ -197,7 +197,7 @@ def dist_docker(
     subprocess.check_call(cmd)
 
 
-def test(image_name):
+def test(image_name, test_cmd):
     """Run unit tests like CI would inside a JAX image."""
 
     gpu_args = [
@@ -215,7 +215,6 @@ def test(image_name):
     cmd = [
         "docker",
         "run",
-        "-it",
         "--rm",
     ]
 
@@ -229,7 +228,7 @@ def test(image_name):
     cmd.extend(mounts)
     cmd.extend(gpu_args)
 
-    container_cmd = "cd /jax && ./build/rocm/build_rocm.sh && ./build/rocm/run_single_gpu.py -c && ./build/rocm/run_multi_gpu.sh"
+    container_cmd = "cd /jax && " + test_cmd 
     cmd.append(image_name)
     cmd.extend(
         [
@@ -292,6 +291,7 @@ def parse_args():
 
     testp = subp.add_parser("test")
     testp.add_argument("image_name")
+    testp.add_argument("--test-cmd", default="./build/rocm/build_rocm.sh && ./build/rocm/run_single_gpu.py -c && ./build/rocm/run_multi_gpu.sh")
 
     ddp = subp.add_parser("dist_docker")
     ddp.add_argument("--dockerfile", default="build/rocm/Dockerfile.ms")
@@ -314,7 +314,7 @@ def main():
         )
 
     elif args.action == "test":
-        test(args.image_name)
+        test(args.image_name, args.test_cmd)
 
     elif args.action == "dist_docker":
         dist_wheels(