Vectorize SGD optimizer step, gradient merging, and gradient normalization loops using SIMD

Author: FFMG

include/neuralnetwork/common/simd_utils.h

@@ -1364,5 +1364,157 @@ class simd
 #endif
     scalar_add_vectors(x, y, n, j);
   }
+
+  // Scalar fallback for scale_vector
+  inline static void scalar_scale_vector(double* y, const double scale, size_t n, size_t start = 0) noexcept
+  {
+    for (size_t j = start; j < n; ++j)
+    {
+      y[j] *= scale;
+    }
+  }
+
+  // Vector-scalar multiplication (y *= scale)
+  inline static void scale_vector(double* y, const double scale, size_t n) noexcept
+  {
+    MYODDWEB_PROFILE_FUNCTION("simd");
+    size_t j = 0;
+#ifdef SIMD_AVX2_ENABLED
+    __m256d vec_scale = _mm256_set1_pd(scale);
+    for (; j + 3 < n; j += 4)
+    {
+      __m256d vec_y = _mm256_loadu_pd(y + j);
+      vec_y = _mm256_mul_pd(vec_y, vec_scale);
+      _mm256_storeu_pd(y + j, vec_y);
+    }
+#endif
+    scalar_scale_vector(y, scale, n, j);
+  }
+
+  // Scalar fallback for sgd_step
+  inline static void scalar_sgd_step(
+    double* values,
+    double* grads,
+    double* velocities,
+    const double* decays,
+    double momentum,
+    double lr,
+    double clipping_scale,
+    bool is_bias,
+    size_t n,
+    size_t start = 0) noexcept
+  {
+    for (size_t i = start; i < n; ++i)
+    {
+      double grad = grads[i] * clipping_scale;
+      if (!is_bias && decays != nullptr && decays[i] > 0.0)
+      {
+        grad += decays[i] * values[i];
+      }
+      double v = momentum * velocities[i] + grad;
+      values[i] -= lr * v;
+      velocities[i] = v;
+      grads[i] = grad;
+    }
+  }
+
+  // Vectorized SGD step
+  inline static void sgd_step(
+    double* values,
+    double* grads,
+    double* velocities,
+    const double* decays,
+    double momentum,
+    double lr,
+    double clipping_scale,
+    bool is_bias,
+    size_t n) noexcept
+  {
+    MYODDWEB_PROFILE_FUNCTION("simd");
+    size_t j = 0;
+#ifdef SIMD_AVX2_ENABLED
+    __m256d vec_clip = _mm256_set1_pd(clipping_scale);
+    __m256d vec_momentum = _mm256_set1_pd(momentum);
+    __m256d vec_lr = _mm256_set1_pd(lr);
+
+    for (; j + 3 < n; j += 4)
+    {
+      __m256d g = _mm256_loadu_pd(&grads[j]);
+      __m256d cur_w = _mm256_loadu_pd(&values[j]);
+      __m256d cur_v = _mm256_loadu_pd(&velocities[j]);
+
+      __m256d grad = _mm256_mul_pd(g, vec_clip);
+
+      if (!is_bias && decays != nullptr)
+      {
+        __m256d d = _mm256_loadu_pd(&decays[j]);
+#ifdef SIMD_FMA_ENABLED
+        grad = _mm256_fmadd_pd(d, cur_w, grad);
+#else
+        grad = _mm256_add_pd(grad, _mm256_mul_pd(d, cur_w));
+#endif
+      }
+
+#ifdef SIMD_FMA_ENABLED
+      __m256d next_v = _mm256_fmadd_pd(vec_momentum, cur_v, grad);
+#else
+      __m256d next_v = _mm256_add_pd(_mm256_mul_pd(vec_momentum, cur_v), grad);
+#endif
+
+      __m256d next_w = _mm256_sub_pd(cur_w, _mm256_mul_pd(vec_lr, next_v));
+
+      _mm256_storeu_pd(&velocities[j], next_v);
+      _mm256_storeu_pd(&values[j], next_w);
+      _mm256_storeu_pd(&grads[j], grad);
+    }
+#endif
+    scalar_sgd_step(values, grads, velocities, decays, momentum, lr, clipping_scale, is_bias, n, j);
+  }
+
+  // Scalar fallback for none_step
+  inline static void scalar_none_step(
+    double* values,
+    double* grads,
+    double lr,
+    double clipping_scale,
+    size_t n,
+    size_t start = 0) noexcept
+  {
+    for (size_t i = start; i < n; ++i)
+    {
+      double grad = grads[i] * clipping_scale;
+      values[i] -= lr * grad;
+      grads[i] = grad;
+    }
+  }
+
+  // Vectorized None step (plain SGD without momentum)
+  inline static void none_step(
+    double* values,
+    double* grads,
+    double lr,
+    double clipping_scale,
+    size_t n) noexcept
+  {
+    MYODDWEB_PROFILE_FUNCTION("simd");
+    size_t j = 0;
+#ifdef SIMD_AVX2_ENABLED
+    __m256d vec_clip = _mm256_set1_pd(clipping_scale);
+    __m256d vec_lr = _mm256_set1_pd(lr);
+
+    for (; j + 3 < n; j += 4)
+    {
+      __m256d g = _mm256_loadu_pd(&grads[j]);
+      __m256d cur_w = _mm256_loadu_pd(&values[j]);
+
+      __m256d grad = _mm256_mul_pd(g, vec_clip);
+      __m256d next_w = _mm256_sub_pd(cur_w, _mm256_mul_pd(vec_lr, grad));
+
+      _mm256_storeu_pd(&values[j], next_w);
+      _mm256_storeu_pd(&grads[j], grad);
+    }
+#endif
+    scalar_none_step(values, grads, lr, clipping_scale, n, j);
+  }
 };
 } // namespace myoddweb::nn

