#pragma once
|
|
#if defined(__AVX__) && !defined(__NVCC__) && \
|
(defined(__x86_64__) || defined(_M_X64) || defined(__i386__))
|
#define CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
#include <immintrin.h>
|
#endif
|
#include <c10/util/Half.h>
|
|
namespace caffe2 {
|
|
namespace internal {
|
|
// The following functions inside internal namespace are inlined because they
|
// are performance critical.
|
|
template <typename T>
|
static inline void adagrad_update_base_inlined(
|
int N,
|
const T* w,
|
const float* g,
|
const T* h,
|
T* nw,
|
T* nh,
|
float decay,
|
float epsilon,
|
float lr) {
|
for (auto i = 0; i < N; ++i) {
|
float gi = g[i];
|
float hi = decay * h[i] + gi * gi;
|
nh[i] = hi;
|
nw[i] = w[i] + lr * gi / (std::sqrt(hi) + epsilon);
|
}
|
}
|
|
// version with prefetching
|
// TODO(msmelyan)
|
// Crux of the computation is computing a / (sqrt(b) + epsilon),
|
// where a and b are vectors and epislon is very small (eg., 10^-5) and does not
|
// change. Today it's computed using two vector sqrt and vector divide simd
|
// instructions. It is slow. We can take advantage of existing fast vector
|
// VRSQRTPS instruction that computes approximate reciprocals of square roots
|
// of the vector. It is 6x faster than vsrt and vdiv combinations. Since the
|
// addition of epislon is just done to avoid division by zero, we approximate a
|
// / (sqrt(b) + epsilon) by a / (sqrt(b + sqrt(epsilon)) If we do that, we can
|
// use VRSQRTPS instead now. VRSQRTPS is not very accurate. Specifically, for
|
// the test on random numbers between 0.1 and 1 the absolute error was about
|
// 10^-3 compared to using slower but more accurate combination of vsqrt and
|
// vdiv. Extend Marat's function with more NR iterations to get more accuracy
|
// for training
|
// TODO(msmelyan)
|
// explore streaming stores, but need to have unique indices (deduplication)
|
inline void adagrad_update_prefetch_inlined(
|
int N,
|
const float* w,
|
#ifdef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
const float* w_n, // prefetch ptr
|
#else
|
const float* /* unused */,
|
#endif
|
|
const float* g,
|
|
const float* h,
|
#ifdef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
const float* h_n, // prefetch ptr
|
#else
|
const float* /* unused */,
|
#endif
|
|
float* nw,
|
#ifdef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
float* nw_n, // prefetch ptr
|
#else
|
float* /* unused */,
|
#endif
|
|
float* nh,
|
#ifdef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
float* nh_n, // prefetch ptr
|
#else
|
float* /* unused */,
|
#endif
|
|
float epsilon,
|
float lr) {
|
auto i = 0;
|
|
#ifdef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
constexpr int kSize = 8;
|
for (; i + kSize <= N; i += kSize) {
|
_mm_prefetch(reinterpret_cast<const char*>(&w_n[i]), _MM_HINT_T0);
|
_mm_prefetch(reinterpret_cast<const char*>(&h_n[i]), _MM_HINT_T0);
|
_mm_prefetch(reinterpret_cast<const char*>(&nw_n[i]), _MM_HINT_T0);
|
_mm_prefetch(reinterpret_cast<const char*>(&nh_n[i]), _MM_HINT_T0);
|
|
__m256 gi = _mm256_loadu_ps(g + i);
|
__m256 hi = _mm256_loadu_ps(h + i);
|
__m256 wi = _mm256_loadu_ps(w + i);
|
|
__m256 nhi = _mm256_add_ps(hi, _mm256_mul_ps(gi, gi));
|
_mm256_storeu_ps(nh + i, nhi);
|
__m256 vtmp = _mm256_div_ps(
|
gi, _mm256_add_ps(_mm256_sqrt_ps(nhi), _mm256_set1_ps(epsilon)));
|
_mm256_storeu_ps(
|
nw + i, _mm256_add_ps(wi, _mm256_mul_ps(_mm256_set1_ps(lr), vtmp)));
|
}
|
#endif
|
|
adagrad_update_base_inlined(
|
N - i, w + i, g + i, h + i, nw + i, nh + i, 1.0f, epsilon, lr);
|
}
|
|
inline void rowwise_adagrad_update_inlined(
|
int N,
|
float* w,
|
#ifdef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
float* w_n, // prefetch ptr
|
#else
|
float* /* unused */,
|
#endif
|
|
const float* g,
|
|
float* h,
|
#ifdef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
float* h_n, // prefetch ptr
|
#else
|
float* /* unused */,
|
#endif
|
|
float epsilon,
|
float lr) {
|
auto i = 0;
|
|
#ifdef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
constexpr int kSize = 8;
|
_mm_prefetch(reinterpret_cast<const char*>(h_n), _MM_HINT_T0);
|
__m256 partial_sum = _mm256_setzero_ps();
|
for (; i + kSize <= N; i += kSize) {
|
__m256 gi = _mm256_loadu_ps(g + i);
|
partial_sum = _mm256_add_ps(partial_sum, _mm256_mul_ps(gi, gi));
|
}
|
// Reduce sum to 1 value
|
__m256 partial_sum_2 = _mm256_hadd_ps(partial_sum, partial_sum);
|
__m256 partial_sum_3 = _mm256_hadd_ps(partial_sum_2, partial_sum_2);
|
float final_sum = _mm_cvtss_f32(_mm256_castps256_ps128(partial_sum_3)) +
|
_mm_cvtss_f32(_mm256_extractf128_ps(partial_sum_3, 1));
|
#else
|
float final_sum = 0.0f;
|
#endif
|
|
for (; i < N; ++i) {
|
final_sum += g[i] * g[i];
|
}
|
final_sum /= N;
|
|
float hi = *h = *h + final_sum;
|
float float_step = lr / (std::sqrt(hi) + epsilon);
|
|
i = 0;
|
#ifdef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
__m256 step = _mm256_set1_ps(float_step);
|
|
for (i = 0; i + kSize <= N; i += kSize) {
|
_mm_prefetch(reinterpret_cast<const char*>(&w_n[i]), _MM_HINT_T0);
|
|
__m256 gi = _mm256_loadu_ps(g + i);
|
__m256 wi = _mm256_loadu_ps(w + i);
|
|
_mm256_storeu_ps(w + i, _mm256_add_ps(wi, _mm256_mul_ps(gi, step)));
|
}
|
#endif
|
|
for (; i < N; ++i) {
|
float gi = g[i];
|
w[i] = w[i] + gi * float_step;
|
}
|
}
|
|
} // namespace internal
|
|
// version with prefetching
|
// TODO(msmelyan)
|
// Crux of the computation is computing a / (sqrt(b) + epsilon),
|
// where a and b are vectors and epislon is very small (eg., 10^-5) and does not
|
// change. Today it's computed using two vector sqrt and vector divide simd
|
// instructions. It is slow. We can take advantage of existing fast vector
|
// VRSQRTPS instruction that computes approximate reciprocals of square roots
|
// of the vector. It is 6x faster than vsrt and vdiv combinations. Since the
|
// addition of epislon is just done to avoid division by zero, we approximate a
|
// / (sqrt(b) + epsilon) by a / (sqrt(b + sqrt(epsilon)) If we do that, we can
|
// use VRSQRTPS instead now. VRSQRTPS is not very accurate. Specifically, for
|
// the test on random numbers between 0.1 and 1 the absolute error was about
|
// 10^-3 compared to using slower but more accurate combination of vsqrt and
|
// vdiv. Extend Marat's function with more NR iterations to get more accuracy
|
// for training
|
// TODO(msmelyan)
|
// explore streaming stores, but need to have inuque indices (deduplication)
|
void adagrad_update_prefetch(
|
int N,
|
const float* w,
|
const float* w_n, // prefetch ptr
|
|
const float* g,
|
|
const float* h,
|
const float* h_n, // prefetch ptr
|
|
float* nw,
|
float* nw_n, // prefetch ptr
|
|
float* nh,
|
float* nh_n, // prefetch ptr
|
|
float epsilon,
|
float lr);
|
|
// Version with prefetching for embeddings and
|
// momentum using fp16
|
void adagrad_fp16_update_prefetch(
|
int N,
|
const at::Half* w,
|
const at::Half* w_n, // prefetch ptr
|
const float* g,
|
const at::Half* h,
|
const at::Half* h_n, // prefetch ptr
|
at::Half* nw,
|
at::Half* nw_n, // prefetch ptr
|
at::Half* nh,
|
at::Half* nh_n, // prefetch ptr
|
float epsilon,
|
float lr);
|
|
void rowwise_adagrad_update(
|
int N,
|
float* w,
|
float* w_n, // prefetch ptr
|
|
const float* g,
|
|
float* h,
|
float* h_n, // prefetch ptr
|
|
float epsilon,
|
float lr);
|
|
// version without prefetching
|
void adagrad_update(
|
int N,
|
const float* w,
|
const float* g,
|
const float* h,
|
float* nw,
|
float* nh,
|
float epsilon,
|
float decay,
|
float lr);
|
|
/**
|
* @return num_rows if succeeds otherwise return the row idx where we pass
|
* the boundary of param_size
|
*/
|
template <typename SIndex>
|
int sparse_adagrad(
|
int num_rows, // number of rows reading
|
int block_size, // number of parameters per rows
|
std::uint64_t param_size, // total number of parameters
|
const float* w, // input parameters
|
const float* g, // input gradients
|
const float* h, // input momentums
|
const SIndex* indices, // indices of each row
|
float* nw, // output parameters
|
float* nh, // output momentums
|
float epsilon,
|
float lr);
|
|
#define SPARSE_ADAGRAD_SPECIALIZATION(SIndex, ISA) \
|
int sparse_adagrad_##SIndex##__##ISA( \
|
int num_rows, \
|
int block_size, \
|
std::uint64_t param_size, \
|
const float* w, \
|
const float* g, \
|
const float* h, \
|
const SIndex* indices, \
|
float* nw, \
|
float* nh, \
|
float epsilon, \
|
float lr) { \
|
for (int i = 0; i < num_rows; ++i) { \
|
std::uint64_t idx = indices[i]; \
|
auto offsetI = i * block_size; \
|
auto offsetIdx = idx * block_size; \
|
\
|
if (block_size + offsetIdx > param_size) { \
|
return i; \
|
} \
|
\
|
if (block_size == 1) { \
|
float gi = g[i]; \
|
float hi = nh[idx] = h[idx] + gi * gi; \
|
nw[idx] = w[idx] + lr * gi / (std::sqrt(hi) + epsilon); \
|
} else { \
|
const int prefdist_T0 = 16; \
|
int i_pref = (i < num_rows - prefdist_T0) ? i + prefdist_T0 : i; \
|
std::uint64_t idx_pref = indices[i_pref]; \
|
\
|
adagrad_update_prefetch__##ISA( \
|
block_size, \
|
w + offsetIdx, \
|
&w[idx_pref * block_size], \
|
g + offsetI, \
|
h + offsetIdx, \
|
&h[idx_pref * block_size], \
|
nw + offsetIdx, \
|
&nw[idx_pref * block_size], \
|
nh + offsetIdx, \
|
&nh[idx_pref * block_size], \
|
epsilon, \
|
lr); \
|
} \
|
} \
|
return num_rows; \
|
};
|
|
} // namespace caffe2
|
|
#ifdef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
#undef CAFFE2_PERFKERNELS_ADAGRAD_H_USE_INTRINSIC
|
#endif
|
