ndarray.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
// Copyright 2024 XCENA Inc.
 
#pragma once
 
#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <initializer_list>
#include <type_traits>
 
namespace mu
{
 
template <typename T>
class NDArray
{
public:
    static constexpr std::int64_t MAX_DIMS = 10;
    using ShapeArray = std::array<std::int64_t, MAX_DIMS>;
    using StridesArray = std::array<std::int64_t, MAX_DIMS>;
 
    NDArray() = default;
    NDArray(T* data, std::initializer_list<int64_t> shape)
        : data_(const_cast<void*>(reinterpret_cast<const void*>(data))),
          dims_(shape.size()),
          elemSize_(sizeof(T))
    {
        std::fill(shape_.begin(), shape_.end(), 0);
        std::fill(strides_.begin(), strides_.end(), 0);
        std::copy(shape.begin(), shape.end(), shape_.begin());
        calculateSizeAndStrides();
    }
 
    NDArray(T* data, const ShapeArray& shape)
        : data_(const_cast<void*>(reinterpret_cast<const void*>(data))),
          dims_(0),
          elemSize_(sizeof(T))
    {
        std::fill(shape_.begin(), shape_.end(), 0);
        std::fill(strides_.begin(), strides_.end(), 0);
        for (std::size_t i = 0; i < MAX_DIMS; ++i)
        {
            if (shape[i] != 0)
            {
                shape_[i] = shape[i];
                dims_++;
            }
        }
        calculateSizeAndStrides();
    }
 
    NDArray(void* data, const ShapeArray& shape, std::int64_t elemSize)
        : data_(data),
          dims_(0),
          elemSize_(elemSize)
    {
        std::fill(shape_.begin(), shape_.end(), 0);
        std::fill(strides_.begin(), strides_.end(), 0);
        for (std::size_t i = 0; i < MAX_DIMS; ++i)
        {
            if (shape[i] != 0)
            {
                shape_[i] = shape[i];
                dims_++;
            }
        }
        calculateSizeAndStrides();
    }
 
    ~NDArray() = default;
 
    T* data()
    {
        return reinterpret_cast<T*>(data_);
    }
    const T* data() const
    {
        return reinterpret_cast<const T*>(data_);
    }
 
    void setData(void* data)
    {
        data_ = data;
    }
 
    const ShapeArray& shape() const
    {
        return shape_;
    }
    const StridesArray& strides() const
    {
        return strides_;
    }
    std::int64_t dims() const
    {
        return dims_;
    }
    std::int64_t elemSize() const
    {
        return elemSize_;
    }
    std::int64_t size() const
    {
        return size_;
    }
    std::int64_t numElements() const
    {
        return numElements_;
    }
 
    template <typename... Idxs>
    T& operator()(Idxs... idxs)
    {
        return const_cast<T&>(static_cast<const NDArray&>(*this)(idxs...));
    }
 
    template <typename... Idxs>
    const T& operator()(Idxs... idxs) const
    {
        return reinterpret_cast<const T*>(data_)[computeOffset({idxs...})];
    }
 
    // Subndarray access with operator[] - always returns a NDArray
    NDArray<T> operator[](std::int64_t idx)
    {
        auto const_result = static_cast<const NDArray&>(*this)[idx];
        return NDArray<T>(
            const_cast<T*>(const_result.data()),
            const_result.shape());
    }
 
    NDArray<const T> operator[](std::int64_t idx) const
    {
        ShapeArray new_shape;
        std::fill(new_shape.begin(), new_shape.end(), 0);
        std::copy(shape_.begin() + 1, shape_.end(), new_shape.begin());
 
        return NDArray<const T>(
            reinterpret_cast<const T*>(data_) + idx * strides_[0],
            new_shape);
    }
 
    // Implicit conversion to scalar for 0-dimensional ndarrays
    template <typename U = T>
    operator typename std::enable_if_t<std::is_same_v<U, std::remove_const_t<T>>, T&>()
    {
        if (dims_ != 0)
        {
            // Return reference to first element as fallback
            return *reinterpret_cast<T*>(data_);
        }
        return *reinterpret_cast<T*>(data_);
    }
 
    operator const T&() const
    {
        if (dims_ != 0)
        {
            // Return reference to first element as fallback
            return *reinterpret_cast<const T*>(data_);
        }
        return *reinterpret_cast<const T*>(data_);
    }
 
    // Explicit scalar() method for clearer intent
    T& scalar()
    {
        return const_cast<T&>(static_cast<const NDArray&>(*this).scalar());
    }
 
    const T& scalar() const
    {
        if (dims_ != 0)
        {
            // Return reference to first element as fallback
            return *reinterpret_cast<const T*>(data_);
        }
        return *reinterpret_cast<const T*>(data_);
    }
 
    // Multi-dimensional access with bounds checking
    template <typename... Idxs>
    T& at(Idxs... idxs)
    {
        return const_cast<T&>(static_cast<const NDArray&>(*this).at(idxs...));
    }
 
    template <typename... Idxs>
    const T& at(Idxs... idxs) const
    {
        std::array<std::int64_t, MAX_DIMS> indices = {idxs...};
        if (sizeof...(Idxs) != static_cast<std::size_t>(dims_))
        {
            // Return reference to first element as fallback
            return *reinterpret_cast<const T*>(data_);
        }
        for (std::size_t i = 0; i < indices.size(); ++i)
        {
            if (indices[i] < 0 || indices[i] >= shape_[i])
            {
                // Return reference to first element as fallback
                return *reinterpret_cast<const T*>(data_);
            }
        }
        return operator()(idxs...);
    }
 
private:
    void* data_;
    ShapeArray shape_;
    StridesArray strides_;
    std::int64_t dims_;
    std::int64_t size_;
    std::int64_t elemSize_;
    std::int64_t numElements_;
 
    void calculateSizeAndStrides()
    {
        // Calculate size
        size_ = elemSize_;
        for (std::int64_t i = 0; i < dims_; ++i)
        {
            size_ *= shape_[i];
        }
        numElements_ = size_ / elemSize_;
 
        // Calculate strides
        std::int64_t stride = 1;
        for (std::int64_t i = dims_ - 1; i >= 0; --i)
        {
            strides_[i] = stride;
            stride *= shape_[i];
        }
    }
 
    std::int64_t computeOffset(const std::array<std::int64_t, MAX_DIMS>& idxs) const
    {
        std::int64_t offset = 0;
        for (std::int64_t i = 0; i < dims_; ++i)
        {
            offset += idxs[i] * strides_[i];
        }
        return offset;
    }
};
 
}  // namespace mu