Skip to content

Commit

Permalink
Update the induction variables as b, m, k1, k2, and n
Browse files Browse the repository at this point in the history
  • Loading branch information
IanNod committed Aug 5, 2024
1 parent ba9ea85 commit 1e65ec5
Show file tree
Hide file tree
Showing 6 changed files with 1,152 additions and 0 deletions.
53 changes: 53 additions & 0 deletions tests/e2e/attention/Build.bazel
Original file line number Diff line number Diff line change
@@ -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 fa2 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_fa2_tests",
srcs = ["generate_e2e_fa2_tests.py"],
)

###########################################################################
##
## LLVMCPU backend
##
###########################################################################

# Default CPU backend.
[iree_generated_e2e_runner_test(
name = "e2e_fa2_cpu_%s_%s_%s_%s" % (dtype, dtype, dtype, size),
generator = ":generate_e2e_fa2_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-fa2-test",
test_type = "fa2",
) for dtype in [
"f16",
] for size in [
"small",
"medium",
"large",
]]
88 changes: 88 additions & 0 deletions tests/e2e/attention/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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_fa2_cpu_f16_f16_f16_small
TEST_TYPE
fa2
GENERATOR
"generate_e2e_fa2_tests.py"
GENERATOR_ARGS
"--query_type=f16"
"--key_type=f16"
"--value_type=f16"
"--shapes=small"
TEST_RUNNER
iree_tools_testing_e2e_iree-e2e-fa2-test
TARGET_BACKENDS
"llvm-cpu"
DRIVERS
"local-task"
LABELS
"hostonly"
"local"
TARGET_CPU_FEATURES_VARIANTS
"default"
)

iree_generated_e2e_runner_test(
NAME
e2e_fa2_cpu_f16_f16_f16_medium
TEST_TYPE
fa2
GENERATOR
"generate_e2e_fa2_tests.py"
GENERATOR_ARGS
"--query_type=f16"
"--key_type=f16"
"--value_type=f16"
"--shapes=medium"
TEST_RUNNER
iree_tools_testing_e2e_iree-e2e-fa2-test
TARGET_BACKENDS
"llvm-cpu"
DRIVERS
"local-task"
LABELS
"hostonly"
"local"
TARGET_CPU_FEATURES_VARIANTS
"default"
)

iree_generated_e2e_runner_test(
NAME
e2e_fa2_cpu_f16_f16_f16_large
TEST_TYPE
fa2
GENERATOR
"generate_e2e_fa2_tests.py"
GENERATOR_ARGS
"--query_type=f16"
"--key_type=f16"
"--value_type=f16"
"--shapes=large"
TEST_RUNNER
iree_tools_testing_e2e_iree-e2e-fa2-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 ###
Loading

0 comments on commit 1e65ec5

Please sign in to comment.