Add 8-bit integer quantization (I8Stream) to Gemma.cpp.

PiperOrigin-RevId: 819787856

Author: pculliton

BUILD.bazel

@@ -349,6 +349,7 @@ cc_library(
         "ops/matmul_static_f32.cc",
         "ops/matmul_static_nuq.cc",
         "ops/matmul_static_sfp.cc",
+        "ops/matmul_static_i8.cc",
     ],
     hdrs = [
         "ops/matmul_static.h",

CMakeLists.txt

@@ -68,6 +68,7 @@ set(SOURCES
   compression/compress.h
   compression/nuq-inl.h
   compression/sfp-inl.h
+  compression/int-inl.h
   compression/types.h
   compression/test_util-inl.h
   evals/benchmark_helper.cc
@@ -109,6 +110,7 @@ set(SOURCES
   ops/matmul_static_f32.cc
   ops/matmul_static_nuq.cc
   ops/matmul_static_sfp.cc
+  ops/matmul_static_i8.cc
   ops/matmul-inl.h
   ops/matmul.cc
   ops/matmul.h

compression/BUILD.bazel

@@ -80,6 +80,37 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "int",
+    textual_hdrs = ["int-inl.h"],
+    deps = [
+        ":types",
+        "//:basics",
+        "@highway//:hwy",
+    ],
+)
+
+cc_test(
+    name = "int_test",
+    size = "small",
+    timeout = "long",
+    srcs = ["int_test.cc"],
+    features = ["fully_static_link"],
+    linkstatic = True,
+    local_defines = ["HWY_IS_TEST"],
+    # for test_suite.
+    tags = ["hwy_ops_test"],
+    deps = [
+        ":distortion",
+        ":int",
+        "@googletest//:gtest_main",  # buildcleaner: keep
+        "//:test_util",
+        "@highway//:hwy",
+        "@highway//:hwy_test_util",
+        "@highway//:nanobenchmark",
+    ],
+)
+
 cc_library(
     name = "test_util",
     textual_hdrs = [
@@ -144,6 +175,7 @@ cc_library(
     textual_hdrs = ["compress-inl.h"],
     deps = [
         ":distortion",
+        ":int",
         ":nuq",
         ":sfp",
         "//:basics",
@@ -182,6 +214,7 @@ cc_library(
     name = "analyze",
     textual_hdrs = ["analyze.h"],
     deps = [
+        ":int",
         ":nuq",
         ":sfp",
         ":types",

compression/compress-inl.h

@@ -47,6 +47,7 @@
 
 #include "hwy/highway.h"
 // After highway.h
+#include "compression/int-inl.h"
 #include "compression/nuq-inl.h"
 #include "compression/sfp-inl.h"
 
@@ -416,6 +417,34 @@ struct CompressTraits<SfpStream> {
   }
 };
 
+// Integer quantization.
+template <>
+struct CompressTraits<I8Stream> {
+  using Packed = I8Stream;
+
+  template <class DF, HWY_IF_F32_D(DF)>
+  static HWY_INLINE void Compress(DF df, const float* HWY_RESTRICT raw,
+                                  size_t num, CompressPerThread& tls,
+                                  const PackedSpan<Packed>& packed,
+                                  const size_t packed_ofs) {
+    IntCodec::Enc(df, raw, num, packed, packed_ofs);
+  }
+
+  template <class D>  // Caller checks this is f32 or bf16
+  static HWY_INLINE void Load2(D d, const PackedSpan<const Packed>& packed,
+                               const size_t packed_ofs, hn::Vec<D>& raw0,
+                               hn::Vec<D>& raw1) {
+    IntCodec::Dec2(d, packed, packed_ofs, raw0, raw1);
+  }
+
+  template <class D, typename Raw>
+  static HWY_INLINE void DecompressAndZeroPad(
+      D d, const PackedSpan<const Packed>& packed, const size_t packed_ofs,
+      Raw* raw, const size_t num) {
+    IntCodec::DecompressAndZeroPad(d, packed, packed_ofs, raw, num);
+  }
+};
+
 // Nonuniform quantization, 4.5 bits per element, two separate streams.
 template <>
 struct CompressTraits<NuqStream> {
@@ -737,9 +766,10 @@ template <class DF, typename T, typename T1, class Func>
 HWY_INLINE void Decompress1AndCompressInplace(DF df, T* HWY_RESTRICT inout,
                                               size_t num,
                                               const T1* HWY_RESTRICT p1,
+                                              const size_t p1_ofs,
                                               Func&& func) {
   const auto packed_inout = MakeSpan(inout, num);
-  const auto packed1 = MakeSpan(p1, num);
+  const auto packed1 = MakeSpan(p1, p1_ofs + num);
 
   using VF = hn::Vec<decltype(df)>;
   HWY_LANES_CONSTEXPR const size_t NF = hn::Lanes(df);
@@ -749,7 +779,7 @@ HWY_INLINE void Decompress1AndCompressInplace(DF df, T* HWY_RESTRICT inout,
       VF v0, v1;
       Decompress2(df, packed_inout, i, v0, v1);
       VF v10, v11;
-      Decompress2(df, packed1, i, v10, v11);
+      Decompress2(df, packed1, p1_ofs + i, v10, v11);
       const VF out0 = func(df, v0, v10);
       const VF out1 = func(df, v1, v11);
       Compress2(df, out0, out1, packed_inout, i);
@@ -765,7 +795,7 @@ HWY_INLINE void Decompress1AndCompressInplace(DF df, T* HWY_RESTRICT inout,
     hn::Store(hn::Zero(df), df, buf_inout + NF);
     hn::Store(hn::Zero(df), df, buf1 + NF);
     DecompressAndZeroPad(df, packed_inout, i, buf_inout, remaining);
-    DecompressAndZeroPad(df, packed1, i, buf1, remaining);
+    DecompressAndZeroPad(df, packed1, p1_ofs + i, buf1, remaining);
     const VF v0 = hn::Load(df, buf_inout);
     const VF v1 = hn::Load(df, buf_inout + NF);
     const VF v10 = hn::Load(df, buf1);
@@ -827,10 +857,10 @@ template <class DF, typename T, typename T1, typename T2, class Func>
 HWY_INLINE void Decompress2AndCompressTo(DF df, T* HWY_RESTRICT out, size_t num,
                                          const T1* HWY_RESTRICT p1,
                                          const T2* HWY_RESTRICT p2,
-                                         Func&& func) {
+                                         const size_t p2_ofs, Func&& func) {
   const auto packed_out = MakeSpan(out, num);
   const auto packed1 = MakeSpan(p1, num);
-  const auto packed2 = MakeSpan(p2, num);
+  const auto packed2 = MakeSpan(p2, p2_ofs + num);
 
   using VF = hn::Vec<decltype(df)>;
   HWY_LANES_CONSTEXPR const size_t NF = hn::Lanes(df);
@@ -839,7 +869,7 @@ HWY_INLINE void Decompress2AndCompressTo(DF df, T* HWY_RESTRICT out, size_t num,
     for (; i <= num - 2 * NF; i += 2 * NF) {
       VF v10, v11, v20, v21;
       Decompress2(df, packed1, i, v10, v11);
-      Decompress2(df, packed2, i, v20, v21);
+      Decompress2(df, packed2, p2_ofs + i, v20, v21);
       const VF out0 = func(df, v10, v20);
       const VF out1 = func(df, v11, v21);
       Compress2(df, out0, out1, packed_out, i);
@@ -856,7 +886,7 @@ HWY_INLINE void Decompress2AndCompressTo(DF df, T* HWY_RESTRICT out, size_t num,
     hn::Store(hn::Zero(df), df, buf1 + NF);
     hn::Store(hn::Zero(df), df, buf2 + NF);
     DecompressAndZeroPad(df, packed1, i, buf1, remaining);
-    DecompressAndZeroPad(df, packed2, i, buf2, remaining);
+    DecompressAndZeroPad(df, packed2, p2_ofs + i, buf2, remaining);
     const VF v10 = hn::Load(df, buf1);
     const VF v11 = hn::Load(df, buf1 + NF);
     const VF v20 = hn::Load(df, buf2);

compression/compress_test.cc

@@ -243,17 +243,17 @@ class TestDecompressAndCompress {
 
         // Uses `out` so as not to overwrite `p`.
         Decompress1AndCompressInplace(
-            df, out.get(), num, p1.get(),
+            df, out.get(), num, p1.get(), /*p1_ofs=*/0,
             [](DF, VF v, VF v1) HWY_ATTR -> VF { return hn::Add(v, v1); });
         HWY_ASSERT_ARRAY_EQ(expected2.get(), out.get(), num);
 
         Decompress1AndCompressTo(df, out.get(), num, p.get(),
                                  [](DF, VF v) HWY_ATTR -> VF { return v; });
         HWY_ASSERT_ARRAY_EQ(expected1.get(), out.get(), num);
 
-        Decompress2AndCompressTo(df, out.get(), num, p.get(), p1.get(),
-                                 [](DF, VF v, VF v1)
-                                     HWY_ATTR -> VF { return hn::Add(v, v1); });
+        Decompress2AndCompressTo(
+            df, out.get(), num, p.get(), p1.get(), /*p2_ofs=*/0,
+            [](DF, VF v, VF v1) HWY_ATTR -> VF { return hn::Add(v, v1); });
         HWY_ASSERT_ARRAY_EQ(expected2.get(), out.get(), num);
 
         Decompress3AndCompressTo(