softmax/softmax_bench.cpp at main · rexes-dev/softmax · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
#include "softmax.h"

#include <cuda_runtime.h>

#include <iomanip>
#include <iostream>

struct CudaTimer {
  CudaTimer() {
    cudaEventCreate(&start);
    cudaEventCreate(&end);
  }
  ~CudaTimer() {
    cudaEventDestroy(start);
    cudaEventDestroy(end);
  }

  template <typename Func> double time(Func &&f) {
    cudaEventRecord(start);
    f();
    cudaEventRecord(end);
    cudaEventSynchronize(end);
    float ms;
    cudaEventElapsedTime(&ms, start, end);
    double us = static_cast<double>(ms) * 1000;
    return us;
  }

  cudaEvent_t start;
  cudaEvent_t end;
};

int getSMCount() {
  static int sm_cnt = [] {
    int device = 0;
    cudaGetDevice(&device);
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, device);
    return prop.multiProcessorCount;
  }();

  return sm_cnt;
}

std::string to_string(SoftmaxType t) {
  switch (t) {
  case SoftmaxType::Naive:
    return "naive";
  case SoftmaxType::Safe:
    return "safe";
  case SoftmaxType::Online:
    return "online";
  }
  return "unknown";
}

void run_bench(int V, int batch_size, SoftmaxType softmax_type) {
  float *x;
  float *y;
  cudaMalloc(&x, V * batch_size * sizeof(float));
  cudaMemset(x, 0, V * batch_size * sizeof(float));
  cudaMalloc(&y, V * batch_size * sizeof(float));

  CudaTimer timer;
  constexpr int WarmUp = 3;
  for (int i = 0; i < WarmUp; ++i)
    softmax_dev(x, y, V, batch_size, softmax_type);
  cudaDeviceSynchronize();

  const int num_iter = 100;
  const auto us = timer.time([&] {
    for (int i = 0; i < num_iter; ++i)
      softmax_dev(x, y, V, batch_size, softmax_type);
  });
  const auto avg_us = us / num_iter;
  std::cout << std::setw(8) << V << ' ' << std::setw(7) << batch_size << ' '
            << std::setw(7) << to_string(softmax_type) << ' ' << std::setw(10)
            << avg_us << std::endl;
  cudaFree(x);
  cudaFree(y);
}

int main() {
  std::cout << "SM count: " << getSMCount() << std::endl;
  std::cout << '\n'
            << "Sweep: batch_size  (V = 128256)" << '\n'
            << std::setw(8) << "V" << ' ' << std::setw(7) << "batch" << ' '
            << std::setw(7) << "type" << ' ' << std::setw(10) << "avg (us)\n";
  for (int batch_size : {1, 64, getSMCount(), 1024, 4096}) {
    const int V = 128256;
    for (const auto t : {SoftmaxType::Safe, SoftmaxType::Online})
      run_bench(V, batch_size, t);
  }

  std::cout << '\n'
            << "Sweep: V  (batch_size = 4096)" << '\n'
            << std::setw(8) << "V" << ' ' << std::setw(7) << "batch" << ' '
            << std::setw(7) << "type" << ' ' << std::setw(10) << "avg (us)\n";
  for (int V : {32000, 50257, 128256, 151936, 256000}) {
    const int batch_size = 4096;
    for (const auto t : {SoftmaxType::Safe, SoftmaxType::Online})
      run_bench(V, batch_size, t);
  }

  return 0;
}