Compiler Explorer

Source code

#include <cooperative_groups.h>
#include <thrust/device_vector.h>
#include <thrust/sequence.h>

#include <stdio.h>

namespace cg = cooperative_groups;

struct sentinel {};

struct strided_iterator_t {
  int i_;
  int size_;
  int group_size_;

__device__ int operator*() const { return i_; }
  __device__ void operator++() { i_ += group_size_; }
  __device__ bool operator==(sentinel) const { return i_ >= size_; }
  __device__ bool operator!=(sentinel s) const { return !(*this == s); }
};

struct federated_sequence_t {
  int group_size_;
  int size_;
  int i_;

__device__ strided_iterator_t begin() { return strided_iterator_t{i_, size_, group_size_}; };
  __device__ sentinel end() { return {}; }
};

template<typename ViewFn>
struct view_closure_t;

struct view_closure_base_t {
  template <typename ProducerT,
            typename ConsumerConstructorT>
  __device__ friend auto
  operator|(ProducerT &&producer,
            view_closure_t<ConsumerConstructorT> consumer) {
    return consumer(producer);
  }
};

template <typename ViewFn>
struct view_closure_t
    : view_closure_base_t
    , ViewFn {
  __device__ explicit view_closure_t(ViewFn fn)
    : ViewFn(fn) {
  }
};

template<typename Fun>
__device__ view_closure_t<Fun> make_view_closure(Fun fun) {
  return view_closure_t<Fun>{fun};
}

// Indirection layer: 
//  allows to abstract data access from subsequent algorithms
template <class G>
__device__ federated_sequence_t 
federated_sequence(G&& group, int size) {
  return federated_sequence_t{
    static_cast<int>(group.size()),
    size,
    static_cast<int>(group.thread_rank())
  };
}

template <typename RangeT,
          typename IteratorT>
struct load_iterator_t {
  using i_t = decltype(std::declval<RangeT>().begin());

RangeT rng_;
  IteratorT in_;

i_t i_;

using value_type = typename std::iterator_traits<IteratorT>::value_type;

__device__ value_type& operator*() { return *(in_ + *i_); }
  __device__ void operator++() { ++i_; }
  __device__ bool operator==(sentinel s) const { return i_ == s; }
  __device__ bool operator!=(sentinel s) const { return !(*this == s); }
};

template <class IteratorT, class Rng>
struct load_t {
  IteratorT in_;
  Rng rng_;

__device__ auto begin() {
    return load_iterator_t<Rng, IteratorT>{rng_, in_, rng_.begin()};
  }

__device__ sentinel end() const { return {}; }
};

template <class IteratorT>
struct load_fn_t {
  IteratorT in_;

template <class Rng>
  __device__ load_t<IteratorT, Rng> operator()(Rng rng) const {
    return load_t<IteratorT, Rng>{in_, rng};
  }
};

template <class IteratorT>
__device__ auto load(IteratorT in) {
  return make_view_closure(load_fn_t<IteratorT>{in});
}

template <typename RangeT,
          typename T>
struct filter_iterator_t {
  using i_t = decltype(std::declval<RangeT>().begin());

RangeT rng_;
  T key_;

i_t i_;

__device__ T& operator*() { return *i_; }
  __device__ void operator++() { 
    ++i_; 
    for (; i_ != rng_.end(); ++i_) {
      if (*i_ == key_) {
        break;
      }
    }
  }
  __device__ bool operator==(sentinel s) const { return i_ == s; }
  __device__ bool operator!=(sentinel s) const { return !(*this == s); }
};

template <class T, class Rng>
struct filter_t {
  T key_;
  Rng rng_;

__device__ auto begin() {
    auto i = rng_.begin();

for (; i != rng_.end(); ++i) {
      if (*i == key_) {
        return filter_iterator_t<Rng, T>{rng_, key_, i};
      }
    }

return filter_iterator_t<Rng, T>{rng_, key_, i};
  }

__device__ sentinel end() const { return {}; }
};

template <class T>
struct filter_fn_t {
  T in_;

template <class Rng>
  __device__ filter_t<T, Rng> operator()(Rng rng) const {
    return filter_t<T, Rng>{in_, rng};
  }
};

template <class T>
__device__ auto filter(T key) {
  return make_view_closure(filter_fn_t<T>{key});
}

template <typename RangeT,
          typename T>
struct sync_iterator_t {
  using i_t = decltype(std::declval<RangeT>().begin());

RangeT rng_;
  T oob_;

i_t i_;

__device__ T& operator*() { 
    if (i_ == sentinel{}) {
      return oob_;
    }
    return *i_; 
  }
  __device__ void operator++() {
    if (i_ != sentinel{}) {
      ++i_; 
    }
  }
  __device__ bool operator==(sentinel s) const {
    // TODO Why there's no ballot in thread group?
    return __syncthreads_and(i_ == s);
  }
  __device__ bool operator!=(sentinel s) const { return !(*this == s); }
};

template <class T, class Rng>
struct sync_t {
  T oob_;
  Rng rng_;

__device__ auto begin() { return sync_iterator_t<Rng, T>{rng_, oob_, rng_.begin()}; }
  __device__ sentinel end() const { return {}; }
};

template <class T>
struct sync_fn_t {
  T oob_;

template <class Rng>
  __device__ sync_t<T, Rng> operator()(Rng rng) const {
    return sync_t<T, Rng>{oob_, rng};
  }
};

template <class T>
__device__ auto sync(T oob) {
  return make_view_closure(sync_fn_t<T>{oob});
}

template <class T>
__global__ void range(T key, int size, int* in, int* out) {
  for (int i: federated_sequence(cg::this_thread_block(), size)) {
    out[i] = threadIdx.x;
  }
}

template <class T>
__global__ void for_each(T key, int size, int* in, int* out) {
  auto group = cg::this_thread_block();

for (int i: federated_sequence(group, size) | load(in)) {
    out[i] = threadIdx.x;
  }
}

template <class T>
__global__ void filter(T key, int size, int* in, bool* out) {
  auto group = cg::this_thread_block();

for (int& i: federated_sequence(group, size) | load(in) | filter(key)) {
    out[(int)(&i - in)] = true;
  }
}

template <class T>
__global__ void contains_unique(T key, int size, int* in, bool* out) {
  auto group = cg::this_thread_block();

for (int& i: federated_sequence(group, size) | load(in) | filter(key)) {
    *out = true;
    break;
  }
}

template <class T>
__global__ void contains_sync(T key, int size, int* in, bool* out) {
  auto group = cg::this_thread_block();
  auto oob = INT_MAX;

for (int& i: federated_sequence(group, size) | load(in) | filter(key) | sync(oob)) {
    if (i != oob) {
      *out = true;
    }
    // Just to illustrate that it's possible
    group.sync();
  }
}

int main() {
  const int size = 64;
  const int key = size / 2;

thrust::device_vector<bool> out(1);
  thrust::device_vector<int> in(size);
  thrust::sequence(in.rbegin(), in.rend());

contains_sync<<<1, 32>>>(
      key, size, 
      thrust::raw_pointer_cast(in.data()),
      thrust::raw_pointer_cast(out.data()));
  cudaDeviceSynchronize();

std::cout << out[0] << std::endl;
}