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

218 lines
8.3 KiB
C++
Raw Blame History

// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <stddef.h>
#include "hwy/base.h"
// Include guard for SIMD code.
#if defined(THIRD_PARTY_GEMMA_CPP_SUM_TOGGLE) == defined(HWY_TARGET_TOGGLE)
#ifdef THIRD_PARTY_GEMMA_CPP_SUM_TOGGLE
#undef THIRD_PARTY_GEMMA_CPP_SUM_TOGGLE
#else
#define THIRD_PARTY_GEMMA_CPP_SUM_TOGGLE
#endif
#include "compression/compress-inl.h"
HWY_BEFORE_NAMESPACE();
namespace gcpp {
namespace HWY_NAMESPACE {
namespace hn = hwy::HWY_NAMESPACE;
// f64 Add, called for f32 inputs promoted to f64. Runs at about half the speed
// of f32 sums.
struct SumKernelDouble {
// Only `CompressTraits<float>` can `Decompress2` to `double`, so both have
// to be `float` in order to have `Raw = double`. Note that if either type is
// smaller than `float`, we may demote the other type from `float` to `BF16`.
template <typename VT, typename WT>
using Raw = hwy::If<IsF32<VT>() && IsF32<WT>(), double, BF16>;
using State = double;
// Raw = double
template <class DRaw, class VR = hn::Vec<DRaw>, HWY_IF_F64_D(DRaw)>
HWY_INLINE void Update4(DRaw /*dd*/, const VR w0, const VR w1, const VR w2,
const VR w3, VR, VR, VR, VR, VR& sum0, VR& sum1,
VR& sum2, VR& sum3, VR&, VR&, VR&, VR&) const {
sum0 = hn::Add(sum0, w0);
sum1 = hn::Add(sum1, w1);
sum2 = hn::Add(sum2, w2);
sum3 = hn::Add(sum3, w3);
}
// Raw = BF16
template <class DRaw, class VR = hn::Vec<DRaw>, HWY_IF_BF16_D(DRaw),
class DS = hn::Repartition<double, DRaw>, class VS = hn::Vec<DS>>
HWY_INLINE void Update4(DRaw dr, const VR w0, const VR w1, const VR w2,
const VR w3, VR, VR, VR, VR, VS& sum0, VS& sum1,
VS& sum2, VS& sum3, VS&, VS&, VS&, VS&) const {
const hn::Repartition<float, DRaw> df;
using VF = hn::Vec<decltype(df)>;
// Reduce to two f32 sums so we can promote them to four f64 vectors.
VF sum02, sum13;
if constexpr (HWY_NATIVE_DOT_BF16) {
const VR k1 = hn::Set(dr, hwy::ConvertScalarTo<BF16>(1.0f));
const VF prod0 = hn::WidenMulPairwiseAdd(df, w0, k1);
const VF prod1 = hn::WidenMulPairwiseAdd(df, w1, k1);
// Fuse WidenMulPairwiseAdd plus Add into ReorderWidenMulAccumulate.
VF unused0 = hn::Zero(df);
VF unused1 = hn::Zero(df);
sum02 = hn::ReorderWidenMulAccumulate(df, w2, k1, prod0, unused0);
sum13 = hn::ReorderWidenMulAccumulate(df, w3, k1, prod1, unused1);
} else {
// If not native, the multiplication costs extra, so convert to f32.
// PromoteEvenTo is cheaper than PromoteUpperTo especially on `SVE`.
const VF fe0 = hn::PromoteEvenTo(df, w0);
const VF fe1 = hn::PromoteEvenTo(df, w1);
const VF fe2 = hn::PromoteEvenTo(df, w2);
const VF fe3 = hn::PromoteEvenTo(df, w3);
const VF fo0 = hn::PromoteOddTo(df, w0);
const VF fo1 = hn::PromoteOddTo(df, w1);
const VF fo2 = hn::PromoteOddTo(df, w2);
const VF fo3 = hn::PromoteOddTo(df, w3);
const VF fe01 = hn::Add(fe0, fe1);
const VF fe23 = hn::Add(fe2, fe3);
const VF fo01 = hn::Add(fo0, fo1);
const VF fo23 = hn::Add(fo2, fo3);
sum02 = hn::Add(fe01, fe23);
sum13 = hn::Add(fo01, fo23);
}
const DS ds;
const VS d0 = hn::PromoteLowerTo(ds, sum02);
const VS d1 = hn::PromoteUpperTo(ds, sum02);
const VS d2 = hn::PromoteLowerTo(ds, sum13);
const VS d3 = hn::PromoteUpperTo(ds, sum13);
sum0 = hn::Add(sum0, d0);
sum1 = hn::Add(sum1, d1);
sum2 = hn::Add(sum2, d2);
sum3 = hn::Add(sum3, d3);
}
// Raw = double
template <class DRaw, class VR = hn::Vec<DRaw>, HWY_IF_F64_D(DRaw)>
HWY_INLINE void Update1(DRaw /*dd*/, const VR w0, const VR v0, VR& sum0,
VR& comp0) const {
sum0 = hn::Add(sum0, w0);
}
// Raw = BF16
template <class DRaw, class VR = hn::Vec<DRaw>, HWY_IF_BF16_D(DRaw),
class DS = hn::Repartition<double, DRaw>, class VS = hn::Vec<DS>>
HWY_INLINE void Update1(DRaw dr, const VR w0, VR, VS& sum0,
VS& extra0) const {
const hn::Repartition<float, DRaw> df;
using VF = hn::Vec<decltype(df)>;
VF f0;
if constexpr (HWY_NATIVE_DOT_BF16) {
const VR k1 = hn::Set(dr, hwy::ConvertScalarTo<BF16>(1.0f));
f0 = hn::WidenMulPairwiseAdd(df, w0, k1);
} else {
const VF fe0 = hn::PromoteEvenTo(df, w0);
const VF fo0 = hn::PromoteOddTo(df, w0);
f0 = hn::Add(fe0, fo0);
}
const DS ds;
const VS d0 = hn::PromoteLowerTo(ds, f0);
const VS d1 = hn::PromoteUpperTo(ds, f0);
sum0 = hn::Add(sum0, d0);
extra0 = hn::Add(extra0, d1);
}
template <class DState, class VS = hn::Vec<DState>>
HWY_INLINE float Reduce(DState dd, VS& sum0, VS& sum1, VS& sum2, VS& sum3,
VS& extra0, VS&, VS&, VS&) const {
// Reduction tree: sum of all accumulators by pairs, then across lanes.
sum0 = hn::Add(sum0, sum1);
sum2 = hn::Add(sum2, sum3);
sum0 = hn::Add(sum0, extra0); // from Update1
sum0 = hn::Add(sum0, sum2);
return static_cast<float>(hn::ReduceSum(dd, sum0));
}
};
// ORO Cascaded Summation, algorithm 6.11 from Handbook of Floating-Point
// Arithmetic. Note that Algorithm 6.7 (KBN) appears erroneous. We use TwoSums
// instead of FastTwoSums because the magnitude of the initial sum is not
// always greater than the next input, and this does actually change the e2e
// generation results. Note that Kahan summation differs in that it first adds
// comp* to w*, so each operation is serially dependent. By contrast, the sum*
// and comp* here have shorter dependency chains.
//
// This about as accurate as SumKernelDouble but slower, hence we only use this
// if f64 is not supported on this target.
struct SumKernelCascaded {
template <typename VT, typename WT>
using Raw = float;
using State = float;
template <class DF, class VF = hn::Vec<DF>, HWY_IF_F32_D(DF)>
HWY_INLINE void Update4(DF df, const VF w0, const VF w1, const VF w2,
const VF w3, VF, VF, VF, VF, VF& sum0, VF& sum1,
VF& sum2, VF& sum3, VF& comp0, VF& comp1, VF& comp2,
VF& comp3) const {
VF serr0, serr1, serr2, serr3;
sum0 = TwoSums(df, sum0, w0, serr0);
sum1 = TwoSums(df, sum1, w1, serr1);
sum2 = TwoSums(df, sum2, w2, serr2);
sum3 = TwoSums(df, sum3, w3, serr3);
comp0 = hn::Add(comp0, serr0);
comp1 = hn::Add(comp1, serr1);
comp2 = hn::Add(comp2, serr2);
comp3 = hn::Add(comp3, serr3);
}
template <class DF, class VF = hn::Vec<DF>, HWY_IF_F32_D(DF)>
HWY_INLINE void Update1(DF df, const VF w0, const VF v0, VF& sum0,
VF& comp0) const {
VF serr0;
sum0 = TwoSums(df, sum0, w0, serr0);
comp0 = hn::Add(comp0, serr0);
}
template <class DF, class VF = hn::Vec<DF>>
HWY_INLINE float Reduce(DF df, VF& sum0, VF& sum1, VF& sum2, VF& sum3,
VF& comp0, VF& comp1, VF& comp2, VF& comp3) const {
// Reduction tree: sum of all accumulators by pairs, then across lanes.
AssimilateCascadedSums(df, sum1, comp1, sum0, comp0);
AssimilateCascadedSums(df, sum3, comp3, sum2, comp2);
AssimilateCascadedSums(df, sum2, comp2, sum0, comp0);
return ReduceCascadedSums(df, sum0, comp0);
}
};
using SumKernelDefault =
hwy::If<HWY_HAVE_FLOAT64, SumKernelDouble, SumKernelCascaded>;
template <class D, typename VT>
HWY_INLINE float Sum(D d, const VT* HWY_RESTRICT vec, size_t num) {
using Raw = hwy::If<HWY_HAVE_FLOAT64, double, float>;
const hn::Repartition<Raw, D> d_raw;
return DecompressAndCall(d_raw, MakeSpan(vec, num), SumKernelDefault());
}
// NOLINTNEXTLINE(google-readability-namespace-comments)
} // namespace HWY_NAMESPACE
} // namespace gcpp
HWY_AFTER_NAMESPACE();
#endif // NOLINT
Reference in New Issue View Git Blame Copy Permalink