diff --git a/backends/aoti/slim/c10/core/SizesAndStrides.h b/backends/aoti/slim/c10/core/SizesAndStrides.h new file mode 100644 index 00000000000..ea1c7150937 --- /dev/null +++ b/backends/aoti/slim/c10/core/SizesAndStrides.h @@ -0,0 +1,415 @@ +/* + * 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 +#include + +/// Maximum number of dimensions that can be stored inline. +/// For tensors with more dimensions, out-of-line storage is used. +#define SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE 5 + +namespace executorch::backends::aoti::slim::c10 { + +using ::executorch::runtime::IntArrayRef; + +/** + * SizesAndStrides - Packed container for tensor sizes and strides. + * + * This class efficiently stores tensor dimension sizes and strides together. + * For tensors with up to 5 dimensions, storage is inline (no heap allocation). + * For larger tensors, heap storage is used. + * + * Memory layout: + * - Inline: 5 int64_t for sizes + 5 int64_t for strides + * - Out-of-line: pointer to heap array [sizes..., strides...] + */ +class SizesAndStrides { + public: + using sizes_iterator = int64_t*; + using sizes_const_iterator = const int64_t*; + using strides_iterator = int64_t*; + using strides_const_iterator = const int64_t*; + + /// Default constructor - creates a 1-dimensional tensor with size 0. + SizesAndStrides() { + size_at_unchecked(0) = 0; + stride_at_unchecked(0) = 1; + } + + ~SizesAndStrides() { + if (SLIMTENSOR_UNLIKELY(!isInline())) { + // NOLINTNEXTLINE(cppcoreguidelines-no-malloc) + free(outOfLineStorage_); + } + } + + SizesAndStrides(const SizesAndStrides& rhs) : size_(rhs.size_) { + if (SLIMTENSOR_LIKELY(rhs.isInline())) { + copyDataInline(rhs); + } else { + allocateOutOfLineStorage(size_); + copyDataOutline(rhs); + } + } + + SizesAndStrides& operator=(const SizesAndStrides& rhs) { + if (this == &rhs) { + return *this; + } + if (SLIMTENSOR_LIKELY(rhs.isInline())) { + if (SLIMTENSOR_UNLIKELY(!isInline())) { + // NOLINTNEXTLINE(cppcoreguidelines-no-malloc) + free(outOfLineStorage_); + } + copyDataInline(rhs); + } else { + if (isInline()) { + allocateOutOfLineStorage(rhs.size_); + } else { + resizeOutOfLineStorage(rhs.size_); + } + copyDataOutline(rhs); + } + size_ = rhs.size_; + return *this; + } + + SizesAndStrides(SizesAndStrides&& rhs) noexcept : size_(rhs.size_) { + if (SLIMTENSOR_LIKELY(isInline())) { + memcpy(inlineStorage_, rhs.inlineStorage_, sizeof(inlineStorage_)); + } else { + outOfLineStorage_ = rhs.outOfLineStorage_; + rhs.outOfLineStorage_ = nullptr; + } + rhs.size_ = 0; + } + + SizesAndStrides& operator=(SizesAndStrides&& rhs) noexcept { + if (this == &rhs) { + return *this; + } + if (SLIMTENSOR_LIKELY(rhs.isInline())) { + if (SLIMTENSOR_UNLIKELY(!isInline())) { + // NOLINTNEXTLINE(cppcoreguidelines-no-malloc) + free(outOfLineStorage_); + } + copyDataInline(rhs); + } else { + if (!isInline()) { + // NOLINTNEXTLINE(cppcoreguidelines-no-malloc) + free(outOfLineStorage_); + } + outOfLineStorage_ = rhs.outOfLineStorage_; + rhs.outOfLineStorage_ = nullptr; + } + size_ = rhs.size_; + rhs.size_ = 0; + return *this; + } + + bool operator==(const SizesAndStrides& other) const { + if (size_ != other.size_) { + return false; + } + return !( + isInline() + ? std::memcmp( + inlineStorage_, other.inlineStorage_, sizeof(inlineStorage_)) + : std::memcmp( + outOfLineStorage_, + other.outOfLineStorage_, + storageBytes(size_))); + } + + /// Returns the number of dimensions. + size_t size() const noexcept { + return size_; + } + + // Size accessors + + const int64_t* sizes_data() const noexcept { + if (SLIMTENSOR_LIKELY(isInline())) { + return &inlineStorage_[0]; + } else { + return &outOfLineStorage_[0]; + } + } + + int64_t* sizes_data() noexcept { + if (SLIMTENSOR_LIKELY(isInline())) { + return &inlineStorage_[0]; + } else { + return &outOfLineStorage_[0]; + } + } + + sizes_const_iterator sizes_begin() const noexcept { + return sizes_data(); + } + + sizes_iterator sizes_begin() noexcept { + return sizes_data(); + } + + sizes_const_iterator sizes_end() const noexcept { + return sizes_begin() + size(); + } + + sizes_iterator sizes_end() noexcept { + return sizes_begin() + size(); + } + + IntArrayRef sizes_arrayref() const noexcept { + return IntArrayRef{sizes_data(), size()}; + } + + void set_sizes(IntArrayRef newSizes) { + resize(newSizes.size()); + std::copy(newSizes.begin(), newSizes.end(), sizes_begin()); + } + + void set_sizes(std::initializer_list newSizes) { + resize(newSizes.size()); + std::copy(newSizes.begin(), newSizes.end(), sizes_begin()); + } + + int64_t size_at(size_t idx) const noexcept { + ET_DCHECK_MSG(idx < size(), "Index out of bounds"); + return sizes_data()[idx]; + } + + int64_t& size_at(size_t idx) noexcept { + ET_DCHECK_MSG(idx < size(), "Index out of bounds"); + return sizes_data()[idx]; + } + + int64_t size_at_unchecked(size_t idx) const noexcept { + return sizes_data()[idx]; + } + + int64_t& size_at_unchecked(size_t idx) noexcept { + return sizes_data()[idx]; + } + + // Stride accessors + + const int64_t* strides_data() const noexcept { + if (SLIMTENSOR_LIKELY(isInline())) { + return &inlineStorage_[SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE]; + } else { + return &outOfLineStorage_[size()]; + } + } + + int64_t* strides_data() noexcept { + if (SLIMTENSOR_LIKELY(isInline())) { + return &inlineStorage_[SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE]; + } else { + return &outOfLineStorage_[size()]; + } + } + + strides_const_iterator strides_begin() const noexcept { + return strides_data(); + } + + strides_iterator strides_begin() noexcept { + return strides_data(); + } + + strides_const_iterator strides_end() const noexcept { + return strides_begin() + size(); + } + + strides_iterator strides_end() noexcept { + return strides_begin() + size(); + } + + IntArrayRef strides_arrayref() const noexcept { + return IntArrayRef{strides_data(), size()}; + } + + void set_strides(IntArrayRef strides) { + ET_DCHECK_MSG( + strides.size() == size(), + "strides size (%zu) must match size (%zu)", + strides.size(), + size()); + std::copy(strides.begin(), strides.end(), strides_begin()); + } + + void set_strides(std::initializer_list strides) { + ET_DCHECK_MSG( + strides.size() == size(), + "strides size (%zu) must match size (%zu)", + strides.size(), + size()); + std::copy(strides.begin(), strides.end(), strides_begin()); + } + + int64_t stride_at(size_t idx) const noexcept { + ET_DCHECK_MSG(idx < size(), "Index out of bounds"); + return strides_data()[idx]; + } + + int64_t& stride_at(size_t idx) noexcept { + ET_DCHECK_MSG(idx < size(), "Index out of bounds"); + return strides_data()[idx]; + } + + int64_t stride_at_unchecked(size_t idx) const noexcept { + return strides_data()[idx]; + } + + int64_t& stride_at_unchecked(size_t idx) noexcept { + return strides_data()[idx]; + } + + /// Resizes to a new number of dimensions. + void resize(size_t newSize) { + const auto oldSize = size(); + if (newSize == oldSize) { + return; + } + if (SLIMTENSOR_LIKELY( + newSize <= SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE && + isInline())) { + if (oldSize < newSize) { + const auto bytesToZero = + (newSize - oldSize) * sizeof(inlineStorage_[0]); + memset(&inlineStorage_[oldSize], 0, bytesToZero); + memset( + &inlineStorage_ + [SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE + oldSize], + 0, + bytesToZero); + } + size_ = newSize; + } else { + resizeSlowPath(newSize, oldSize); + } + } + + private: + bool isInline() const noexcept { + return size_ <= SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE; + } + + void copyDataInline(const SizesAndStrides& rhs) { + ET_DCHECK_MSG(rhs.isInline(), "rhs must be inline"); + memcpy(inlineStorage_, rhs.inlineStorage_, sizeof(inlineStorage_)); + } + + void copyDataOutline(const SizesAndStrides& rhs) noexcept { + memcpy(outOfLineStorage_, rhs.outOfLineStorage_, storageBytes(rhs.size_)); + } + + static size_t storageBytes(size_t size) noexcept { + return size * 2 * sizeof(int64_t); + } + + void allocateOutOfLineStorage(size_t size) { + // NOLINTNEXTLINE(cppcoreguidelines-no-malloc) + outOfLineStorage_ = static_cast(malloc(storageBytes(size))); + ET_CHECK_MSG( + outOfLineStorage_, + "Could not allocate memory for Tensor SizesAndStrides!"); + } + + void resizeOutOfLineStorage(size_t newSize) { + ET_DCHECK_MSG(!isInline(), "must not be inline"); + outOfLineStorage_ = static_cast( + // NOLINTNEXTLINE(cppcoreguidelines-no-malloc) + realloc(outOfLineStorage_, storageBytes(newSize))); + ET_CHECK_MSG( + outOfLineStorage_, + "Could not allocate memory for Tensor SizesAndStrides!"); + } + + void resizeSlowPath(size_t newSize, size_t oldSize) { + if (newSize <= SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE) { + ET_DCHECK_MSG( + !isInline(), + "resizeSlowPath called when fast path should have been hit!"); + int64_t* tempStorage = outOfLineStorage_; + memcpy( + &inlineStorage_[0], + &tempStorage[0], + SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE * + sizeof(inlineStorage_[0])); + memcpy( + &inlineStorage_[SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE], + &tempStorage[oldSize], + SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE * + sizeof(inlineStorage_[0])); + // NOLINTNEXTLINE(cppcoreguidelines-no-malloc) + free(tempStorage); + } else { + if (isInline()) { + int64_t* tempStorage = + // NOLINTNEXTLINE(cppcoreguidelines-no-malloc) + static_cast(malloc(storageBytes(newSize))); + ET_CHECK_MSG( + tempStorage, + "Could not allocate memory to change Tensor SizesAndStrides!"); + const auto bytesToCopy = oldSize * sizeof(inlineStorage_[0]); + const auto bytesToZero = (newSize > oldSize) + ? (newSize - oldSize) * sizeof(tempStorage[0]) + : 0; + memcpy(&tempStorage[0], &inlineStorage_[0], bytesToCopy); + if (bytesToZero) { + memset(&tempStorage[oldSize], 0, bytesToZero); + } + memcpy( + &tempStorage[newSize], + &inlineStorage_[SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE], + bytesToCopy); + if (bytesToZero) { + memset(&tempStorage[newSize + oldSize], 0, bytesToZero); + } + outOfLineStorage_ = tempStorage; + } else { + const bool isGrowing = oldSize < newSize; + if (isGrowing) { + resizeOutOfLineStorage(newSize); + } + memmove( + outOfLineStorage_ + newSize, + outOfLineStorage_ + oldSize, + std::min(oldSize, newSize) * sizeof(outOfLineStorage_[0])); + if (!isGrowing) { + resizeOutOfLineStorage(newSize); + } else { + const auto bytesToZero = + (newSize - oldSize) * sizeof(outOfLineStorage_[0]); + memset(&outOfLineStorage_[oldSize], 0, bytesToZero); + memset(&outOfLineStorage_[newSize + oldSize], 0, bytesToZero); + } + } + } + size_ = newSize; + } + + size_t size_{1}; + union { + int64_t* outOfLineStorage_; + // NOLINTNEXTLINE(*c-array*) + int64_t inlineStorage_[SLIMTENSOR_SIZES_AND_STRIDES_MAX_INLINE_SIZE * 2]{}; + }; +}; + +} // namespace executorch::backends::aoti::slim::c10 diff --git a/backends/aoti/slim/c10/core/targets.bzl b/backends/aoti/slim/c10/core/targets.bzl index 9b7d1259df0..c421081f095 100644 --- a/backends/aoti/slim/c10/core/targets.bzl +++ b/backends/aoti/slim/c10/core/targets.bzl @@ -40,6 +40,20 @@ def define_common_targets(): ], ) + # Header-only library for SizesAndStrides + runtime.cxx_library( + name = "sizes_and_strides", + headers = [ + "SizesAndStrides.h", + ], + visibility = ["@EXECUTORCH_CLIENTS"], + exported_deps = [ + "//executorch/backends/aoti/slim/c10/macros:macros", + "//executorch/runtime/core:core", + "//executorch/runtime/platform:platform", + ], + ) + # Combined c10 core library runtime.cxx_library( name = "core", @@ -48,5 +62,6 @@ def define_common_targets(): ":device", ":device_type", ":scalar_type", + ":sizes_and_strides", ], ) diff --git a/backends/aoti/slim/c10/core/test/targets.bzl b/backends/aoti/slim/c10/core/test/targets.bzl index f7abf59a273..87352c619d6 100644 --- a/backends/aoti/slim/c10/core/test/targets.bzl +++ b/backends/aoti/slim/c10/core/test/targets.bzl @@ -23,3 +23,13 @@ def define_common_targets(): "//executorch/backends/aoti/slim/c10/core:scalar_type", ], ) + + runtime.cxx_test( + name = "test_sizes_and_strides", + srcs = [ + "test_sizes_and_strides.cpp", + ], + deps = [ + "//executorch/backends/aoti/slim/c10/core:sizes_and_strides", + ], + ) diff --git a/backends/aoti/slim/c10/core/test/test_sizes_and_strides.cpp b/backends/aoti/slim/c10/core/test/test_sizes_and_strides.cpp new file mode 100644 index 00000000000..4a40963a5bb --- /dev/null +++ b/backends/aoti/slim/c10/core/test/test_sizes_and_strides.cpp @@ -0,0 +1,316 @@ +/* + * 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 + +namespace executorch::backends::aoti::slim::c10 { + +// ============================================================================= +// Default Construction Tests +// ============================================================================= + +TEST(SizesAndStridesTest, DefaultConstruct) { + SizesAndStrides ss; + + EXPECT_EQ(ss.size(), 1u); + EXPECT_EQ(ss.size_at(0), 0); + EXPECT_EQ(ss.stride_at(0), 1); +} + +// ============================================================================= +// Set Sizes and Strides Tests +// ============================================================================= + +TEST(SizesAndStridesTest, SetSizes) { + SizesAndStrides ss; + ss.set_sizes({2, 3, 4}); + + EXPECT_EQ(ss.size(), 3u); + EXPECT_EQ(ss.sizes_arrayref()[0], 2); + EXPECT_EQ(ss.sizes_arrayref()[1], 3); + EXPECT_EQ(ss.sizes_arrayref()[2], 4); +} + +TEST(SizesAndStridesTest, SetSizesAndStrides) { + SizesAndStrides ss; + ss.set_sizes({2, 3, 4}); + ss.set_strides({12, 4, 1}); + + EXPECT_EQ(ss.size(), 3u); + EXPECT_EQ(ss.sizes_arrayref()[0], 2); + EXPECT_EQ(ss.sizes_arrayref()[1], 3); + EXPECT_EQ(ss.sizes_arrayref()[2], 4); + EXPECT_EQ(ss.strides_arrayref()[0], 12); + EXPECT_EQ(ss.strides_arrayref()[1], 4); + EXPECT_EQ(ss.strides_arrayref()[2], 1); +} + +TEST(SizesAndStridesTest, SizeAt) { + SizesAndStrides ss; + ss.set_sizes({10, 20, 30}); + + EXPECT_EQ(ss.size_at(0), 10); + EXPECT_EQ(ss.size_at(1), 20); + EXPECT_EQ(ss.size_at(2), 30); +} + +TEST(SizesAndStridesTest, StrideAt) { + SizesAndStrides ss; + ss.set_sizes({2, 3}); + ss.set_strides({3, 1}); + + EXPECT_EQ(ss.stride_at(0), 3); + EXPECT_EQ(ss.stride_at(1), 1); +} + +TEST(SizesAndStridesTest, SizeAtUnchecked) { + SizesAndStrides ss; + ss.set_sizes({5, 6, 7}); + + EXPECT_EQ(ss.size_at_unchecked(0), 5); + EXPECT_EQ(ss.size_at_unchecked(1), 6); + EXPECT_EQ(ss.size_at_unchecked(2), 7); +} + +TEST(SizesAndStridesTest, ModifySizeAt) { + SizesAndStrides ss; + ss.set_sizes({2, 3, 4}); + + ss.size_at(1) = 10; + + EXPECT_EQ(ss.size_at(1), 10); +} + +TEST(SizesAndStridesTest, ModifyStrideAt) { + SizesAndStrides ss; + ss.set_sizes({2, 3}); + ss.set_strides({3, 1}); + + ss.stride_at(0) = 100; + + EXPECT_EQ(ss.stride_at(0), 100); +} + +// ============================================================================= +// Inline vs Out-of-Line Storage Tests +// ============================================================================= + +TEST(SizesAndStridesTest, InlineStorage) { + SizesAndStrides ss; + ss.set_sizes({1, 2, 3, 4, 5}); + ss.set_strides({120, 60, 20, 5, 1}); + + EXPECT_EQ(ss.size(), 5u); + EXPECT_EQ(ss.sizes_arrayref()[4], 5); + EXPECT_EQ(ss.strides_arrayref()[4], 1); +} + +TEST(SizesAndStridesTest, OutOfLineStorage) { + SizesAndStrides ss; + ss.set_sizes({1, 2, 3, 4, 5, 6, 7}); + ss.set_strides({5040, 2520, 840, 210, 42, 7, 1}); + + EXPECT_EQ(ss.size(), 7u); + EXPECT_EQ(ss.sizes_arrayref()[6], 7); + EXPECT_EQ(ss.strides_arrayref()[6], 1); +} + +TEST(SizesAndStridesTest, ResizeFromInlineToOutOfLine) { + SizesAndStrides ss; + ss.set_sizes({1, 2, 3}); + + EXPECT_EQ(ss.size(), 3u); + + ss.set_sizes({1, 2, 3, 4, 5, 6, 7, 8}); + ss.set_strides({40320, 20160, 6720, 1680, 336, 56, 8, 1}); + + EXPECT_EQ(ss.size(), 8u); + EXPECT_EQ(ss.sizes_arrayref()[7], 8); +} + +TEST(SizesAndStridesTest, ResizeFromOutOfLineToInline) { + SizesAndStrides ss; + ss.set_sizes({1, 2, 3, 4, 5, 6, 7}); + ss.set_strides({5040, 2520, 840, 210, 42, 7, 1}); + + EXPECT_EQ(ss.size(), 7u); + + ss.set_sizes({2, 3}); + ss.set_strides({3, 1}); + + EXPECT_EQ(ss.size(), 2u); + EXPECT_EQ(ss.sizes_arrayref()[0], 2); + EXPECT_EQ(ss.strides_arrayref()[0], 3); +} + +// ============================================================================= +// Copy and Move Tests +// ============================================================================= + +TEST(SizesAndStridesTest, CopyConstructInline) { + SizesAndStrides original; + original.set_sizes({2, 3, 4}); + original.set_strides({12, 4, 1}); + + SizesAndStrides copy = original; + + EXPECT_EQ(copy.size(), 3u); + EXPECT_EQ(copy.sizes_arrayref()[0], 2); + EXPECT_EQ(copy.strides_arrayref()[0], 12); +} + +TEST(SizesAndStridesTest, CopyConstructOutOfLine) { + SizesAndStrides original; + original.set_sizes({1, 2, 3, 4, 5, 6, 7}); + original.set_strides({5040, 2520, 840, 210, 42, 7, 1}); + + SizesAndStrides copy = original; + + EXPECT_EQ(copy.size(), 7u); + EXPECT_EQ(copy.sizes_arrayref()[6], 7); + EXPECT_EQ(copy.strides_arrayref()[6], 1); +} + +TEST(SizesAndStridesTest, CopyAssignInline) { + SizesAndStrides original; + original.set_sizes({2, 3}); + original.set_strides({3, 1}); + + SizesAndStrides copy; + copy = original; + + EXPECT_EQ(copy.size(), 2u); + EXPECT_EQ(copy.sizes_arrayref()[1], 3); +} + +TEST(SizesAndStridesTest, MoveConstructInline) { + SizesAndStrides original; + original.set_sizes({2, 3, 4}); + original.set_strides({12, 4, 1}); + + SizesAndStrides moved = std::move(original); + + EXPECT_EQ(moved.size(), 3u); + EXPECT_EQ(moved.sizes_arrayref()[0], 2); + EXPECT_EQ(original.size(), 0u); +} + +TEST(SizesAndStridesTest, MoveConstructOutOfLine) { + SizesAndStrides original; + original.set_sizes({1, 2, 3, 4, 5, 6, 7}); + original.set_strides({5040, 2520, 840, 210, 42, 7, 1}); + + SizesAndStrides moved = std::move(original); + + EXPECT_EQ(moved.size(), 7u); + EXPECT_EQ(moved.sizes_arrayref()[6], 7); + EXPECT_EQ(original.size(), 0u); +} + +TEST(SizesAndStridesTest, MoveAssignOutOfLine) { + SizesAndStrides original; + original.set_sizes({1, 2, 3, 4, 5, 6, 7}); + original.set_strides({5040, 2520, 840, 210, 42, 7, 1}); + + SizesAndStrides target; + target = std::move(original); + + EXPECT_EQ(target.size(), 7u); + EXPECT_EQ(target.sizes_arrayref()[6], 7); + EXPECT_EQ(original.size(), 0u); +} + +// ============================================================================= +// Equality Tests +// ============================================================================= + +TEST(SizesAndStridesTest, EqualityTrue) { + SizesAndStrides ss1; + ss1.set_sizes({2, 3, 4}); + ss1.set_strides({12, 4, 1}); + + SizesAndStrides ss2; + ss2.set_sizes({2, 3, 4}); + ss2.set_strides({12, 4, 1}); + + EXPECT_TRUE(ss1 == ss2); +} + +TEST(SizesAndStridesTest, EqualityFalseDifferentSizes) { + SizesAndStrides ss1; + ss1.set_sizes({2, 3, 4}); + ss1.set_strides({12, 4, 1}); + + SizesAndStrides ss2; + ss2.set_sizes({2, 3, 5}); + ss2.set_strides({15, 5, 1}); + + EXPECT_FALSE(ss1 == ss2); +} + +TEST(SizesAndStridesTest, EqualityFalseDifferentDims) { + SizesAndStrides ss1; + ss1.set_sizes({2, 3}); + + SizesAndStrides ss2; + ss2.set_sizes({2, 3, 4}); + + EXPECT_FALSE(ss1 == ss2); +} + +// ============================================================================= +// Iterator Tests +// ============================================================================= + +TEST(SizesAndStridesTest, SizesIterator) { + SizesAndStrides ss; + ss.set_sizes({2, 3, 4}); + + int64_t sum = 0; + for (auto it = ss.sizes_begin(); it != ss.sizes_end(); ++it) { + sum += *it; + } + EXPECT_EQ(sum, 9); +} + +TEST(SizesAndStridesTest, StridesIterator) { + SizesAndStrides ss; + ss.set_sizes({2, 3, 4}); + ss.set_strides({12, 4, 1}); + + int64_t sum = 0; + for (auto it = ss.strides_begin(); it != ss.strides_end(); ++it) { + sum += *it; + } + EXPECT_EQ(sum, 17); +} + +// ============================================================================= +// Edge Cases +// ============================================================================= + +TEST(SizesAndStridesTest, ScalarTensor) { + SizesAndStrides ss; + ss.resize(0); + + EXPECT_EQ(ss.size(), 0u); +} + +TEST(SizesAndStridesTest, OneDimensional) { + SizesAndStrides ss; + ss.set_sizes({10}); + ss.set_strides({1}); + + EXPECT_EQ(ss.size(), 1u); + EXPECT_EQ(ss.size_at(0), 10); + EXPECT_EQ(ss.stride_at(0), 1); +} + +} // namespace executorch::backends::aoti::slim::c10 diff --git a/backends/aoti/slim/c10/macros/Macros.h b/backends/aoti/slim/c10/macros/Macros.h new file mode 100644 index 00000000000..0c9f53c58a4 --- /dev/null +++ b/backends/aoti/slim/c10/macros/Macros.h @@ -0,0 +1,26 @@ +/* + * 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 + +// SlimTensor Macros Header +// +// This header provides common compiler hint macros for SlimTensor code. + +// Compiler hint macros for branch prediction +#if defined(__GNUC__) || defined(__clang__) +#define SLIMTENSOR_LIKELY(x) __builtin_expect(!!(x), 1) +#define SLIMTENSOR_UNLIKELY(x) __builtin_expect(!!(x), 0) +#else +#define SLIMTENSOR_LIKELY(x) (x) +#define SLIMTENSOR_UNLIKELY(x) (x) +#endif + +namespace executorch::backends::aoti::slim::c10 { +// Empty namespace - macros are defined above +} // namespace executorch::backends::aoti::slim::c10 diff --git a/backends/aoti/slim/c10/macros/TARGETS b/backends/aoti/slim/c10/macros/TARGETS new file mode 100644 index 00000000000..77871de4469 --- /dev/null +++ b/backends/aoti/slim/c10/macros/TARGETS @@ -0,0 +1,3 @@ +load("targets.bzl", "define_common_targets") + +define_common_targets() diff --git a/backends/aoti/slim/c10/macros/targets.bzl b/backends/aoti/slim/c10/macros/targets.bzl new file mode 100644 index 00000000000..417590ca651 --- /dev/null +++ b/backends/aoti/slim/c10/macros/targets.bzl @@ -0,0 +1,13 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + """Define targets for SlimTensor c10 macros module.""" + + # Header-only library for Macros + runtime.cxx_library( + name = "macros", + headers = [ + "Macros.h", + ], + visibility = ["@EXECUTORCH_CLIENTS"], + ) diff --git a/backends/aoti/slim/util/ArrayRefUtil.h b/backends/aoti/slim/util/ArrayRefUtil.h new file mode 100644 index 00000000000..5daec3f8337 --- /dev/null +++ b/backends/aoti/slim/util/ArrayRefUtil.h @@ -0,0 +1,43 @@ +/* + * 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 + +// Utilities for working with ExecuTorch's ArrayRef in SlimTensor code. + +#include + +#include +#include + +namespace executorch::backends::aoti::slim { + +// Bring ExecuTorch's ArrayRef types into the SlimTensor namespace +using ::executorch::runtime::ArrayRef; +using ::executorch::runtime::IntArrayRef; +using ::executorch::runtime::makeArrayRef; + +/// Helper function to construct an ArrayRef from a std::vector. +template +inline ArrayRef makeArrayRef(const std::vector& Vec) { + return ArrayRef(Vec.data(), Vec.size()); +} + +/// Helper function to construct an ArrayRef from a std::initializer_list. +template +inline ArrayRef makeArrayRef(std::initializer_list list) { + return ArrayRef(list.begin(), list.size()); +} + +/// Helper function to convert ArrayRef to std::vector. +template +inline std::vector toVec(ArrayRef arr) { + return std::vector(arr.begin(), arr.end()); +} + +} // namespace executorch::backends::aoti::slim diff --git a/backends/aoti/slim/util/SizeUtil.h b/backends/aoti/slim/util/SizeUtil.h new file mode 100644 index 00000000000..faf6e9c914f --- /dev/null +++ b/backends/aoti/slim/util/SizeUtil.h @@ -0,0 +1,90 @@ +/* + * 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 + +namespace executorch::backends::aoti::slim { + +/// Computes the number of elements (numel) from the sizes array. +/// @param sizes Array of dimension sizes. +/// @return The total number of elements (product of all sizes). +inline int64_t compute_numel(IntArrayRef sizes) { + int64_t n = 1; + for (auto s : sizes) { + n *= s; + } + return n; +} + +/// Computes the storage size in bytes for a contiguous tensor. +/// @param sizes Array of dimension sizes. +/// @param itemsize_bytes Size of each element in bytes. +/// @param storage_offset Offset into the storage in elements. +/// @return The required storage size in bytes. +inline size_t compute_storage_nbytes_contiguous( + IntArrayRef sizes, + size_t itemsize_bytes, + size_t storage_offset) { + const auto numel = compute_numel(sizes); + return itemsize_bytes * (storage_offset + numel); +} + +/// Computes the storage size in bytes for a strided tensor. +/// @param sizes Array of dimension sizes. +/// @param strides Array of dimension strides. +/// @param itemsize_bytes Size of each element in bytes. +/// @param storage_offset Offset into the storage in elements. +/// @return The required storage size in bytes. +inline size_t compute_storage_nbytes( + IntArrayRef sizes, + IntArrayRef strides, + size_t itemsize_bytes, + size_t storage_offset) { + ET_CHECK_MSG( + sizes.size() == strides.size(), + "dimensionality of sizes (%zu) must match dimensionality of strides (%zu)", + sizes.size(), + strides.size()); + + // Size of the underlying storage is 1 bigger than the offset + // of the last element according to stride. + size_t size = 1; + for (size_t i = 0; i < sizes.size(); ++i) { + if (sizes[i] == 0) { + return 0; + } + size += strides[i] * (sizes[i] - 1); + } + return itemsize_bytes * (storage_offset + size); +} + +/// Computes contiguous strides from sizes. +/// @param sizes Array of dimension sizes. +/// @return Vector of strides for a contiguous tensor. +inline std::vector compute_contiguous_strides(IntArrayRef sizes) { + int64_t ndim = static_cast(sizes.size()); + std::vector strides(ndim); + if (ndim > 0) { + int64_t stride = 1; + for (int64_t i = ndim - 1; i >= 0; i--) { + strides[i] = stride; + if (sizes[i] != 0) { + stride *= sizes[i]; + } + } + } + return strides; +} + +} // namespace executorch::backends::aoti::slim diff --git a/backends/aoti/slim/util/targets.bzl b/backends/aoti/slim/util/targets.bzl index 13f49168a0f..60a7933a861 100644 --- a/backends/aoti/slim/util/targets.bzl +++ b/backends/aoti/slim/util/targets.bzl @@ -14,3 +14,28 @@ def define_common_targets(): "//executorch/runtime/platform:platform", ], ) + + # Header-only library for ArrayRefUtil + runtime.cxx_library( + name = "array_ref_util", + headers = [ + "ArrayRefUtil.h", + ], + visibility = ["@EXECUTORCH_CLIENTS"], + exported_deps = [ + "//executorch/runtime/core:core", + ], + ) + + # Header-only library for SizeUtil + runtime.cxx_library( + name = "size_util", + headers = [ + "SizeUtil.h", + ], + visibility = ["@EXECUTORCH_CLIENTS"], + exported_deps = [ + ":array_ref_util", + "//executorch/runtime/platform:platform", + ], + ) diff --git a/backends/aoti/slim/util/test/TARGETS b/backends/aoti/slim/util/test/TARGETS new file mode 100644 index 00000000000..77871de4469 --- /dev/null +++ b/backends/aoti/slim/util/test/TARGETS @@ -0,0 +1,3 @@ +load("targets.bzl", "define_common_targets") + +define_common_targets() diff --git a/backends/aoti/slim/util/test/targets.bzl b/backends/aoti/slim/util/test/targets.bzl new file mode 100644 index 00000000000..7334a2d4da1 --- /dev/null +++ b/backends/aoti/slim/util/test/targets.bzl @@ -0,0 +1,14 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + """Define test targets for SlimTensor util module.""" + + runtime.cxx_test( + name = "test_size_util", + srcs = [ + "test_size_util.cpp", + ], + deps = [ + "//executorch/backends/aoti/slim/util:size_util", + ], + ) diff --git a/backends/aoti/slim/util/test/test_size_util.cpp b/backends/aoti/slim/util/test/test_size_util.cpp new file mode 100644 index 00000000000..fd1deed1909 --- /dev/null +++ b/backends/aoti/slim/util/test/test_size_util.cpp @@ -0,0 +1,232 @@ +/* + * 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 + +namespace executorch::backends::aoti::slim { + +// ============================================================================= +// compute_numel Tests +// ============================================================================= + +TEST(SizeUtilTest, ComputeNumel1D) { + std::vector sizes = {10}; + EXPECT_EQ(compute_numel(makeArrayRef(sizes)), 10); +} + +TEST(SizeUtilTest, ComputeNumel2D) { + std::vector sizes = {3, 4}; + EXPECT_EQ(compute_numel(makeArrayRef(sizes)), 12); +} + +TEST(SizeUtilTest, ComputeNumel3D) { + std::vector sizes = {2, 3, 4}; + EXPECT_EQ(compute_numel(makeArrayRef(sizes)), 24); +} + +TEST(SizeUtilTest, ComputeNumel4D) { + std::vector sizes = {2, 3, 4, 5}; + EXPECT_EQ(compute_numel(makeArrayRef(sizes)), 120); +} + +TEST(SizeUtilTest, ComputeNumelEmpty) { + std::vector sizes = {0, 3, 4}; + EXPECT_EQ(compute_numel(makeArrayRef(sizes)), 0); +} + +TEST(SizeUtilTest, ComputeNumelScalar) { + std::vector sizes = {}; + EXPECT_EQ(compute_numel(makeArrayRef(sizes)), 1); +} + +TEST(SizeUtilTest, ComputeNumelWithOnes) { + std::vector sizes = {1, 1, 5, 1}; + EXPECT_EQ(compute_numel(makeArrayRef(sizes)), 5); +} + +// ============================================================================= +// compute_contiguous_strides Tests +// ============================================================================= + +TEST(SizeUtilTest, ComputeContiguousStrides1D) { + std::vector sizes = {10}; + auto strides = compute_contiguous_strides(makeArrayRef(sizes)); + + EXPECT_EQ(strides.size(), 1u); + EXPECT_EQ(strides[0], 1); +} + +TEST(SizeUtilTest, ComputeContiguousStrides2D) { + std::vector sizes = {3, 4}; + auto strides = compute_contiguous_strides(makeArrayRef(sizes)); + + EXPECT_EQ(strides.size(), 2u); + EXPECT_EQ(strides[0], 4); + EXPECT_EQ(strides[1], 1); +} + +TEST(SizeUtilTest, ComputeContiguousStrides3D) { + std::vector sizes = {2, 3, 4}; + auto strides = compute_contiguous_strides(makeArrayRef(sizes)); + + EXPECT_EQ(strides.size(), 3u); + EXPECT_EQ(strides[0], 12); + EXPECT_EQ(strides[1], 4); + EXPECT_EQ(strides[2], 1); +} + +TEST(SizeUtilTest, ComputeContiguousStrides4D) { + std::vector sizes = {2, 3, 4, 5}; + auto strides = compute_contiguous_strides(makeArrayRef(sizes)); + + EXPECT_EQ(strides.size(), 4u); + EXPECT_EQ(strides[0], 60); + EXPECT_EQ(strides[1], 20); + EXPECT_EQ(strides[2], 5); + EXPECT_EQ(strides[3], 1); +} + +TEST(SizeUtilTest, ComputeContiguousStridesScalar) { + std::vector sizes = {}; + auto strides = compute_contiguous_strides(makeArrayRef(sizes)); + + EXPECT_EQ(strides.size(), 0u); +} + +TEST(SizeUtilTest, ComputeContiguousStridesWithZero) { + std::vector sizes = {0, 3, 4}; + auto strides = compute_contiguous_strides(makeArrayRef(sizes)); + + EXPECT_EQ(strides.size(), 3u); + EXPECT_EQ(strides[0], 12); + EXPECT_EQ(strides[1], 4); + EXPECT_EQ(strides[2], 1); +} + +// ============================================================================= +// compute_storage_nbytes_contiguous Tests +// ============================================================================= + +TEST(SizeUtilTest, ComputeStorageNbytesContiguousFloat) { + std::vector sizes = {2, 3}; + size_t nbytes = + compute_storage_nbytes_contiguous(makeArrayRef(sizes), sizeof(float), 0); + EXPECT_EQ(nbytes, 6 * sizeof(float)); +} + +TEST(SizeUtilTest, ComputeStorageNbytesContiguousDouble) { + std::vector sizes = {2, 3, 4}; + size_t nbytes = + compute_storage_nbytes_contiguous(makeArrayRef(sizes), sizeof(double), 0); + EXPECT_EQ(nbytes, 24 * sizeof(double)); +} + +TEST(SizeUtilTest, ComputeStorageNbytesContiguousWithOffset) { + std::vector sizes = {2, 3}; + size_t nbytes = + compute_storage_nbytes_contiguous(makeArrayRef(sizes), sizeof(float), 10); + EXPECT_EQ(nbytes, (10 + 6) * sizeof(float)); +} + +TEST(SizeUtilTest, ComputeStorageNbytesContiguousEmpty) { + std::vector sizes = {0, 3, 4}; + size_t nbytes = + compute_storage_nbytes_contiguous(makeArrayRef(sizes), sizeof(float), 0); + EXPECT_EQ(nbytes, 0u); +} + +// ============================================================================= +// compute_storage_nbytes (strided) Tests +// ============================================================================= + +TEST(SizeUtilTest, ComputeStorageNbytesContiguousTensor) { + std::vector sizes = {2, 3}; + std::vector strides = {3, 1}; + size_t nbytes = compute_storage_nbytes( + makeArrayRef(sizes), makeArrayRef(strides), sizeof(float), 0); + EXPECT_EQ(nbytes, 6 * sizeof(float)); +} + +TEST(SizeUtilTest, ComputeStorageNbytesTransposedTensor) { + std::vector sizes = {2, 3}; + std::vector strides = {1, 2}; + size_t nbytes = compute_storage_nbytes( + makeArrayRef(sizes), makeArrayRef(strides), sizeof(float), 0); + EXPECT_EQ(nbytes, 6 * sizeof(float)); +} + +TEST(SizeUtilTest, ComputeStorageNbytesWithOffset) { + std::vector sizes = {2, 3}; + std::vector strides = {3, 1}; + size_t nbytes = compute_storage_nbytes( + makeArrayRef(sizes), makeArrayRef(strides), sizeof(float), 5); + EXPECT_EQ(nbytes, (5 + 6) * sizeof(float)); +} + +TEST(SizeUtilTest, ComputeStorageNbytesStrided3D) { + std::vector sizes = {2, 3, 4}; + std::vector strides = {12, 4, 1}; + size_t nbytes = compute_storage_nbytes( + makeArrayRef(sizes), makeArrayRef(strides), sizeof(float), 0); + EXPECT_EQ(nbytes, 24 * sizeof(float)); +} + +TEST(SizeUtilTest, ComputeStorageNbytesEmpty) { + std::vector sizes = {0, 3, 4}; + std::vector strides = {12, 4, 1}; + size_t nbytes = compute_storage_nbytes( + makeArrayRef(sizes), makeArrayRef(strides), sizeof(float), 0); + EXPECT_EQ(nbytes, 0u); +} + +TEST(SizeUtilTest, ComputeStorageNbytesNonContiguous) { + std::vector sizes = {2, 2}; + std::vector strides = {4, 1}; + size_t nbytes = compute_storage_nbytes( + makeArrayRef(sizes), makeArrayRef(strides), sizeof(float), 0); + EXPECT_EQ(nbytes, 6 * sizeof(float)); +} + +// ============================================================================= +// ArrayRefUtil Tests +// ============================================================================= + +TEST(ArrayRefUtilTest, MakeArrayRefFromVector) { + std::vector vec = {1, 2, 3, 4, 5}; + IntArrayRef ref = makeArrayRef(vec); + + EXPECT_EQ(ref.size(), 5u); + EXPECT_EQ(ref[0], 1); + EXPECT_EQ(ref[4], 5); +} + +TEST(ArrayRefUtilTest, MakeArrayRefFromVectorConstruction) { + std::vector values = {10, 20, 30}; + IntArrayRef ref = makeArrayRef(values); + + EXPECT_EQ(ref.size(), 3u); + EXPECT_EQ(ref[0], 10); + EXPECT_EQ(ref[2], 30); +} + +TEST(ArrayRefUtilTest, ToVec) { + std::vector original = {1, 2, 3}; + IntArrayRef ref = makeArrayRef(original); + std::vector copy = toVec(ref); + + EXPECT_EQ(copy.size(), 3u); + EXPECT_EQ(copy[0], 1); + EXPECT_EQ(copy[2], 3); + + copy[0] = 100; + EXPECT_EQ(original[0], 1); +} + +} // namespace executorch::backends::aoti::slim