diff --git a/kernels/portable/cpu/util/delinearized_indexes_range.h b/kernels/portable/cpu/util/delinearized_indexes_range.h new file mode 100644 index 00000000000..c6a9cc91bfb --- /dev/null +++ b/kernels/portable/cpu/util/delinearized_indexes_range.h @@ -0,0 +1,108 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include + +#include +#include + +namespace torch::executor { + +namespace internal { +class DelinearizedIndexesIterator { + public: + using difference_type = ssize_t; + using value_type = std::array; + using reference = const value_type&; + using pointer = const value_type*; + using iterator_category = std::forward_iterator_tag; + + DelinearizedIndexesIterator() = default; + + explicit DelinearizedIndexesIterator(const Tensor& t) + : idx_(0), dim_(t.dim()), shape_(t.sizes()) { + } + + struct make_end_t { + explicit constexpr make_end_t() = default; + }; + + DelinearizedIndexesIterator(make_end_t, const Tensor& t) + : idx_(t.numel()) {} + + bool operator==(const DelinearizedIndexesIterator& rhs) const { + return idx_ == rhs.idx_; + } + + bool operator!=(const DelinearizedIndexesIterator& rhs) const { + return !operator==(rhs); + } + + reference operator*() const { + return repr_; + } + + pointer operator->() const { + return &repr_; + } + + DelinearizedIndexesIterator& operator++() { + idx_++; + for (auto ii = dim_ - 1; ii >= 0; --ii) { + repr_[ii]++; + ET_DCHECK(repr_[ii] <= shape_[ii]); + if ET_LIKELY (repr_[ii] < shape_[ii]) { + break; + } else { + repr_[ii] = 0; + } + } + return *this; + } + + DelinearizedIndexesIterator operator++(int) { + auto it = *this; + operator++(); + return it; + } + + difference_type operator-(const DelinearizedIndexesIterator& rhs) const { + return difference_type(idx_ - rhs.idx_); + } + + private: + std::size_t idx_ = 0; + value_type repr_ = {0,}; + ssize_t dim_; + ArrayRef shape_; +}; +} // namespace internal + +class DelinearizedIndexesRange { + public: + using iterator = internal::DelinearizedIndexesIterator; + + DelinearizedIndexesRange(const Tensor& t) : + tensor_(t) {} + + iterator begin() const { + return iterator(tensor_); + } + + iterator end() { + return iterator(iterator::make_end_t(), tensor_); + } + private: + const Tensor& tensor_; +}; +} // namespace torch::executor diff --git a/kernels/portable/cpu/util/targets.bzl b/kernels/portable/cpu/util/targets.bzl index eef765d5eec..95970b600ef 100644 --- a/kernels/portable/cpu/util/targets.bzl +++ b/kernels/portable/cpu/util/targets.bzl @@ -280,6 +280,19 @@ def define_common_targets(): visibility = ["//executorch/kernels/portable/cpu/..."], ) + runtime.cxx_library( + name = "delinearized_indexes_range", + exported_headers = ["delinearized_indexes_range.h"], + deps = [ + "//executorch/runtime/core/exec_aten:lib", + "//executorch/runtime/core/exec_aten/util:tensor_dimension_limit", + ], + visibility = [ + "//executorch/...", + "@EXECUTORCH_CLIENTS", + ], + ) + # Utility functions that can be used by operators that perform reduction for aten_mode in get_aten_mode_options(): suffix = "_aten" if aten_mode else "" diff --git a/kernels/portable/cpu/util/test/CMakeLists.txt b/kernels/portable/cpu/util/test/CMakeLists.txt index 5f81e4b6aec..76c53ea8448 100644 --- a/kernels/portable/cpu/util/test/CMakeLists.txt +++ b/kernels/portable/cpu/util/test/CMakeLists.txt @@ -19,7 +19,9 @@ set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..) include(${EXECUTORCH_ROOT}/build/Test.cmake) -set(_test_srcs broadcast_test.cpp reduce_test.cpp) +set(_test_srcs broadcast_test.cpp delinearized_indexes_range_test.cpp + reduce_test.cpp +) et_cxx_test( kernels_portable_cpu_util_test SOURCES ${_test_srcs} EXTRA_LIBS diff --git a/kernels/portable/cpu/util/test/delinearized_indexes_range_test.cpp b/kernels/portable/cpu/util/test/delinearized_indexes_range_test.cpp new file mode 100644 index 00000000000..395e1a74d98 --- /dev/null +++ b/kernels/portable/cpu/util/test/delinearized_indexes_range_test.cpp @@ -0,0 +1,69 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include + +#include + +using executorch::aten::ScalarType; +using executorch::aten::Tensor; +using executorch::runtime::testing::TensorFactory; +using torch::executor::DelinearizedIndexesRange; + +TEST(DelinearizedIndexesRangeTest, Empty) { + TensorFactory tf; + + Tensor a = tf.make({0}, {}); + ASSERT_EQ(a.numel(), 0); + bool loop_entered = false; + for (auto _ : DelinearizedIndexesRange(a)) { + loop_entered = true; + } + EXPECT_FALSE(loop_entered); +} + +TEST(DelinearizedIndexesRangeTest, OneD) { + TensorFactory tf; + + Tensor a = tf.zeros({5}); + DelinearizedIndexesRange r(a); + std::vector v(r.begin(), r.end()); + int idx = 0; + for (const auto& elem: v) { + EXPECT_EQ(elem[0], idx++); + } +} + +TEST(DelinearizedIndexesRangeTest, ThreeD) { + TensorFactory tf; + Tensor a = tf.zeros({3, 2, 3}); + DelinearizedIndexesRange r(a); + std::vector v(r.begin(), r.end()); + std::vector expected = { + {0, 0, 0}, + {0, 0, 1}, + {0, 0, 2}, + {0, 1, 0}, + {0, 1, 1}, + {0, 1, 2}, + {1, 0, 0}, + {1, 0, 1}, + {1, 0, 2}, + {1, 1, 0}, + {1, 1, 1}, + {1, 1, 2}, + {2, 0, 0}, + {2, 0, 1}, + {2, 0, 2}, + {2, 1, 0}, + {2, 1, 1}, + {2, 1, 2}, + }; + EXPECT_EQ(v, expected); +} diff --git a/kernels/portable/cpu/util/test/targets.bzl b/kernels/portable/cpu/util/test/targets.bzl index 28988b90dcc..25e16237361 100644 --- a/kernels/portable/cpu/util/test/targets.bzl +++ b/kernels/portable/cpu/util/test/targets.bzl @@ -12,6 +12,16 @@ def define_common_targets(): ], ) + runtime.cxx_test( + name = "delinearized_indexes_range_test", + srcs = ["delinearized_indexes_range_test.cpp"], + deps = [ + "//executorch/kernels/portable/cpu/util:delinearized_indexes_range", + "//executorch/runtime/core/exec_aten:lib", + "//executorch/runtime/core/exec_aten/testing_util:tensor_util", + ], + ) + runtime.cxx_test( name = "reduce_test", srcs = ["reduce_test.cpp"], diff --git a/test/utils/OSSTestConfig.json b/test/utils/OSSTestConfig.json index 70cb2d2e44f..b94f11ea94e 100644 --- a/test/utils/OSSTestConfig.json +++ b/test/utils/OSSTestConfig.json @@ -7,7 +7,8 @@ "op_fast_hadamard_transform_test.cpp" ], "additional_libs": [ - "custom_ops" + "custom_ops", + "dumb_fht" ] }, { @@ -62,6 +63,7 @@ "directory": "kernels/portable/cpu/util/test", "sources": [ "broadcast_test.cpp", + "delinearized_indexes_range_test.cpp", "reduce_test.cpp" ], "additional_libs": [