Using thrust to perform lookup on device

28 Views Asked by At

I want to use Thrust to make a lookup operation faster (using GPU). I have a relative small lookup table (122 elements) saved as array in a header file for that. The idea is:

  • *float4 molecule_device is the pointer to the array holding the key to look up (in 4th component)
  • after the lookup I want to replace the key with additional information from the lookup table
  • int colour_device* is where I want the looked up information to be stored
  • a parameter (int) to determine what information/column is to be returned

There are three points in particular I'm struggling with:

  1. The array are already malloced and populated on the device. To use the array with thrust I need to cast it to a thrust device pointer. Can I use this pointer in a thrust for each as begin iterator or do I have to take additional steps there?
  2. What combination of zip iterator and tuples do I need to loop over the first array and store the data in both the molecule array (replace key) and the colour array?
  3. How can I add the column integer to the function I defined for thrust?

Below is a basic version I'd like to get working (shortened at some places)

#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <fstream>

#include <thrust/device_vector.h>
#include <thrust/for_each.h>
#include <thrust/tuple.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/transform_reduce.h>

#include <atomdata.cuh>
#include <reader.hpp>

template <typename T>
struct lookup
{
    __host__ __device__
        thrust::tuple<T, T>
        operator()(thrust::tuple<float4 &, int &> a) const
    {
        unsigned int symbol = (int)thrust::get<0>(a).w; // the key
        float x;
        uint y;
        findEntry(symbol, 1, &x, &y); // the 1 needs to be replaced by the column integer
        return thrust::make_tuple(x, y);
    }
};

int main(void)
{
    int numAtoms = 400;
    float4 *molecule;
    float4 *molecule_device;
    uint *colors_device;
    int column = 1;

    molecule = new float4[numAtoms];
    std::string filePath = ".../3i40.pdb";  // path shortend
    readFile(filePath, molecule, numAtoms); // this populates the molecule array; every datapoint is an atom with coordinates (x,y,z) and a radius (w)
    checkCudaErrors(cudaMalloc((void **)&(molecule_device), numAtoms * sizeof(float4)));
    checkCudaErrors(cudaMemcpy(molecule_device, molecule, numAtoms * sizeof(float4), cudaMemcpyHostToDevice));
    checkCudaErrors(cudaMalloc((void **)&(colors_device), numAtoms * sizeof(uint)));

    thrust::device_ptr<float4> dev_molecule_ptr = thrust::device_pointer_cast(molecule_device);
    thrust::device_ptr<int> dev_colour_ptr = thrust::device_pointer_cast(colour);
    auto start_zip = thrust::make_zip_iterator(thrust::make_tuple(dev_molecule_ptr, dev_colour_ptr));

    thrust::for_each_n(start_zip, atom_count, lookup()); // here I'm totally lost
    
    return 0;
}

The header file with the lookup table:

#ifndef ATOMDATA_CUH
#define ATOMDATA_CUH

#ifndef ELEMENTS
#define ELEMENTS 112
#endif

typedef unsigned int uint;

struct entry
{
    uint symbol;

    float rSingleBonds1;
    float rSingleBonds2;
    float rVanDerWaals;

    uint cCorey;
    uint cKoltun;
    uint cJmol;
    uint cRasmolOld;
    uint cRasmolNew;
    uint cPubChem;
};

const entry pse[ELEMENTS] = {
#ifndef ATOMDATA_CUH
#define ATOMDATA_CUH

#include <string>

#ifndef ELEMENTS
#define ELEMENTS 112
#endif

typedef unsigned int uint;

struct entry
{
    uint symbol;

    float rSingleBonds1;
    float rSingleBonds2;
    float rVanDerWaals;

    uint cCorey;
    uint cKoltun;
    uint cJmol;
    uint cRasmolOld;
    uint cRasmolNew;
    uint cPubChem;
};

const entry pse[ELEMENTS] = {
    {2384, .31f, .32f, 1.2f, (uint)0xffffff, (uint)0xffffff, (uint)0xffffff, (uint)0xffffff, (uint)0xffffff, (uint)0x638c8c},
    {5324, .28f, .46f, 1.4f, (uint)0xffc0cb, (uint)0xffc0cb, (uint)0xd9ffff, (uint)0xffc0cb, (uint)0xffc0cb, (uint)0xd593a1},
    // ... remaining lines similar; cut for size purposes ...   
    {6219, .0f, 1.22f, .0f, (uint)0xffc0cb, (uint)0xffc0cb, (uint)0xffc0cb, (uint)0xffb6c1, (uint)0xffb6c1, (uint)0xffb6c1}};

__host__ __device__ void findEntry(uint symbol, uint case, float *radius_out, uint *color_out)
{
    uint n = 0;
    for (uint i = 0; i < ELEMENTS; i++)
    {
        if (symbol == pse[i].symbol)
        {
            *radius_out = pse[n].rVanDerWaals;
            int color;
            // ... some switch statements
            *color_out = color;
            break;
        }
    }
}
#endif

I searched online for similar thrust implementations/problems. Most use some combination of zip iterators and tuples to solve this (hence my code above). However I haven't found an implementation, that changes values in both the first (search) vector as well as the output vector.

The code above is minimal example. I didn't include the population of the array, as it's quite complex and not interesting (for testing any random floats do the trick).

I didn't compile the minimal example because my code setup isn't very flexible; likewise I don't expect any full coded answer and would be happy about any help about my approach.

0

There are 0 best solutions below