Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
192 changes: 131 additions & 61 deletions runtime/compute/cker/include/cker/Shape.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,11 @@
#define __NNFW_CKER_SHAPE_H__

#include <algorithm>
#include <cstring>
#include <array>
#include <cassert>
#include <cstring>
#include <iterator>
#include <variant>
#include <vector>

namespace nnfw
Expand All @@ -35,18 +38,28 @@ class Shape
// larger shapes are separately allocated.
static constexpr int kMaxSmallSize = 6;

// Delete copy assignment operator.
Shape &operator=(Shape const &) = delete;

Shape() : _size(0) {}
// Default constructor: initializes an empty shape (size = 0) with small storage.
Shape() : _size(0), dims_(std::array<int32_t, kMaxSmallSize>{}) {}

// Constructor that takes a dimension count.
// If dimensions_count <= kMaxSmallSize, it uses a fixed-size array.
// Otherwise, it uses a dynamic vector.
explicit Shape(int dimensions_count) : _size(dimensions_count)
{
if (dimensions_count > kMaxSmallSize)
if (dimensions_count <= kMaxSmallSize)
{
dims_ = std::array<int32_t, kMaxSmallSize>{};
}
else
{
_dims_pointer = new int32_t[dimensions_count];
dims_ = std::vector<int32_t>(dimensions_count);
Comment thread
ragmani marked this conversation as resolved.
}
}

// Constructor that creates a shape of given size and fills all dimensions with "value".
Shape(int shape_size, int32_t value) : _size(0)
{
Resize(shape_size);
Expand All @@ -56,136 +69,190 @@ class Shape
}
}

// Constructor that creates a shape from an array of dimension data.
Shape(int dimensions_count, const int32_t *dims_data) : _size(0)
{
ReplaceWith(dimensions_count, dims_data);
}

// Initializer list constructor.
// Marked explicit to avoid unintended overload resolution.
Shape(const std::initializer_list<int> init_list) : _size(0) { BuildFrom(init_list); }

// Avoid using this constructor. We should be able to delete it when C++17
// rolls out.
Shape(Shape const &other) : _size(other.DimensionsCount())
// Copy constructor
Shape(const Shape &other) : _size(other._size)
{
if (_size > kMaxSmallSize)
if (_size <= kMaxSmallSize)
{
// When the number of dimensions is small, copy the fixed array.
dims_ = std::get<std::array<int32_t, kMaxSmallSize>>(other.dims_);
}
else
{
_dims_pointer = new int32_t[_size];
// Otherwise, copy the dynamically allocated vector.
dims_ = std::get<std::vector<int32_t>>(other.dims_);
}
std::memcpy(DimsData(), other.DimsData(), sizeof(int32_t) * _size);
}
Shape(Shape &&other) = default;
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I enabled Move Semantics deleted by the custom copy constructor:
It allows instances of Shape to be moved efficiently. Instead of copying potentially large amounts of data (especially for dynamically allocated storage when dimensions exceed kMaxSmallSize), the object’s resources can be transferred (or "moved") to a new object with minimal overhead.

For Performance Benefits:
In modern C++ (C++11 and later), move semantics are essential to optimize performance and resource management. By marking the move constructor as default, we ensure that move operations are available and implemented in the most efficient way possible by the compiler.


bool operator==(const Shape &comp) const
{
return this->_size == comp._size &&
std::memcmp(DimsData(), comp.DimsData(), _size * sizeof(int32_t)) == 0;
}

~Shape()
~Shape() = default;

// Returns the number of dimensions.
inline int32_t DimensionsCount() const { return _size; }

// Returns the dimension size at index i.
inline int32_t Dims(int i) const
{
if (_size > kMaxSmallSize)
assert(i >= 0 && i < _size);
if (_size <= kMaxSmallSize)
{
delete[] _dims_pointer;
return std::get<std::array<int32_t, kMaxSmallSize>>(dims_)[i];
}
else
{
return std::get<std::vector<int32_t>>(dims_)[i];
}
}

