diff --git a/.github/workflows/pytest_rocm_abort.yml b/.github/workflows/pytest_rocm_abort.yml new file mode 100644 index 000000000000..854b6681978b --- /dev/null +++ b/.github/workflows/pytest_rocm_abort.yml @@ -0,0 +1,175 @@ +# CI - Pytest ROCm (Abort Support) +# +# This workflow runs the ROCm tests with Pytest in ROCm GHCR containers, +# using the ROCm `pytest-abort` retry wrapper to detect/retry aborts/crashes. +# +# It can be triggered manually via workflow_dispatch or called by other workflows +# via workflow_call. +# +# It consists of the following job: +# run-tests: +# - Runs in ROCm container (ghcr.io/rocm/jax-base-ubu24-rocm*:latest) +# - Downloads the JAX and jaxlib wheels from GCS, and ROCm plugins from latest release. +# - Executes the `run_pytest_rocm_abort.sh` script, which installs wheel artifacts and +# runs the ROCm tests with Pytest under `pytest-abort-retry`. +name: CI - Pytest ROCm (Abort Support) + +on: + workflow_dispatch: + inputs: + runner: + description: "Which runner should the workflow run on?" + type: choice + default: "linux-x86-64-4gpu-amd" + options: + - "linux-x86-64-1gpu-amd" + - "linux-x86-64-4gpu-amd" + - "linux-x86-64-8gpu-amd" + python: + description: "Which Python version to use?" + type: choice + default: "3.11" + options: + - "3.11" + - "3.12" + rocm-version: + description: "Which ROCm version to test?" + type: choice + default: "7.2.0" + options: + - "7.2.0" + rocm-tag: + description: "ROCm tag for container image (e.g., rocm720)" + type: string + default: "rocm720" + jaxlib-version: + description: "Which jaxlib version to use? (head/pypi_latest)" + type: choice + default: "head" + options: + - "head" + - "pypi_latest" + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)" + type: choice + default: '0' + options: + - '0' + - '1' + gcs_download_uri: + description: "GCS location prefix from where the artifacts should be downloaded" + type: string + default: 'gs://jax-nightly-artifacts/latest' + halt-for-connection: + description: 'Should this workflow run wait for a remote connection?' + type: string + default: 'no' + max-worker-restart: + description: "Max xdist worker restarts (passed to pytest --max-worker-restart)" + type: string + default: '50' + workflow_call: + inputs: + runner: + description: "Which runner should the workflow run on?" + type: string + default: "linux-x86-64-4gpu-amd" + python: + description: "Which Python version to use?" + type: string + default: "3.11" + rocm-version: + description: "Which ROCm version to test?" + type: string + default: "7.2.0" + rocm-tag: + description: "ROCm tag for container image (e.g., rocm720)" + type: string + default: "rocm720" + jaxlib-version: + description: "Which jaxlib version to use? (head/pypi_latest)" + type: string + default: "head" + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)" + default: '0' + type: string + gcs_download_uri: + description: "GCS location prefix from where the artifacts should be downloaded" + default: 'gs://jax-nightly-artifacts/latest' + type: string + halt-for-connection: + description: 'Should this workflow run wait for a remote connection?' + type: string + default: 'no' + max-worker-restart: + description: "Max xdist worker restarts (passed to pytest --max-worker-restart)" + type: string + default: '50' + +permissions: {} + +env: + UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple" + +jobs: + run-tests: + defaults: + run: + # Set the shell to bash as GitHub actions run with /bin/sh by default + shell: bash + runs-on: ${{ inputs.runner }} + continue-on-error: true + # Run in ROCm GHCR container with GPU access + container: + image: ghcr.io/rocm/jax-base-ubu24.${{ inputs.rocm-tag }}:latest + credentials: + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + options: --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --shm-size 64G --env-file /etc/podinfo/gha-gpu-isolation-settings + name: "${{ (contains(inputs.runner, '1gpu') && '1gpu') || + (contains(inputs.runner, '4gpu') && '4gpu') || + (contains(inputs.runner, '8gpu') && '8gpu') }}, ROCm ${{ inputs.rocm-version }}, py${{ inputs.python }}" + + env: + JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}" + JAXCI_PYTHON: "python${{ inputs.python }}" + JAXCI_ENABLE_X64: "0" + MAX_WORKER_RESTART: "${{ inputs['max-worker-restart'] }}" + + steps: + - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0 + with: + persist-credentials: false + - name: Download JAX ROCm wheels + uses: ./.github/actions/download-jax-rocm-wheels + with: + python: ${{ inputs.python }} + rocm-version: ${{ inputs.rocm-version }} + jaxlib-version: ${{ inputs.jaxlib-version }} + skip-download-jaxlib-and-plugins-from-gcs: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }} + gcs_download_uri: ${{ inputs.gcs_download_uri }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + - name: Install Python dependencies + run: | + $JAXCI_PYTHON -m pip install uv~=0.5.30 + $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt + # Halt for testing + - name: Wait For Connection + uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c + with: + halt-dispatch-input: ${{ inputs.halt-for-connection }} + - name: Run Pytest ROCm tests (abort support) + timeout-minutes: 180 + run: ./ci/run_pytest_rocm_abort.sh + - name: Upload pytest results to artifact + if: always() + uses: actions/upload-artifact@v4 + with: + name: logs_abort + path: | + logs_abort/ + if-no-files-found: warn + retention-days: 2 + overwrite: true diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index 657eccf4c248..4d0255e3b135 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.github/workflows/wheel_tests_continuous.yml @@ -35,8 +35,8 @@ permissions: actions: read on: - schedule: - - cron: "0 */3 * * *" # Run once every 3 hours + # schedule: + # - cron: "0 */3 * * *" # Run once every 3 hours workflow_dispatch: # allows triggering the workflow run manually concurrency: diff --git a/.github/workflows/wheel_tests_nightly_release_abort.yml b/.github/workflows/wheel_tests_nightly_release_abort.yml new file mode 100644 index 000000000000..518306956eff --- /dev/null +++ b/.github/workflows/wheel_tests_nightly_release_abort.yml @@ -0,0 +1,53 @@ +# CI - Wheel Tests (Nightly/Release) (ROCm abort only) +# +# This workflow runs only the ROCm wheel tests using the abort/retry wrapper workflow. +name: CI - Wheel Tests (Nightly/Release) (ROCm abort only) + +on: + workflow_dispatch: + inputs: + gcs_download_uri: + description: "GCS location URI from where the artifacts should be downloaded" + required: true + default: 'gs://jax-nightly-artifacts/latest' + type: string + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to download only the jax wheel from GCS (e.g for testing a jax only release)" + required: true + default: '0' + type: string + halt-for-connection: + description: 'Should this workflow run wait for a remote connection? (yes/no)' + required: false + default: 'no' + type: string + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} + cancel-in-progress: true +permissions: {} + +env: + UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple" + +jobs: + run-pytest-rocm: + uses: ./.github/workflows/pytest_rocm_abort.yml + strategy: + fail-fast: false # don't cancel all jobs on failure + matrix: + runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"] + python: ["3.11", "3.12", "3.13", "3.14"] + rocm: [ + {version: "7.2.0", tag: "rocm720"}, + ] + name: "Pytest ROCm abort (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})" + with: + runner: ${{ matrix.runner }} + python: ${{ matrix.python }} + rocm-version: ${{ matrix.rocm.version }} + rocm-tag: ${{ matrix.rocm.tag }} + jaxlib-version: "head" + skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}} + gcs_download_uri: ${{inputs.gcs_download_uri}} + halt-for-connection: ${{inputs.halt-for-connection}} diff --git a/.gitignore b/.gitignore index d30c019b31e8..e55a9a5de101 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,7 @@ jax.iml /include/ /lib/ /share/ + +/compile_commands.json +/strace.txt +/external diff --git a/build/build.py b/build/build.py index 32a7b0016639..994e2df222ae 100755 --- a/build/build.py +++ b/build/build.py @@ -638,6 +638,10 @@ async def main(): ) if "rocm" in args.wheels: + if not args.configure_only: + print("ERROR: This repo is not used for building the ROCm JAX plugins. Please use the new plugin repo: https://github.com/ROCm/rocm-jax") + exit(1) + wheel_build_command_base.append("--config=rocm_base") wheel_build_command_base.append("--config=rocm") if clang_local: diff --git a/build/rocm-test-requirements.txt b/build/rocm-test-requirements.txt new file mode 100644 index 000000000000..399237175957 --- /dev/null +++ b/build/rocm-test-requirements.txt @@ -0,0 +1,24 @@ +absl-py +build +cloudpickle +colorama>=0.4.4 +filelock +flatbuffers +hypothesis +mpmath>=1.3 +pillow>=10.4.0 +# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t +portpicker; python_version<"3.13" +pytest-xdist +pytest-json-report +pytest-html +pytest-csv +pytest-rerunfailures +pytest-html-merger +pytest-reportlog +wheel +rich +setuptools +matplotlib +opt-einsum +auditwheel diff --git a/ci/run_pytest_rocm_abort.sh b/ci/run_pytest_rocm_abort.sh new file mode 100755 index 000000000000..6c6ab6d18281 --- /dev/null +++ b/ci/run_pytest_rocm_abort.sh @@ -0,0 +1,182 @@ +#!/bin/bash +# Copyright 2024 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# Runs Pytest ROCm tests (with ROCm pytest-abort retry wrapper). +# Requires the jaxlib and ROCm plugin wheels to be present inside $JAXCI_OUTPUT_DIR (../dist) +# +# -e: abort script if one command fails +# -u: error if undefined variable used +# -x: log all commands +# -o history: record shell history +# -o allexport: export all functions and variables to be available to subscripts +set -exu -o history -o allexport + +source ci/envs/default.env + +# Install jaxlib and ROCm plugin wheels inside the $JAXCI_OUTPUT_DIR directory +echo "Installing wheels locally..." +source ./ci/utilities/install_wheels_locally.sh + +# Print all the installed packages +echo "Installed packages:" +"$JAXCI_PYTHON" -m uv pip freeze + +"$JAXCI_PYTHON" -c "import jax; print(jax.default_backend()); print(jax.devices()); print(len(jax.devices()))" + +rocm-smi + +# ============================================================================== +# Set up the generic test environment variables +# ============================================================================== +export PY_COLORS=1 +export JAX_SKIP_SLOW_TESTS=true +export NCCL_DEBUG=WARN +export TF_CPP_MIN_LOG_LEVEL=0 +export JAX_ENABLE_X64="$JAXCI_ENABLE_X64" + +# ============================================================================== +# Calculate the optimal number of parallel processes for pytest +# This will be the minimum of: GPU capacity, CPU core count, and a system RAM limit. +# ============================================================================== + +export gpu_count=$(rocminfo | egrep -c "Device Type:\\s+GPU") +echo "Number of GPUs detected: $gpu_count" + +# Query GPU 0 memory using rocm-smi +export memory_per_gpu_mib=$(rocm-smi -d 0 --showmeminfo vram | grep -i "vram total" | awk '{print int($NF/1024/1024)}' | head -1) +echo "Reported memory per GPU: $memory_per_gpu_mib MiB" + +# Convert effective memory from MiB to GiB. +export memory_per_gpu_gib=$((memory_per_gpu_mib / 1024)) +echo "Effective memory per GPU: $memory_per_gpu_gib GiB" + +# Allow 2 GiB of GPU RAM per test. +export max_tests_per_gpu=$((memory_per_gpu_gib / 2)) +echo "Max tests per GPU (assuming 2GiB/test): $max_tests_per_gpu" + +export num_processes=$((gpu_count * max_tests_per_gpu)) +echo "Initial number of processes based on GPU capacity: $num_processes" + +export num_cpu_cores=$(nproc) +echo "Number of CPU cores available: $num_cpu_cores" + +# Reads total memory from /proc/meminfo (in KiB) and converts to GiB. +export total_ram_gib=$(awk '/MemTotal/ {printf "%.0f", $2/1048576}' /proc/meminfo) +echo "Total system RAM: $total_ram_gib GiB" + +# Set a safety limit for system RAM usage, e.g., 1/6th of total. +export host_memory_limit=$((total_ram_gib / 6)) +echo "Host memory process limit (1/6th of total RAM): $host_memory_limit" + +if [[ $num_cpu_cores -lt $num_processes ]]; then + num_processes=$num_cpu_cores + echo "Adjusting num_processes to match CPU core count: $num_processes" +fi + +if [[ $host_memory_limit -lt $num_processes ]]; then + num_processes=$host_memory_limit + echo "Adjusting num_processes to match host memory limit: $num_processes" +fi + +if [[ 16 -lt $num_processes ]]; then + num_processes=16 + echo "Reducing num_processes to $num_processes" +fi + +echo "Final number of processes to run: $num_processes" + +export JAX_ENABLE_ROCM_XDIST="$gpu_count" +export XLA_PYTHON_CLIENT_ALLOCATOR=platform +export XLA_FLAGS="--xla_gpu_force_compilation_parallelism=1 --xla_gpu_enable_nccl_comm_splitting=false --xla_gpu_enable_command_buffer=" + +# Disable core dumps just in case +ulimit -c 0 + +# Keep deselected tests in one place for the abort wrapper. +ROCM_PYTEST_DESELECT_ARGS=( + --deselect=tests/multi_device_test.py::MultiDeviceTest::test_computation_follows_data + --deselect=tests/multiprocess_gpu_test.py::MultiProcessGpuTest::test_distributed_jax_visible_devices + --deselect=tests/compilation_cache_test.py::CompilationCacheTest::test_task_using_cache_metric +) + +# --max-runs: retry the entire pytest run up to N times on abort/crash. +# --max-worker-restart: restart crashed xdist workers up to N times. +# --maxfail: stop the run after N test failures. +rocm_test_cmd() { + local abort_flag="${1:-0}" + shift + if [[ "$abort_flag" == "1" ]]; then + pytest-abort-retry --max-runs 3 --clear-crash-log -- "$@" + else + "$@" + fi +} + +rocm_log_tail_on_failure() { + local logfile="$1" + local status="$2" + if [[ "$status" -ne 0 ]]; then + echo "Pytest failed (exit=$status). Showing last 200 lines of $logfile:" + tail -n 200 "$logfile" || true + else + echo "Pytest output saved to $logfile (uploaded as artifact)." + fi +} + +rocm_install_extra_requirements() { + if [[ -n "${GITHUB_WORKSPACE:-}" ]]; then + cd "$GITHUB_WORKSPACE" + fi + + # Install extra requirements. + "$JAXCI_PYTHON" -m uv pip install pytest-timeout pytest-html pytest-csv pytest-json-report pytest-abort +} + +rocm_install_extra_requirements + +echo "Running ROCm tests (with abort/retry wrapper)..." +mkdir -p logs_abort +logfile="logs_abort/jax_ToT_UT_abort.log" + +# Allow the workflow to override worker restart limit. +max_worker_restart="${MAX_WORKER_RESTART:-50}" + +# pytest-abort output directories (must be set before running pytest). +export PYTEST_ABORT_LAST_RUNNING_DIR="logs_abort/last_running" +export PYTEST_ABORT_CRASHED_TESTS_LOG="logs_abort/crashed_tests.jsonl" +mkdir -p "$PYTEST_ABORT_LAST_RUNNING_DIR" + +set +e +rocm_test_cmd 1 "$JAXCI_PYTHON" -m pytest -n "$num_processes" --max-worker-restart="$max_worker_restart" --tb=short --timeout=1200 --timeout-method=thread tests \ + "${ROCM_PYTEST_DESELECT_ARGS[@]}" \ + --json-report \ + --json-report-file=logs_abort/tests-report-abort.json \ + --csv=logs_abort/tests-report-abort.csv \ + --html=logs_abort/tests-report-abort.html \ + --self-contained-html \ + >"$logfile" 2>&1 +pytest_status=$? +set -e +rocm_log_tail_on_failure "$logfile" "$pytest_status" + +echo "Postprocessing reports with crashed tests..." +pytest-abort-postprocess \ + --crash-log "$PYTEST_ABORT_CRASHED_TESTS_LOG" \ + --json-report logs_abort/tests-report-abort.json \ + --html-report logs_abort/tests-report-abort.html \ + --csv-report logs_abort/tests-report-abort.csv \ + >>"$logfile" 2>&1 + +exit "$pytest_status" diff --git a/conftest.py b/conftest.py index 1e6e17ab966a..c313ecaf31c1 100644 --- a/conftest.py +++ b/conftest.py @@ -16,7 +16,6 @@ import os import pytest - @pytest.fixture(autouse=True) def add_imports(doctest_namespace): import jax diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py index fe3ab7f910f0..5ea7f584e3f2 100644 --- a/jax/_src/interpreters/mlir.py +++ b/jax/_src/interpreters/mlir.py @@ -640,9 +640,8 @@ def make_ir_context() -> ir.Context: # multi threaded execution aborts the process if we try to register a new # dialect after this point. The dialect registry in a context is not thread # safe, and a fatal error is much better than a data race. - # jax_mlir_ext.enter_multi_threaded_execution(context) - # TODO(phawkins): clean up users who add their own dialects to JAX's contexts - # and enable this. + # if jaxlib_version >= (0, 8): + # jax_mlir_ext.enter_multi_threaded_execution(context) return context diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py index d688221e7881..0dee9c26e552 100644 --- a/jax/_src/test_util.py +++ b/jax/_src/test_util.py @@ -25,6 +25,7 @@ import math import os +from pathlib import Path import platform import re import sys @@ -389,6 +390,15 @@ def supported_dtypes() -> set[DTypeLike]: def is_device_rocm() -> bool: return 'rocm' in xla_bridge.get_backend().platform_version +def get_rocm_version(): + rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm") + version_path = Path(rocm_path) / ".info" / "version" + if not version_path.exists(): + raise FileNotFoundError(f"Expected ROCm version file at {version_path}") + version_str = version_path.read_text().strip() + major, minor, *_ = version_str.split(".") + return int(major), int(minor) + def is_device_cuda() -> bool: return 'cuda' in xla_bridge.get_backend().platform_version @@ -610,7 +620,14 @@ def skip_on_devices(*disabled_devices, skip_reason=None): skip_reason: Optional custom skip message when test is skipped. """ if skip_reason is None: - skip_reason = "Skipped on devices with tags: " + ", ".join(disabled_devices) + skip_messages = { + ("gpu",): "Skipped on all GPUs.", + ("cpu",): "Skipped on CPU.", + ("tpu",): "Skipped on TPU.", + ("cuda",): "Skipped on CUDA GPUs.", + ("rocm",): "Skipped on ROCm GPUs.", + } + skip_reason = skip_messages.get(disabled_devices) return _device_filter(lambda: not test_device_matches(disabled_devices), skip_reason) def run_on_devices(*enabled_devices, skip_reason=None): @@ -621,9 +638,14 @@ def run_on_devices(*enabled_devices, skip_reason=None): skip_reason: Optional custom skip message when test is skipped. """ if skip_reason is None: - skip_reason = ( - "Skipped unless running on devices with tags: " + ", ".join(enabled_devices) - ) + device_specific_skip_reasons = { + ("cpu",): "Skipped: CPU-only test.", + ("tpu",): "Skipped: TPU-only test.", + ("gpu",): "Skipped: GPU-only test.", + ("rocm",): "Skipped: ROCm-only test.", + ("cuda",): "Skipped: CUDA-only test.", + } + skip_reason = device_specific_skip_reasons.get(enabled_devices) return _device_filter(lambda: test_device_matches(enabled_devices), skip_reason) def device_supports_buffer_donation(): diff --git a/jaxlib/BUILD b/jaxlib/BUILD index a434e7ebc119..134f364ffde6 100644 --- a/jaxlib/BUILD +++ b/jaxlib/BUILD @@ -41,7 +41,7 @@ licenses(["notice"]) package( default_applicable_licenses = [], - default_visibility = ["//jax:internal"], + default_visibility = ["//visibility:public"], ) package_group( diff --git a/jaxlib/gpu/solver_interface.cc b/jaxlib/gpu/solver_interface.cc index 7ebcf7483652..e530e5bc1f2f 100644 --- a/jaxlib/gpu/solver_interface.cc +++ b/jaxlib/gpu/solver_interface.cc @@ -63,8 +63,8 @@ JAX_GPU_DEFINE_GETRF(gpuDoubleComplex, gpusolverDnZgetrf); JAX_GPU_DEFINE_GETRF_BATCHED(float, gpublasSgetrfBatched); JAX_GPU_DEFINE_GETRF_BATCHED(double, gpublasDgetrfBatched); -JAX_GPU_DEFINE_GETRF_BATCHED(gpublasComplex, gpublasCgetrfBatched); -JAX_GPU_DEFINE_GETRF_BATCHED(gpublasDoubleComplex, gpublasZgetrfBatched); +JAX_GPU_DEFINE_GETRF_BATCHED(gpuComplex, gpublasCgetrfBatched); +JAX_GPU_DEFINE_GETRF_BATCHED(gpuDoubleComplex, gpublasZgetrfBatched); #undef JAX_GPU_DEFINE_GETRF_BATCHED // QR decomposition: geqrf @@ -102,8 +102,8 @@ JAX_GPU_DEFINE_GEQRF(gpuDoubleComplex, gpusolverDnZgeqrf); JAX_GPU_DEFINE_GEQRF_BATCHED(float, gpublasSgeqrfBatched); JAX_GPU_DEFINE_GEQRF_BATCHED(double, gpublasDgeqrfBatched); -JAX_GPU_DEFINE_GEQRF_BATCHED(gpublasComplex, gpublasCgeqrfBatched); -JAX_GPU_DEFINE_GEQRF_BATCHED(gpublasDoubleComplex, gpublasZgeqrfBatched); +JAX_GPU_DEFINE_GEQRF_BATCHED(gpuComplex, gpublasCgeqrfBatched); +JAX_GPU_DEFINE_GEQRF_BATCHED(gpuDoubleComplex, gpublasZgeqrfBatched); #undef JAX_GPU_DEFINE_GEQRF_BATCHED // Householder transformations: orgqr @@ -302,8 +302,8 @@ JAX_GPU_DEFINE_SYEVD(gpuDoubleComplex, gpusolverDnZheevd); JAX_GPU_DEFINE_SYRK(float, gpublasSsyrk); JAX_GPU_DEFINE_SYRK(double, gpublasDsyrk); -JAX_GPU_DEFINE_SYRK(gpublasComplex, gpublasCsyrk); -JAX_GPU_DEFINE_SYRK(gpublasDoubleComplex, gpublasZsyrk); +JAX_GPU_DEFINE_SYRK(gpuComplex, gpublasCsyrk); +JAX_GPU_DEFINE_SYRK(gpuDoubleComplex, gpublasZsyrk); #undef JAX_GPU_DEFINE_SYRK // Singular Value Decomposition: gesvd diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h index 9adb1f11776f..fb9c1f407145 100644 --- a/jaxlib/gpu/vendor.h +++ b/jaxlib/gpu/vendor.h @@ -457,6 +457,8 @@ inline constexpr uint32_t kNumThreadsPerWarp = 32; #elif defined(JAX_GPU_HIP) +#define HIPBLAS_V2 1 + // IWYU pragma: begin_exports #include "rocm/include/hip/hip_cooperative_groups.h" #include "rocm/include/hip/hip_runtime_api.h" @@ -477,17 +479,11 @@ inline constexpr uint32_t kNumThreadsPerWarp = 32; // MIOpen lib. Remove when MIOpen support is complete. #define MIOPEN_STATUS_SUCCESS 0 -typedef hipFloatComplex gpuComplex; +typedef hipComplex gpuComplex; typedef hipDoubleComplex gpuDoubleComplex; -#if TF_ROCM_VERSION >= 70000 -typedef hipFloatComplex gpublasComplex; +typedef hipComplex gpublasComplex; typedef hipDoubleComplex gpublasDoubleComplex; -#else -typedef hipblasComplex gpublasComplex; -typedef hipblasDoubleComplex gpublasDoubleComplex; -#endif // TF_ROCM_VERSION >= 70000 - typedef struct hipsolverHandle_* gpusolverDnHandle_t; typedef hipblasFillMode_t gpublasFillMode_t; typedef hipsolverFillMode_t gpusolverFillMode_t; @@ -545,6 +541,7 @@ inline hipblasStatus_t gpublasCreate(gpublasHandle_t* handle) { return hipblasCreate(reinterpret_cast(handle)); } } // namespace jax::hip + #define gpublasCreate ::jax::hip::gpublasCreate #define gpublasSetStream hipblasSetStream #define gpublasSgeqrfBatched hipblasSgeqrfBatched @@ -601,6 +598,7 @@ inline hipsolverStatus_t gpusolverDnCreate(gpusolverDnHandle_t* handle) { return hipsolverCreate(reinterpret_cast(handle)); } } // namespace jax::hip + #define gpusolverDnCreate ::jax::hip::gpusolverDnCreate #define gpusolverDnSetStream hipsolverSetStream #define gpusolverDnCreateSyevjInfo hipsolverCreateSyevjInfo @@ -780,10 +778,10 @@ inline hipsparseStatus_t gpusparseCreate(gpusparseHandle_t* handle) { #define GPUSPARSE_INDEX_32I HIPSPARSE_INDEX_32I #define GPUSPARSE_INDEX_64I HIPSPARSE_INDEX_64I #define GPUSPARSE_DENSETOSPARSE_ALG_DEFAULT HIPSPARSE_DENSETOSPARSE_ALG_DEFAULT -#define GPUSPARSE_SPMV_COO_ALG HIPSPARSE_MV_ALG_DEFAULT -#define GPUSPARSE_SPMV_CSR_ALG HIPSPARSE_MV_ALG_DEFAULT -#define GPUSPARSE_SPMM_COO_ALG HIPSPARSE_SPMM_ALG_DEFAULT -#define GPUSPARSE_SPMM_CSR_ALG HIPSPARSE_SPMM_ALG_DEFAULT +#define GPUSPARSE_SPMV_COO_ALG HIPSPARSE_COOMV_ALG +#define GPUSPARSE_SPMV_CSR_ALG HIPSPARSE_CSRMV_ALG1 +#define GPUSPARSE_SPMM_COO_ALG HIPSPARSE_SPMM_COO_ALG1 +#define GPUSPARSE_SPMM_CSR_ALG HIPSPARSE_SPMM_CSR_ALG1 #define GPUSPARSE_INDEX_BASE_ZERO HIPSPARSE_INDEX_BASE_ZERO #define GPUSPARSE_OPERATION_NON_TRANSPOSE HIPSPARSE_OPERATION_NON_TRANSPOSE #define GPUSPARSE_OPERATION_TRANSPOSE HIPSPARSE_OPERATION_TRANSPOSE @@ -796,7 +794,7 @@ inline hipsparseStatus_t gpusparseCreate(gpusparseHandle_t* handle) { #define GPU_STREAM_NON_BLOCKING hipStreamNonBlocking #define gpuMalloc hipMalloc -#define gpuGetLastError hipGetLastError +#define gpuGetLastError hipExtGetLastError #define gpuGetErrorString hipGetErrorString #define gpuMemcpyAsync hipMemcpyAsync #define gpuMemcpyDeviceToDevice hipMemcpyDeviceToDevice diff --git a/jaxlib/jax.bzl b/jaxlib/jax.bzl index 924400c6fe61..cb6d497c4618 100644 --- a/jaxlib/jax.bzl +++ b/jaxlib/jax.bzl @@ -22,7 +22,6 @@ load("@jax_wheel//:wheel.bzl", "WHEEL_VERSION") load("@jax_wheel_version_suffix//:wheel_version_suffix.bzl", "WHEEL_VERSION_SUFFIX") load("@local_config_cuda//cuda:build_defs.bzl", _cuda_library = "cuda_library", _if_cuda_is_configured = "if_cuda_is_configured") load("@local_config_rocm//rocm:build_defs.bzl", _if_rocm_is_configured = "if_rocm_is_configured", _rocm_library = "rocm_library") -load("@nvidia_wheel_versions//:versions.bzl", "NVIDIA_WHEEL_VERSIONS") load("@python_version_repo//:py_version.bzl", "HERMETIC_PYTHON_VERSION", "HERMETIC_PYTHON_VERSION_KIND") load("@rocm_external_test_deps//:external_deps.bzl", "EXTERNAL_DEPS") load("@rocm_prebuilt_test_deps//:external_deps.bzl", PREBUILT_EXTERNAL_DEPS = "EXTERNAL_DEPS") @@ -485,7 +484,6 @@ def _jax_wheel_impl(ctx): if ctx.attr.platform_version == "": fail("platform_version must be set to a valid cuda version for cuda wheels") args.add("--platform_version", ctx.attr.platform_version) # required for gpu wheels - args.add("--nvidia_wheel_versions_data", NVIDIA_WHEEL_VERSIONS) # required for gpu wheels if ctx.attr.enable_rocm: args.add("--enable-rocm", "True") if ctx.attr.platform_version == "": diff --git a/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc b/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc index fc04f9ce1718..d6515d146e00 100644 --- a/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc +++ b/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc @@ -208,7 +208,6 @@ NB_MODULE(_jax_mlir_ext, m) { unwrap(registry)->insert(); unwrap(registry)->insert(); unwrap(registry)->insert(); - // For Mosaic GPU REGISTER_DIALECT(cf); REGISTER_DIALECT(gpu); diff --git a/jaxlib/rocm/BUILD b/jaxlib/rocm/BUILD index 65a844a2e99f..3c0c93844487 100644 --- a/jaxlib/rocm/BUILD +++ b/jaxlib/rocm/BUILD @@ -28,7 +28,7 @@ licenses(["notice"]) package( default_applicable_licenses = [], - default_visibility = ["//:__subpackages__"], + default_visibility = ["//visibility:public"], ) cc_library( @@ -425,6 +425,10 @@ rocm_nanobind_extension( "@xla//xla/ffi/api:ffi", "@xla//xla/python:safe_static_init", ], + linkopts = [ + "-L/opt/rocm/lib", + "-lamdhip64", + ], ) cc_library( @@ -554,6 +558,7 @@ rocm_nanobind_extension( enable_stub_generation = False, module_name = "rocm_plugin_extension", deps = [ + ":hip_gpu_kernel_helpers", ":py_client_gpu", "//jaxlib:kernel_nanobind_helpers", "//jaxlib/gpu:gpu_plugin_extension", diff --git a/jaxlib/rocm/rocm_plugin_extension.cc b/jaxlib/rocm/rocm_plugin_extension.cc index 05be1e81c858..e28e4927b81f 100644 --- a/jaxlib/rocm/rocm_plugin_extension.cc +++ b/jaxlib/rocm/rocm_plugin_extension.cc @@ -23,6 +23,7 @@ limitations under the License. #include "jaxlib/gpu/gpu_plugin_extension.h" #include "jaxlib/gpu/py_client_gpu.h" #include "jaxlib/kernel_nanobind_helpers.h" +#include "jaxlib/gpu/gpu_kernel_helpers.h" namespace nb = nanobind; @@ -96,6 +97,13 @@ nb::dict FfiHandlers() { return dict; } +int ROCmDeviceCount() { + int device_count = -1; + JAX_THROW_IF_ERROR(JAX_AS_STATUS(hipInit(0))); + JAX_THROW_IF_ERROR(JAX_AS_STATUS(hipGetDeviceCount(&device_count))); + return device_count; +} + } // namespace NB_MODULE(rocm_plugin_extension, m) { @@ -122,5 +130,6 @@ NB_MODULE(rocm_plugin_extension, m) { return device_ordinal; }, nb::arg("data_value")); + m.def("get_device_count", &ROCmDeviceCount); } } // namespace jax diff --git a/jaxlib/rocm/rocm_rpath.bzl b/jaxlib/rocm/rocm_rpath.bzl index 354082c0d200..1da2d640ca42 100644 --- a/jaxlib/rocm/rocm_rpath.bzl +++ b/jaxlib/rocm/rocm_rpath.bzl @@ -35,8 +35,11 @@ _ROCM_LINK_ONLY = "@local_config_rocm//rocm:link_only" _WHEEL_RPATHS = [ "-Wl,-rpath,$$ORIGIN/../rocm/lib", + "-Wl,-rpath,$$ORIGIN/../rocm/lib/rocm_sysdeps/lib", "-Wl,-rpath,$$ORIGIN/../../rocm/lib", + "-Wl,-rpath,$$ORIGIN/../../rocm/lib/rocm_sysdeps/lib", "-Wl,-rpath,/opt/rocm/lib", + "-Wl,-rpath,/opt/rocm/lib/rocm_sysdeps/lib", ] def _wheel_features(): diff --git a/tests/array_interoperability_test.py b/tests/array_interoperability_test.py index d033a490f547..b14e9171f952 100644 --- a/tests/array_interoperability_test.py +++ b/tests/array_interoperability_test.py @@ -312,7 +312,7 @@ def testCudaArrayInterfaceOnNonCudaFails(self): self.assertFalse(hasattr(x, "__cuda_array_interface__")) with self.assertRaisesRegex( AttributeError, - "__cuda_array_interface__ is only defined for .*GPU buffers.", + "__cuda_array_interface__ is only defined for GPU buffers.", ): _ = x.__cuda_array_interface__ diff --git a/tests/array_test.py b/tests/array_test.py index 5e46ba709fee..fd67a31e28e5 100644 --- a/tests/array_test.py +++ b/tests/array_test.py @@ -1541,7 +1541,7 @@ class RngShardingTest(jtu.JaxTestCase): # tests that the PRNGs are automatically sharded as expected @parameterized.named_parameters(("3", 3), ("4", 4), ("5", 5)) - @jtu.skip_on_devices("gpu") + @jtu.skip_on_devices("cuda") def test_random_bits_is_pure_map_1d(self, num_devices): @jax.jit def f(x): @@ -1575,7 +1575,7 @@ def f(x): "mesh_shape": mesh_shape, "pspec": pspec} for mesh_shape in [(3, 2), (4, 2), (2, 3)] for pspec in [P('x', None), P(None, 'y'), P('x', 'y')]) - @jtu.skip_on_devices("gpu") + @jtu.skip_on_devices("cuda") def test_random_bits_is_pure_map_2d(self, mesh_shape, pspec): @jax.jit def f(x): diff --git a/tests/core_test.py b/tests/core_test.py index e7e090f28d21..a99f09501037 100644 --- a/tests/core_test.py +++ b/tests/core_test.py @@ -256,6 +256,13 @@ def bar(y): assert jvp(foo, (1.0,), (2.0,)) == (1.0, 2.0) def test_simple_jit(self): + # XLA reduction autotuner produces no valid tile configs for gfx1151 (RDNA3.5, + # 32-thread wavefronts). JIT compilation triggers autotuning which fails with + # "NOT_FOUND: No valid config found!". See ROCM-21429. + if (jtu.is_device_rocm() + and "Radeon" in jax.local_devices()[0].device_kind): + self.skipTest( + "XLA reduction autotuner has no valid config for Navi/RDNA GPUs (gfx11xx)") def foo(x): if x.shape == (): return x + 1. @@ -403,6 +410,13 @@ def f(y): self.assertEqual(e1.outvars[0].aval.shape, (3, 3)) # only primal out shape def test_tracer_reprs(self): + # XLA reduction autotuner produces no valid tile configs for gfx1151 (RDNA3.5, + # 32-thread wavefronts). JIT compilation triggers autotuning which fails with + # "NOT_FOUND: No valid config found!". See ROCM-21429. + if (jtu.is_device_rocm() + and "Radeon" in jax.local_devices()[0].device_kind): + self.skipTest( + "XLA reduction autotuner has no valid config for Navi/RDNA GPUs (gfx11xx)") def f(x): nonlocal x_repr x_repr = repr(x) @@ -425,6 +439,13 @@ def f(x): self.assertEqual(x_repr, "JVPTracer(primal=[0. 1. 2.], tangent=VmapTracer(aval=float32[3], batched=float32[3,3]))") def test_verbose_tracer_reprs(self): + # XLA reduction autotuner produces no valid tile configs for gfx1151 (RDNA3.5, + # 32-thread wavefronts). JIT compilation triggers autotuning which fails with + # "NOT_FOUND: No valid config found!". See ROCM-21429. + if (jtu.is_device_rocm() + and "Radeon" in jax.local_devices()[0].device_kind): + self.skipTest( + "XLA reduction autotuner has no valid config for Navi/RDNA GPUs (gfx11xx)") # Verbose reprs, avaiable via tracer._pretty_print() def f(x): nonlocal x_repr diff --git a/tests/lax_vmap_test.py b/tests/lax_vmap_test.py index 57fdfc4dda88..de260ec4ed93 100644 --- a/tests/lax_vmap_test.py +++ b/tests/lax_vmap_test.py @@ -759,6 +759,8 @@ def testSort(self, shape, dimension, arity, bdims, is_stable): # TODO Collapse # TODO Scatter + # b/183233858: variadic reduce-window not implemented on XLA:CUDA + @jtu.skip_on_devices("cuda") def test_variadic_reduce_window(self): # https://github.com/jax-ml/jax/discussions/9818 and # https://github.com/jax-ml/jax/issues/9837 diff --git a/tests/linalg_test.py b/tests/linalg_test.py index 5776aec6cb10..56f2f91c5381 100644 --- a/tests/linalg_test.py +++ b/tests/linalg_test.py @@ -585,11 +585,12 @@ def testEighZeroDiagonal(self): np.linalg.norm(np.matmul(a, v) - w * v), 2.5 * eps * np.linalg.norm(a) ) + def testEighTinyNorm(self): - # Skip test on ROCm due to numerical error. Issue #34711 - # TODO(GulsumGudukbay): Unskip once fixed. if jtu.is_device_rocm(): - self.skipTest("Skipped on ROCm due to numerical error.") + # numerical errors seen as of ROCm 7.2 due to hipSolver issue + # TODO: re-enable the test once the hipSolver issue is fixed + self.skipTest("testEighNorm not supported on ROCm due to hipSOLVER issue") rng = jtu.rand_default(self.rng()) a = rng((300, 300), dtype=np.float32) eps = jnp.finfo(a.dtype).eps diff --git a/tests/pallas/gpu_paged_attention_test.py b/tests/pallas/gpu_paged_attention_test.py index 1b778c787a6d..6a1d8de22a78 100644 --- a/tests/pallas/gpu_paged_attention_test.py +++ b/tests/pallas/gpu_paged_attention_test.py @@ -112,6 +112,64 @@ class PagedAttentionKernelTest(PallasBaseTest): def setUp(self): super().setUp() + def _estimate_shared_memory_bytes(self, block_h, pages_per_compute_block, + page_size, head_dim, dtype): + """Estimate shared memory usage for paged attention kernel.""" + dtype_size = jnp.dtype(dtype).itemsize + # Approximate calculation based on kernel's memory usage + # Q block: block_h * head_dim + # K/V blocks: pages_per_compute_block * page_size * head_dim + # Plus accumulators and intermediate values + block_k = pages_per_compute_block * page_size + estimated = dtype_size * ( + block_h * head_dim + # Q + 2 * block_k * head_dim + # K and V + block_h * block_k + # logits/attention weights + block_h * 8 # accumulators (m, l, etc.) in float32 + ) + return estimated + + def _adjust_params_for_shared_memory(self, block_h, pages_per_compute_block, + page_size, head_dim, dtype): + """Adjust parameters to fit within device shared memory limits. + + Uses XLA's DeviceDescription.shared_memory_per_block_optin() to query + the actual device capability rather than hardcoding values. + """ + try: + device = jax.local_devices()[0] + # Query XLA DeviceDescription for max shared memory per block + # This is exposed from stream_executor::DeviceDescription::shared_memory_per_block_optin() + max_smem = device.shared_memory_per_block_optin + except (AttributeError, IndexError): + # Fallback if XLA doesn't expose shared_memory_per_block_optin (older versions) + # or if no devices are available. Use conservative 48KB (safe for most GPUs). + max_smem = 48 * 1024 + + estimated = self._estimate_shared_memory_bytes( + block_h, pages_per_compute_block, page_size, head_dim, dtype) + + # If within limits, no adjustment needed + if estimated <= max_smem: + return block_h, pages_per_compute_block, page_size + + # Try to reduce parameters to fit + while estimated > max_smem: + if pages_per_compute_block > 2: + pages_per_compute_block = pages_per_compute_block // 2 + elif page_size > 8: + page_size = page_size // 2 + elif block_h > 8: + block_h = block_h // 2 + else: + # Can't reduce further, will need to skip + return None, None, None + + estimated = self._estimate_shared_memory_bytes( + block_h, pages_per_compute_block, page_size, head_dim, dtype) + + return block_h, pages_per_compute_block, page_size + @jtu.sample_product( dtype=(jnp.float16,), page_size=(8, 16, 32), @@ -201,6 +259,17 @@ def test_quantized_paged_attention( if (quant_dtype == jnp.float8_e4m3fn and not jtu.is_cuda_compute_capability_at_least("8.9")): self.skipTest("Skipping since float8_e4m3fn is not supported on < sm89") + + # Check and adjust parameters if needed to fit device limits for ROCm + if jtu.is_device_rocm(): + adjusted = self._adjust_params_for_shared_memory( + block_h, pages_per_compute_block, page_size, head_dim, dtype) + + if adjusted == (None, None, None): + self.skipTest("Cannot adjust parameters to fit ROCm device shared memory limits") + + block_h, pages_per_compute_block, page_size = adjusted + max_kv_len = 2048 seq_lens = np.asarray([3, 256, 513, 1023, 2048], dtype=jnp.int32) q, k_pages, v_pages, block_tables = _generate_qkv( @@ -218,7 +287,7 @@ def test_quantized_paged_attention( k_, k_scales = (_quantize(k_pages, quant_dtype) if quantize_k else (k_pages, None)) - v_, v_scales = (_quantize(k_pages, quant_dtype) + v_, v_scales = (_quantize(v_pages, quant_dtype) if quantize_v else (v_pages, None)) o = paged_attention.paged_attention( diff --git a/tests/pallas/gpu_pallas_distributed_test.py b/tests/pallas/gpu_pallas_distributed_test.py index 2f174b21ca23..7137e19d9ec0 100644 --- a/tests/pallas/gpu_pallas_distributed_test.py +++ b/tests/pallas/gpu_pallas_distributed_test.py @@ -126,14 +126,17 @@ def setUp(self): if jtu.test_device_matches(["rocm"]): self.skipTest("Mosaic GPU is not supported on ROCm.") + # Check mosaic support first (before GPU capability check) + if not mgpu.supports_cross_device_collectives(): + if jtu.test_device_matches(["rocm"]): + self.skipTest("Mosaic not supported on ROCm currently.") + else: + self.skipTest("NVSHMEM library unavailable.") if (not jtu.test_device_matches(["cuda"]) or not jtu.is_cuda_compute_capability_at_least("9.0")): self.skipTest("Only works on GPU with capability >= sm90") - if not mgpu.supports_cross_device_collectives(): - self.skipTest( - "Skip test since cross-device collectives are not supported" - " (either NVSHMEM is not available in multi-process mode, or mixed" - " mode is used).") + if jax.process_count() == 1: + self.skipTest("Test requires multiple processes.") if os.environ.get("XLA_PYTHON_CLIENT_ALLOCATOR", "") == "platform": self.skipTest("NVSHMEM doesn't work with the platform allocator.") diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py index 7b2b670a0db1..1031f72c80b5 100644 --- a/tests/pallas/ops_test.py +++ b/tests/pallas/ops_test.py @@ -1047,6 +1047,13 @@ def test_is_finite(self, dtype): if jtu.is_device_rocm() and jtu.parse_version(jax.__version__) == (0, 8, 0): self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + if jtu.test_device_matches(["cuda"]): + self.skipTest("Not tested on CUDA") # set this b/c this how the test was + # originally configured. Have no way to test cuda. + + if jtu.is_device_rocm(): + self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + size = len(self.IS_FINITE_TEST_VALUES) @functools.partial( @@ -1099,6 +1106,9 @@ def test_is_finite_scalar(self, dtype): if jtu.is_device_rocm() and jtu.parse_version(jax.__version__) == (0, 8, 0): self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + if jtu.is_device_rocm(): + self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + size = len(self.IS_FINITE_TEST_VALUES) @functools.partial( @@ -2160,12 +2170,20 @@ def test_dot(self, lhs_and_rhs_shape, dtype, trans_x, trans_y): self.skipTest("bfloat16 type are not supported on GPU") # Check shared memory limit: Triton loads lhs + rhs into shared memory if jtu.is_device_rocm(): + # ROCm: use correct formula with dynamic limit from rocminfo dtype_size = jnp.dtype(dtype).itemsize - if (math.prod(lhs_shape) + math.prod(rhs_shape)) * dtype_size > get_rocm_shared_memory_limit(): + shared_mem_bytes = (math.prod(lhs_shape) + math.prod(rhs_shape)) * dtype_size + shared_mem_limit = get_rocm_shared_memory_limit() + if shared_mem_bytes > shared_mem_limit: self.skipTest("Shared memory size limit exceeded") - elif math.prod(lhs_shape) + math.prod(rhs_shape) + math.prod(out_shape) > (256 * 256) * 2: - self.skipTest("Shared memory size limit exceeded") - if (jax.local_devices()[0].shared_memory_per_block_optin == 99 * 1024 and + else: + # NVIDIA: keep original check + if ( + math.prod(lhs_shape) + math.prod(rhs_shape) + math.prod(out_shape) + > (256 * 256) * 2 + ): + self.skipTest("Shared memory size limit exceeded") + if (jax.local_devices()[0].device_kind == "NVIDIA L4" and dtype == jnp.float32 and lhs_and_rhs_shape in [ ((128, 16), (128, 256)), diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py index 0f317d81330a..ef0af2ec9ce3 100644 --- a/tests/random_lax_test.py +++ b/tests/random_lax_test.py @@ -968,7 +968,7 @@ def testMultivariateNormalCovariance(self): check_dtypes=False) @jtu.sample_product(method=['cholesky', 'eigh', 'svd']) - @jtu.skip_on_devices('cuda', 'tpu') # Some NaNs on accelerators. + @jtu.skip_on_devices('cuda', 'tpu') # Some NaNs on accelerators. ROCm supported def testMultivariateNormalSingularCovariance(self, method): # Singular covariance matrix https://github.com/jax-ml/jax/discussions/13293 mu = jnp.zeros((2,)) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index affcaf058c00..fa91f4f18185 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -139,8 +139,6 @@ def test_csr_fromdense_ad(self, shape, dtype): @jax.default_matmul_precision("float32") def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") csr_matmul = sparse_csr._csr_matvec if len(bshape) == 1 else sparse_csr._csr_matmat tol = {np.float32: 2E-5, np.float64: 1E-12, np.complex64: 1E-5, @@ -221,8 +219,6 @@ def test_csr_fromdense(self, shape, dtype): ) def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") op = lambda M: M.T if transpose else M @@ -594,8 +590,6 @@ def test_coo_spmm(self, shape, dtype, transpose): @jtu.run_on_devices("gpu") def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") tol = {np.float32: 2E-5, np.float64: 2E-14} @@ -1052,8 +1046,6 @@ def test_transpose(self, shape, dtype, Obj): @jax.default_matmul_precision("float32") def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") rng = sptu.rand_sparse(self.rng(), post=jnp.array) rng_b = jtu.rand_default(self.rng())