ceph/src/arrow/cpp/src/arrow/gpu/cuda_benchmark.cc

   1 // Licensed to the Apache Software Foundation (ASF) under one
   2 // or more contributor license agreements.  See the NOTICE file
   3 // distributed with this work for additional information
   4 // regarding copyright ownership.  The ASF licenses this file
   5 // to you under the Apache License, Version 2.0 (the
   6 // "License"); you may not use this file except in compliance
   7 // with the License.  You may obtain a copy of the License at
   8 //
   9 //   http://www.apache.org/licenses/LICENSE-2.0
  10 //
  11 // Unless required by applicable law or agreed to in writing,
  12 // software distributed under the License is distributed on an
  13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 // KIND, either express or implied.  See the License for the
  15 // specific language governing permissions and limitations
  16 // under the License.
  17
  18 #include "benchmark/benchmark.h"
  19
  20 #include <cstdint>
  21 #include <memory>
  22 #include <vector>
  23
  24 #include "arrow/array.h"
  25 #include "arrow/memory_pool.h"
  26 #include "arrow/testing/gtest_util.h"
  27 #include "arrow/testing/util.h"
  28
  29 #include "arrow/gpu/cuda_api.h"
  30
  31 namespace arrow {
  32 namespace cuda {
  33
  34 constexpr int64_t kGpuNumber = 0;
  35
  36 static void CudaBufferWriterBenchmark(benchmark::State& state, const int64_t total_bytes,
  37                                       const int64_t chunksize,
  38                                       const int64_t buffer_size) {
  39   CudaDeviceManager* manager;
  40   ABORT_NOT_OK(CudaDeviceManager::Instance().Value(&manager));
  41   std::shared_ptr<CudaContext> context;
  42   ABORT_NOT_OK(manager->GetContext(kGpuNumber).Value(&context));
  43
  44   std::shared_ptr<CudaBuffer> device_buffer;
  45   ABORT_NOT_OK(context->Allocate(total_bytes).Value(&device_buffer));
  46   CudaBufferWriter writer(device_buffer);
  47
  48   if (buffer_size > 0) {
  49     ABORT_NOT_OK(writer.SetBufferSize(buffer_size));
  50   }
  51
  52   std::shared_ptr<ResizableBuffer> buffer;
  53   ASSERT_OK(MakeRandomByteBuffer(total_bytes, default_memory_pool(), &buffer));
  54
  55   const uint8_t* host_data = buffer->data();
  56   while (state.KeepRunning()) {
  57     int64_t bytes_written = 0;
  58     ABORT_NOT_OK(writer.Seek(0));
  59     while (bytes_written < total_bytes) {
  60       int64_t bytes_to_write = std::min(chunksize, total_bytes - bytes_written);
  61       ABORT_NOT_OK(writer.Write(host_data + bytes_written, bytes_to_write));
  62       bytes_written += bytes_to_write;
  63     }
  64   }
  65   state.SetBytesProcessed(int64_t(state.iterations()) * total_bytes);
  66 }
  67
  68 static void Writer_Buffered(benchmark::State& state) {
  69   // 128MB
  70   const int64_t kTotalBytes = 1 << 27;
  71
  72   // 8MB
  73   const int64_t kBufferSize = 1 << 23;
  74
  75   CudaBufferWriterBenchmark(state, kTotalBytes, state.range(0), kBufferSize);
  76 }
  77
  78 static void Writer_Unbuffered(benchmark::State& state) {
  79   // 128MB
  80   const int64_t kTotalBytes = 1 << 27;
  81   CudaBufferWriterBenchmark(state, kTotalBytes, state.range(0), 0);
  82 }
  83
  84 // Vary chunk write size from 256 bytes to 64K
  85 BENCHMARK(Writer_Buffered)->RangeMultiplier(16)->Range(1 << 8, 1 << 16)->UseRealTime();
  86
  87 BENCHMARK(Writer_Unbuffered)
  88     ->RangeMultiplier(4)
  89     ->RangeMultiplier(16)
  90     ->Range(1 << 8, 1 << 16)
  91     ->UseRealTime();
  92
  93 }  // namespace cuda
  94 }  // namespace arrow