Major cleanup of profiler zones, add Caller annotation for all pool.Run

Pass ThreadingContext instead of Pools/Profiler individually, for access to Zones
Add GCPP_ZONE helper
Add Caller argument to pool.Run to enable new stats
Remove most direct dependencies on ThreadPool, prefer ParallelFor

PiperOrigin-RevId: 822934530

Author: jan-wassenberg

BUILD.bazel

@@ -95,7 +95,6 @@ cc_library(
         ":topology",
         # Placeholder for container detection, do not remove
         "@highway//:hwy",
-        "@highway//:profiler",
         "@highway//:thread_pool",
         "@highway//:topology",
     ],
@@ -124,7 +123,9 @@ cc_library(
     srcs = ["util/zones.cc"],
     hdrs = ["util/zones.h"],
     deps = [
+        "@highway//:hwy",
         "@highway//:profiler",
+        "@highway//:thread_pool",
     ],
 )
 
@@ -258,7 +259,6 @@ cc_library(
         "//io:fields",
         "@highway//:hwy",
         "@highway//:profiler",
-        "@highway//:thread_pool",
     ],
 )
 
@@ -278,7 +278,6 @@ cc_library(
         "//io:blob_store",
         "@highway//:hwy",
         "@highway//:profiler",
-        "@highway//:thread_pool",
     ],
 )
 
@@ -309,7 +308,6 @@ cc_library(
     deps = [
         ":allocator",
         ":basics",
-        ":configs",
         ":mat",
         ":threading",
         ":threading_context",
@@ -397,7 +395,6 @@ cc_library(
         "@highway//:hwy",
         "@highway//:math",
         "@highway//:profiler",
-        "@highway//:thread_pool",
         "@highway//hwy/contrib/sort:vqsort",
     ],
 )
@@ -424,7 +421,6 @@ cc_test(
         "@highway//:nanobenchmark",  #buildcleaner: keep
         "@highway//:profiler",
         "@highway//:stats",
-        "@highway//:thread_pool",
     ],
 )
 
@@ -507,7 +503,6 @@ cc_test(
         "@highway//:hwy_test_util",
         "@highway//:nanobenchmark",
         "@highway//:profiler",
-        "@highway//:thread_pool",
     ],
 )
 

CMakeLists.txt

@@ -22,7 +22,7 @@ set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
-FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 9781a1698ee0756ef1eaaf96930113ed7cb6d3ee EXCLUDE_FROM_ALL)
+FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 2a16a50ff61071bb25ddef0ce35d92b0e2b9c579 EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(highway)
 
 ## Note: absl needs to be installed by sentencepiece. This will only happen if

MODULE.bazel

@@ -18,7 +18,7 @@ bazel_dep(name = "google_benchmark", version = "1.8.5")
 # Require a more recent version.
 git_override(
     module_name = "highway",
-    commit = "9781a1698ee0756ef1eaaf96930113ed7cb6d3ee",
+    commit = "2a16a50ff61071bb25ddef0ce35d92b0e2b9c579",
     remote = "https://github.com/google/highway",
 )
 

README.md