inline int32_t DimensionsCount() const { return _size; }
inline int32_t Dims(int i) const
// Sets the dimension at index i.
inline void SetDim(int i, int32_t val)
{
assert(i >= 0);
assert(i < _size);
return _size > kMaxSmallSize ? _dims_pointer[i] : _dims[i];
assert(i >= 0 && i < _size);
if (_size <= kMaxSmallSize)
{
std::get<std::array<int32_t, kMaxSmallSize>>(dims_)[i] = val;
}
else
{
std::get<std::vector<int32_t>>(dims_)[i] = val;
}
}
inline void SetDim(int i, int32_t val)

// Returns a pointer to the dimension data (mutable).
inline int32_t *DimsData()
{
if (_size <= kMaxSmallSize)
{
return std::get<std::array<int32_t, kMaxSmallSize>>(dims_).data();
}
else
{
return std::get<std::vector<int32_t>>(dims_).data();
}
}

// Returns a pointer to the dimension data (const).
inline const int32_t *DimsData() const
{
assert(i >= 0);
assert(i < _size);
if (_size > kMaxSmallSize)
if (_size <= kMaxSmallSize)
{
_dims_pointer[i] = val;
return std::get<std::array<int32_t, kMaxSmallSize>>(dims_).data();
}
else
{
_dims[i] = val;
return std::get<std::vector<int32_t>>(dims_).data();
}
}

inline int32_t *DimsData() { return _size > kMaxSmallSize ? _dims_pointer : _dims; }
inline const int32_t *DimsData() const { return _size > kMaxSmallSize ? _dims_pointer : _dims; }
// The caller must ensure that the shape is no bigger than 6D.
inline const int32_t *DimsDataUpTo6D() const { return _dims; }
// The caller must ensure that the shape is no larger than 6D.
inline const int32_t *DimsDataUpTo6D() const
{
return std::get<std::array<int32_t, kMaxSmallSize>>(dims_).data();
}

// Resizes the shape to dimensions_count while preserving existing data.
inline void Resize(int dimensions_count)
{
if (_size > kMaxSmallSize)
std::vector<int32_t> oldDims;
oldDims.reserve(_size);
if (_size <= kMaxSmallSize)
{
delete[] _dims_pointer;
const auto &arr = std::get<std::array<int32_t, kMaxSmallSize>>(dims_);
oldDims.assign(arr.begin(), arr.begin() + _size);
}
_size = dimensions_count;
if (dimensions_count > kMaxSmallSize)
else
{
oldDims = std::get<std::vector<int32_t>>(dims_);
}

int count = std::min(_size, dimensions_count);

if (dimensions_count <= kMaxSmallSize)
{
_dims_pointer = new int32_t[dimensions_count];
std::array<int32_t, kMaxSmallSize> dims = {};
std::copy_n(oldDims.begin(), count, dims.begin());
dims_ = dims;
}
else
{
std::vector<int32_t> dims(dimensions_count, 0);
std::copy_n(oldDims.begin(), count, dims.begin());
dims_ = dims;
}

_size = dimensions_count;
}

// Replaces the current shape with a new one defined by dimensions_count and dims_data.
inline void ReplaceWith(int dimensions_count, const int32_t *dims_data)
{
Resize(dimensions_count);
int32_t *dst_dims = DimsData();
std::memcpy(dst_dims, dims_data, dimensions_count * sizeof(int32_t));
std::memcpy(DimsData(), dims_data, dimensions_count * sizeof(int32_t));
}

// Replaces the current shape with another shape.
inline void ReplaceWith(const Shape &other)
{
ReplaceWith(other.DimensionsCount(), other.DimsData());
}

// Replaces the current shape with another shape using move semantics.
inline void ReplaceWith(Shape &&other)
{
Resize(0);
std::swap(_size, other._size);
if (_size <= kMaxSmallSize)
std::copy(other._dims, other._dims + kMaxSmallSize, _dims);
else
_dims_pointer = other._dims_pointer;
dims_ = std::move(other.dims_);
}

