Skip to content

Commit

Permalink
WIP: support for making shared arrays from deques
Browse files Browse the repository at this point in the history
iraikov committed Jan 8, 2025
1 parent 111aed0 commit b66d14a
Showing 1 changed file with 236 additions and 19 deletions.
255 changes: 236 additions & 19 deletions python/neuroh5/shared_array.cc
Original file line number Diff line number Diff line change
@@ -20,6 +20,12 @@ template<> struct NumpyTypeMap<uint32_t> { static constexpr int type_num = NPY_U
template<> struct NumpyTypeMap<uint16_t> { static constexpr int type_num = NPY_UINT16; };
template<> struct NumpyTypeMap<uint8_t> { static constexpr int type_num = NPY_UINT8; };

// Structure to hold type information
struct TypeInfo {
int numpy_type;
size_t item_size;
};

// Template class to manage array memory
template<typename T>
class SharedArrayHolder {
@@ -35,29 +41,97 @@ class SharedArrayHolder {
data[i] = static_cast<T>(i);
}
}

// Approach 1: Constructor that takes ownership of vector's data
explicit SharedArrayHolder(std::vector<T>&& vec) : size(vec.size())
{
// Get the vector's allocator
auto alloc = vec.get_allocator();

// Get pointer to vector's data
T* ptr = vec.data();
size = vec.size();

// Release ownership from vector (C++17)
vec.release();

// Take ownership of the pointer
data.reset(ptr);
}

T* get_data() { return data.get(); }
size_t get_size() const { return size; }
};
// Approach 2: Construct array from pointer and size
SharedArrayHolder(T* ptr, size_t n) : data(ptr), size(n) {}

// Structure to hold type information
struct TypeInfo {
int numpy_type;
size_t item_size;
// Approach 3: Constructor that copies vector's buffer
explicit SharedArrayHolder(const std::vector<T>& vec) : size(vec.size())
{
// Allocate new buffer
data = std::make_unique<T[]>(size);
// Use memcpy for POD types (more efficient than std::copy)
if (std::is_trivially_copyable_v<T>)
{
std::memcpy(data.get(), vec.data(), size * sizeof(T));
} else
{
std::copy(vec.begin(), vec.end(), data.get());
}
}

// Constructor for deque input
explicit SharedArrayHolder(const std::deque<T>& deq) : size(deq.size())
{
data = std::make_unique<T[]>(size);

if constexpr (std::is_trivially_copyable_v<T>) {
// For POD types, copy each chunk efficiently
size_t pos = 0;
for (auto it = deq.begin(); it != deq.end(); ++it) {
data[pos++] = *it;
}
} else {
// For non-POD types, use standard copy
std::copy(deq.begin(), deq.end(), data.get());
}
}

// Move constructor for deque
explicit SharedArrayHolder(std::deque<T>&& deq) : size(deq.size())
{
data = std::make_unique<T[]>(size);

if constexpr (std::is_trivially_copyable_v<T>) {
// For POD types, move each element efficiently
size_t pos = 0;
while (!deq.empty()) {
data[pos++] = std::move(deq.front());
deq.pop_front();
}
} else {
// For non-POD types, use move iterator
std::move(deq.begin(), deq.end(), data.get());
deq.clear();
}
}

T* get_data() { return data.get(); }
size_t get_size() const { return size; }
};



