Add e2e test suite for the Attention - CPU Backend (iree-org#17751)

Add the e2e test suite for the Attention. For now, it only checks CPU FP16, and the reference implementation is FP32. --------- Signed-off-by: ERMAN GURSES <[email protected]> Co-authored-by: ian <[email protected]>
nirvedhmeshram · Aug 19, 2024 · 2d629c6 · 2d629c6
1 parent b637a30
commit 2d629c6
Show file tree

Hide file tree

Showing 6 changed files with 1,165 additions and 0 deletions.
diff --git a/tests/e2e/attention/BUILD.bazel b/tests/e2e/attention/BUILD.bazel
@@ -0,0 +1,53 @@
+# Copyright 2024 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# End-to-end attention tests.
+
+load("//build_tools/bazel:iree_e2e_generated_runner_test.bzl", "iree_generated_e2e_runner_test")
+
+package(
+    features = ["layering_check"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+py_binary(
+    name = "generate_e2e_attention_tests",
+    srcs = ["generate_e2e_attention_tests.py"],
+)
+
+###########################################################################
+##
+## LLVMCPU backend
+##
+###########################################################################
+
+# Default CPU backend.
+[iree_generated_e2e_runner_test(
+    name = "e2e_attention_cpu_%s_%s_%s_%s" % (dtype, dtype, dtype, size),
+    generator = ":generate_e2e_attention_tests",
+    generator_args = [
+        "--query_type=%s" % dtype,
+        "--key_type=%s" % dtype,
+        "--value_type=%s" % dtype,
+        "--shapes=%s" % size,
+    ],
+    tags = [
+        "hostonly",
+        "local",
+    ],
+    target_backends_and_drivers = [
+        ("llvm-cpu", "local-task"),
+    ],
+    target_cpu_features_variants = ["default"],
+    test_runner = "//tools/testing/e2e:iree-e2e-attention-test",
+    test_type = "attention",
+) for dtype in [
+    "f16",
+] for size in [
+    "small",
+    "medium",
+    "large",
+]]
diff --git a/tests/e2e/attention/CMakeLists.txt b/tests/e2e/attention/CMakeLists.txt
@@ -0,0 +1,88 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from           #
+# tests/e2e/attention/BUILD.bazel                                              #
+#                                                                              #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary   #
+# CMake-only content.                                                          #
+#                                                                              #
+# To disable autogeneration for this file entirely, delete this header.        #
+################################################################################
+
+iree_add_all_subdirs()
+
+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"
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###