diff options
| author | Stan Seibert <stan@mtrr.org> | 2011-12-29 21:48:05 -0500 |
|---|---|---|
| committer | tlatorre <tlatorre@uchicago.edu> | 2021-05-09 08:42:38 -0700 |
| commit | 4cc95716febf04dcbb926d056db3ffa4f8c3ab4a (patch) | |
| tree | 80a42c9c0bc0fc6638b30492ce4c44c38849ecc5 | |
| parent | 56718834a085b6dd12cc2cf66336e057d5363990 (diff) | |
| download | chroma-4cc95716febf04dcbb926d056db3ffa4f8c3ab4a.tar.gz chroma-4cc95716febf04dcbb926d056db3ffa4f8c3ab4a.tar.bz2 chroma-4cc95716febf04dcbb926d056db3ffa4f8c3ab4a.zip | |
Allow make_gpu_struct to pack structs containing the CUDA vector types.
| -rw-r--r-- | chroma/gpu/tools.py | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/chroma/gpu/tools.py b/chroma/gpu/tools.py index 38d67a1..98df7ce 100644 --- a/chroma/gpu/tools.py +++ b/chroma/gpu/tools.py @@ -142,6 +142,8 @@ def create_cuda_context(device_id=None): return context +vec_dtypes = set([ x for x in ga.vec.__dict__.values() if type(x) == np.dtype ]) + def make_gpu_struct(size, members): struct = cuda.mem_alloc(size) @@ -158,7 +160,7 @@ def make_gpu_struct(size, members): cuda.memcpy_htod(int(struct)+i, np.intp(int(member))) i += 8 - elif np.isscalar(member): + elif np.isscalar(member) or hasattr(member, 'dtype') and member.dtype in vec_dtypes: cuda.memcpy_htod(int(struct)+i, member) i += member.nbytes else: |