[MLIR][KernelGen] Explicitly specify special inputs in all binary tests.
PiperOrigin-RevId: 350725301
Change-Id: I516dde4a6496af7a7419c3f877e63ce4a5c8f086
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 29a9fec..f5bc895 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc
@@ -489,13 +489,19 @@
return static_cast<Eigen::half>(std::floor(static_cast<float>(lhs / rhs)));
}
-GENERATE_DEFAULT_TESTS(FloorDiv,
- /*test_name=*/Half, Eigen::half, Eigen::half,
- baseline_floor_div)
-GENERATE_DEFAULT_TESTS(FloorDiv,
- /*test_name=*/Float, float, float, baseline_floor_div)
-GENERATE_DEFAULT_TESTS(FloorDiv,
- /*test_name=*/Double, double, double, baseline_floor_div)
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ FloorDiv,
+ /*test_name=*/Half, Eigen::half, Eigen::half,
+ test::DefaultInput<Eigen::half>("Div"),
+ test::DefaultInputNonZero<Eigen::half>(), baseline_floor_div);
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ FloorDiv,
+ /*test_name=*/Float, float, float, test::DefaultInput<float>("Div"),
+ test::DefaultInputNonZero<float>(), baseline_floor_div);
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ FloorDiv,
+ /*test_name=*/Double, double, double, test::DefaultInput<double>("Div"),
+ test::DefaultInputNonZero<double>(), baseline_floor_div);
/// Test `tf.Greater`.
@@ -544,14 +550,22 @@
return lhs << rhs;
}
-GENERATE_DEFAULT_TESTS(LeftShift, /*test_name=*/Int8, int8, int8,
- baseline_left_shift)
-GENERATE_DEFAULT_TESTS(LeftShift, /*test_name=*/Int16, int16, int16,
- baseline_left_shift)
-GENERATE_DEFAULT_TESTS(LeftShift, /*test_name=*/Int32, int32, int32,
- baseline_left_shift)
-GENERATE_DEFAULT_TESTS(LeftShift, /*test_name=*/Int64, int64, int64,
- baseline_left_shift)
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ LeftShift, /*test_name=*/Int8, int8, int8,
+ test::DefaultInput<int8>("LeftShift"),
+ test::DefaultInputLessThanBitwidth<int8>(), baseline_left_shift)
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ LeftShift, /*test_name=*/Int16, int16, int16,
+ test::DefaultInput<int16>("LeftShift"),
+ test::DefaultInputLessThanBitwidth<int16>(), baseline_left_shift)
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ LeftShift, /*test_name=*/Int32, int32, int32,
+ test::DefaultInput<int32>("LeftShift"),
+ test::DefaultInputLessThanBitwidth<int32>(), baseline_left_shift)
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ LeftShift, /*test_name=*/Int64, int64, int64,
+ test::DefaultInput<int64>("LeftShift"),
+ test::DefaultInputLessThanBitwidth<int64>(), baseline_left_shift)
/// Test `tf.Less`.
@@ -656,14 +670,22 @@
return lhs >> rhs;
}
-GENERATE_DEFAULT_TESTS(RightShift,
- /*test_name=*/Int8, int8, int8, baseline_right_shift)
-GENERATE_DEFAULT_TESTS(RightShift,
- /*test_name=*/Int16, int16, int16, baseline_right_shift)
-GENERATE_DEFAULT_TESTS(RightShift,
- /*test_name=*/Int32, int32, int32, baseline_right_shift)
-GENERATE_DEFAULT_TESTS(RightShift,
- /*test_name=*/Int64, int64, int64, baseline_right_shift)
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ RightShift,
+ /*test_name=*/Int8, int8, int8, test::DefaultInput<int8>("RightShift"),
+ test::DefaultInputLessThanBitwidth<int8>(), baseline_right_shift)
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ RightShift,
+ /*test_name=*/Int16, int16, int16, test::DefaultInput<int16>("RightShift"),
+ test::DefaultInputLessThanBitwidth<int16>(), baseline_right_shift)
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ RightShift,
+ /*test_name=*/Int32, int32, int32, test::DefaultInput<int32>("RightShift"),
+ test::DefaultInputLessThanBitwidth<int32>(), baseline_right_shift)
+GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
+ RightShift,
+ /*test_name=*/Int64, int64, int64, test::DefaultInput<int64>("RightShift"),
+ test::DefaultInputLessThanBitwidth<int64>(), baseline_right_shift)
/// Test `tf.Sub`.
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h b/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h
index 62e5b81..b33041f 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h
+++ b/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h
@@ -139,14 +139,17 @@
template <typename T,
std::enable_if_t<llvm::is_one_of<T, int8, int16, int32, int64>::value,
bool> = true>
+absl::InlinedVector<T, 10> DefaultInputLessThanBitwidth() {
+ auto max_shift = sizeof(T) * 8 - 1;
+ absl::InlinedVector<T, 10> v(max_shift);
+ for (auto i = 0; i < max_shift; ++i) v.push_back(i);
+ return v;
+}
+
+template <typename T,
+ std::enable_if_t<llvm::is_one_of<T, int8, int16, int32, int64>::value,
+ bool> = true>
absl::InlinedVector<T, 10> DefaultInput(absl::string_view op_name) {
- // Only generate values less than the bitwidth of the data type.
- if (op_name == "LeftShift" || op_name == "RightShift") {
- auto max_shift = sizeof(T) * 8 - 1;
- absl::InlinedVector<T, 10> v(max_shift);
- for (auto i = 0; i < max_shift; ++i) v.push_back(i);
- return v;
- }
return InputAsVector<T, int>({-18, -9, -1, 0, 0, 1, 1, 2, 3, 5, 7, 9, 9, 18});
}
@@ -154,10 +157,6 @@
llvm::is_one_of<T, Eigen::half, float, double>::value,
bool> = true>
absl::InlinedVector<T, 10> DefaultInput(absl::string_view op_name) {
- if (op_name == "FloorDiv") {
- return InputAsVector<T, double>({-18.0, -9.0, -1e-6, -0.1, 0.1, 1e-6, 0.1,
- 0.2, 0.3, 0.5, 0.7, 0.9, 9.0, 18.0});
- }
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});
}