include/neuralnetwork/layers/elmanrnnlayer.cpp

@@ -875,38 +875,20 @@ void ElmanRNNLayer::calculate_and_store_gradients(const std::vector<GradientsAnd
   zero_gradients();
   for (unsigned int t = 0; t < num_threads; ++t)
   {
-    for (size_t i = 0; i < _w_grads.size(); ++i)
-    {
-      _w_grads[i] += _thread_w_grads[t][i];
-    }
-    for (size_t i = 0; i < _rw_grads.size(); ++i)
-    {
-      _rw_grads[i] += _thread_rw_grads[t][i];
-    }
+    simd::add_vectors(_thread_w_grads[t].data(), _w_grads.data(), _w_grads.size());
+    simd::add_vectors(_thread_rw_grads[t].data(), _rw_grads.data(), _rw_grads.size());
     if (has_bias())
     {
-      for (size_t i = 0; i < _b_grads.size(); ++i)
-      {
-        _b_grads[i] += _thread_b_grads[t][i];
-      }
+      simd::add_vectors(_thread_b_grads[t].data(), _b_grads.data(), _b_grads.size());
     }
   }
 
   const double inv_batch = 1.0 / static_cast<double>(batch_size);
-  for (double& x : _w_grads)
-  {
-    x *= inv_batch;
-  }
-  for (double& x : _rw_grads)
-  {
-    x *= inv_batch;
-  }
+  simd::scale_vector(_w_grads.data(), inv_batch, _w_grads.size());
+  simd::scale_vector(_rw_grads.data(), inv_batch, _rw_grads.size());
   if (has_bias())
   {
-    for (double& x : _b_grads)
-    {
-      x *= inv_batch;
-    }
+    simd::scale_vector(_b_grads.data(), inv_batch, _b_grads.size());
   }
 }
 

include/neuralnetwork/layers/fflayer.cpp

