Add f16 e2e tests for vulkan (#5257)

build_tools/bazel/build_core.sh

@@ -48,6 +48,9 @@ declare -a test_env_args=(
 declare -a default_build_tag_filters=("-nokokoro")
 declare -a default_test_tag_filters=("-nokokoro" "-driver=metal")
 
+# The VK_KHR_shader_float16_int8 extension is optional prior to Vulkan 1.2.
+# We test on SwiftShader, which does not support this extension.
+default_test_tag_filters+=("-vulkan_uses_vk_khr_shader_float16_int8")
 # CUDA CI testing disabled until we setup a target for it.
 default_test_tag_filters+=("-driver=cuda")
 

build_tools/cmake/test.sh

@@ -30,6 +30,9 @@ 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}
+# The VK_KHR_shader_float16_int8 extension is optional prior to Vulkan 1.2.
+# We test on SwiftShader, which does not support this extension.
+export IREE_VULKAN_F16_DISABLE=${IREE_VULKAN_F16_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
@@ -62,6 +65,9 @@ fi
 if [[ "${IREE_CUDA_DISABLE?}" == 1 ]]; then
   label_exclude_args+=("^driver=cuda$")
 fi
+if [[ "${IREE_VULKAN_F16_DISABLE?}" == 1 ]]; then
+  label_exclude_args+=("^vulkan_uses_vk_khr_shader_float16_int8$")
+fi
 
 # Join on "|"
 label_exclude_regex="($(IFS="|" ; echo "${label_exclude_args[*]?}"))"

build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-swiftshader/build.sh

@@ -71,4 +71,4 @@ export CTEST_PARALLEL_LEVEL=${CTEST_PARALLEL_LEVEL:-$(nproc)}
 echo "Testing with CTest"
 ctest --output-on-failure \
    --tests-regex "^integrations/tensorflow/|^bindings/python/" \
-   --label-exclude "^nokokoro$"
+   --label-exclude "^nokokoro$|^vulkan_uses_vk_khr_shader_float16_int8$"

build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build.sh

@@ -44,6 +44,7 @@ echo "Initializing submodules"
 echo "Building with cmake"
 ./build_tools/cmake/clean_build.sh
 
+export IREE_VULKAN_F16_DISABLE=0
 export IREE_CUDA_DISABLE=0
 echo "Testing with ctest"
 ./build_tools/cmake/test.sh

iree/test/e2e/xla_ops/partial/BUILD

@@ -0,0 +1,51 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Tests of end-to-end IREE support for support of specific subsets of ops in the
+# XLA HLO dialect. This is for cases where some variation (e.g. tensor element
+# types) of an op is not supported on all backends. When the test is supported
+# on all backends it should be moved into the corresponding op test in the main
+# xla_ops direcotry. Each test file should have a name matching the
+# corresponding XLA HLO op with a suffix indicating the subset it tests. Only
+# the functionality of that op should be tessted (though it may make use of
+# other ops where necessary). Tests should be written using the IREE Check
+# framework and should always pass on the reference VMLA backend. See
+# https://google.github.io/iree/TestingGuide#iree-core-end-to-end-tests.
+
+load("//build_tools/bazel:iree_check_test.bzl", "iree_check_single_backend_test_suite")
+
+package(
+    default_visibility = ["//visibility:public"],
+    features = ["layering_check"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+iree_check_single_backend_test_suite(
+    name = "check_linalg_on_tensors_vulkan-spirv_vulkan_f16",
+    srcs = [
+        "add_f16.mlir",
+        "dot_f16.mlir",
+    ],
+    compiler_flags = [
+        "-iree-flow-dispatch-linalg-on-tensors",
+        "-iree-codegen-spirv-experimental-linalg-on-tensors",
+        "-iree-vulkan-target-triple=valhall-g77-unknown-android10",
+    ],
+    driver = "vulkan",
+    tags = [
+        "notap",
+        "vulkan_uses_vk_khr_shader_float16_int8",
+    ],
+    target_backend = "vulkan-spirv",
+)

iree/test/e2e/xla_ops/partial/CMakeLists.txt

@@ -0,0 +1,32 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from           #
+# iree/test/e2e/xla_ops/partial/BUILD                                          #
+#                                                                              #
+# 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_check_single_backend_test_suite(
+  NAME
+    check_linalg_on_tensors_vulkan-spirv_vulkan_f16
+  SRCS
+    "add_f16.mlir"
+    "dot_f16.mlir"
+  TARGET_BACKEND
+    "vulkan-spirv"
+  DRIVER
+    "vulkan"
+  COMPILER_FLAGS
+    "-iree-flow-dispatch-linalg-on-tensors"
+    "-iree-codegen-spirv-experimental-linalg-on-tensors"
+    "-iree-vulkan-target-triple=valhall-g77-unknown-android10"
+  LABELS
+    "notap"
+    "vulkan_uses_vk_khr_shader_float16_int8"
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###

iree/test/e2e/xla_ops/partial/add_f16.mlir

@@ -0,0 +1,7 @@
+func @add_f16() attributes { iree.module.export } {
+  %0 = iree.unfoldable_constant dense<[1.5, 2.0, 3.0, 4.2]> : tensor<4xf16>
+  %1 = iree.unfoldable_constant dense<[5.0, 6.2, 7.0, 8.1]> : tensor<4xf16>
+  %result = "mhlo.add"(%0, %1) : (tensor<4xf16>, tensor<4xf16>) -> tensor<4xf16>
+  check.expect_almost_eq_const(%result, dense<[6.5, 8.2, 10.0, 12.3]> : tensor<4xf16>) : tensor<4xf16>
+  return
+}

iree/test/e2e/xla_ops/partial/dot_f16.mlir

@@ -0,0 +1,20 @@
+func @f16() attributes { iree.module.export } {
+  %lhs = iree.unfoldable_constant dense<[
+    [15.0, 14.0, 13.0],
+    [12.0, 11.0, 10.0],
+    [09.0, 08.0, 07.0],
+    [06.0, 05.0, 04.0],
+    [03.0, 02.0, 01.0]]> : tensor<5x3xf16>
+  %rhs = iree.unfoldable_constant dense<[
+    [15.0, 14.0, 13.0, 12.0, 11.0],
+    [10.0, 09.0, 08.0, 07.0, 06.0],
+    [05.0, 04.0, 03.0, 02.0, 01.0]]> : tensor<3x5xf16>
+  %res = "mhlo.dot"(%lhs, %rhs) : (tensor<5x3xf16>, tensor<3x5xf16>) -> tensor<5x5xf16>
+  check.expect_almost_eq_const(%res, dense<[
+    [430.0, 388.0, 346.0, 304.0, 262.0],
+    [340.0, 307.0, 274.0, 241.0, 208.0],
+    [250.0, 226.0, 202.0, 178.0, 154.0],
+    [160.0, 145.0, 130.0, 115.0, 100.0],
+    [70.0, 64.0, 58.0, 52.0, 46.0]]> : tensor<5x5xf16>) : tensor<5x5xf16>
+  return
+}