Create a test for every tf_to_kernel invocation.
This additionally will create a nice reproducer command on any failure.
PiperOrigin-RevId: 344310249
Change-Id: Id076d95a80430f4f71f3fa091c4e7943e981b83f
diff --git a/tensorflow/core/kernels/mlir_generated/build_defs.bzl b/tensorflow/core/kernels/mlir_generated/build_defs.bzl
index ec1bbdf..e20257d 100644
--- a/tensorflow/core/kernels/mlir_generated/build_defs.bzl
+++ b/tensorflow/core/kernels/mlir_generated/build_defs.bzl
@@ -366,6 +366,22 @@
static_library = "{name}_{type}.a".format(name = name, type = type),
)
+ # We have to use a sh_test instead of build_test because it doesn't properly find the dependent targets.
+ native.sh_test(
+ name = "{name}_{type}_gen_test".format(name = name, type = type),
+ srcs = ["build_test.sh"],
+ tags = ["no_rocm"],
+ args = [
+ "$(location //tensorflow/compiler/mlir/tools/kernel_gen:tf_to_kernel)",
+ "$(location {name}_{type}.mlir)".format(name = name, type = type),
+ ],
+ size = "small",
+ data = [
+ ":{name}_{type}.mlir".format(name = name, type = type),
+ "//tensorflow/compiler/mlir/tools/kernel_gen:tf_to_kernel",
+ ],
+ )
+
native.cc_library(
name = name + "_kernels",
compatible_with = get_compatible_with_cloud(),
diff --git a/tensorflow/core/kernels/mlir_generated/build_test.sh b/tensorflow/core/kernels/mlir_generated/build_test.sh
new file mode 100755
index 0000000..a0748a9
--- /dev/null
+++ b/tensorflow/core/kernels/mlir_generated/build_test.sh
@@ -0,0 +1,32 @@
+#!/bin/bash
+#
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+# http://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.
+# ==============================================================================
+
+# A script to catch kernel_gen failures as a test rather than build failure.
+set -e
+
+# KernelGen binary
+TF_TO_KERNEL="$1"
+OUTPUT_FILE="${TEST_TMPDIR}/output.mlir"
+INPUT="$2"
+
+# Do something
+${TF_TO_KERNEL} --input=${INPUT} --output=${OUTPUT_FILE} --unroll_factors=4 --tile_sizes=256 --arch=sm_70,compute_75 || die "Failed to generate kernel"
+
+# Check something
+[ -s ${OUTPUT_FILE} ] || die "output file was empty"
+
+echo "PASS"
