From b50e11b52c73705daa5b24b2673f6fba79e8084d Mon Sep 17 00:00:00 2001 From: Scott Wolchok Date: Thu, 23 Jan 2025 14:13:01 -0800 Subject: [PATCH] Update [ghstack-poisoned] --- kernels/portable/cpu/op_var.cpp | 10 ++--- kernels/test/op_var_test.cpp | 65 ++++++++++++++++++++++----------- 2 files changed, 49 insertions(+), 26 deletions(-) diff --git a/kernels/portable/cpu/op_var.cpp b/kernels/portable/cpu/op_var.cpp index e5499fa064e..c3627281481 100644 --- a/kernels/portable/cpu/op_var.cpp +++ b/kernels/portable/cpu/op_var.cpp @@ -38,7 +38,7 @@ void compute_variance( in, dim_list, out_ix); - CTYPE_OUT mean = sum / num; + CTYPE_OUT mean = sum / static_cast(num); CTYPE_OUT sum2 = map_reduce_over_dim_list( [mean](CTYPE_IN v) { return ( @@ -90,8 +90,8 @@ Tensor& var_out( constexpr auto name = "var.out"; - ET_SWITCH_FLOAT_TYPES(in.scalar_type(), ctx, name, CTYPE_IN, [&] { - ET_SWITCH_FLOAT_TYPES(out.scalar_type(), ctx, name, CTYPE_OUT, [&] { + ET_SWITCH_FLOATHBF16_TYPES(in.scalar_type(), ctx, name, CTYPE_IN, [&] { + ET_SWITCH_FLOATHBF16_TYPES(out.scalar_type(), ctx, name, CTYPE_OUT, [&] { compute_variance(in, out, dim_list, num, denom); }); }); @@ -135,8 +135,8 @@ Tensor& var_correction_out( const size_t num = get_reduced_dim_product(in, dim_list); const double denom = num - correction_val; - ET_SWITCH_FLOAT_TYPES(in.scalar_type(), ctx, name, CTYPE_IN, [&] { - ET_SWITCH_FLOAT_TYPES(out.scalar_type(), ctx, name, CTYPE_OUT, [&] { + ET_SWITCH_FLOATHBF16_TYPES(in.scalar_type(), ctx, name, CTYPE_IN, [&] { + ET_SWITCH_FLOATHBF16_TYPES(out.scalar_type(), ctx, name, CTYPE_OUT, [&] { compute_variance(in, out, dim_list, num, denom); }); }); diff --git a/kernels/test/op_var_test.cpp b/kernels/test/op_var_test.cpp index 5501635c229..749fd62d1ce 100644 --- a/kernels/test/op_var_test.cpp +++ b/kernels/test/op_var_test.cpp @@ -25,6 +25,19 @@ using exec_aten::ScalarType; using exec_aten::Tensor; using torch::executor::testing::TensorFactory; +namespace { +void expect_tensor_close_with_increased_tol( + const Tensor& actual, + const Tensor& expected) { + if (actual.scalar_type() == ScalarType::BFloat16 || + actual.scalar_type() == ScalarType::Half) { + EXPECT_TENSOR_CLOSE_WITH_TOL(expected, actual, 1e-2, 1e-2); + } else { + EXPECT_TENSOR_CLOSE(expected, actual); + } +} +} // namespace + class OpVarOutTest : public OperatorTest { protected: Tensor& op_var_out( @@ -142,7 +155,7 @@ class OpVarOutTest : public OperatorTest { op_var_out( self, optional_dim_list, /*unbiased=*/true, /*keepdim=*/true, out); // clang-format off - EXPECT_TENSOR_CLOSE(out, tf_out.make( + expect_tensor_close_with_increased_tol(out, tf_out.make( {2, 3, 1}, { 1.666667, @@ -160,7 +173,7 @@ class OpVarOutTest : public OperatorTest { op_var_out( self, optional_dim_list, /*unbiased=*/true, /*keepdim=*/false, out); // clang-format off - EXPECT_TENSOR_CLOSE(out, tf_out.make( + expect_tensor_close_with_increased_tol(out, tf_out.make( {2, 3}, { 1.666667, 1.666667, 1.666667, @@ -174,12 +187,14 @@ class OpVarOutTest : public OperatorTest { optional_dim_list = ArrayRef{dims_2, 2}; op_var_out( self, optional_dim_list, /*unbiased=*/true, /*keepdim=*/true, out); - EXPECT_TENSOR_CLOSE(out, tf_out.make({1, 1, 4}, {56.0, 56.0, 56.0, 56.0})); + expect_tensor_close_with_increased_tol( + out, tf_out.make({1, 1, 4}, {56.0, 56.0, 56.0, 56.0})); out = tf_out.zeros({4}); op_var_out( self, optional_dim_list, /*unbiased=*/true, /*keepdim=*/false, out); - EXPECT_TENSOR_CLOSE(out, tf_out.make({4}, {56.0, 56.0, 56.0, 56.0})); + expect_tensor_close_with_increased_tol( + out, tf_out.make({4}, {56.0, 56.0, 56.0, 56.0})); // dim list with negative dimensions should work out = tf_out.zeros({2, 1, 4}); @@ -188,7 +203,7 @@ class OpVarOutTest : public OperatorTest { op_var_out( self, optional_dim_list, /*unbiased=*/false, /*keepdim=*/true, out); // clang-format off - EXPECT_TENSOR_CLOSE(out, tf_out.make( + expect_tensor_close_with_increased_tol(out, tf_out.make( {2, 1, 4}, { 10.666667, 10.666667, 10.666667, 10.666667, @@ -201,18 +216,19 @@ class OpVarOutTest : public OperatorTest { out = tf_out.zeros({1, 1, 1}); optional> null_dim_list; op_var_out(self, null_dim_list, /*unbiased=*/true, /*keepdim=*/true, out); - EXPECT_TENSOR_CLOSE(out, tf_out.make({1, 1, 1}, {50.0})); + expect_tensor_close_with_increased_tol(out, tf_out.make({1, 1, 1}, {50.0})); optional> empty_dim_list{ArrayRef{}}; op_var_out(self, empty_dim_list, /*unbiased=*/false, /*keepdim=*/true, out); - EXPECT_TENSOR_CLOSE(out, tf_out.make({1, 1, 1}, {47.916668})); + expect_tensor_close_with_increased_tol( + out, tf_out.make({1, 1, 1}, {47.916668})); out = tf_out.zeros({}); op_var_out(self, null_dim_list, /*unbiased=*/false, /*keepdim=*/false, out); - EXPECT_TENSOR_CLOSE(out, tf_out.make({}, {47.916668})); + expect_tensor_close_with_increased_tol(out, tf_out.make({}, {47.916668})); op_var_out(self, empty_dim_list, /*unbiased=*/true, /*keepdim=*/false, out); - EXPECT_TENSOR_CLOSE(out, tf_out.make({}, {50.0})); + expect_tensor_close_with_increased_tol(out, tf_out.make({}, {50.0})); } }; @@ -227,6 +243,20 @@ class OpVarCorrectionOutTest : public OperatorTest { return torch::executor::aten::var_outf( context_, self, dim, correction, keepdim, out); } + + template + void test_dtype() { + TensorFactory tf; + + Tensor x = tf.make({2, 3}, {4.9, 4.0, 5.6, 3.8, 4.9, 5.6}); + Tensor expected = tf.make({2}, {0.72693, 0.93032}); + optional correction(1.23); + Tensor out = tf.zeros({2}); + + op_var_correction_out( + x, ArrayRef{1}, correction, /*keepdim=*/false, out); + expect_tensor_close_with_increased_tol(out, expected); + } }; TEST_F(OpVarOutTest, InvalidDimensionListDies) { @@ -303,9 +333,9 @@ TEST_F(OpVarOutTest, AllFloatInputFloatOutputPasses) { test_var_out_dtype(); #define TEST_ENTRY(INPUT_CTYPE, INPUT_DTYPE) \ - ET_FORALL_FLOAT_TYPES_WITH2(INPUT_CTYPE, INPUT_DTYPE, TEST_KERNEL); + ET_FORALL_FLOATHBF16_TYPES_WITH2(INPUT_CTYPE, INPUT_DTYPE, TEST_KERNEL); - ET_FORALL_FLOAT_TYPES(TEST_ENTRY); + ET_FORALL_FLOATHBF16_TYPES(TEST_ENTRY); #undef TEST_ENTRY #undef TEST_KERNEL } @@ -387,14 +417,7 @@ TEST_F(OpVarOutTest, DynamicShapeUnbound) { } TEST_F(OpVarCorrectionOutTest, SmokeTest) { - TensorFactory tf; - - Tensor x = tf.make({2, 3}, {4.9, 4.0, 5.6, 3.8, 4.9, 5.6}); - Tensor expected = tf.make({2}, {0.72693, 0.93032}); - optional correction(1.23); - Tensor out = tf.zeros({2}); - - op_var_correction_out( - x, ArrayRef{1}, correction, /*keepdim=*/false, out); - EXPECT_TENSOR_CLOSE(out, expected); +#define TEST_ENTRY(ctype, dtype) test_dtype(); + ET_FORALL_FLOATHBF16_TYPES(TEST_ENTRY); +#undef TEST_ENTRY }