Re-land e2e attention test suite. (#18320)

Add the e2e test suite for the Attention. It only checks CPU FP16, and the reference implementation is FP32. #17751 #18302 --------- Signed-off-by: erman-gurses <[email protected]>
iree-org · Aug 22, 2024 · 588732c · 588732c
1 parent c44d29b
commit 588732c
Show file tree

Hide file tree

Showing 5 changed files with 1,104 additions and 0 deletions.
diff --git a/tests/e2e/attention/CMakeLists.txt b/tests/e2e/attention/CMakeLists.txt
@@ -0,0 +1,80 @@
+# TODO: (#17751) Add the arm_64 tests when the bug resolved. See:
+#   https://github.com/iree-org/iree/actions/runs/10468944505/job/28990909321#step:4:9815
+if(IREE_ARCH STREQUAL "arm_64")
+  return()
+endif()
+
+iree_generated_e2e_runner_test(
+  NAME
+    e2e_attention_cpu_f16_f16_f16_small
+  TEST_TYPE
+    attention
+  GENERATOR
+    "generate_e2e_attention_tests.py"
+  GENERATOR_ARGS
+    "--query_type=f16"
+    "--key_type=f16"
+    "--value_type=f16"
+    "--shapes=small"
+  TEST_RUNNER
+    iree_tools_testing_e2e_iree-e2e-attention-test
+  TARGET_BACKENDS
+    "llvm-cpu"
+  DRIVERS
+    "local-task"
+  LABELS
+    "hostonly"
+    "local"
+  TARGET_CPU_FEATURES_VARIANTS
+    "default"
+)
+
+iree_generated_e2e_runner_test(
+  NAME
+    e2e_attention_cpu_f16_f16_f16_medium
+  TEST_TYPE
+    attention
+  GENERATOR
+    "generate_e2e_attention_tests.py"
+  GENERATOR_ARGS
+    "--query_type=f16"
+    "--key_type=f16"
+    "--value_type=f16"
+    "--shapes=medium"
+  TEST_RUNNER
+    iree_tools_testing_e2e_iree-e2e-attention-test
+  TARGET_BACKENDS
+    "llvm-cpu"
+  DRIVERS
+    "local-task"
+  LABELS
+    "hostonly"
+    "local"
+  TARGET_CPU_FEATURES_VARIANTS
+    "default"
+)
+
+iree_generated_e2e_runner_test(
+  NAME
+    e2e_attention_cpu_f16_f16_f16_large
+  TEST_TYPE
+    attention
+  GENERATOR
+    "generate_e2e_attention_tests.py"
+  GENERATOR_ARGS
+    "--query_type=f16"
+    "--key_type=f16"
+    "--value_type=f16"
+    "--shapes=large"
+  TEST_RUNNER
+    iree_tools_testing_e2e_iree-e2e-attention-test
+  TARGET_BACKENDS
+    "llvm-cpu"
+  DRIVERS
+    "local-task"
+  LABELS
+    "hostonly"
+    "local"
+  TARGET_CPU_FEATURES_VARIANTS
+    "default"
+)