@@ -700,30 +700,18 @@ void FFLayer::calculate_and_store_gradients(const std::vector<GradientsAndOutput
 
   for (unsigned int t = 0; t < num_threads; ++t)
   {
-    for (size_t i = 0; i < _w_grads.size(); ++i)
-    {
-      _w_grads[i] += _thread_w_grads[t][i];
-    }
+    simd::add_vectors(_thread_w_grads[t].data(), _w_grads.data(), _w_grads.size());
     if (has_bias())
     {
-      for (size_t i = 0; i < _b_grads.size(); ++i)
-      {
-        _b_grads[i] += _thread_b_grads[t][i];
-      }
+      simd::add_vectors(_thread_b_grads[t].data(), _b_grads.data(), _b_grads.size());
     }
   }
 
   const double inv_batch = 1.0 / static_cast<double>(batch_size);
-  for (double& grad : _w_grads)
-  {
-    grad *= inv_batch;
-  }
+  simd::scale_vector(_w_grads.data(), inv_batch, _w_grads.size());
   if (has_bias())
   {
-    for (double& grad : _b_grads)
-    {
-      grad *= inv_batch;
-    }
+    simd::scale_vector(_b_grads.data(), inv_batch, _b_grads.size());
   }
 }
 

include/neuralnetwork/layers/grurnnlayer.cpp

