Function CudaBufferWriterBenchmark

cpp/src/arrow/gpu/cuda_benchmark.cc:36–66

constexpr int64_t kGpuNumber = 0;

static void CudaBufferWriterBenchmark(benchmark::State& state, const int64_t total_bytes,
                                      const int64_t chunksize,
                                      const int64_t buffer_size) {
  CudaDeviceManager* manager = nullptr;
  ABORT_NOT_OK(CudaDeviceManager::Instance().Value(&manager));
  std::shared_ptr<CudaContext> context;
  ABORT_NOT_OK(manager->GetContext(kGpuNumber).Value(&context));

  std::shared_ptr<CudaBuffer> device_buffer;
  ABORT_NOT_OK(context->Allocate(total_bytes).Value(&device_buffer));
  CudaBufferWriter writer(device_buffer);

  if (buffer_size > 0) {
    ABORT_NOT_OK(writer.SetBufferSize(buffer_size));
  }

  std::shared_ptr<ResizableBuffer> buffer;
  ASSERT_OK(MakeRandomByteBuffer(total_bytes, default_memory_pool(), &buffer));

  const uint8_t* host_data = buffer->data();
  while (state.KeepRunning()) {
    int64_t bytes_written = 0;
    ABORT_NOT_OK(writer.Seek(0));
    while (bytes_written < total_bytes) {
      int64_t bytes_to_write = std::min(chunksize, total_bytes - bytes_written);
      ABORT_NOT_OK(writer.Write(host_data + bytes_written, bytes_to_write));
      bytes_written += bytes_to_write;
    }
  }
  state.SetBytesProcessed(int64_t(state.iterations()) * total_bytes);
}

static void Writer_Buffered(benchmark::State& state) {
  // 128MB