ci(blas): derive host-ASAN thread budgets from the CPU allocation

.github/scripts/tests/test_blas_test_wrappers.py

@@ -0,0 +1,78 @@
+# Copyright Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+"""Check CPU budgets and smoke coverage at the BLAS test launch boundary."""
+
+import fnmatch
+import os
+import runpy
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+
+
+class BlasTestWrappersTest(unittest.TestCase):
+    def launch(self, component, artifact_group, budget=None):
+        env = {
+            "ARTIFACT_GROUP": artifact_group,
+            "AMDGPU_FAMILIES": "gfx950",
+            "RUNNER_OS": "Linux",
+            "THEROCK_BIN_DIR": "/unused/bin",
+            "THEROCK_DIR": str(REPO_ROOT),
+            "TEST_TYPE": "quick",
+            # The first shard also performs unrelated layout and tuning-store checks.
+            "SHARD_INDEX": "2",
+            "TOTAL_SHARDS": "2",
+            "HSA_XNACK": "0",
+        }
+        if budget is not None:
+            env.update(OMP_NUM_THREADS=budget, OPENBLAS_NUM_THREADS=budget)
+        with patch.dict(os.environ, env, clear=True), patch("subprocess.run") as run:
+            runpy.run_path(str(REPO_ROOT / "test/therock" / f"test_{component}.py"))
+        run.assert_called_once()
+        args, kwargs = run.call_args
+        self.assertEqual(args[0][0], f"/unused/bin/{component}-test")
+        self.assertTrue(kwargs["check"])
+        return args[0], kwargs["env"]
+
+    def test_host_asan_preserves_cpu_budget(self):
+        for component in ("hipblaslt", "hipsparselt"):
+            for budget in (None, "1", "25"):
+                with self.subTest(component=component, budget=budget):
+                    _, env = self.launch(component, "gfx950-host-asan", budget)
+                    self.assertEqual(env.get("OMP_NUM_THREADS"), budget)
+                    self.assertEqual(env.get("OPENBLAS_NUM_THREADS"), budget)
+                    self.assertEqual(env["HSA_XNACK"], "0")
+
+    def test_device_asan_keeps_hipblaslt_workarounds(self):
+        _, env = self.launch("hipblaslt", "gfx950-asan", "25")
+        self.assertEqual(env["OMP_NUM_THREADS"], "1")
+        self.assertEqual(env["OPENBLAS_NUM_THREADS"], "25")
+        self.assertEqual(env["HSA_XNACK"], "1")
+
+    def test_sparse_smoke_includes_128_by_128_under_asan(self):
+        for artifact in ("gfx950", "gfx950-host-asan", "gfx950-asan"):
+            with self.subTest(artifact=artifact):
+                cmd, _ = self.launch("hipsparselt", artifact, "25")
+                filter_arg = next(
+                    arg for arg in cmd if arg.startswith("--gtest_filter=")
+                )
+                positive, _, negative = filter_arg.split("=", 1)[1].partition("-")
+
+                def selected(name):
+                    return any(
+                        fnmatch.fnmatchcase(name, p) for p in positive.split(":")
+                    ) and not any(
+                        fnmatch.fnmatchcase(name, p) for p in negative.split(":")
+                    )
+
+                for transpose in ("NN", "NT", "TN", "TT"):
+                    name = f"spmm/smoke_spmm_strided_batched_smoke_f16_r_{transpose}_128_128_128"
+                    self.assertTrue(selected(name), name)
+                    self.assertFalse(selected(name + "_known_bug"))
+
+
+if __name__ == "__main__":
+    unittest.main()

.github/workflows/hipblaslt-asan-ci.yml

@@ -14,6 +14,7 @@ on:
       - "test/therock/lsan.supp"
       - ".github/workflows/hipblaslt-asan-ci.yml"
       - ".github/actions/ci-env/**"
+      - ".github/scripts/tests/test_blas_test_wrappers.py"
   workflow_dispatch:
     inputs:
       test_type:
@@ -60,6 +61,11 @@ jobs:
         id: ci-env
         uses: ./.github/actions/ci-env
 
+      - name: Check BLAS test budgets and smoke coverage
+        run: |
+          git sparse-checkout add test/therock
+          python3 .github/scripts/tests/test_blas_test_wrappers.py
+
       - name: Compute matrices and select timeouts
         id: matrix
         run: |
@@ -278,6 +284,7 @@ jobs:
         --group-add 992
         --group-add 110
         --env-file /etc/podinfo/gha-gpu-isolation-settings
