Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

1
#include <ATen/ATen.h>
2
#include <ATen/AccumulateType.h>
3
#include <ATen/cuda/CUDAContext.h>
4
#include <ATen/cuda/Exceptions.h>
5
// Another possibility:
6
// #include <torch/all.h>
7

8
#include <assert.h>
9

10
#include "multi_tensor_apply.cuh"
11
#include "type_shim.h"
12

13
#define BLOCK_SIZE 1024
14
#define ILP 4
15

16
typedef enum {
17
  ADAGRAD_MODE_0 = 0, // L2 regularization mode.
18
  ADAGRAD_MODE_1 = 1, // AdamW-style weight decay.
19

20
} adagradMode_t;
21

22
using MATH_T = float;
23

24
template <typename T> struct AdagradFunctor {
25
  __device__ __forceinline__ void
26
  operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<3> &tl,
27
             const float epsilon, const float lr, adagradMode_t mode,
28
             const float weight_decay) {
29
    int tensor_loc = tl.block_to_tensor[blockIdx.x];
30
    int chunk_idx = tl.block_to_chunk[blockIdx.x];
31
    int n = tl.sizes[tensor_loc];
32

33
    T *g = (T *)tl.addresses[0][tensor_loc];
34
    g += chunk_idx * chunk_size;
35

36
    T *p = (T *)tl.addresses[1][tensor_loc];
37
    p += chunk_idx * chunk_size;
38

39
    T *h = (T *)tl.addresses[2][tensor_loc];
40
    h += chunk_idx * chunk_size;
41

42
    n -= chunk_idx * chunk_size;
43

44
    // see note in multi_tensor_scale_kernel.cu
45
    for (int i_start = 0; i_start < n && i_start < chunk_size;
46
         i_start += blockDim.x * ILP) {
47
      MATH_T r_g[ILP];
48
      MATH_T r_p[ILP];
49
      MATH_T r_h[ILP];
50
#pragma unroll
51
      for (int ii = 0; ii < ILP; ii++) {
52
        int i = i_start + threadIdx.x + ii * blockDim.x;
53
        if (i < n && i < chunk_size) {
54
          r_g[ii] = g[i];
55
          r_p[ii] = p[i];
56
          r_h[ii] = h[i];
57
        } else {
58
          r_g[ii] = MATH_T(0);
59
          r_p[ii] = MATH_T(0);
60
          r_h[ii] = MATH_T(0);
61
        }
62
      }
63
#pragma unroll
64
      for (int ii = 0; ii < ILP; ii++) {
65
        if (mode == ADAGRAD_MODE_0) { // L2
66
          r_g[ii] = r_g[ii] + weight_decay * r_p[ii];
67
          r_h[ii] = r_h[ii] + r_g[ii] * r_g[ii];
68
          r_p[ii] = r_p[ii] - lr * (r_g[ii] / (sqrtf(r_h[ii]) + epsilon));
69
        } else { // AdamW-style
70
          r_h[ii] = r_h[ii] + r_g[ii] * r_g[ii];
71
          r_p[ii] = r_p[ii] - lr * (r_g[ii] / (sqrtf(r_h[ii]) + epsilon) + weight_decay * r_p[ii]);
72
        }
73
      }
74
#pragma unroll
75
      for (int ii = 0; ii < ILP; ii++) {
76
        int i = i_start + threadIdx.x + ii * blockDim.x;
77
        if (i < n && i < chunk_size) {
78
          p[i] = r_p[ii];
79
          h[i] = r_h[ii];
80
        }
81
      }
82
    }
83
  }
84
};
85

86
void multi_tensor_adagrad_cuda(
87
    int chunk_size, at::Tensor noop_flag,
88
    std::vector<std::vector<at::Tensor>> tensor_lists, const float lr,
89
    const float epsilon, const int mode, const float weight_decay) {
90
  using namespace at;
91

92
  // Assume single type across p,g,h now
93
  DISPATCH_DOUBLE_FLOAT_AND_HALF(
94
      tensor_lists[0][0].scalar_type(), 0, "adagrad",
95
      multi_tensor_apply<3>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
96
                            AdagradFunctor<scalar_t_0>(), epsilon, lr,
97
                            (adagradMode_t)mode, weight_decay);)
98

99
  AT_CUDA_CHECK(cudaGetLastError());
100
}
101

102