@@ -452,7 +452,7 @@ FetchContent_MakeAvailable(sentencepiece)
 FetchContent_Declare(gemma GIT_REPOSITORY https://github.com/google/gemma.cpp GIT_TAG origin/main)
 FetchContent_MakeAvailable(gemma)
 
-FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 9781a1698ee0756ef1eaaf96930113ed7cb6d3ee)
+FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 2a16a50ff61071bb25ddef0ce35d92b0e2b9c579)
 FetchContent_MakeAvailable(highway)
 ```
 

compression/BUILD.bazel

@@ -120,9 +120,9 @@ cc_library(
         ":compress",
         ":distortion",
         "//:mat",
+        "//:threading_context",
         "@highway//:hwy",
         "@highway//:hwy_test_util",
-        "@highway//:thread_pool",
     ],
 )
 
@@ -180,6 +180,7 @@ cc_library(
         ":sfp",
         "//:basics",
         "//:mat",
+        "//:threading_context",
         "@highway//:hwy",
         "@highway//:nanobenchmark",
         "@highway//:profiler",
@@ -203,9 +204,9 @@ cc_test(
         ":test_util",
         "@googletest//:gtest_main",  # buildcleaner: keep
         "//:test_util",
+        "//:threading_context",
         "@highway//:hwy",
         "@highway//:hwy_test_util",
-        "@highway//:thread_pool",
     ],
 )
 

compression/compress-inl.h

@@ -26,10 +26,10 @@
 
 #include "compression/compress.h"  // IWYU pragma: export
 #include "compression/distortion.h"
+#include "util/threading_context.h"
 #include "hwy/aligned_allocator.h"
 #include "hwy/base.h"
 #include "hwy/contrib/thread_pool/thread_pool.h"
-#include "hwy/timer.h"
 
 #if COMPRESS_STATS
 #include <cmath>  // lroundf
@@ -493,13 +493,13 @@ struct CompressTraits<NuqStream> {
   }
 };
 
-// Compresses `num` elements of `raw` to `packed` starting at `packed_ofs`,
-// which is useful for compressing sub-regions of an array.
+// DEPRECATED: Use the overload with ThreadingContext instead.
 template <typename Packed>
 HWY_NOINLINE void Compress(const float* HWY_RESTRICT raw, size_t num,
                            CompressWorkingSet& work,
                            const PackedSpan<Packed>& packed,
-                           const size_t packed_ofs, hwy::ThreadPool& pool) {
+                           const size_t packed_ofs, hwy::ThreadPool& pool,
+                           hwy::pool::Caller caller = hwy::pool::Caller()) {
   packed.BoundsCheck(packed_ofs, num);
   work.tls.resize(pool.NumWorkers());
   if constexpr (COMPRESS_STATS) {
@@ -511,7 +511,7 @@ HWY_NOINLINE void Compress(const float* HWY_RESTRICT raw, size_t num,
   using Traits = CompressTraits<hwy::RemoveConst<Packed>>;
   constexpr size_t kBatch = 8192;
   const size_t num_batches = hwy::DivCeil(num, kBatch);
-  pool.Run(0, num_batches,
+  pool.Run(0, num_batches, caller,
            [&](const uint32_t idx_batch, size_t thread) HWY_ATTR {
              const hn::ScalableTag<float> df;
 
@@ -530,6 +530,17 @@ HWY_NOINLINE void Compress(const float* HWY_RESTRICT raw, size_t num,
   }
 }
 
+// Compresses `num` elements of `raw` to `packed` starting at `packed_ofs`,
+// which is useful for compressing sub-regions of an array.
+template <typename Packed>
+HWY_NOINLINE void Compress(const float* HWY_RESTRICT raw, size_t num,
+                           CompressWorkingSet& work,
+                           const PackedSpan<Packed>& packed,
+                           const size_t packed_ofs, ThreadingContext& ctx) {
+  Compress(raw, num, work, packed, packed_ofs, ctx.pools.Pool(),
+           ctx.pool_callers.Get(Callers::kCompress));
+}
+
 // Same as above, but without parallelization nor benchmarking.
 template <typename Packed>
 HWY_NOINLINE void Compress(const float* HWY_RESTRICT raw, size_t num,

compression/compress_test.cc

@@ -24,9 +24,9 @@
 #include "compression/compress.h"
 #include "compression/distortion.h"
 #include "util/test_util.h"
+#include "util/threading_context.h"
 #include "hwy/aligned_allocator.h"
 #include "hwy/base.h"
-#include "hwy/contrib/thread_pool/thread_pool.h"
 #include "hwy/tests/hwy_gtest.h"
 // clang-format off
 #undef HWY_TARGET_INCLUDE
@@ -42,14 +42,27 @@ namespace gcpp {
 namespace HWY_NAMESPACE {
 namespace hn = hwy::HWY_NAMESPACE;
 
+static ThreadingArgs SingleThreadArgs() {
+  ThreadingArgs args;
+  args.max_lps = 1;
+  return args;
+}
+
+static ThreadingContext& Ctx() {
+  static ThreadingContext* ctx = new ThreadingContext(SingleThreadArgs());
+  return *ctx;
+}
+
 // Calls Compress and Decompress2 and verifies the distortion/error.
 template <typename Packed>
 struct TestDecompress2 {
   template <typename T, class D>
   HWY_INLINE void operator()(T /*unused*/, D d) {
     const size_t N = hn::Lanes(d);
     CompressWorkingSet work;
-    hwy::ThreadPool pool(0);
+    ThreadingArgs args;
+    args.max_lps = 1;
+    ThreadingContext ctx(args);
     hwy::RandomState rng;
 
     const size_t num = 2 * N;
@@ -68,7 +81,7 @@ struct TestDecompress2 {
     // Short inputs fail VerifyGaussian.
 
     const size_t packed_ofs = 0;
-    Compress(raw.get(), num, work, packed_span, packed_ofs, pool);
+    Compress(raw.get(), num, work, packed_span, packed_ofs, ctx);
     hn::Vec<D> raw0, raw1;
     Decompress2(d, MakeConst(packed_span), packed_ofs, raw0, raw1);
     hn::Store(raw0, d, dec.get());
@@ -129,7 +142,6 @@ struct TestShortLengths {
   HWY_INLINE void operator()(T /*unused*/, D d) {
     const size_t N = hn::Lanes(d);
     CompressWorkingSet work;
-    hwy::ThreadPool pool(0);
     hwy::RandomState rng;
 
     for (size_t num = 1; num < 5 * hn::Lanes(d); ++num) {
@@ -149,7 +161,7 @@ struct TestShortLengths {
       // Short inputs fail VerifyGaussian.
 
       const size_t packed_ofs = 0;
-      Compress(raw.get(), num, work, packed_span, packed_ofs, pool);
+      Compress(raw.get(), num, work, packed_span, packed_ofs, Ctx());
       DecompressAndZeroPad(d, MakeConst(packed_span), packed_ofs, dec.get(),
                            num);
 

compression/python/compression_clif_aux.cc

@@ -37,7 +37,6 @@
 #include "util/basics.h"
 #include "util/mat.h"
 #include "util/threading_context.h"
-#include "hwy/contrib/thread_pool/thread_pool.h"
 
 #undef HWY_TARGET_INCLUDE
 #define HWY_TARGET_INCLUDE \
@@ -57,9 +56,6 @@ class SbsWriterImpl : public ISbsWriter {
   template <typename Packed>
   void InsertT(const char* name, F32Span weights,
                const TensorInfo& tensor_info) {
-    // TODO(janwas): 1D parallel-for.
-    hwy::ThreadPool& pool = ctx_.pools.Pool();
-
     MatPtrT<Packed> mat(name, ExtentsFromInfo(&tensor_info));
     // SFP and NUQ (which uses SFP for cluster centers) have a limited range
     // and depending on the input values may require rescaling. Scaling is
@@ -82,21 +78,20 @@ class SbsWriterImpl : public ISbsWriter {
     // succeeds, but we only have 10 floats, not the full tensor.
     if (weights.size() == 10 && mat.Extents().Area() != 10) {
       Compress(weights.data(), weights.size(), working_set_, mat.Span(),
-               /*packed_ofs=*/0, pool);
+               /*packed_ofs=*/0, ctx_);
       writer_.Add(name, mat.Packed(), mat.ElementBytes() * 10);
       return;
     }
 
     HWY_ASSERT(weights.size() == mat.Extents().Area());
     Compress(weights.data(), weights.size(), working_set_, mat.Span(),
-             /*packed_ofs=*/0, pool);
+             /*packed_ofs=*/0, ctx_);
     writer_.Add(name, mat.Packed(), mat.PackedBytes());
   }
 
  public:
   SbsWriterImpl(const std::string& sbs_path)
-      : ctx_(ThreadingArgs()),
-        writer_(gcpp::Path(sbs_path), ctx_.pools.Pool()) {}
+      : ctx_(ThreadingArgs()), writer_(gcpp::Path(sbs_path), ctx_) {}
 
   void Insert(const char* name, F32Span weights, Type type,
               const TensorInfo& tensor_info) override {

compression/test_util-inl.h

@@ -23,7 +23,7 @@
 // IWYU pragma: end_exports
 
 #include "compression/compress.h"
-#include "hwy/contrib/thread_pool/thread_pool.h"
+#include "util/threading_context.h"
 
 #endif  // THIRD_PARTY_GEMMA_CPP_COMPRESSION_TEST_UTIL_INL_H_
 
@@ -98,25 +98,26 @@ void ForeachActivationType3(D d) {
 
 // Generates inputs: deterministic, within max SfpStream range.
 template <typename MatT>
-MatStorageT<MatT> GenerateMat(const Extents2D& extents,
-                              const Allocator& allocator, MatPadding padding,
-                              hwy::ThreadPool& pool) {
+MatStorageT<MatT> GenerateMat(const Extents2D& extents, MatPadding padding,
+                              ThreadingContext& ctx) {
   gcpp::CompressWorkingSet ws;
-  ws.tls.resize(pool.NumWorkers());
-  MatStorageT<float> raw("raw", extents, allocator, MatPadding::kPacked);
-  MatStorageT<MatT> compressed("mat", extents, allocator, padding);
+  ws.tls.resize(ctx.pools.MaxWorkers());
+  MatStorageT<float> raw("raw", extents, ctx.allocator, MatPadding::kPacked);
+  MatStorageT<MatT> compressed("mat", extents, ctx.allocator, padding);
   const float scale = SfpStream::kMax / extents.Area();
-  pool.Run(0, extents.rows, [&](const size_t r, size_t thread) {
-    float* HWY_RESTRICT row = raw.Row(r);
-    for (size_t c = 0; c < extents.cols; c++) {
-      float f = static_cast<float>(r * extents.cols + c) * scale;
-      if ((r + c) & 1) f = -f;  // Also generate some negative values.
-      row[c] = f;
-    }
-    Compress(raw.Row(r), raw.Cols(), ws.tls[thread],
-             MakeSpan(compressed.Row(r), extents.cols),
-             /*packed_ofs=*/0);
-  });
+  ParallelFor(ParallelismStrategy::kFlat, extents.rows, ctx, /*cluster_idx=*/0,
+              Callers::kTest, [&](size_t r, size_t thread) {
+                float* HWY_RESTRICT row = raw.Row(r);
+                for (size_t c = 0; c < extents.cols; c++) {
+                  float f = static_cast<float>(r * extents.cols + c) * scale;
+                  if ((r + c) & 1)
+                    f = -f;  // Also generate some negative values.
+                  row[c] = f;
+                }
+                Compress(raw.Row(r), raw.Cols(), ws.tls[thread],
+                         MakeSpan(compressed.Row(r), extents.cols),
+                         /*packed_ofs=*/0);
+              });
 
   compressed.SetScale(0.6f);  // Arbitrary value, different from 1.
   return compressed;
@@ -126,25 +127,26 @@ MatStorageT<MatT> GenerateMat(const Extents2D& extents,
 // `f` swaps `r` and `c`.
 template <typename MatT>
 MatStorageT<MatT> GenerateTransposedMat(const Extents2D extents,
-                                        const Allocator& allocator,
                                         MatPadding padding,
-                                        hwy::ThreadPool& pool) {
+                                        ThreadingContext& ctx) {
   gcpp::CompressWorkingSet ws;
-  ws.tls.resize(pool.NumWorkers());
-  MatStorageT<float> raw("raw", extents, allocator, MatPadding::kPacked);
-  MatStorageT<MatT> compressed("trans", extents, allocator, padding);
+  ws.tls.resize(ctx.pools.MaxWorkers());
+  MatStorageT<float> raw("raw", extents, ctx.allocator, MatPadding::kPacked);
+  MatStorageT<MatT> compressed("trans", extents, ctx.allocator, padding);
   const float scale = SfpStream::kMax / extents.Area();
-  pool.Run(0, extents.rows, [&](const size_t r, size_t thread) {
-    float* HWY_RESTRICT row = raw.Row(r);
-    for (size_t c = 0; c < extents.cols; c++) {
-      float f = static_cast<float>(c * extents.rows + r) * scale;
-      if ((r + c) & 1) f = -f;  // Also generate some negative values.
-      row[c] = f;
-    }
-    Compress(raw.Row(r), raw.Cols(), ws.tls[thread],
-             MakeSpan(compressed.Row(r), extents.cols),
-             /*packed_ofs=*/0);
-  });
+  ParallelFor(ParallelismStrategy::kFlat, extents.rows, ctx, /*cluster_idx=*/0,
+              Callers::kTest, [&](size_t r, size_t thread) {
+                float* HWY_RESTRICT row = raw.Row(r);
+                for (size_t c = 0; c < extents.cols; c++) {
+                  float f = static_cast<float>(c * extents.rows + r) * scale;
+                  if ((r + c) & 1)
+                    f = -f;  // Also generate some negative values.
+                  row[c] = f;
+                }
+                Compress(raw.Row(r), raw.Cols(), ws.tls[thread],
+                         MakeSpan(compressed.Row(r), extents.cols),
+                         /*packed_ofs=*/0);
+              });
 
   // Arbitrary value, different from 1, must match `GenerateMat`.
   compressed.SetScale(0.6f);