+        -e KUBE_CPU_REQUEST
     strategy:
       fail-fast: false
       matrix: ${{ fromJSON(needs.setup.outputs.test_matrix) }}
@@ -295,6 +302,22 @@ jobs:
       HSA_XNACK: "0"
       TEST_COMPONENT: ${{ matrix.component }}
     steps:
+      # Match the allocation-derived budget used by TheRock component jobs.
+      - name: Set thread limits
+        run: |
+          if [[ "${KUBE_CPU_REQUEST:-}" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
+            CPU_COUNT="${KUBE_CPU_REQUEST%%.*}"
+            CPU_COUNT=$((10#${CPU_COUNT}))
+            if (( CPU_COUNT < 1 )); then
+              CPU_COUNT=1
+            fi
+            echo "OPENBLAS_NUM_THREADS=${CPU_COUNT}" >> "$GITHUB_ENV"
+            echo "OMP_NUM_THREADS=${CPU_COUNT}" >> "$GITHUB_ENV"
+            echo "Thread limits set: OPENBLAS_NUM_THREADS=${CPU_COUNT}, OMP_NUM_THREADS=${CPU_COUNT} (from KUBE_CPU_REQUEST=${KUBE_CPU_REQUEST})"
+          else
+            echo "KUBE_CPU_REQUEST absent or invalid; preserving inherited thread limits"
+          fi
+
       - name: Checkout TheRock repository
         uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:

projects/hipblaslt/TESTING.md

@@ -965,6 +965,11 @@ manually dispatched runs.
 
 **Runtime configuration.** The lane sets a large ASAN quarantine, a LeakSanitizer suppression file
 at `test/therock/lsan.supp` in the repository root, an explicit symbolizer path, and `HSA_XNACK=0`.
+Both libraries use `OMP_NUM_THREADS` and `OPENBLAS_NUM_THREADS` derived from the
+runner's `KUBE_CPU_REQUEST`: the integer part, with a minimum of one. If the allocation
+is absent or invalid, the workflow preserves inherited thread settings. Host-ASAN
+hipBLASLt preserves this budget; full device ASAN retains its one-thread OpenMP override.
+The hipSPARSELt quick selection includes the 128x128 strided-batched smoke cases under ASAN.
 Host-only AddressSanitizer does not require GPU page-fault retry (XNACK), so this workflow disables
 it. The shared test launcher preserves the caller's XNACK setting for host-ASAN artifacts and sets
 `HSA_XNACK=1` for full-ASAN artifacts. The suppression file is the thing to look at first when a

test/therock/test_hipblaslt.py

@@ -36,7 +36,7 @@ def is_asan():
     # Only device ASAN requires XNACK; host ASAN preserves the caller's setting.
     if "host-asan" not in os.getenv("ARTIFACT_GROUP", ""):
         environ_vars["HSA_XNACK"] = "1"
-    environ_vars["OMP_NUM_THREADS"] = "1"
+        environ_vars["OMP_NUM_THREADS"] = "1"
 
 # ---------------------------------------------------------------------------
 # Pre-flight: validate the installed library tree's physical layout BEFORE

test/therock/test_hipsparselt.py

@@ -15,11 +15,6 @@
 ).resolve()
 
 
-def is_asan():
-    """Return True for ASAN builds, detected via the ARTIFACT_GROUP env var."""
-    return "asan" in os.getenv("ARTIFACT_GROUP", "")
-
-
 logging.basicConfig(level=logging.INFO)
 
 # GTest sharding
@@ -36,17 +31,7 @@ def is_asan():
 
 test_filter = []
 if test_type == "quick":
-    quick_filter = "--gtest_filter=*smoke*-*known_bug*"
-    if is_asan():
-        # Under ASAN, the M=N=128 spmm_strided_batched smoke cases take ~40-360 ms each instead
-        # of ~7 ms, adding 43-67 minutes on gfx942/gfx950. The standard tests still run them.
-        quick_filter += (
-            ":*smoke_spmm_strided_batched*_NN_128_128_*"
-            ":*smoke_spmm_strided_batched*_NT_128_128_*"
-            ":*smoke_spmm_strided_batched*_TN_128_128_*"
-            ":*smoke_spmm_strided_batched*_TT_128_128_*"
-        )
-    test_filter.append(quick_filter)
+    test_filter.append("--gtest_filter=*smoke*-*known_bug*")
 elif test_type == "standard":
     test_filter.append("--gtest_filter=*smoke*:*quick*:*pre_checkin*-*known_bug*")