Attention body support for blas / eigen / mkl

src/neural/blas/encoder.h

@@ -30,15 +30,22 @@ namespace lczero {
 
 void LayerNorm2DWithSkipConnection(const size_t batch_size,
                                    const size_t channels, float* data,
-                                   const float* skip, const float* gammas,
-                                   const float* betas, float epsilon) {
+                                   const float alpha, const float* skip,
+                                   const float* gammas, const float* betas,
+                                   float epsilon) {
   for (size_t i = 0; i < batch_size; i++) {
 #ifndef USE_ISPC
     // Mean taken in dimension C.
     float mean = 0;
-    for (size_t c = 0; c < channels; ++c) {
-      data[i * channels + c] += skip[i * channels + c];
-      mean += data[i * channels + c];
+    if (skip != nullptr) {
+      for (size_t c = 0; c < channels; ++c) {
+        data[i * channels + c] += alpha * skip[i * channels + c];
+        mean += data[i * channels + c];
+      }
+    } else {
+      for (size_t c = 0; c < channels; ++c) {
+        mean += data[i * channels + c];
+      }
     }
     mean /= channels;
 
@@ -57,9 +64,15 @@ void LayerNorm2DWithSkipConnection(const size_t batch_size,
           betas[c] + gammas[c] * (data[i * channels + c] - mean) * den;
     }
 #else
-    ispc::LayerNorm2DWithSkipConnection(channels, data + i * channels,
-                                        skip + i * channels, gammas, betas,
-                                        epsilon);
+    if (skip != nullptr) {
+      ispc::LayerNorm2DWithSkipConnection(channels, data + i * channels, alpha,
+                                          skip + i * channels, gammas, betas,
+                                          epsilon);
+    } else {
+      ispc::LayerNorm2DWithSkipConnection(channels, data + i * channels, 0.0f,
+                                          nullptr, gammas, betas, epsilon);
+    }
+
 #endif
   }
 }

src/neural/blas/fully_connected_layer.cc

@@ -103,7 +103,9 @@ void FullyConnectedLayer<false>::Forward1D(
                 outputs,            // C
                 (int)output_size);  // ldc, leading rank of C
   }
-  ApplyBias(batch_size, output_size, biases, activation, outputs);
+  if (biases != nullptr) {
+    ApplyBias(batch_size, output_size, biases, activation, outputs);
+  }
 }
 
 template <>
@@ -134,7 +136,9 @@ void FullyConnectedLayer<true>::Forward1D(
             .transpose() *
         ConstEigenMatrixMap<float>(inputs, input_size, batch_size);
   }
-  ApplyBias(batch_size, output_size, biases, activation, outputs);
+  if (biases != nullptr) {
+    ApplyBias(batch_size, output_size, biases, activation, outputs);
+  }
 }
 
 template <>

src/neural/blas/layer_norm.ispc

@@ -18,6 +18,7 @@
 
 export void LayerNorm2DWithSkipConnection(uniform const size_t channels,
                                           uniform float data[],
+                                          const uniform float alpha,
                                           const uniform float skip[],
                                           const uniform float gammas[],
                                           const uniform float betas[],
@@ -27,23 +28,39 @@ export void LayerNorm2DWithSkipConnection(uniform const size_t channels,
     // One pass mean and variance taken in dimension C. Uses shifted variance calculation.
     float imean = 0;
     float ivar = 0;
-    float k = data[0] + skip[0];
-    foreach (c = 0 ... channels) {
-      float t = data[c] + skip[c];
-      data[c] = t;
-      t -= k;
-      imean += t;
-      ivar += t * t;
+    float k = data[0];
+    if (skip != NULL) {
+      k += alpha * skip[0];
+      foreach (c = 0 ... channels) {
+        float t = data[c] + alpha * skip[c];
+        data[c] = t;
+        t -= k;
+        imean += t;
+        ivar += t * t;
+      }
+    } else {
+      foreach (c = 0 ... channels) {
+        float t = data[c];
+        t -= k;
+        imean += t;
+        ivar += t * t;
+      }
     }
     float mean = reduce_add(imean) / channels;
     float var = (reduce_add(ivar) - channels * mean * mean) / channels;
     mean += k;
 #else
   // Mean taken in dimension C.
   float imean = 0;
-  foreach (c = 0 ... channels) {
-    data[c] += skip[c];
-    imean += data[c];
+  if (skip != NULL) {
+    foreach (c = 0 ... channels) {
+      data[c] += alpha * skip[c];
+      imean += data[c];
+    }
+  } else {
+    foreach (c = 0 ... channels) {
+      imean += data[c];
+    }
   }
   float mean = reduce_add(imean) / channels;