diff --git a/compression/compress-inl.h b/compression/compress-inl.h index 7fd097c..28eb497 100644 --- a/compression/compress-inl.h +++ b/compression/compress-inl.h @@ -386,7 +386,7 @@ struct CompressTraits { size_t num, CompressPerThread& tls, const PackedSpan& packed, const size_t packed_ofs) { - NuqCodec::Enc(df, raw, num, tls.buf, packed, packed_ofs); + NuqCodec::EncInterleaved(df, raw, num, tls.buf, packed, packed_ofs); if (COMPRESS_STATS) { for (size_t i = 0; i < num; ++i) { @@ -396,8 +396,8 @@ struct CompressTraits { const hn::Repartition dbf; const size_t N16 = hn::Lanes(dbf); auto distorted = hwy::AllocateAligned(hwy::RoundUpTo(num, N16)); - NuqCodec::DecompressAndZeroPad(dbf, MakeConst(packed), packed_ofs, - distorted.get(), num); + NuqCodec::DecompressAndZeroPadInterleaved( + dbf, MakeConst(packed), packed_ofs, distorted.get(), num); DistortionStats stats; for (size_t i = 0; i < num; ++i) { stats.Notify(raw[i], hwy::F32FromBF16(distorted[i])); @@ -410,7 +410,7 @@ struct CompressTraits { static HWY_INLINE void Load2(D d, const PackedSpan& packed, const size_t packed_ofs, hn::Vec& raw0, hn::Vec& raw1) { - NuqCodec::Dec2(d, packed, packed_ofs, raw0, raw1); + NuqCodec::Dec2Interleaved(d, packed, packed_ofs, raw0, raw1); } // Store2 is not yet implemented. @@ -419,7 +419,7 @@ struct CompressTraits { static HWY_INLINE void DecompressAndZeroPad( D d, const PackedSpan& packed, const size_t packed_ofs, Raw* raw, const size_t num) { - NuqCodec::DecompressAndZeroPad(d, packed, packed_ofs, raw, num); + NuqCodec::DecompressAndZeroPadInterleaved(d, packed, packed_ofs, raw, num); } }; diff --git a/compression/compress.h b/compression/compress.h index 9050d53..c7067d6 100644 --- a/compression/compress.h +++ b/compression/compress.h @@ -17,6 +17,7 @@ #ifndef THIRD_PARTY_GEMMA_CPP_COMPRESSION_COMPRESS_H_ #define THIRD_PARTY_GEMMA_CPP_COMPRESSION_COMPRESS_H_ +#include "hwy/base.h" #define COMPRESS_STATS 0 #include @@ -134,7 +135,7 @@ class MatPtr { size_t NumElements() const { return num_elements_; } // Returns the number of bytes in the array. - size_t SizeBytes() const { return num_elements_ * element_size_; } + virtual size_t SizeBytes() const { return num_elements_ * element_size_; } // Returns the number of rows in the 2-d array (outer dimension). size_t Rows() const { return rows_; } @@ -240,10 +241,13 @@ class MatPtrT : public MatPtr { return name; } - // Sets the number of elements in the array. For use when the number of - // elements is != rows * cols ONLY. - void SetNumElements(size_t num_elements) { - num_elements_ = CompressedArrayElements(num_elements); + // Returns the number of bytes in the array. Overrides MatPtr::SizeBytes() + // to account for NUQ's differing packed size. + size_t SizeBytes() const override { + if (hwy::IsSame, NuqStream>()) { + return NuqStream::PackedEnd(num_elements_); + } + return num_elements_ * element_size_; } // 2-d Accessor for a specific type but with a dynamic inner dimension. @@ -333,6 +337,12 @@ class MatStorageT : public MatPtrT { // from the current num_elements_ which was set by the constructor from the // rows and cols. void Allocate(size_t num_elements = 0) { + // size_t num_elements = 0; + // TODO: optimize this check or obviate it. + if (hwy::IsSame, NuqStream>()) { + HWY_DASSERT(num_elements == 0); + } + if (num_elements == 0) { num_elements = hwy::DivCeil(this->SizeBytes(), sizeof(MatT)); } else { diff --git a/compression/nuq-inl.h b/compression/nuq-inl.h index 63c4255..7c2327c 100644 --- a/compression/nuq-inl.h +++ b/compression/nuq-inl.h @@ -40,6 +40,7 @@ // After highway.h #include "compression/sfp-inl.h" #include "hwy/contrib/sort/vqsort-inl.h" +#include "hwy/profiler.h" // uses SIMD HWY_BEFORE_NAMESPACE(); namespace gcpp { @@ -529,12 +530,21 @@ class NuqCodec { return (!HWY_HAVE_SCALABLE && du.MaxBytes() >= 32) ? 1 : 2; } + // Offset (in bytes) of a group's table for packed_ofs (in elements) within a + // set of groups. + static constexpr size_t TableByteOffset(size_t packed_ofs) { + const size_t kBytesPerGroup = + (kClusters * sizeof(SfpStream)) + kGroupSize / 2; + return (packed_ofs / kGroupSize) * kBytesPerGroup; + } + // Unpacks `centers` from SFP into bf16 and loads them into one or two vectors // for use by [Two]TableLookups. Returns as u16 because TableLookupLanes might // not be available for bf16. template static HWY_INLINE hn::Vec LoadTable(DU du, const uint8_t* centers, hn::Vec* HWY_RESTRICT tbl1) { + PROFILER_FUNC; // Cap to the table size (kClusters) for decoding SFP - sufficient, and may // be faster than a large vector. const hn::CappedTag d_table; @@ -606,6 +616,81 @@ class NuqCodec { } public: + // Encodes `num` floats from `raw` into `packed`. `packed` points to + // compressed storage and `packed_ofs` indicates the destination offset within + // it, in number of elements. Tables are interleaved with indices (clustered + // elements) to allow for easier unpacking. Returns the total number of + // unused clusters, which is typically zero. + template + static HWY_INLINE size_t EncInterleaved(DF df, const float* HWY_RESTRICT raw, + const size_t num, + NuqStream::ClusterBuf& buf, + const PackedSpan& packed, + size_t packed_ofs) { + const hn::Repartition d16; + const hn::Repartition d8; + using V16 = hn::Vec; + using V8 = hn::Vec; + const size_t N16 = hn::Lanes(d16); + + HWY_ASSERT(packed_ofs % kGroupSize == 0); + + const size_t num_groups = hwy::DivCeil(num, kGroupSize); + // TODO: dynamic resize should be removed; it is no longer necessary as + // interleaved encoding uses only a single buffer of the same size. + buf.Resize(1); + + size_t unused_clusters = 0; + size_t current_offset = packed_ofs; + for (size_t g = 0; g < num_groups; ++g) { + const size_t g_num = HWY_MIN(num - g * kGroupSize, kGroupSize); + const float* HWY_RESTRICT g_in = raw + g * kGroupSize; + + float* HWY_RESTRICT g_centers = buf.centers.get(); + uint16_t* HWY_RESTRICT g_idx = buf.idx.get(); + + unused_clusters += + NuqClustering::ClusterExactL2(df, g_in, g_num, buf, g_centers, g_idx); + + uint8_t* centers = &packed.ptr->byte + TableByteOffset(current_offset); + SfpCodec::Enc(df, buf.centers.get(), kClusters, + reinterpret_cast(centers)); + uint8_t* packed_start = centers + kClusters; + + current_offset += g_num; + + HWY_DASSERT(g_num % (4 * N16) == 0); + + size_t i = 0; + HWY_UNROLL(1) + for (; i < g_num; i += 4 * N16) { + const V16 idx0 = hn::LoadU(d16, g_idx + i + 0 * N16); + const V16 idx1 = hn::LoadU(d16, g_idx + i + 1 * N16); + const V16 idx2 = hn::LoadU(d16, g_idx + i + 2 * N16); + const V16 idx3 = hn::LoadU(d16, g_idx + i + 3 * N16); + const V8 nibbles = + NibbleCodec::OrderedPackU16(d16, idx0, idx1, idx2, idx3); + hn::StoreU(nibbles, d8, packed_start + i / 2); + } + + const size_t remaining = g_num - i; + + HWY_DASSERT(remaining < 4 * N16); + if (HWY_UNLIKELY(remaining != 0)) { + const V16 idx0 = hn::LoadU(d16, g_idx + i + 0 * N16); + const V16 idx1 = hn::LoadU(d16, g_idx + i + 1 * N16); + const V16 idx2 = hn::LoadU(d16, g_idx + i + 2 * N16); + const V16 idx3 = hn::LoadU(d16, g_idx + i + 3 * N16); + const V8 nibbles = + NibbleCodec::OrderedPackU16(d16, idx0, idx1, idx2, idx3); + // i is even, but remaining might not be. + hn::StoreN(nibbles, d8, packed_start + i / 2, + hwy::DivCeil(remaining, 2)); + } + } + return unused_clusters; + } + // Encodes `num` floats from `raw`. `packed` points to compressed storage and // `packed_ofs` indicates the destination offset within it, in units of float // values, for parallel encoding by multiple threads. Returns the total @@ -733,6 +818,8 @@ class NuqCodec { raw1 = BitCast(dbf, c1); } + // TODO(philculliton): Remove non-interleaved function versions now that + // interleaved is working / the default. // Decompresses to two f32 vectors. `packed_ofs` must be a multiple of two // vectors so that we only have to load one group's table. template @@ -765,6 +852,107 @@ class NuqCodec { raw1 = hn::PromoteUpperTo(df, BitCast(dbf, c0)); } + // Decompresses to two bf16 vectors. `packed_ofs` must be a multiple of two + // vectors so that we only have to load one group's table. + template + static HWY_INLINE void Dec2Interleaved( + DBF dbf, const PackedSpan& packed, + const size_t packed_ofs, hn::Vec& raw0, hn::Vec& raw1) { + PROFILER_FUNC; + const hn::RebindToUnsigned d16; + const D8HFromD16 d8h; + using V16 = hn::Vec; + using V8H = hn::Vec; + + const size_t within_group = packed_ofs % kGroupSize; + HWY_DASSERT(within_group % (2 * hn::Lanes(d16)) == 0); + const uint8_t* table = &packed.ptr->byte + TableByteOffset(packed_ofs); + const uint8_t* indices = table + kClusters + hwy::DivCeil(within_group, 2); + + V16 tbl1 = Zero(d16); + const V16 tbl0 = LoadTable(d16, table, &tbl1); + + const V8H nibbles = hn::LoadU(d8h, indices); + + V16 c0, c1; + TableLookups(d16, tbl0, tbl1, nibbles, c0, c1); + raw0 = BitCast(dbf, c0); + raw1 = BitCast(dbf, c1); + } + + // Decompresses to two f32 vectors. `packed_ofs` must be a multiple of two + // vectors so that we only have to load one group's table. + template + static HWY_INLINE void Dec2Interleaved( + DF df, const PackedSpan& packed, const size_t packed_ofs, + hn::Vec& raw0, hn::Vec& raw1) { + const hn::Repartition dbf; + const hn::RebindToUnsigned d16; + const hn::Half> d8q; + using V8Q = hn::Vec; + using V16 = hn::Vec; + + const size_t within_group = packed_ofs % kGroupSize; + HWY_DASSERT(within_group % (2 * hn::Lanes(df)) == 0); + const uint8_t* table = &packed.ptr->byte + TableByteOffset(packed_ofs); + const uint8_t* indices = table + kClusters + hwy::DivCeil(within_group, 2); + + V16 tbl1 = Zero(d16); + const V16 tbl0 = LoadTable(d16, table, &tbl1); + + // The single-vector TableLookups overload only calls OrderedUnpackU16<0>, + // which expects a quarter vector of bytes. + const V8Q nibbles = hn::LoadU(d8q, indices); + + // TODO(janwas): From janwas: on AVX-512 I imagine we can get a + // bit more speed for this function by changing LoadTable to return floats, + // then we could have a single lookup here instead of PromoteUpperTo which + // is not cheap. + const V16 c0 = TableLookups(d16, tbl0, tbl1, nibbles); + raw0 = hn::PromoteLowerTo(df, BitCast(dbf, c0)); + raw1 = hn::PromoteUpperTo(df, BitCast(dbf, c0)); + } + + template > + static HWY_INLINE void DecompressAndZeroPadInterleaved( + D d, const PackedSpan& packed, size_t packed_ofs, + Raw* HWY_RESTRICT raw, size_t num) { + // If unaligned, load elements from the first group and update the args, + // from which we compute new tables/indices below. + size_t current_offset = packed_ofs; + if (size_t within_group = packed_ofs % kGroupSize; within_group != 0) { + const uint8_t* tables = + &packed.ptr->byte + TableByteOffset(current_offset); + const uint8_t* indices = tables + kClusters; + const size_t remaining = HWY_MIN(num, kGroupSize - within_group); + + DecPartialGroup(d, tables, indices, raw, remaining); + packed_ofs += remaining; + current_offset += remaining; + raw += remaining; + num -= remaining; + if (num == 0) return; + } + + HWY_DASSERT(packed_ofs % kGroupSize == 0); + + const size_t num_groups = hwy::DivCeil(num, kGroupSize); + HWY_UNROLL(1) + for (size_t g = 0; g < num_groups - 1; ++g) { + const uint8_t* tables = + &packed.ptr->byte + TableByteOffset(current_offset); + const uint8_t* indices = tables + kClusters; + DecWholeGroup(d, tables, indices, raw + g * kGroupSize); + current_offset += kGroupSize; + } + + const size_t g = num_groups - 1; + const uint8_t* tables = &packed.ptr->byte + TableByteOffset(current_offset); + const uint8_t* indices = tables + kClusters; + DecPartialGroup(d, tables, indices, raw + g * kGroupSize, + num - g * kGroupSize); + } + // Decompresses from `packed`, starting at (any) `packed_ofs`, to (any) `num` // elements in `raw`, then appends `[0, hn::Lanes(d))` zeroes as required to // round `num` up to one vector, if it is not already. @@ -955,6 +1143,7 @@ class NuqCodec { } const size_t remaining = num - i; + HWY_DASSERT(remaining < 4 * NF); if (HWY_UNLIKELY(remaining != 0)) { // i is even, but remaining might not be. diff --git a/compression/nuq_test.cc b/compression/nuq_test.cc index 8cbce6c..2536816 100644 --- a/compression/nuq_test.cc +++ b/compression/nuq_test.cc @@ -277,6 +277,189 @@ struct TestOffset { void TestOffsetBF16() { hn::ForGEVectors<128, TestOffset>()(BF16()); } void TestOffsetF32() { hn::ForGEVectors<128, TestOffset>()(float()); } +// Can encode and decode sub-regions. +struct TestUnalignedOffset { + template + HWY_INLINE void operator()(T /*unused*/, D d) { + const hn::Repartition df; + const size_t total = 10 * kGroupSize; // already padded + + const int num_unaligned_offsets = 4; + const std::array unaligned_offsets = { + 4, kGroupSize + 100, 2 * kGroupSize + 100, 3 * kGroupSize + 100}; + const std::array num = {4, 16, 32, 64}; + + for (int i = 0; i < num_unaligned_offsets; ++i) { + const size_t unaligned_offset = unaligned_offsets[i]; + const size_t num_decompressed = num[i]; + + auto in = hwy::AllocateAligned(total); // Enc() requires f32 + auto dec1 = hwy::AllocateAligned(total); + auto nuq = hwy::AllocateAligned(NuqStream::PackedEnd(total)); + auto dec2 = hwy::AllocateAligned(num_decompressed); + HWY_ASSERT(in && dec1 && dec2 && nuq); + const auto nuq_span = MakeSpan(nuq.get(), total); + + hwy::RandomState rng; + for (size_t i = 0; i < total; ++i) { + in[i] = static_cast(RandomGaussian(rng)); + } + + // Encode + decode everything + NuqStream::ClusterBuf buf; + (void)NuqCodec::Enc(df, in.get(), total, buf, nuq_span, 0); + NuqCodec::DecompressAndZeroPad(d, MakeConst(nuq_span), 0, dec1.get(), + total); + + NuqCodec::DecompressAndZeroPad(d, MakeConst(nuq_span), unaligned_offset, + dec2.get(), num_decompressed); + + for (size_t i = 0; i < num_decompressed; ++i) { + T expected = hwy::ConvertScalarTo(dec1[unaligned_offset + i]); + T actual = hwy::ConvertScalarTo(dec2[i]); + + HWY_ASSERT_EQ(f_should_be_correct, f_might_be_wrong); + } + } + } +}; + +void TestUnalignedOffsetBF16() { + hn::ForGEVectors<128, TestUnalignedOffset>()(BF16()); +} +void TestUnalignedOffsetF32() { + hn::ForGEVectors<128, TestUnalignedOffset>()(float()); +} + +// Can encode and decode sub-regions. +// Uses Dec2Interleaved to decode all elements in the packed buffer, then +// compares against the non-interleaved decode. +struct TestDec2Interleaved { + template + HWY_INLINE void operator()(T /*unused*/, D d) { + const hn::Repartition df; + // incl. partial group to test partial group handling + const size_t total = 1 * kGroupSize + (kGroupSize / 2); + const size_t kMidLen = 2 * kGroupSize; // length of middle piece + + auto in = hwy::AllocateAligned(total); // Enc() requires f32 + auto dec0 = hwy::AllocateAligned(total); + auto dec1 = hwy::AllocateAligned(total); + auto dec2 = hwy::AllocateAligned(kMidLen); + auto nuq = hwy::AllocateAligned(NuqStream::PackedEnd(total)); + HWY_ASSERT(in && dec0 && dec1 && dec2 && nuq); + const auto nuq_span = MakeSpan(nuq.get(), total); + + hwy::RandomState rng; + for (size_t i = 0; i < total; ++i) { + in[i] = static_cast(RandomGaussian(rng)); + } + + // Non-interleaved encode + decode for comparison + NuqStream::ClusterBuf buf0; + (void)NuqCodec::Enc(df, in.get(), total, buf0, nuq_span, 0); + NuqCodec::DecompressAndZeroPad(d, MakeConst(nuq_span), 0, dec0.get(), + total); + + // Encode + decode everything + NuqStream::ClusterBuf buf; + (void)NuqCodec::EncInterleaved(df, in.get(), total, buf, nuq_span, 0); + + using V = hn::Vec; + const size_t N = Lanes(d); + + for (size_t i = 0; i < total; i += 2 * N) { + V f0, f1; + NuqCodec::Dec2Interleaved(d, MakeConst(nuq_span), i, f0, f1); + + hn::StoreU(f0, d, dec1.get() + i + 0 * N); + hn::StoreU(f1, d, dec1.get() + i + 1 * N); + } + + for (size_t i = 0; i < total; ++i) { + if (dec0[i] != dec1[i]) { + fprintf(stderr, "dec0[%zu] = %g, dec1[%zu] = %g\n", i, (float)dec0[i], + i, (float)dec1[i]); + } + + HWY_ASSERT(dec0[i] == dec1[i]); + } + } +}; + +void TestDec2BF16Interleaved() { + hn::ForGEVectors<128, TestDec2Interleaved>()(BF16()); +} +void TestDec2F32Interleaved() { + hn::ForGEVectors<128, TestDec2Interleaved>()(float()); +} + +// Can encode and decode sub-regions. +struct TestOffsetInterleaved { + template + HWY_INLINE void operator()(T /*unused*/, D d) { + const hn::Repartition df; + const size_t total = + 10 * kGroupSize + + (kGroupSize / + 2); // adding a partial group to test... partial group handling! + const size_t kMidLen = 2 * kGroupSize; // length of middle piece + + auto in = hwy::AllocateAligned(total); // Enc() requires f32 + auto dec0 = hwy::AllocateAligned(total); + auto dec1 = hwy::AllocateAligned(total); + auto dec2 = hwy::AllocateAligned(kMidLen); + auto nuq = hwy::AllocateAligned(NuqStream::PackedEnd(total)); + HWY_ASSERT(in && dec0 && dec1 && dec2 && nuq); + const auto nuq_span = MakeSpan(nuq.get(), total); + + hwy::RandomState rng; + for (size_t i = 0; i < total; ++i) { + in[i] = static_cast(RandomGaussian(rng)); + } + + // Non-interleaved encode + decode for comparison + NuqStream::ClusterBuf buf0; + (void)NuqCodec::Enc(df, in.get(), total, buf0, nuq_span, 0); + NuqCodec::DecompressAndZeroPad(d, MakeConst(nuq_span), 0, dec0.get(), + total); + + // Encode + decode everything + NuqStream::ClusterBuf buf; + (void)NuqCodec::EncInterleaved(df, in.get(), total, buf, nuq_span, 0); + NuqCodec::DecompressAndZeroPadInterleaved(d, MakeConst(nuq_span), 0, + dec1.get(), total); + + for (size_t i = 0; i < total; ++i) { + if (dec0[i] != dec1[i]) { + fprintf(stderr, "dec0[%zu] = %g, dec1[%zu] = %g\n", i, (float)dec0[i], + i, (float)dec1[i]); + } + + HWY_ASSERT(dec0[i] == dec1[i]); + } + + // Overwrite middle with first inputs + const size_t offset = 5 * kGroupSize; + (void)NuqCodec::EncInterleaved(df, in.get(), kMidLen, buf, nuq_span, + offset); + + // Decoded middle now matches previously decoded first + NuqCodec::DecompressAndZeroPadInterleaved(d, MakeConst(nuq_span), offset, + dec2.get(), kMidLen); + for (size_t i = 0; i < kMidLen; ++i) { + HWY_ASSERT(dec1[i] == dec2[i]); + } + } +}; + +void TestOffsetBF16Interleaved() { + hn::ForGEVectors<128, TestOffsetInterleaved>()(BF16()); +} +void TestOffsetF32Interleaved() { + hn::ForGEVectors<128, TestOffsetInterleaved>()(float()); +} + struct TestNibble { template HWY_INLINE void operator()(T /*unused*/, D d) { @@ -409,6 +592,12 @@ HWY_EXPORT_AND_TEST_P(NuqTest, TestAllRamp); HWY_EXPORT_AND_TEST_P(NuqTest, TestAllNormal); HWY_EXPORT_AND_TEST_P(NuqTest, TestOffsetBF16); HWY_EXPORT_AND_TEST_P(NuqTest, TestOffsetF32); +HWY_EXPORT_AND_TEST_P(NuqTest, TestDec2BF16Interleaved); +HWY_EXPORT_AND_TEST_P(NuqTest, TestDec2F32Interleaved); +HWY_EXPORT_AND_TEST_P(NuqTest, TestUnalignedOffsetBF16); +HWY_EXPORT_AND_TEST_P(NuqTest, TestUnalignedOffsetF32); +HWY_EXPORT_AND_TEST_P(NuqTest, TestOffsetBF16Interleaved); +HWY_EXPORT_AND_TEST_P(NuqTest, TestOffsetF32Interleaved); HWY_EXPORT_AND_TEST_P(NuqTest, TestAllNibble); HWY_EXPORT_AND_TEST_P(NuqTest, TestEncDecBF16); HWY_EXPORT_AND_TEST_P(NuqTest, TestEncDecF32); diff --git a/compression/python/compression_clif_aux.cc b/compression/python/compression_clif_aux.cc index ba91781..c466ef1 100644 --- a/compression/python/compression_clif_aux.cc +++ b/compression/python/compression_clif_aux.cc @@ -55,8 +55,7 @@ class SbsWriterImpl : public WriterInterface { template void AllocateAndCompress(const std::string& name, absl::Span weights) { - const size_t num_packed = CompressedArrayElements(weights.size()); - MatPtrT storage(name, 1, num_packed); + MatPtrT storage(name, 1, weights.size()); model_memory_.push_back(storage); model_memory_.back().Allocate(); storage.SetPtr(model_memory_.back()); diff --git a/compression/shared.h b/compression/shared.h index 40c8f1c..81f5954 100644 --- a/compression/shared.h +++ b/compression/shared.h @@ -159,16 +159,20 @@ struct NuqStream { // Returns offset of packed indices from the start of the stream. This matches // the (padded) total table size because table entries are bytes. + // TODO(philculliton): Remove when removing non-interleaved functions. static constexpr size_t PackedStart(size_t capacity) { // Round up to avoid cache-line splits when loading indices. No effect on // size as long as capacity / kGroupSize is a multiple of 4. - return hwy::RoundUpTo(hwy::DivCeil(capacity, kGroupSize) * kClusters, 64); + return kClusters; // hwy::RoundUpTo(hwy::DivCeil(capacity, kGroupSize) * + // kClusters, 64); } // Returns number of NuqStream to allocate for the stream, which matches its // size in bytes. static constexpr size_t PackedEnd(size_t capacity) { - return PackedStart(capacity) + hwy::DivCeil(capacity, 2); // 2x 4-bit/byte + const size_t num_groups = hwy::DivCeil(capacity, kGroupSize); + return (kClusters * num_groups) + + hwy::DivCeil(capacity, 2); // 2x 4-bit/byte } uint8_t byte; diff --git a/gemma/weights.h b/gemma/weights.h index ce2df43..280068b 100644 --- a/gemma/weights.h +++ b/gemma/weights.h @@ -26,6 +26,9 @@ #include #include +// TODO: remove compress-inl.h and highway.h when no longer required - i.e. +// necessary functionality in Rehape() is moved to weights.cc. +#include "compression/compress-inl.h" #include "compression/compress.h" #include "compression/shared.h" #include "gemma/common.h" @@ -33,6 +36,7 @@ #include "hwy/aligned_allocator.h" #include "hwy/base.h" #include "hwy/contrib/thread_pool/thread_pool.h" +#include "hwy/highway.h" namespace gcpp { @@ -203,6 +207,7 @@ struct LayerWeightsPtrs { // Initializes att_weights from attn_vec_einsum_w, hence this must be called // after loading weights via ForEachTensor. // TODO: update compression/convert_weights to bake this in. + // TODO(janwas): shift to weights.cc. void Reshape(MatStorage* storage) { if (attn_vec_einsum_w.data() == nullptr) return; @@ -218,6 +223,41 @@ struct LayerWeightsPtrs { storage->Allocate(); att_weights.SetPtr(*storage); } + + if (hwy::IsSame()) { + namespace hn = hwy::HWY_NAMESPACE; + const hn::ScalableTag df; + + hwy::AlignedFreeUniquePtr attn_vec_einsum_w_tmp = + hwy::AllocateAligned(model_dim * heads * qkv_dim); + hwy::AlignedFreeUniquePtr att_weights_tmp = + hwy::AllocateAligned(model_dim * heads * qkv_dim); + + HWY_NAMESPACE::DecompressAndZeroPad( + df, MakeSpan(attn_vec_einsum_w.data(), model_dim * heads * qkv_dim), + 0, attn_vec_einsum_w_tmp.get(), model_dim * heads * qkv_dim); + + for (size_t m = 0; m < model_dim; ++m) { + float* HWY_RESTRICT out_row = + att_weights_tmp.get() + m * heads * qkv_dim; + for (size_t h = 0; h < heads; ++h) { + hwy::CopyBytes(attn_vec_einsum_w_tmp.get() + h * model_dim * qkv_dim + + m * qkv_dim, + out_row + h * qkv_dim, qkv_dim * sizeof(float)); + } + } + + CompressWorkingSet work; + hwy::ThreadPool pool(0); + + HWY_NAMESPACE::Compress( + att_weights_tmp.get(), model_dim * heads * qkv_dim, work, + MakeSpan(att_weights.data(), model_dim * heads * qkv_dim), + /*packed_ofs=*/0, pool); + + return; + } + for (size_t m = 0; m < model_dim; ++m) { Weight* HWY_RESTRICT out_row = att_weights.data() + m * heads * qkv_dim; for (size_t h = 0; h < heads; ++h) {