@@ -1487,37 +1487,28 @@ void GRURNNLayer::calculate_and_store_gradients(
   zero_gradients();
   for (unsigned int t = 0; t < num_threads; ++t)
   {
-    for (size_t i = 0; i < _w_grads.size(); ++i)
-    {
-      _w_grads[i] += _thread_w_grads[t][i];
-      _z_w_grads[i] += _thread_z_w_grads[t][i];
-      _r_w_grads[i] += _thread_r_w_grads[t][i];
-    }
-    for (size_t i = 0; i < _rw_grads.size(); ++i)
-    {
-      _rw_grads[i] += _thread_rw_grads[t][i];
-      _z_rw_grads[i] += _thread_z_rw_grads[t][i];
-      _r_rw_grads[i] += _thread_r_rw_grads[t][i];
-    }
+    simd::add_vectors(_thread_w_grads[t].data(), _w_grads.data(), _w_grads.size());
+    simd::add_vectors(_thread_z_w_grads[t].data(), _z_w_grads.data(), _z_w_grads.size());
+    simd::add_vectors(_thread_r_w_grads[t].data(), _r_w_grads.data(), _r_w_grads.size());
+
+    simd::add_vectors(_thread_rw_grads[t].data(), _rw_grads.data(), _rw_grads.size());
+    simd::add_vectors(_thread_z_rw_grads[t].data(), _z_rw_grads.data(), _z_rw_grads.size());
+    simd::add_vectors(_thread_r_rw_grads[t].data(), _r_rw_grads.data(), _r_rw_grads.size());
+
     if (has_bias())
     {
-      for (size_t i = 0; i < _b_grads.size(); ++i)
-      {
-        _b_grads[i] += _thread_b_grads[t][i];
-        _z_b_grads[i] += _thread_z_b_grads[t][i];
-        _r_b_grads[i] += _thread_r_b_grads[t][i];
-      }
+      simd::add_vectors(_thread_b_grads[t].data(), _b_grads.data(), _b_grads.size());
+      simd::add_vectors(_thread_z_b_grads[t].data(), _z_b_grads.data(), _z_b_grads.size());
+      simd::add_vectors(_thread_r_b_grads[t].data(), _r_b_grads.data(), _r_b_grads.size());
     }
   }
 
   const double denom = static_cast<double>(batch_size);
+  const double inv_batch = 1.0 / denom;
 
-  const auto normalize = [&denom](std::vector<double>& grads)
+  const auto normalize = [inv_batch](std::vector<double>& grads)
   {
-    for (double& g : grads)
-    {
-      g /= denom;
-    }
+    simd::scale_vector(grads.data(), inv_batch, grads.size());
   };
 
   normalize(_w_grads);

include/neuralnetwork/layers/layer.cpp

@@ -531,28 +531,14 @@ void Layer::apply_update_to_vector(
   switch (optimiser_type)
   {
   case OptimiserType::None:
-    for (size_t i = 0; i < n; ++i)
-    {
-      double grad = grads[i] * clipping_scale;
-      values[i] -= learning_rate * grad;
-      grads[i] = grad;
-    }
+    simd::none_step(values.data(), grads.data(), learning_rate, clipping_scale, n);
     break;
 
   case OptimiserType::SGD:
   {
-    const auto& momentum = get_momentum();
-    for (size_t i = 0; i < n; ++i)
-    {
-      double grad = grads[i] * clipping_scale;
-      if (!is_bias && i < decays.size() && decays[i] > 0.0)
-      {
-        grad += decays[i] * values[i];
-      }
-      velocities[i] = momentum * velocities[i] + grad;
-      values[i] -= learning_rate * velocities[i];
-      grads[i] = grad;
-    }
+    const double momentum = get_momentum();
+    const double* decay_ptr = (!is_bias && decays.size() >= n) ? decays.data() : nullptr;
+    simd::sgd_step(values.data(), grads.data(), velocities.data(), decay_ptr, momentum, learning_rate, clipping_scale, is_bias, n);
   }
   break;
 

include/neuralnetwork/layers/lstmlayer.cpp

@@ -952,36 +952,26 @@ const std::vector<GradientsAndOutputs>& batch_gradients_and_outputs, const std::
   zero_gradients();
   for (unsigned int t = 0; t < num_threads; ++t)
   {
-    for (size_t i = 0; i < _w_grads.size(); ++i)
-    {
-      _w_grads[i] += _thread_w_grads[t][i];
-      _f_w_grads[i] += _thread_f_w_grads[t][i];
-      _i_w_grads[i] += _thread_i_w_grads[t][i];
-      _o_w_grads[i] += _thread_o_w_grads[t][i];
-    }
-    for (size_t i = 0; i < _rw_grads.size(); ++i)
-    {
-      _rw_grads[i] += _thread_rw_grads[t][i];
-      _f_rw_grads[i] += _thread_f_rw_grads[t][i];
-      _i_rw_grads[i] += _thread_i_rw_grads[t][i];
-      _o_rw_grads[i] += _thread_o_rw_grads[t][i];
-    }
-    for (size_t i = 0; i < _b_grads.size(); ++i)
-    {
-      _b_grads[i] += _thread_b_grads[t][i];
-      _f_b_grads[i] += _thread_f_b_grads[t][i];
-      _i_b_grads[i] += _thread_i_b_grads[t][i];
-      _o_b_grads[i] += _thread_o_b_grads[t][i];
-    }
+    simd::add_vectors(_thread_w_grads[t].data(), _w_grads.data(), _w_grads.size());
+    simd::add_vectors(_thread_f_w_grads[t].data(), _f_w_grads.data(), _f_w_grads.size());
+    simd::add_vectors(_thread_i_w_grads[t].data(), _i_w_grads.data(), _i_w_grads.size());
+    simd::add_vectors(_thread_o_w_grads[t].data(), _o_w_grads.data(), _o_w_grads.size());
+
+    simd::add_vectors(_thread_rw_grads[t].data(), _rw_grads.data(), _rw_grads.size());
+    simd::add_vectors(_thread_f_rw_grads[t].data(), _f_rw_grads.data(), _f_rw_grads.size());
+    simd::add_vectors(_thread_i_rw_grads[t].data(), _i_rw_grads.data(), _i_rw_grads.size());
+    simd::add_vectors(_thread_o_rw_grads[t].data(), _o_rw_grads.data(), _o_rw_grads.size());
+
+    simd::add_vectors(_thread_b_grads[t].data(), _b_grads.data(), _b_grads.size());
+    simd::add_vectors(_thread_f_b_grads[t].data(), _f_b_grads.data(), _f_b_grads.size());
+    simd::add_vectors(_thread_i_b_grads[t].data(), _i_b_grads.data(), _i_b_grads.size());
+    simd::add_vectors(_thread_o_b_grads[t].data(), _o_b_grads.data(), _o_b_grads.size());
   }
 
   const double inv_batch = 1.0 / static_cast<double>(batch_size);
-  auto norm = [&inv_batch](std::vector<double>& v)
+  auto norm = [inv_batch](std::vector<double>& v)
   {
-    for (double& x : v)
-    {
-      x *= inv_batch;
-    }
+    simd::scale_vector(v.data(), inv_batch, v.size());
   };
   norm(_w_grads);
   norm(_b_grads);