template <typename T> inline void BuildFrom(const T &src_iterable)
// Builds the shape from an iterable sequence.
template <typename Iterable> inline void BuildFrom(const Iterable &src_iterable)
{
const int dimensions_count = std::distance(src_iterable.begin(), src_iterable.end());
const int dimensions_count =
static_cast<int>(std::distance(src_iterable.begin(), src_iterable.end()));
Resize(dimensions_count);
int32_t *data = DimsData();
for (auto &&it : src_iterable)
for (auto it = src_iterable.begin(); it != src_iterable.end(); ++it)
{
*data = it;
++data;
*data++ = static_cast<int32_t>(*it);
}
}

// This will probably be factored out. Old code made substantial use of 4-D
// shapes, and so this function is used to extend smaller shapes. Note that
// (a) as Dims<4>-dependent code is eliminated, the reliance on this should be
// reduced, and (b) some kernels are stricly 4-D, but then the shapes of their
// inputs should already be 4-D, so this function should not be needed.
// Returns the total count of elements, that is the size when flattened into a
// vector.
inline static Shape ExtendedShape(int new_shape_size, const Shape &shape)
{
return Shape(new_shape_size, shape, 1);
}

// Overload for initializer list building.
inline void BuildFrom(const std::initializer_list<int> init_list)
{
BuildFrom<const std::initializer_list<int>>(init_list);
}

// Returns the total count of elements, that is the size when flattened into a
// vector.
// Returns the total count of elements (flattened size).
inline int FlatSize() const
{
int buffer_size = 1;
const int *dims_data = DimsData();
for (int i = 0; i < _size; i++)
{
const int dim = dims_data[i];
buffer_size *= dim;
buffer_size *= dims_data[i];
}
return buffer_size;
}
Expand All @@ -211,12 +278,13 @@ class Shape
}

int32_t _size;
union {
int32_t _dims[kMaxSmallSize];
int32_t *_dims_pointer{nullptr};
};
// Internal storage: use std::array for shapes with dimensions up to kMaxSmallSize,
// and std::vector for larger shapes.
std::variant<std::array<int32_t, kMaxSmallSize>, std::vector<int32_t>> dims_;
};

// Utility functions below.

inline int MatchingDim(const Shape &shape1, int index1, [[maybe_unused]] const Shape &shape2,
[[maybe_unused]] int index2)
{
Expand All @@ -232,7 +300,10 @@ int MatchingDim(const Shape &shape1, int index1, [[maybe_unused]] const Shape &s
return MatchingDim(shape1, index1, args...);
}

inline Shape GetShape(const std::vector<int32_t> &data) { return Shape(data.size(), data.data()); }
inline Shape GetShape(const std::vector<int32_t> &data)
{
return Shape(static_cast<int>(data.size()), data.data());
}

inline int Offset(const Shape &shape, int i0, int i1, int i2, int i3)
{
Expand Down Expand Up @@ -278,8 +349,7 @@ inline int FlatSizeSkipDim(const Shape &shape, int skip_dim)
return flat_size;
}

// Flat size calculation, checking that dimensions match with one or more other
// arrays.
// Flat size calculation, checking that dimensions match with one or more other shapes.
template <typename... Ts> inline bool checkMatching(const Shape &shape, Ts... check_shapes)
{
auto match = [&shape](const Shape &s) -> bool {
Expand Down
2 changes: 1 addition & 1 deletion runtime/compute/cker/include/cker/operation/Einsum.h
Original file line number Diff line number Diff line change
Expand Up @@ -871,7 +871,7 @@ class Einsum
Tensor rhs;
reshapeToRank3(inputs[1], bcast.y_batch_size(), &rhs);
Shape old_output_shape = bcast.output_batch_shape();
Shape output_shape(old_output_shape.DimensionsCount() + inputs.size());
Shape output_shape(static_cast<int>(old_output_shape.DimensionsCount() + inputs.size()));
for (int i = 0; i < old_output_shape.DimensionsCount(); i++)
{
output_shape.SetDim(i, old_output_shape.Dims(i));
Expand Down