From be42013b9c8a4fbf374a0eb16f400993f8195c83 Mon Sep 17 00:00:00 2001 From: Scott Wolchok Date: Fri, 17 Jan 2025 16:04:15 -0800 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- kernels/portable/cpu/op_topk.cpp | 2 +- kernels/test/CMakeLists.txt | 1 + kernels/test/op_topk_test.cpp | 53 ++++++++++++++++++-------------- 3 files changed, 32 insertions(+), 24 deletions(-) diff --git a/kernels/portable/cpu/op_topk.cpp b/kernels/portable/cpu/op_topk.cpp index caaf7e033db..987e974bbf5 100644 --- a/kernels/portable/cpu/op_topk.cpp +++ b/kernels/portable/cpu/op_topk.cpp @@ -186,7 +186,7 @@ std::tuple topk_values( bool temp_mem_allocated = false; - ET_SWITCH_REALH_TYPES(in.scalar_type(), ctx, name, CTYPE, [&]() { + ET_SWITCH_REALHBF16_TYPES(in.scalar_type(), ctx, name, CTYPE, [&]() { using elem_t = std::pair; size_t temp_mem_size = nonempty_size(in, dim) * sizeof(elem_t); diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt index f1f51e73fbb..b22bf772db3 100644 --- a/kernels/test/CMakeLists.txt +++ b/kernels/test/CMakeLists.txt @@ -220,6 +220,7 @@ set(all_test_sources "op_tan_test.cpp" "op_tanh_test.cpp" "op_to_copy_test.cpp" + "op_topk_test.cpp" "op_transpose_copy_test.cpp" "op_tril_test.cpp" "op_trunc_test.cpp" diff --git a/kernels/test/op_topk_test.cpp b/kernels/test/op_topk_test.cpp index 9e77e9ad5a8..8baef548a5b 100644 --- a/kernels/test/op_topk_test.cpp +++ b/kernels/test/op_topk_test.cpp @@ -118,32 +118,39 @@ class OpTopkValuesTest : public ::testing::Test { // first. torch::executor::runtime_init(); } + + template + void run_smoke_test() { + TensorFactory tfDtype; + TensorFactory tfLong; + + Tensor input = + tfDtype.make({3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + int64_t k = 2; + int64_t dim = 0; + bool largest = true; + bool sorted = true; + Tensor values = tfDtype.zeros({2, 2, 2}); + Tensor indices = tfLong.zeros({2, 2, 2}); + Tensor values_expected = tfDtype.make({2, 2, 2}, {9, 10, 11, 12, 5, 6, 7, 8}); + Tensor indices_expected = tfLong.make({2, 2, 2}, {2, 2, 2, 2, 1, 1, 1, 1}); + op_topk_values(input, k, dim, largest, sorted, values, indices); + EXPECT_TENSOR_CLOSE(values, values_expected); + EXPECT_TENSOR_EQ(indices, indices_expected); + + largest = false; + values_expected = tfDtype.make({2, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8}); + indices_expected = tfLong.make({2, 2, 2}, {0, 0, 0, 0, 1, 1, 1, 1}); + op_topk_values(input, k, dim, largest, sorted, values, indices); + EXPECT_TENSOR_CLOSE(values, values_expected); + EXPECT_TENSOR_EQ(indices, indices_expected); + + } }; TEST_F(OpTopkValuesTest, SmokeTest) { - TensorFactory tfFloat; - TensorFactory tfLong; - - Tensor input = - tfFloat.make({3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); - int64_t k = 2; - int64_t dim = 0; - bool largest = true; - bool sorted = true; - Tensor values = tfFloat.zeros({2, 2, 2}); - Tensor indices = tfLong.zeros({2, 2, 2}); - Tensor values_expected = tfFloat.make({2, 2, 2}, {9, 10, 11, 12, 5, 6, 7, 8}); - Tensor indices_expected = tfLong.make({2, 2, 2}, {2, 2, 2, 2, 1, 1, 1, 1}); - op_topk_values(input, k, dim, largest, sorted, values, indices); - EXPECT_TENSOR_CLOSE(values, values_expected); - EXPECT_TENSOR_EQ(indices, indices_expected); - - largest = false; - values_expected = tfFloat.make({2, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8}); - indices_expected = tfLong.make({2, 2, 2}, {0, 0, 0, 0, 1, 1, 1, 1}); - op_topk_values(input, k, dim, largest, sorted, values, indices); - EXPECT_TENSOR_CLOSE(values, values_expected); - EXPECT_TENSOR_EQ(indices, indices_expected); +#define RUN_SMOKE_TEST(ctype, dtype) run_smoke_test(); + ET_FORALL_REALHBF16_TYPES(RUN_SMOKE_TEST); } TEST_F(OpTopkValuesTest, NonPartialSort) { From 7686c7e42d21054184a37851d84c5033ec1ff784 Mon Sep 17 00:00:00 2001 From: Scott Wolchok Date: Wed, 22 Jan 2025 09:35:29 -0800 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- kernels/test/op_topk_test.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/kernels/test/op_topk_test.cpp b/kernels/test/op_topk_test.cpp index 8baef548a5b..54bdde9edaa 100644 --- a/kernels/test/op_topk_test.cpp +++ b/kernels/test/op_topk_test.cpp @@ -125,14 +125,15 @@ class OpTopkValuesTest : public ::testing::Test { TensorFactory tfLong; Tensor input = - tfDtype.make({3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + tfDtype.make({3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); int64_t k = 2; int64_t dim = 0; bool largest = true; bool sorted = true; Tensor values = tfDtype.zeros({2, 2, 2}); Tensor indices = tfLong.zeros({2, 2, 2}); - Tensor values_expected = tfDtype.make({2, 2, 2}, {9, 10, 11, 12, 5, 6, 7, 8}); + Tensor values_expected = + tfDtype.make({2, 2, 2}, {9, 10, 11, 12, 5, 6, 7, 8}); Tensor indices_expected = tfLong.make({2, 2, 2}, {2, 2, 2, 2, 1, 1, 1, 1}); op_topk_values(input, k, dim, largest, sorted, values, indices); EXPECT_TENSOR_CLOSE(values, values_expected); @@ -144,7 +145,6 @@ class OpTopkValuesTest : public ::testing::Test { op_topk_values(input, k, dim, largest, sorted, values, indices); EXPECT_TENSOR_CLOSE(values, values_expected); EXPECT_TENSOR_EQ(indices, indices_expected); - } };