How to combine thrust comparisons using placeholders? - cuda

Question

I have a device_vector of float values A of size N. I also have a float value V for comparison. Depending on an input value I need to extract indices of A for which the values are > OR < OR == V.
I use the following code but it seems cumbersome. Is there a more concise way to do it?
void detect_indices_lesser_greater_equal_to_value(thrust::device_vector<float> S, float value,
int criterion, thrust::device_vector<int>& indices)
{
int N=S.size();
int size=N;
if(criterion==0) // criterion =0 => equal
{
thrust::device_vector<int>::iterator end = thrust::copy_if(thrust::device,thrust::make_counting_iterator(0),
thrust::make_counting_iterator(N),
S.begin(),
indices.begin(),
thrust::placeholders::_1 == value);
size = end-indices.begin();
}
if(criterion==1) // criterion =1 => less
{
thrust::device_vector<int>::iterator end = thrust::copy_if(thrust::device,thrust::make_counting_iterator(0),
thrust::make_counting_iterator(N),
S.begin(),
indices.begin(),
thrust::placeholders::_1 < value);
size = end-indices.begin();
}
if(criterion==2) // criterion =2 => greater
{
thrust::device_vector<int>::iterator end = thrust::copy_if(thrust::device,thrust::make_counting_iterator(0),
thrust::make_counting_iterator(N),
S.begin(),
indices.begin(),
thrust::placeholders::_1 > value);
size = end-indices.begin();
}
indices.resize(size);
}

Answer 1

This can be done with two thrust::partition operations. Partitioning is pretty simple: everything that results in a true predicate is moved to the left side of the input vector. Everything else is moved to the right. Here's a simple example:
$ cat t22.cu
#include <thrust/partition.h>
#include <thrust/copy.h>
#include <thrust/device_vector.h>
typedef float mt;
using namespace thrust::placeholders;
int main(){
const mt pval = 4;
mt data[] = {1,3,7,4,5,2,4,3,9};
const int ds = sizeof(data)/sizeof(data[0]);
thrust::device_vector<mt> d(data, data+ds);
auto end1 = thrust::partition(d.begin(), d.end(), _1<pval);
auto end2 = thrust::partition(end1, d.end(), _1==pval);
std::cout << "less than pval:" << std::endl;
thrust::copy(d.begin(), end1, std::ostream_iterator<mt>(std::cout,","));
std::cout << std::endl << "equal to pval:" << std::endl;
thrust::copy(end1, end2, std::ostream_iterator<mt>(std::cout,","));
std::cout << std::endl << "greater than pval:" << std::endl;
thrust::copy(end2, d.end(), std::ostream_iterator<mt>(std::cout,","));
std::cout << std::endl;
}
$ nvcc -o t22 t22.cu
$ ./t22
less than pval:
1,3,2,3,
equal to pval:
4,4,
greater than pval:
7,5,9,
$
If you require that the ordering in the 3 resultant sub-vectors be the same as the original input ordering, you could use the thrust::stable_partition variant.
(Note that in your question you refer to float quantities, but your example code uses <int> iterators. However the above code can work with either by modifying the typedef).