Handling Lack of Support for Cupy Structured Arrays

[fstrug]

Problem

Support for structured arrays in cupy has been an open request since 2019 cupy/cupy#2031. While it is possible to create a structured cupy.ndarray, once it is created you cannot interact with it in the same way as a cupy.ndarray. It is possible to copy back to the host as a numpy.ndarray and get the expected behavior.

zeros_cupy = cp.zeros(shape=5, dtype=np.dtype([('index', '<i8'), ('tag', '<i4')]))
# index = zeros_cupy["index"] # ValueError: invalid literal for int() with base 10: 'index'

zeros_cupy_tonumpy = cupy.asnumpy(zeros_cupy)
index = zeros_cupy["index"]

In awkward, one can create an analogous awkward array with the cuda backend by calling

ak_zeros_fromnumpy = ak.Array(np.zeros(shape=5, dtype=np.dtype([('index', '<i8'), ('tag', '<i4')])))
ak_zeros_fromnumpy_tocuda = ak.to_backend(ak_zeros_numpy, backend = "cuda") 

while calling

ak_zeros_cupy = ak.Array(cp.zeros(shape=5, dtype=np.dtype([('index', '<i8'), ('tag', '<i4')])))

returns an error TypeError: unsupported dtype: dtype([('index', '<i8'), ('tag', '<i4')]).

Potential Workaround

Since it is unclear what cupy will do in regards to this issue, it may be useful to provide a work around.

import cupy as cp
import numpy as np
import awkward as ak

def get_RecordForm(array):
    dtypes = zeros_cupy.dtype
    fields = array.dtype.fields

    form = ak.forms.RecordForm(
        [ak.forms.NumpyForm(str(value[0]), form_key = "node{}".format(i+1)) 
         for i, value in enumerate(fields.values())],
        fields.keys(),
        form_key = "node0",
    )
    return form

def get_Container(array, form):
    container = {}
    fields = array.dtype.fields
    reshaped = array.view("uint8").reshape(-1,array.dtype.itemsize)
    tracker = 0
    for i, key in enumerate(zeros_cupy.dtype.fields.values()):

        form_key = form.contents[i].form_key
        data = reshaped[:,tracker:tracker+key[0].itemsize].view(key[0])
        tracker += key[0].itemsize
        container[f"{form_key}-data"] = data
    return(container)
        
        

def structured_cupy_toak(array):
    length = array.shape[0]
    form = get_RecordForm(array)
    container = get_Container(array, form)
    return ak.from_buffers(form, length, container, backend = "cuda")

zeros_cupy = cp.zeros(shape=5, dtype=np.dtype([('index', '<i8'), ('tag', '<i4')]))
ak_zeros_cupy = structured_cupy_toak(zeros_cupy)
print(ak_zeros_cupy)

This appears to give the same results as

zeros_numpy = np.zeros(shape=5, dtype=np.dtype([('index', '<i8'), ('tag', '<i4')]))
ak_zeros_numpy = ak.Array(zeros_numpy)
ak_zeros_cuda_fromnumpy = ak.to_backend(ak_zeros_numpy, backend = "cuda")
print(ak_zeros_cuda_fromnumpy)

but you get CUDARuntimeError: cudaErrorMisalignedAddress: misaligned address when you try to operate on the awkward.Array constructed directly from the cupy.ndarray slices. It would be nice to avoid bouncing the buffer back to the CPU if possible here, but that would require making contiguous copies from the slices? What are your thoughts here?

[fstrug]

@nsmith- @lgray

[ianna]

@fstrug - could you please try to

import cupy as cp
import numpy as np
import awkward as ak

def get_RecordForm(array):
    fields = array.dtype.fields

    form = ak.forms.RecordForm(
        [ak.forms.NumpyForm(str(value[0]), form_key=f"node{i+1}") 
         for i, value in enumerate(fields.values())],
        fields.keys(),
        form_key="node0",
    )
    return form

def get_Container(array, form):
    container = {}
    fields = array.dtype.fields
    reshaped = cp.ascontiguousarray(array.view("uint8")).reshape(-1, array.dtype.itemsize)
    tracker = 0

    for i, key in enumerate(fields.values()):
        form_key = form.contents[i].form_key
        data = reshaped[:, tracker:tracker + key[0].itemsize].view(key[0])
        tracker += key[0].itemsize
        container[f"{form_key}-data"] = cp.ascontiguousarray(data)  # Enforce alignment

    return container

def structured_cupy_to_ak(array):
    length = array.shape[0]
    form = get_RecordForm(array)
    container = get_Container(array, form)
    return ak.from_buffers(form, length, container, backend="cuda")

[fstrug]

It fails at cp.ascontiguousarray(data): CUDADriverError: CUDA_ERROR_MISALIGNED_ADDRESS. This is for cupy version 13.4.1.

[fstrug]

It seems that we must call cp.ascontiguousarray() before taking an alternate view on the slice of reshaped. Below works for me as expected and we get the contiguous buffers needed.

import cupy as cp
import numpy as np
import awkward as ak

def get_RecordForm(array):
    fields = array.dtype.fields

    form = ak.forms.RecordForm(
        [ak.forms.NumpyForm(str(value[0]), form_key=f"node{i+1}") 
         for i, value in enumerate(fields.values())],
        fields.keys(),
        form_key="node0",
    )
    return form

def get_Container(array, form):
    container = {}
    fields = array.dtype.fields
    reshaped = cp.ascontiguousarray(array.view("uint8")).reshape(-1, array.dtype.itemsize)
    tracker = 0

    for i, key in enumerate(fields.values()):
        form_key = form.contents[i].form_key
        data = cp.ascontiguousarray(reshaped[:, tracker:tracker + key[0].itemsize]).view(key[0]) # Enforce alignment
        tracker += key[0].itemsize
        container[f"{form_key}-data"] = data  
    return container

def structured_cupy_to_ak(array):
    length = array.shape[0]
    form = get_RecordForm(array)
    container = get_Container(array, form)
    return ak.from_buffers(form, length, container, backend="cuda")

[ianna]

indeed, .view() slicing can lead to misaligned addresses...