Created
March 3, 2023 21:28
-
-
Save jzrake/78d0e47d2dbe90fa11f9743d1efc6011 to your computer and use it in GitHub Desktop.
Implements a simple ndarray class template on CPU/GPU with C++17 standard (Apple Clang 14.0 or nvcc CUDA 11.6)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#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; | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment