Implements a simple ndarray class template on CPU/GPU with C++17 standard (Apple Clang 14.0 or nvcc CUDA 11.6)

gpu_array_example.cpp

#include <cstdlib>
#include <utility> // swap

#ifdef GPU
#define HD __host__ __device__
#else
#define HD
#endif

using uint = unsigned int;

template<uint Rank>
struct uivec
{
    HD uint& operator[](uint i)
    {
        return data[i];
    }
    HD const uint& operator[](uint i) const
    {
        return data[i];
    }
    uint data[Rank];
};

template<class... Args>
HD auto make_uivec(Args... args)
{
    return uivec<sizeof...(Args)>{uint(args)...};
}

template<uint Rank>
HD uint product(uivec<Rank> t)
{
    uint n = 1;

    for (uint i = 0; i < Rank; ++i)
        n *= t[i];
    return n;
}

template<uint Rank>
HD uint dot(uivec<Rank> t, uivec<Rank> u)
{
    uint n = 0;

    for (uint i = 0; i < Rank; ++i)
        n += t[i] * u[i];
    return n;
}

template<uint Rank>
HD uivec<Rank> strides_row_major(uivec<Rank> shape)
{
    auto result = uivec<Rank>{};

    if constexpr (Rank > 0)
        result[Rank - 1] = 1;

    if constexpr (Rank > 1)
        for (int n = Rank - 2; n >= 0; --n)
            result[n] = result[n + 1] * shape[n + 1];

    return result;
}

template<class ValueType, uint Rank>
class array
{
public:
    using value_type = ValueType;

    array& operator=(const array&) = delete;

    HD array(const array& other)
        : memory(other.memory)
        , owned(false)
        , shape(other.shape)
        , jumps(other.jumps)
    {
    }

    HD array(uivec<Rank> shape)
        : memory(new value_type[product(shape)])
        , owned(true)
        , shape(shape)
        , jumps(strides_row_major(shape))
    {
    }

    HD array(array&& other) noexcept
    {
        memory = other.memory;
        owned = other.owned;
        shape = other.shape;
        jumps = other.jumps;

        other.memory = nullptr;
        other.owned = false;
        other.shape = {0};
        other.jumps = {0};
    }

    HD ~array()
    {
        if (owned)
        {
            delete [] memory;
        }
    }

    HD const value_type& operator[](std::size_t i) const
    {
        return memory[i];
    }

    HD value_type& operator[](std::size_t i)
    {
        return memory[i];
    }

    template<typename... I>
    HD const value_type& operator()(I... i) const
    {
        return memory[dot(make_uivec(i...), jumps)];
    }

    template<typename... I>
    HD value_type& operator()(I... i)
    {
        return memory[dot(make_uivec(i...), jumps)];
    }

    HD const value_type* data() const
    {
        return memory;
    }

    HD value_type* data()
    {
        return memory;
    }

    friend void swap(array& a, array& b) noexcept
    {
        std::swap(a.memory, b.memory);
        std::swap(a.owned, b.owned);
        std::swap(a.shape, b.shape);
        std::swap(a.jumps, b.jumps);
    }

private:
    ValueType* memory = nullptr;
    bool owned = false;
    uivec<Rank> shape = {0};
    uivec<Rank> jumps = {0};
};


#ifdef GPU
template<class F>
__global__ void my_kernel(F f)
{
    f();
}
#endif


int main()
{
    auto a = array<double, 3>(make_uivec(1, 2, 3));
    auto b = array<double, 3>(make_uivec(1, 2, 3));

    swap(a, b);

    b(0, 0, 0) = 2.0;

    auto f = [a] HD () mutable {
        a[0] = 7.0;
    };

    #ifdef GPU
    my_kernel<<<1, 1>>>(f);
    #else
    f();
    #endif

    printf("%lf\n", a(0, 0, 0));
    printf("%lf\n", b(0, 0, 0));

    return 0;
}