[MLIR][KernelGen] Extract test util that can be shared among kernel tests

This concerns helper functions to generate default input values and helper
functions to derive inputs of the right type and size.

PiperOrigin-RevId: 348022802
Change-Id: If84c06bde49de780206beae9c832cfab746b3999
This commit is contained in:
A. Unique TensorFlower 2020-12-17 08:06:02 -08:00 committed by TensorFlower Gardener
parent 6ce6b611f9
commit 687056c351
4 changed files with 211 additions and 120 deletions

View File

@ -219,6 +219,7 @@ tf_cuda_cc_test(
"no_cuda_asan", # b/173033461
],
deps = [
":gpu_ops_test_util",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:tensorflow",
@ -237,6 +238,23 @@ tf_cuda_cc_test(
],
)
cc_library(
name = "gpu_ops_test_util",
testonly = 1,
srcs = [
"gpu_ops_test_util.cc",
"gpu_ops_test_util.h",
],
hdrs = [
"gpu_ops_test_util.h",
],
deps = [
"//tensorflow/core:tensorflow",
"@com_google_absl//absl/container:inlined_vector",
"@llvm-project//llvm:Support",
],
)
# TODO(b/160731748): Re-enable when it works again.
# gen_kernel_library(
# name = "bias_add",

View File

@ -26,6 +26,7 @@ limitations under the License.
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
@ -108,9 +109,9 @@ class GpuBinaryOpTest : public OpsTestBase {
TensorShape lhs_shape{3};
TensorShape rhs_shape{2};
auto repeated_lhs_input =
RepeatInputToMatchShape(lhs_input, lhs_shape.num_elements());
test::RepeatInputToMatchShape(lhs_input, lhs_shape.num_elements());
auto repeated_rhs_input =
RepeatInputToMatchShape(rhs_input, rhs_shape.num_elements());
test::RepeatInputToMatchShape(rhs_input, rhs_shape.num_elements());
RunAndExpectInvalidArgument<T, OutT>(op_name, lhs_shape, repeated_lhs_input,
rhs_shape, repeated_rhs_input,
@ -126,8 +127,10 @@ class GpuBinaryOpTest : public OpsTestBase {
bool use_constraint = true) {
// Prepare inputs.
int input_size = shape.num_elements();
auto repeated_lhs_input = RepeatInputToMatchShape(lhs_input, input_size);
auto repeated_rhs_input = RepeatInputToMatchShape(rhs_input, input_size);
auto repeated_lhs_input =
test::RepeatInputToMatchShape(lhs_input, input_size);
auto repeated_rhs_input =
test::RepeatInputToMatchShape(rhs_input, input_size);
// Compute expected results.
absl::InlinedVector<OutT, 10> expected_output;
@ -156,7 +159,7 @@ class GpuBinaryOpTest : public OpsTestBase {
// Prepare inputs.
TensorShape scalar_shape{};
auto repeated_other_input =
RepeatInputToMatchShape(other_input, other_shape.num_elements());
test::RepeatInputToMatchShape(other_input, other_shape.num_elements());
// Compute expected results.
absl::InlinedVector<OutT, 10> expected_output;
@ -169,7 +172,7 @@ class GpuBinaryOpTest : public OpsTestBase {
expected_output.push_back(result);
}
auto scalar_input_vector = InputAsVector<T>({scalar_input});
auto scalar_input_vector = test::InputAsVector<T>({scalar_input});
RunAndExpectResult<T, OutT>(op_name, scalar_shape, scalar_input_vector,
other_shape, repeated_other_input,
/*expected_shape=*/other_shape, expected_output,
@ -188,9 +191,9 @@ class GpuBinaryOpTest : public OpsTestBase {
TensorShape lhs_shape{1};
TensorShape rhs_shape{6};
auto repeated_lhs_input =
RepeatInputToMatchShape(lhs_input, lhs_shape.num_elements());
test::RepeatInputToMatchShape(lhs_input, lhs_shape.num_elements());
auto repeated_rhs_input =
RepeatInputToMatchShape(rhs_input, rhs_shape.num_elements());
test::RepeatInputToMatchShape(rhs_input, rhs_shape.num_elements());
// Compute expected results.
std::vector<int> lhs_indices = {0, 0, 0, 0, 0, 0};
@ -217,9 +220,9 @@ class GpuBinaryOpTest : public OpsTestBase {
TensorShape lhs_shape{3};
TensorShape rhs_shape{2, 3};
auto repeated_lhs_input =
RepeatInputToMatchShape(lhs_input, lhs_shape.num_elements());
test::RepeatInputToMatchShape(lhs_input, lhs_shape.num_elements());
auto repeated_rhs_input =
RepeatInputToMatchShape(rhs_input, rhs_shape.num_elements());
test::RepeatInputToMatchShape(rhs_input, rhs_shape.num_elements());
// Compute expected results.
std::vector<int> lhs_indices = {0, 1, 2, 0, 1, 2};
@ -245,9 +248,9 @@ class GpuBinaryOpTest : public OpsTestBase {
TensorShape lhs_shape{2, 1};
TensorShape rhs_shape{3};
auto repeated_lhs_input =
RepeatInputToMatchShape(lhs_input, lhs_shape.num_elements());
test::RepeatInputToMatchShape(lhs_input, lhs_shape.num_elements());
auto repeated_rhs_input =
RepeatInputToMatchShape(rhs_input, rhs_shape.num_elements());
test::RepeatInputToMatchShape(rhs_input, rhs_shape.num_elements());
// Compute expected results.
TensorShape expected_shape{2, 3};
@ -284,17 +287,6 @@ class GpuBinaryOpTest : public OpsTestBase {
}
private:
template <typename T>
absl::InlinedVector<T, 10> RepeatInputToMatchShape(
absl::InlinedVector<T, 10> input, int size) {
absl::InlinedVector<T, 10> result;
for (int i = 0; i < size; i++) {
auto value = input[i % input.size()];
result.push_back(value);
}
return result;
}
template <typename T, typename BaselineT, typename OutT,
typename BaselineOutT>
absl::InlinedVector<OutT, 10> ComputeExpectedOutput(
@ -310,108 +302,61 @@ class GpuBinaryOpTest : public OpsTestBase {
}
return expected_output;
}
// Helper functions to get default input values.
protected:
TensorShape DefaultInputShape() { return TensorShape{3, 4}; }
template <typename T>
T DefaultScalarInput() {
return static_cast<T>(2.0);
}
template <typename T>
absl::InlinedVector<T, 10> InfZeroInput() {
return InputAsVector<T, double>({-std::numeric_limits<double>::infinity(),
-0.1, -0.0, 0.0, 0.1,
std::numeric_limits<float>::infinity()});
}
template <
typename T,
std::enable_if_t<llvm::is_one_of<T, int8, int16, int32, int64>::value,
bool> = true>
absl::InlinedVector<T, 10> DefaultInput() {
return InputAsVector<T, int>(
{-18, -9, -1, 0, 0, 1, 1, 2, 3, 5, 7, 9, 9, 18});
}
template <
typename T,
std::enable_if_t<llvm::is_one_of<T, Eigen::half, float, double>::value,
bool> = true>
absl::InlinedVector<T, 10> DefaultInput() {
return InputAsVector<T, double>({-18.0, -9.0, -1e-6, -0.0, 0.0, 1e-6, 0.1,
0.2, 0.3, 0.5, 0.7, 0.9, 9.0, 18.0});
}
template <typename T,
std::enable_if_t<llvm::is_one_of<T, bool>::value, bool> = true>
absl::InlinedVector<T, 10> DefaultInput() {
return InputAsVector<T, bool>({true, false, true, true, false});
}
template <typename T, typename LiteralT>
absl::InlinedVector<T, 10> InputAsVector(
std::initializer_list<LiteralT> input) {
absl::InlinedVector<T, 10> result;
result.reserve(input.size());
for (const LiteralT& value : input) {
result.push_back(static_cast<T>(value));
}
return result;
}
};
// Macros to easily generate common test cases. For specific inputs, please
// define your own test fixtures.
#define GENERATE_DEFAULT_TESTS_2(op_name, test_name, T, BaselineT, OutT, \
BaselineOutT, baseline_callback, \
use_constraint) \
TEST_F(GpuBinaryOpTest, op_name##EqShapes##test_name) { \
TestEqualShapes<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*shape=*/DefaultInputShape(), \
/*lhs_input=*/DefaultInput<T>(), /*rhs_input=*/DefaultInput<T>(), \
baseline_callback, use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##OneScalar##test_name) { \
TestOneScalar<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*scalar_input=*/DefaultScalarInput<T>(), \
/*other_shape=*/DefaultInputShape(), \
/*other_input=*/DefaultInput<T>(), baseline_callback, use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##IncompatibleShapes##test_name) { \
TestIncompatibleShapes<T, OutT>(#op_name, /*lhs_input=*/DefaultInput<T>(), \
/*rhs_input=*/DefaultInput<T>(), \
use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##BroadcastingExpand##test_name) { \
TestBroadcastingExpand<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*lhs_input=*/DefaultInput<T>(), \
/*rhs_input=*/DefaultInput<T>(), baseline_callback, use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##BroadcastingInDim##test_name) { \
TestBroadcastingInDim<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*lhs_input=*/DefaultInput<T>(), \
/*rhs_input=*/DefaultInput<T>(), baseline_callback, use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##Broadcasting##test_name) { \
TestBroadcasting<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*lhs_input=*/DefaultInput<T>(), \
/*rhs_input=*/DefaultInput<T>(), baseline_callback, use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##EmptyShapeBroadcasting##test_name) { \
TestEmptyShapeBroadcasting<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*lhs_input=*/DefaultInput<T>(), \
/*rhs_input=*/DefaultInput<T>(), use_constraint); \
#define GENERATE_DEFAULT_TESTS_2(op_name, test_name, T, BaselineT, OutT, \
BaselineOutT, baseline_callback, \
use_constraint) \
TEST_F(GpuBinaryOpTest, op_name##EqShapes##test_name) { \
TestEqualShapes<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*shape=*/test::DefaultInputShape(), \
/*lhs_input=*/test::DefaultInput<T>(), \
/*rhs_input=*/test::DefaultInput<T>(), baseline_callback, \
use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##OneScalar##test_name) { \
TestOneScalar<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*scalar_input=*/test::DefaultScalarInput<T>(), \
/*other_shape=*/test::DefaultInputShape(), \
/*other_input=*/test::DefaultInput<T>(), baseline_callback, \
use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##IncompatibleShapes##test_name) { \
TestIncompatibleShapes<T, OutT>( \
#op_name, /*lhs_input=*/test::DefaultInput<T>(), \
/*rhs_input=*/test::DefaultInput<T>(), use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##BroadcastingExpand##test_name) { \
TestBroadcastingExpand<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*lhs_input=*/test::DefaultInput<T>(), \
/*rhs_input=*/test::DefaultInput<T>(), baseline_callback, \
use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##BroadcastingInDim##test_name) { \
TestBroadcastingInDim<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*lhs_input=*/test::DefaultInput<T>(), \
/*rhs_input=*/test::DefaultInput<T>(), baseline_callback, \
use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##Broadcasting##test_name) { \
TestBroadcasting<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*lhs_input=*/test::DefaultInput<T>(), \
/*rhs_input=*/test::DefaultInput<T>(), baseline_callback, \
use_constraint); \
} \
\
TEST_F(GpuBinaryOpTest, op_name##EmptyShapeBroadcasting##test_name) { \
TestEmptyShapeBroadcasting<T, BaselineT, OutT, BaselineOutT>( \
#op_name, /*lhs_input=*/test::DefaultInput<T>(), \
/*rhs_input=*/test::DefaultInput<T>(), use_constraint); \
}
#define GENERATE_DEFAULT_TESTS(op_name, test_name, T, OutT, baseline_callback) \

View File

@ -0,0 +1,24 @@
/* 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.
==============================================================================*/
#include "tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h"
namespace tensorflow {
namespace test {
TensorShape DefaultInputShape() { return TensorShape{3, 4}; }
} // namespace test
} // namespace tensorflow

View File

@ -0,0 +1,104 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_MLIR_GENERATED_GPU_OPS_TEST_UTIL_H_
#define TENSORFLOW_CORE_KERNELS_MLIR_GENERATED_GPU_OPS_TEST_UTIL_H_
#include "absl/container/inlined_vector.h"
#include "llvm/ADT/STLExtras.h"
#include "tensorflow/core/framework/tensor_shape.h"
namespace tensorflow {
namespace test {
/// Helper functions to create or derive inputs of the right type and size.
template <typename T, typename LiteralT>
absl::InlinedVector<T, 10> InputAsVector(
std::initializer_list<LiteralT> input) {
absl::InlinedVector<T, 10> result;
result.reserve(input.size());
for (const LiteralT& value : input) {
result.push_back(static_cast<T>(value));
}
return result;
}
template <typename T>
absl::InlinedVector<T, 10> RepeatInputToMatchShape(
absl::InlinedVector<T, 10> input, int size) {
absl::InlinedVector<T, 10> result;
for (int i = 0; i < size; i++) {
auto value = input[i % input.size()];
result.push_back(value);
}
return result;
}
/// Helper functions to get default input values.
TensorShape DefaultInputShape();
template <typename T,
std::enable_if_t<llvm::is_one_of<T, int8, int16, int32, int64>::value,
bool> = true>
T DefaultScalarInput() {
return static_cast<T>(3);
}
template <typename T, std::enable_if_t<
llvm::is_one_of<T, Eigen::half, float, double>::value,
bool> = true>
T DefaultScalarInput() {
return static_cast<T>(2.0);
}
template <typename T,
std::enable_if_t<llvm::is_one_of<T, bool>::value, bool> = true>
T DefaultScalarInput() {
return static_cast<T>(true);
}
template <typename T>
absl::InlinedVector<T, 10> InfZeroInput() {
return InputAsVector<T, double>({-std::numeric_limits<double>::infinity(),
-0.1, -0.0, 0.0, 0.1,
std::numeric_limits<float>::infinity()});
}
template <typename T,
std::enable_if_t<llvm::is_one_of<T, int8, int16, int32, int64>::value,
bool> = true>
absl::InlinedVector<T, 10> DefaultInput() {
return InputAsVector<T, int>({-18, -9, -1, 0, 0, 1, 1, 2, 3, 5, 7, 9, 9, 18});
}
template <typename T, std::enable_if_t<
llvm::is_one_of<T, Eigen::half, float, double>::value,
bool> = true>
absl::InlinedVector<T, 10> DefaultInput() {
return InputAsVector<T, double>({-18.0, -9.0, -1e-6, -0.0, 0.0, 1e-6, 0.1,
0.2, 0.3, 0.5, 0.7, 0.9, 9.0, 18.0});
}
template <typename T,
std::enable_if_t<llvm::is_one_of<T, bool>::value, bool> = true>
absl::InlinedVector<T, 10> DefaultInput() {
return InputAsVector<T, bool>({true, false, true, true, false});
}
} // namespace test
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_MLIR_GENERATED_GPU_OPS_TEST_UTIL_H_