commit c1d608f74c48ac9c52fc2bd1f0c0fa83de740c48
author Scott Todd <scott.todd0@gmail.com> Wed Feb 14 13:59:46 2024 -0800
committer GitHub <noreply@github.com> Wed Feb 14 13:59:46 2024 -0800
tree 2a65313d960c9797b95d44da6b03915cfa7cd103
parent c066ceb0b9c0de6608d8ac33c4fd960af73d9ee2

Add host cpu tests for regression_suite/.../ukernel. (#16413)

We shouldn't just be testing this on ROCm. Even if the specific coverage
is only going to be looked at there, developers without ROCm on their
machines should still be able to run and debug this test.

Progress on https://github.com/openxla/iree/issues/16372 (tidying up
`regression_suite` while I evaluating extending it)

ci-extra: build_packages, regression_test_cpu,
regression_test_amdgpu_rocm

experimental/regression_suite/tests/pregenerated/test_ukernel.py

diff --git a/experimental/regression_suite/tests/pregenerated/test_ukernel.py b/experimental/regression_suite/tests/pregenerated/test_ukernel.py
index 1fc5015..22a5aa2 100644
--- a/experimental/regression_suite/tests/pregenerated/test_ukernel.py
+++ b/experimental/regression_suite/tests/pregenerated/test_ukernel.py
@@ -13,9 +13,6 @@
 
 COMMON_FLAGS = [
     "--iree-input-type=none",
-    "--iree-stream-resource-index-bits=64",
-    "--iree-vm-target-index-bits=64",
-    "--iree-stream-resource-max-allocation-size=3221225472",
 ]
 
 argmax_ukernel_source = fetch_source_fixture(
@@ -25,6 +22,19 @@
 
 
 @pytest.fixture
+def argmax_ukernel_host_cpu_vmfb(argmax_ukernel_source):
+    return iree_compile(
+        argmax_ukernel_source,
+        "host_cpu",
+        flags=COMMON_FLAGS
+        + [
+            "--iree-hal-target-backends=llvm-cpu",
+            "--iree-llvmcpu-target-cpu-features=host",
+        ],
+    )
+
+
+@pytest.fixture
 def argmax_ukernel_gfx90a_rocm_vmfb(argmax_ukernel_source):
     return iree_compile(
         argmax_ukernel_source,
@@ -89,6 +99,55 @@
 
 @pytest.mark.presubmit
 @pytest.mark.unstable_linalg
+@pytest.mark.plat_host_cpu
+def test_correctness_host_cpu(
+    argmax_ukernel_host_cpu_vmfb,
+    argmax_input_f16,
+    argmax_output_f16,
+    argmax_input_f32,
+    argmax_output_f32,
+):
+    iree_run_module(
+        argmax_ukernel_host_cpu_vmfb,
+        device="local-task",
+        function="argmax_3d_dyn_f16i32",
+        args=[
+            f"--input=@{argmax_input_f16.path}",
+            f"--expected_output=@{argmax_output_f16.path}",
+        ],
+    )
+    iree_run_module(
+        argmax_ukernel_host_cpu_vmfb,
+        device="local-task",
+        function="argmax_3d_dyn_f16i64",
+        args=[
+            f"--input=@{argmax_input_f16.path}",
+            f"--expected_output=@{argmax_output_f16.path}",
+        ],
+    )
+
+    iree_run_module(
+        argmax_ukernel_host_cpu_vmfb,
+        device="local-task",
+        function="argmax_3d_dyn_f32i32",
+        args=[
+            f"--input=@{argmax_input_f32.path}",
+            f"--expected_output=@{argmax_output_f32.path}",
+        ],
+    )
+    iree_run_module(
+        argmax_ukernel_host_cpu_vmfb,
+        device="local-task",
+        function="argmax_3d_dyn_f32i64",
+        args=[
+            f"--input=@{argmax_input_f32.path}",
+            f"--expected_output=@{argmax_output_f32.path}",
+        ],
+    )
+
+
+@pytest.mark.presubmit
+@pytest.mark.unstable_linalg
 @pytest.mark.plat_gfx90a_rocm
 def test_correctness_gfx90a_rocm(
     argmax_ukernel_gfx90a_rocm_vmfb,