// Cleanup function template
template<typename T>
static void array_dealloc(PyObject* capsule) {
static void shared_array_dealloc(PyObject* capsule) {
auto* holder = static_cast<SharedArrayHolder<T>*>(
PyCapsule_GetPointer(capsule, "array_memory")
);
delete holder;
}

// Helper function to create array of specific type
// Helper function to create shared numpy array of specific type
template<typename T>
static PyObject* create_typed_array(Py_ssize_t size) {
static PyObject* create_typed_shared_array(Py_ssize_t size)
{
// Create the array holder
auto holder = new SharedArrayHolder<T>(size);

@@ -68,7 +142,7 @@ static PyObject* create_typed_array(Py_ssize_t size) {
PyObject* capsule = PyCapsule_New(
holder,
"array_memory",
array_dealloc<T>
shared_array_dealloc<T>
);

if (!capsule) {
@@ -94,6 +168,101 @@ static PyObject* create_typed_array(Py_ssize_t size) {
}

// Set the array's base object to our capsule
if (PyArray_SetBaseObject((PyArrayObject*)array, capsule) < 0)
{
Py_DECREF(capsule);
Py_DECREF(array);
return nullptr;
}

return array;
}


// Create shared array from deque (by const reference)
template<typename T>
static PyObject* create_shared_array_from_deque(const std::deque<T>& deq)
{
auto holder = new SharedArrayHolder<T>(deq);

npy_intp dims[1] = {static_cast<npy_intp>(holder->get_size())};

PyObject* capsule = PyCapsule_New(
holder,
"array_memory",
shared_array_dealloc<T>
);

if (!capsule) {
delete holder;
return nullptr;
}

PyObject* array = PyArray_NewFromDescr(
&PyArray_Type,
PyArray_DescrFromType(NumpyTypeMap<T>::type_num),
1,
dims,
nullptr,
holder->get_data(),
NPY_ARRAY_WRITEABLE,
nullptr
);

if (!array)
{
Py_DECREF(capsule);
return nullptr;
}

if (PyArray_SetBaseObject((PyArrayObject*)array, capsule) < 0)
{
Py_DECREF(capsule);
Py_DECREF(array);
return nullptr;
}

return array;
}



// Create shared array from deque (by rvalue reference)
template<typename T>
static PyObject* create_shared_array_from_deque(std::deque<T>&& deq)
{
auto holder = new SharedArrayHolder<T>(std::move(deq));

npy_intp dims[1] = {static_cast<npy_intp>(holder->get_size())};

PyObject* capsule = PyCapsule_New(
holder,
"array_memory",
shared_array_dealloc<T>
);

if (!capsule)
{
delete holder;
return nullptr;
}

PyObject* array = PyArray_NewFromDescr(
&PyArray_Type,
PyArray_DescrFromType(NumpyTypeMap<T>::type_num),
1,
dims,
nullptr,
holder->get_data(),
NPY_ARRAY_WRITEABLE,
nullptr
);

if (!array) {
Py_DECREF(capsule);
return nullptr;
}

if (PyArray_SetBaseObject((PyArrayObject*)array, capsule) < 0) {
Py_DECREF(capsule);
Py_DECREF(array);
@@ -103,7 +272,55 @@ static PyObject* create_typed_array(Py_ssize_t size) {
return array;
}

// Function to create the shared numpy array
// Function to create shared numpy array from existing C++ vector using move semantics
template<typename T>
static PyObject* create_shared_array_from_vector(std::vector<T>&& vec)
{

// Create the array holder using move constructor
auto holder = new SharedArrayHolder<T>(std::move(vec));

npy_intp dims[1] = {static_cast<npy_intp>(holder->get_size())};

PyObject* capsule = PyCapsule_New(
holder,
"array_memory",
shared_array_dealloc<T>
);

if (!capsule) {
delete holder;
return nullptr;
}

PyObject* array = PyArray_NewFromDescr(
&PyArray_Type,
PyArray_DescrFromType(NumpyTypeMap<T>::type_num),
1,
dims,
nullptr,
holder->get_data(),
NPY_ARRAY_WRITEABLE,
nullptr
);

if (!array) {
Py_DECREF(capsule);
return nullptr;
}

if (PyArray_SetBaseObject((PyArrayObject*)array, capsule) < 0) {
Py_DECREF(capsule);
Py_DECREF(array);
return nullptr;
}

return array;
}



// Function to create shared numpy array
static PyObject* create_shared_array(PyObject* self, PyObject* args) {
Py_ssize_t size;
const char* dtype_str;
@@ -116,25 +333,25 @@ static PyObject* create_shared_array(PyObject* self, PyObject* args) {
// Map dtype string to appropriate creation function
std::string dtype(dtype_str);
if (dtype == "float32" || dtype == "float") {
return create_typed_array<float>(size);
return create_typed_shared_array<float>(size);
}
else if (dtype == "int32") {
return create_typed_array<int32_t>(size);
return create_typed_shared_array<int32_t>(size);
}
else if (dtype == "int16") {
return create_typed_array<int16_t>(size);
return create_typed_shared_array<int16_t>(size);
}
else if (dtype == "int8") {
return create_typed_array<int8_t>(size);
return create_typed_shared_array<int8_t>(size);
}
else if (dtype == "uint32") {
return create_typed_array<uint32_t>(size);
return create_typed_shared_array<uint32_t>(size);
}
else if (dtype == "uint16") {
return create_typed_array<uint16_t>(size);
return create_typed_shared_array<uint16_t>(size);
}
else if (dtype == "uint8") {
return create_typed_array<uint8_t>(size);
return create_typed_shared_array<uint8_t>(size);
}

PyErr_SetString(PyExc_ValueError, "Unsupported dtype");

0 comments on commit b66d14a

Please sign in to comment.