123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335 |
- #ifndef CERES_INTERNAL_CUDA_STREAMED_BUFFER_H_
- #define CERES_INTERNAL_CUDA_STREAMED_BUFFER_H_
- #include "ceres/internal/config.h"
- #ifndef CERES_NO_CUDA
- #include "ceres/cuda_buffer.h"
- namespace ceres::internal {
- template <typename T>
- class CERES_NO_EXPORT CudaStreamedBuffer {
- public:
-
-
-
-
-
-
-
-
-
-
-
- static constexpr int kNumBatches = 2;
-
-
-
-
-
- CudaStreamedBuffer(ContextImpl* context, const int max_buffer_size)
- : kValuesPerBatch(max_buffer_size / kNumBatches),
- context_(context),
- values_gpu_(context, kValuesPerBatch * kNumBatches) {
- static_assert(ContextImpl::kNumCudaStreams >= kNumBatches);
- CHECK_GE(max_buffer_size, kNumBatches);
-
-
-
- CHECK_EQ(cudaSuccess,
- cudaHostAlloc(&values_cpu_pinned_,
- sizeof(T) * kValuesPerBatch * kNumBatches,
- cudaHostAllocWriteCombined));
- for (auto& e : copy_finished_) {
- CHECK_EQ(cudaSuccess,
- cudaEventCreateWithFlags(&e, cudaEventDisableTiming));
- }
- }
- CudaStreamedBuffer(const CudaStreamedBuffer&) = delete;
- ~CudaStreamedBuffer() {
- CHECK_EQ(cudaSuccess, cudaFreeHost(values_cpu_pinned_));
- for (auto& e : copy_finished_) {
- CHECK_EQ(cudaSuccess, cudaEventDestroy(e));
- }
- }
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
- template <typename Fun>
- void CopyToGpu(const T* from, const int num_values, Fun&& callback) {
-
-
-
- CHECK_EQ(cudaSuccess, cudaStreamSynchronize(context_->DefaultStream()));
-
-
-
-
- const bool copy_to_pinned_memory = MemoryTypeResultsInSynchronousCopy(from);
- T* batch_values_gpu[kNumBatches];
- T* batch_values_cpu[kNumBatches];
- auto streams = context_->streams_;
- for (int i = 0; i < kNumBatches; ++i) {
- batch_values_gpu[i] = values_gpu_.data() + kValuesPerBatch * i;
- batch_values_cpu[i] = values_cpu_pinned_ + kValuesPerBatch * i;
- }
- int batch_id = 0;
- for (int offset = 0; offset < num_values; offset += kValuesPerBatch) {
- const int num_values_batch =
- std::min(num_values - offset, kValuesPerBatch);
- const T* batch_from = from + offset;
- T* batch_to = batch_values_gpu[batch_id];
- auto stream = streams[batch_id];
- auto copy_finished = copy_finished_[batch_id];
- if (copy_to_pinned_memory) {
-
-
- CHECK_EQ(cudaSuccess, cudaEventSynchronize(copy_finished));
- std::copy_n(batch_from, num_values_batch, batch_values_cpu[batch_id]);
- batch_from = batch_values_cpu[batch_id];
- }
- CHECK_EQ(cudaSuccess,
- cudaMemcpyAsync(batch_to,
- batch_from,
- sizeof(T) * num_values_batch,
- cudaMemcpyHostToDevice,
- stream));
- if (copy_to_pinned_memory) {
-
-
-
-
-
-
-
- CHECK_EQ(cudaSuccess, cudaEventRecord(copy_finished, stream));
- }
- callback(batch_to, num_values_batch, offset, stream);
- batch_id = (batch_id + 1) % kNumBatches;
- }
-
- for (int i = 0; i < kNumBatches; ++i) {
- CHECK_EQ(cudaSuccess, cudaStreamSynchronize(streams[i]));
- }
- }
- private:
-
-
-
- static bool MemoryTypeResultsInSynchronousCopy(const void* ptr) {
- cudaPointerAttributes attributes;
- auto status = cudaPointerGetAttributes(&attributes, ptr);
- #if CUDART_VERSION < 11000
-
-
- if (status == cudaErrorInvalidValue) {
- return true;
- }
- #endif
- CHECK_EQ(status, cudaSuccess);
-
- CHECK_NE(attributes.type, cudaMemoryTypeDevice);
-
-
-
-
-
- return attributes.type == cudaMemoryTypeUnregistered;
- }
- const int kValuesPerBatch;
- ContextImpl* context_ = nullptr;
- CudaBuffer<T> values_gpu_;
- T* values_cpu_pinned_ = nullptr;
- cudaEvent_t copy_finished_[kNumBatches] = {nullptr};
- };
- }
- #endif
- #endif
|