Enable some E2E tests for CUDA (#4998)
Those are disabled by default in CI until a CUDA system is setup.
diff --git a/build_tools/bazel/build_core.sh b/build_tools/bazel/build_core.sh
index b7cd47c..4c7a54a 100755
--- a/build_tools/bazel/build_core.sh
+++ b/build_tools/bazel/build_core.sh
@@ -48,6 +48,9 @@
declare -a default_build_tag_filters=("-nokokoro")
declare -a default_test_tag_filters=("-nokokoro" "-driver=metal")
+# CUDA CI testing disabled until we setup a target for it.
+default_test_tag_filters+=("-driver=cuda")
+
if [[ "${IREE_VULKAN_DISABLE?}" == 1 ]]; then
default_test_tag_filters+=("-driver=vulkan")
fi
diff --git a/build_tools/cmake/test.sh b/build_tools/cmake/test.sh
index b38df1c..dc3ab5c 100755
--- a/build_tools/cmake/test.sh
+++ b/build_tools/cmake/test.sh
@@ -28,6 +28,8 @@
# and turning on the llvmaot ones.
export IREE_VULKAN_DISABLE=${IREE_VULKAN_DISABLE:-0}
export IREE_LLVMAOT_DISABLE=${IREE_LLVMAOT_DISABLE:-0}
+# CUDA is off by default.
+export IREE_CUDA_DISABLE=${IREE_CUDA_DISABLE:-1}
# Tests to exclude by label. In addition to any custom labels (which are carried
# over from Bazel tags), every test should be labeled with the directory it is
@@ -57,6 +59,9 @@
if [[ "${IREE_LLVMAOT_DISABLE?}" == 1 ]]; then
label_exclude_args+=("^driver=dylib$")
fi
+if [[ "${IREE_CUDA_DISABLE?}" == 1 ]]; then
+ label_exclude_args+=("^driver=cuda$")
+fi
# Join on "|"
label_exclude_regex="($(IFS="|" ; echo "${label_exclude_args[*]?}"))"
diff --git a/iree/test/e2e/xla_ops/BUILD b/iree/test/e2e/xla_ops/BUILD
index 2db32d2..f3e3f3a 100644
--- a/iree/test/e2e/xla_ops/BUILD
+++ b/iree/test/e2e/xla_ops/BUILD
@@ -256,6 +256,32 @@
target_backend = "vulkan-spirv",
)
+iree_check_single_backend_test_suite(
+ name = "check_linalg_on_tensors_cuda",
+ srcs = [
+ "add.mlir",
+ "broadcast.mlir",
+ "broadcast_add.mlir",
+ "broadcast_in_dim.mlir",
+ "clamp.mlir",
+ "compare.mlir",
+ "constant.mlir",
+ "divide.mlir",
+ "iota.mlir",
+ "maximum.mlir",
+ "minimum.mlir",
+ "multiply.mlir",
+ "negate.mlir",
+ "select.mlir",
+ "subtract.mlir",
+ ],
+ compiler_flags = [
+ "-iree-flow-dispatch-linalg-on-tensors",
+ ],
+ driver = "cuda",
+ target_backend = "cuda",
+)
+
test_suite(
name = "check",
tests = [
diff --git a/iree/test/e2e/xla_ops/CMakeLists.txt b/iree/test/e2e/xla_ops/CMakeLists.txt
index 6d2b15e..e1c0dc0 100644
--- a/iree/test/e2e/xla_ops/CMakeLists.txt
+++ b/iree/test/e2e/xla_ops/CMakeLists.txt
@@ -221,4 +221,31 @@
"-iree-codegen-spirv-experimental-linalg-on-tensors"
)
+iree_check_single_backend_test_suite(
+ NAME
+ check_linalg_on_tensors_cuda
+ SRCS
+ "add.mlir"
+ "broadcast.mlir"
+ "broadcast_add.mlir"
+ "broadcast_in_dim.mlir"
+ "clamp.mlir"
+ "compare.mlir"
+ "constant.mlir"
+ "divide.mlir"
+ "iota.mlir"
+ "maximum.mlir"
+ "minimum.mlir"
+ "multiply.mlir"
+ "negate.mlir"
+ "select.mlir"
+ "subtract.mlir"
+ TARGET_BACKEND
+ "cuda"
+ DRIVER
+ "cuda"
+ COMPILER_FLAGS
+ "-iree-flow-dispatch-linalg-on-tensors"
+)
+
### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
