add DelinearizedIndexesRange

ghstack-source-id: c8e09aa
ghstack-comment-id: 2691804929
Pull Request resolved: pytorch/executorch#8859

Author: swolchok

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 <algorithm>
+#include <array>
+#include <cstdint>
+#include <iterator>
+
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/util/tensor_dimension_limit.h>
+
+namespace torch::executor {
+
+namespace internal {
+class DelinearizedIndexesIterator {
+ public:
+  using difference_type = ssize_t;
+  using value_type = std::array<std::size_t, ::executorch::runtime::kTensorDimensionLimit>;
+  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<exec_aten::SizesType> 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

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 ""

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

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 <executorch/kernels/portable/cpu/util/delinearized_indexes_range.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
+
+#include <gtest/gtest.h>
+
+using executorch::aten::ScalarType;
+using executorch::aten::Tensor;
+using executorch::runtime::testing::TensorFactory;
+using torch::executor::DelinearizedIndexesRange;
+
+TEST(DelinearizedIndexesRangeTest, Empty) {
+  TensorFactory<ScalarType::Int> 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<ScalarType::Int> tf;
+
+  Tensor a = tf.zeros({5});
+  DelinearizedIndexesRange r(a);
+  std::vector<typename DelinearizedIndexesRange::iterator::value_type> v(r.begin(), r.end());
+  int idx = 0;
+  for (const auto& elem: v) {
+    EXPECT_EQ(elem[0], idx++);
+  }
+}
+
+TEST(DelinearizedIndexesRangeTest, ThreeD) {
+  TensorFactory<ScalarType::Int> tf;
+  Tensor a = tf.zeros({3, 2, 3});
+  DelinearizedIndexesRange r(a);
+  std::vector<typename DelinearizedIndexesRange::iterator::value_type> v(r.begin(), r.end());
+  std::vector<typename DelinearizedIndexesRange::iterator::value_type> 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);
+}

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"],

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": [