THCThrustAllocator.cuh

#ifndef THC_THRUST_ALLOCATOR_INC
#define THC_THRUST_ALLOCATOR_INC

#include <cstddef>

/// Allocator for Thrust to re-route its internal device allocations
/// to the THC allocator
class THCThrustAllocator {
 public:
  typedef char value_type;

  THCThrustAllocator(THCState* state)
      : state_(state) {
  }

  ~THCThrustAllocator() {
  }

  char* allocate(std::ptrdiff_t size) {
    return static_cast<char*>(THCudaMalloc(state_, size));
  }

  void deallocate(char* p, size_t size) {
    THCudaFree(state_, p);
  }

 private:
  THCState* state_;
};

#endif // THC_THRUST_ALLOCATOR_INC