Introduce local skip function to work-around issue with CPU driver

Introduce pytest option  --skip-known-top-k-failures-on-cpu to skip
tests known to fail with CPU device on architectures that support
AVX2 but not AVX512 (reproduced with
`CL_CONFIG_CPU_TARGET_ARCH=core-avx2`).

Use the option in pytest command call in conda-packages worfklow,
and in scripts/gen_coverage.py

Author: oleksandr-pavlyk

.github/workflows/conda-package.yml

@@ -228,7 +228,7 @@ jobs:
         run: |
           . $CONDA/etc/profile.d/conda.sh
           conda activate ${{ env.TEST_ENV_NAME }}
-          python -m pytest -v --pyargs $MODULE_NAME
+          python -m pytest -v --skip-known-top-k-failures-on-cpu --pyargs $MODULE_NAME
 
   test_windows:
     needs: build_windows
@@ -399,7 +399,7 @@ jobs:
           SYCL_CACHE_PERSISTENT: 1
         working-directory: ${{ env.workdir }}\test_tmp
         run: >-
-          conda activate ${{ env.TEST_ENV_NAME }} && python -m pytest -v -s --pyargs ${{ env.MODULE_NAME }}
+          conda activate ${{ env.TEST_ENV_NAME }} && python -m pytest -v -s --skip-known-top-k-failures-on-cpu --pyargs ${{ env.MODULE_NAME }}
 
   upload_linux:
     needs: test_linux

dpctl/tests/conftest.py

@@ -62,6 +62,14 @@ def pytest_addoption(parser):
         default=False,
         help="skip test_syclinterface",
     )
+    parser.addoption(
+        "--skip-known-top-k-failures-on-cpu",
+        action="store_true",
+        default=False,
+        help="skip top_k tests for types int8/int16 known to "
+        "fail on CPU devices on older architectures with "
+        "DPC++ 2025.0",
+    )
 
 
 def pytest_collection_modifyitems(config, items):

dpctl/tests/test_usm_ndarray_top_k.py

@@ -20,6 +20,11 @@
 from dpctl.tests.helper import get_queue_or_skip, skip_if_dtype_not_supported
 
 
+@pytest.fixture
+def skip_known_failues_on_cpu(request):
+    return request.config.getoption("--skip-known-top-k-failures-on-cpu")
+
+
 def _expected_largest_inds(inp, n, shift, k):
     "Computed expected top_k indices for mode='largest'"
     assert k < n
@@ -52,10 +57,17 @@ def _expected_largest_inds(inp, n, shift, k):
     return expected_inds
 
 
+def _skip_if_workaround_is_needed(q, dtype, n, enabled):
+    if enabled:
+        dev = q.sycl_device
+        if dev.is_cpu and dtype in ["i1", "i2"] and n > 128:
+            pytest.skip(reason="CPU driver bug")
+
+
 @pytest.mark.parametrize(
     "dtype",
     [
-        pytest.param("i1", marks=pytest.mark.skip(reason="CPU bug")),
+        "i1",
         "u1",
         "i2",
         "u2",
@@ -71,11 +83,10 @@ def _expected_largest_inds(inp, n, shift, k):
     ],
 )
 @pytest.mark.parametrize("n", [33, 43, 255, 511, 1021, 8193])
-def test_top_k_1d_largest(dtype, n):
+def test_top_k_1d_largest(dtype, n, skip_known_failues_on_cpu):
     q = get_queue_or_skip()
     skip_if_dtype_not_supported(dtype, q)
-    if dtype == "i1":
-        pytest.skip()
+    _skip_if_workaround_is_needed(q, dtype, n, skip_known_failues_on_cpu)
 
     shift, k = 734, 5
     o = dpt.ones(n, dtype=dtype)
@@ -128,7 +139,7 @@ def _expected_smallest_inds(inp, n, shift, k):
 @pytest.mark.parametrize(
     "dtype",
     [
-        pytest.param("i1", marks=pytest.mark.skip(reason="CPU bug")),
+        "i1",
         "u1",
         "i2",
         "u2",
@@ -144,10 +155,12 @@ def _expected_smallest_inds(inp, n, shift, k):
     ],
 )
 @pytest.mark.parametrize("n", [37, 39, 61, 255, 257, 513, 1021, 8193])
