From 687056c351d2760665ef8b9b52eb8bb3491bd2c9 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 17 Dec 2020 08:06:02 -0800 Subject: [PATCH] [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 --- tensorflow/core/kernels/mlir_generated/BUILD | 18 ++ .../mlir_generated/gpu_binary_ops_test.cc | 185 ++++++------------ .../mlir_generated/gpu_ops_test_util.cc | 24 +++ .../mlir_generated/gpu_ops_test_util.h | 104 ++++++++++ 4 files changed, 211 insertions(+), 120 deletions(-) create mode 100644 tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.cc create mode 100644 tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD index c07c62ac7de..5d6c922f72d 100644 --- a/tensorflow/core/kernels/mlir_generated/BUILD +++ b/tensorflow/core/kernels/mlir_generated/BUILD @@ -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", diff --git a/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc b/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc index c97fbc159ca..06bd9fa1348 100644 --- a/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc +++ b/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc @@ -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(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 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 expected_output; @@ -169,7 +172,7 @@ class GpuBinaryOpTest : public OpsTestBase { expected_output.push_back(result); } - auto scalar_input_vector = InputAsVector({scalar_input}); + auto scalar_input_vector = test::InputAsVector({scalar_input}); RunAndExpectResult(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 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 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 - absl::InlinedVector RepeatInputToMatchShape( - absl::InlinedVector input, int size) { - absl::InlinedVector result; - for (int i = 0; i < size; i++) { - auto value = input[i % input.size()]; - result.push_back(value); - } - return result; - } - template absl::InlinedVector 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 - T DefaultScalarInput() { - return static_cast(2.0); - } - - template - absl::InlinedVector InfZeroInput() { - return InputAsVector({-std::numeric_limits::infinity(), - -0.1, -0.0, 0.0, 0.1, - std::numeric_limits::infinity()}); - } - - template < - typename T, - std::enable_if_t::value, - bool> = true> - absl::InlinedVector DefaultInput() { - return InputAsVector( - {-18, -9, -1, 0, 0, 1, 1, 2, 3, 5, 7, 9, 9, 18}); - } - - template < - typename T, - std::enable_if_t::value, - bool> = true> - absl::InlinedVector DefaultInput() { - return InputAsVector({-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 ::value, bool> = true> - absl::InlinedVector DefaultInput() { - return InputAsVector({true, false, true, true, false}); - } - - template - absl::InlinedVector InputAsVector( - std::initializer_list input) { - absl::InlinedVector result; - result.reserve(input.size()); - for (const LiteralT& value : input) { - result.push_back(static_cast(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( \ - #op_name, /*shape=*/DefaultInputShape(), \ - /*lhs_input=*/DefaultInput(), /*rhs_input=*/DefaultInput(), \ - baseline_callback, use_constraint); \ - } \ - \ - TEST_F(GpuBinaryOpTest, op_name##OneScalar##test_name) { \ - TestOneScalar( \ - #op_name, /*scalar_input=*/DefaultScalarInput(), \ - /*other_shape=*/DefaultInputShape(), \ - /*other_input=*/DefaultInput(), baseline_callback, use_constraint); \ - } \ - \ - TEST_F(GpuBinaryOpTest, op_name##IncompatibleShapes##test_name) { \ - TestIncompatibleShapes(#op_name, /*lhs_input=*/DefaultInput(), \ - /*rhs_input=*/DefaultInput(), \ - use_constraint); \ - } \ - \ - TEST_F(GpuBinaryOpTest, op_name##BroadcastingExpand##test_name) { \ - TestBroadcastingExpand( \ - #op_name, /*lhs_input=*/DefaultInput(), \ - /*rhs_input=*/DefaultInput(), baseline_callback, use_constraint); \ - } \ - \ - TEST_F(GpuBinaryOpTest, op_name##BroadcastingInDim##test_name) { \ - TestBroadcastingInDim( \ - #op_name, /*lhs_input=*/DefaultInput(), \ - /*rhs_input=*/DefaultInput(), baseline_callback, use_constraint); \ - } \ - \ - TEST_F(GpuBinaryOpTest, op_name##Broadcasting##test_name) { \ - TestBroadcasting( \ - #op_name, /*lhs_input=*/DefaultInput(), \ - /*rhs_input=*/DefaultInput(), baseline_callback, use_constraint); \ - } \ - \ - TEST_F(GpuBinaryOpTest, op_name##EmptyShapeBroadcasting##test_name) { \ - TestEmptyShapeBroadcasting( \ - #op_name, /*lhs_input=*/DefaultInput(), \ - /*rhs_input=*/DefaultInput(), 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( \ + #op_name, /*shape=*/test::DefaultInputShape(), \ + /*lhs_input=*/test::DefaultInput(), \ + /*rhs_input=*/test::DefaultInput(), baseline_callback, \ + use_constraint); \ + } \ + \ + TEST_F(GpuBinaryOpTest, op_name##OneScalar##test_name) { \ + TestOneScalar( \ + #op_name, /*scalar_input=*/test::DefaultScalarInput(), \ + /*other_shape=*/test::DefaultInputShape(), \ + /*other_input=*/test::DefaultInput(), baseline_callback, \ + use_constraint); \ + } \ + \ + TEST_F(GpuBinaryOpTest, op_name##IncompatibleShapes##test_name) { \ + TestIncompatibleShapes( \ + #op_name, /*lhs_input=*/test::DefaultInput(), \ + /*rhs_input=*/test::DefaultInput(), use_constraint); \ + } \ + \ + TEST_F(GpuBinaryOpTest, op_name##BroadcastingExpand##test_name) { \ + TestBroadcastingExpand( \ + #op_name, /*lhs_input=*/test::DefaultInput(), \ + /*rhs_input=*/test::DefaultInput(), baseline_callback, \ + use_constraint); \ + } \ + \ + TEST_F(GpuBinaryOpTest, op_name##BroadcastingInDim##test_name) { \ + TestBroadcastingInDim( \ + #op_name, /*lhs_input=*/test::DefaultInput(), \ + /*rhs_input=*/test::DefaultInput(), baseline_callback, \ + use_constraint); \ + } \ + \ + TEST_F(GpuBinaryOpTest, op_name##Broadcasting##test_name) { \ + TestBroadcasting( \ + #op_name, /*lhs_input=*/test::DefaultInput(), \ + /*rhs_input=*/test::DefaultInput(), baseline_callback, \ + use_constraint); \ + } \ + \ + TEST_F(GpuBinaryOpTest, op_name##EmptyShapeBroadcasting##test_name) { \ + TestEmptyShapeBroadcasting( \ + #op_name, /*lhs_input=*/test::DefaultInput(), \ + /*rhs_input=*/test::DefaultInput(), use_constraint); \ } #define GENERATE_DEFAULT_TESTS(op_name, test_name, T, OutT, baseline_callback) \ diff --git a/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.cc b/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.cc new file mode 100644 index 00000000000..39b23c7ad33 --- /dev/null +++ b/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.cc @@ -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 diff --git a/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h b/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h new file mode 100644 index 00000000000..d40387b18ea --- /dev/null +++ b/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h @@ -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 +absl::InlinedVector InputAsVector( + std::initializer_list input) { + absl::InlinedVector result; + result.reserve(input.size()); + for (const LiteralT& value : input) { + result.push_back(static_cast(value)); + } + return result; +} + +template +absl::InlinedVector RepeatInputToMatchShape( + absl::InlinedVector input, int size) { + absl::InlinedVector 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 ::value, + bool> = true> +T DefaultScalarInput() { + return static_cast(3); +} + +template ::value, + bool> = true> +T DefaultScalarInput() { + return static_cast(2.0); +} + +template ::value, bool> = true> +T DefaultScalarInput() { + return static_cast(true); +} + +template +absl::InlinedVector InfZeroInput() { + return InputAsVector({-std::numeric_limits::infinity(), + -0.1, -0.0, 0.0, 0.1, + std::numeric_limits::infinity()}); +} + +template ::value, + bool> = true> +absl::InlinedVector DefaultInput() { + return InputAsVector({-18, -9, -1, 0, 0, 1, 1, 2, 3, 5, 7, 9, 9, 18}); +} + +template ::value, + bool> = true> +absl::InlinedVector DefaultInput() { + return InputAsVector({-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 ::value, bool> = true> +absl::InlinedVector DefaultInput() { + return InputAsVector({true, false, true, true, false}); +} +} // namespace test +} // namespace tensorflow + +#endif // TENSORFLOW_CORE_KERNELS_MLIR_GENERATED_GPU_OPS_TEST_UTIL_H_