happyz
/
gemma.cpp
mirror of https://github.com/google/gemma.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #77 from enum-class:rmsnorm 

PiperOrigin-RevId: 614642173
This commit is contained in:
Copybara-Service 2024-03-11 06:25:57 -07:00
parent e577198fc0 6712f07ee7
commit a9be065c28
1 changed files with 14 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
ops.h
View File

@ -340,11 +340,21 @@ static HWY_NOINLINE HWY_MAYBE_UNUSED float Dot(const float* HWY_RESTRICT a,
// = Dot(a, a, size), but that is not allowed due to HWY_RESTRICT. // = Dot(a, a, size), but that is not allowed due to HWY_RESTRICT.
static HWY_NOINLINE HWY_MAYBE_UNUSED float SquaredL2( static HWY_NOINLINE HWY_MAYBE_UNUSED float SquaredL2(
const float* HWY_RESTRICT a, size_t size) { const float* HWY_RESTRICT a, size_t size) {
float total = 0.f; const hn::ScalableTag<float> d;
for (size_t i = 0; i < size; ++i) { const size_t N = hn::Lanes(d);
total += a[i] * a[i]; HWY_DASSERT(size >= 2 * N);
HWY_DASSERT(size % (2 * N) == 0);
auto sum0 = hn::Zero(d);
auto sum1 = hn::Zero(d);
for (size_t i = 0; i <= size - 2 * N; i += 2 * N) {
const auto a0 = LoadU(d, a + i);
sum0 = MulAdd(a0, a0, sum0);
const auto a1 = LoadU(d, a + i + N);
sum1 = MulAdd(a1, a1, sum1);
} }
return total;
return ReduceSum(d, Add(sum0, sum1));
} }
static HWY_NOINLINE HWY_MAYBE_UNUSED void RMSNorm( static HWY_NOINLINE HWY_MAYBE_UNUSED void RMSNorm(