-def test_top_k_1d_smallest(dtype, n):
+def test_top_k_1d_smallest(dtype, n, skip_known_failues_on_cpu):
     q = get_queue_or_skip()
     skip_if_dtype_not_supported(dtype, q)
 
+    _skip_if_workaround_is_needed(q, dtype, n, skip_known_failues_on_cpu)
+
     shift, k = 734, 5
     o = dpt.ones(n, dtype=dtype)
     z = dpt.zeros(n, dtype=dtype)
@@ -163,3 +176,91 @@ def test_top_k_1d_smallest(dtype, n):
     assert dpt.all(s.indices == expected_inds)
     assert dpt.all(s.values == dpt.zeros(k, dtype=dtype)), s.values
     assert dpt.all(s.values == inp[s.indices]), s.indices
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        # skip short types to ensure that m*n can be represented
+        # in the type
+        "i4",
+        "u4",
+        "i8",
+        "u8",
+        "f2",
+        "f4",
+        "f8",
+        "c8",
+        "c16",
+    ],
+)
+@pytest.mark.parametrize("n", [37, 39, 61, 255, 257, 513, 1021, 8193])
+def test_top_k_2d_largest(dtype, n):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dtype, q)
+
+    m, k = 8, 3
+    if dtype == "f2" and m * n > 2000:
+        pytest.skip(
+            "f2 can not distinguish between large integers used in this test"
+        )
+
+    x = dpt.reshape(dpt.arange(m * n, dtype=dtype), (m, n))
+
+    r = dpt.top_k(x, k, axis=1)
+
+    assert r.values.shape == (m, k)
+    assert r.indices.shape == (m, k)
+    expected_inds = dpt.reshape(dpt.arange(n, dtype=r.indices.dtype), (1, n))[
+        :, -k:
+    ]
+    assert expected_inds.shape == (1, k)
+    assert dpt.all(
+        dpt.sort(r.indices, axis=1) == dpt.sort(expected_inds, axis=1)
+    ), (r.indices, expected_inds)
+    expected_vals = x[:, -k:]
+    assert dpt.all(
+        dpt.sort(r.values, axis=1) == dpt.sort(expected_vals, axis=1)
+    )
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        # skip short types to ensure that m*n can be represented
+        # in the type
+        "i4",
+        "u4",
+        "i8",
+        "u8",
+        "f2",
+        "f4",
+        "f8",
+        "c8",
+        "c16",
+    ],
+)
+@pytest.mark.parametrize("n", [37, 39, 61, 255, 257, 513, 1021, 8193])
+def test_top_k_2d_smallest(dtype, n):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dtype, q)
+
+    m, k = 8, 3
+    if dtype == "f2" and m * n > 2000:
+        pytest.skip(
+            "f2 can not distinguish between large integers used in this test"
+        )
+
+    x = dpt.reshape(dpt.arange(m * n, dtype=dtype), (m, n))
+
+    r = dpt.top_k(x, k, axis=1, mode="smallest")
+
+    assert r.values.shape == (m, k)
+    assert r.indices.shape == (m, k)
+    expected_inds = dpt.reshape(dpt.arange(n, dtype=r.indices.dtype), (1, n))[
+        :, :k
+    ]
+    assert dpt.all(
+        dpt.sort(r.indices, axis=1) == dpt.sort(expected_inds, axis=1)
+    )
+    assert dpt.all(dpt.sort(r.values, axis=1) == dpt.sort(x[:, :k], axis=1))

scripts/gen_coverage.py

@@ -102,6 +102,7 @@ def run(
             "dpctl",
             "-vv",
             "--ignore=dpctl/tensor/libtensor/tests",
+            "--skip-known-top-k-failures-on-cpu",
             "--no-sycl-interface-test",
         ],
         cwd=setup_dir,