Minor cleanup, Windows+Bazel build fixes

add app.h comment compress-inl: remove unused typedef gemma-inl: add missing HWY_ATTR separate sum-inl.h and basics.h headers update include pragmas update dot_test thresholds update Highway version in Bazel for HWY_RCAST_ALIGNED fix PiperOrigin-RevId: 684398826
google · Oct 10, 2024 · e14c98b · e14c98b
1 parent a570e3f
commit e14c98b
Show file tree

Hide file tree

Showing 16 changed files with 314 additions and 242 deletions.
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -20,6 +20,14 @@ licenses(["notice"])
 
 exports_files(["LICENSE"])
 
+cc_library(
+    name = "basics",
+    hdrs = ["util/basics.h"],
+    deps = [
+        "@hwy//:hwy",
+    ],
+)
+
 cc_library(
     name = "allocator",
     hdrs = ["util/allocator.h"],
@@ -66,13 +74,15 @@ cc_library(
     ],
     textual_hdrs = [
         "ops/dot-inl.h",
+        "ops/sum-inl.h",
         "ops/fp_arith-inl.h",
         "ops/matmul-inl.h",
         "ops/matvec-inl.h",
         "ops/ops-inl.h",
     ],
     deps = [
         ":allocator",
+        ":basics",
         ":threading",
         "//compression:compress",
         "//compression:sfp",
@@ -271,6 +281,7 @@ cc_library(
     ],
     deps = [
         ":allocator",
+        ":basics",
         ":common",
         ":ops",
         ":tokenizer",

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -106,11 +106,13 @@ set(SOURCES
   ops/matmul-inl.h
   ops/matvec-inl.h
   ops/ops-inl.h
+  ops/sum-inl.h
   paligemma/image.cc
   paligemma/image.h
   util/allocator.h
   util/app.h
   util/args.h
+  util/basics.h
   util/test_util.h
   util/threading.h
   )

diff --git a/MODULE.bazel b/MODULE.bazel
@@ -3,37 +3,33 @@ module(
     version = "0.1.0",
 )
 
-bazel_dep(name = "rules_license", version = "0.0.7")
-bazel_dep(name = "googletest", version = "1.14.0")
-
-# Copied from Highway because Bazel does not load them transitively
+bazel_dep(name = "abseil-cpp", version = "20240722.0")
 bazel_dep(name = "bazel_skylib", version = "1.4.1")
-bazel_dep(name = "rules_cc", version = "0.0.9")
+bazel_dep(name = "googletest", version = "1.14.0")
+bazel_dep(name = "highway", version = "1.1.0")
+bazel_dep(name = "nlohmann_json", version = "3.11.3")
 bazel_dep(name = "platforms", version = "0.0.7")
+bazel_dep(name = "rules_cc", version = "0.0.9")
+bazel_dep(name = "rules_license", version = "0.0.7")
+bazel_dep(name = "google_benchmark", version = "1.8.5")
 
-http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
-
-http_archive(
-    name = "hwy",
-    urls = ["https://github.com/google/highway/archive/refs/tags/1.2.0.zip"],
-    integrity = "sha256-fbtKAGj5hhhBr5Bggtsrj4aIodC2OHb1njB8LGfom8A=",    strip_prefix = "highway-1.2.0",
+# Require a more recent version for HWY_RCAST_ALIGNED
+git_override(
+    "highway",
+    "https://github.com/google/highway",
+    commit = "bb6c3f36b0c8dde8a8ef98b0f0884f4de820a7ca",
 )
 
-http_archive(
-    name = "nlohmann_json",
-    urls = ["https://github.com/nlohmann/json/archive/refs/tags/v3.11.3.zip"],
-    integrity = "sha256-BAIrBdgG61/3MCPCgLaGl9Erk+G3JnoLIqGjnsdXgGk=",
-    strip_prefix = "json-3.11.3",
-)
+http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
 http_archive(
     name = "com_google_sentencepiece",
+    build_file = "@//bazel:sentencepiece.bazel",
+    patch_args = ["-p1"],
+    patches = ["@//bazel:sentencepiece.patch"],
     sha256 = "8409b0126ebd62b256c685d5757150cf7fcb2b92a2f2b98efb3f38fc36719754",
     strip_prefix = "sentencepiece-0.1.96",
     urls = ["https://github.com/google/sentencepiece/archive/refs/tags/v0.1.96.zip"],
-    build_file = "@//bazel:sentencepiece.bazel",
-    patches = ["@//bazel:sentencepiece.patch"],
-    patch_args = ["-p1"],
 )
 
 # For sentencepiece
@@ -56,17 +52,3 @@ cc_library(
         "https://github.com/s-yata/darts-clone/archive/e40ce4627526985a7767444b6ed6893ab6ff8983.zip",
     ],
 )
-# ABSL on 2023-10-18
-http_archive(
-    name = "com_google_absl",
-    sha256 = "f841f78243f179326f2a80b719f2887c38fe226d288ecdc46e2aa091e6aa43bc",
-    strip_prefix = "abseil-cpp-9687a8ea750bfcddf790372093245a1d041b21a3",
-    urls = ["https://github.com/abseil/abseil-cpp/archive//9687a8ea750bfcddf790372093245a1d041b21a3.tar.gz"],
-)
-# Benchmark
-http_archive(
-    name = "benchmark",
-    urls = ["https://github.com/google/benchmark/archive/refs/tags/v1.8.2.tar.gz"],
-    integrity = "sha256-KqspgNA3YTf5adkoSPu2gharsHYzA0U0/IxlzE56DpM=",
-    strip_prefix = "benchmark-1.8.2",
-)
diff --git a/compression/BUILD.bazel b/compression/BUILD.bazel
@@ -89,7 +89,7 @@ cc_library(
     textual_hdrs = ["nuq-inl.h"],
     deps = [
         ":sfp",
-        "//:allocator",
+        "//:basics",
         "@hwy//:hwy",
         "@hwy//hwy/contrib/sort:vqsort",
     ],

diff --git a/compression/compress-inl.h b/compression/compress-inl.h
@@ -24,7 +24,7 @@
 #include <cmath>  // lroundf, only if COMPRESS_STATS
 
 #include "compression/blob_store.h"
-#include "compression/compress.h"
+#include "compression/compress.h"  // IWYU pragma: export
 #include "compression/distortion.h"
 #include "hwy/aligned_allocator.h"
 #include "hwy/base.h"

diff --git a/compression/nuq-inl.h b/compression/nuq-inl.h
@@ -22,7 +22,7 @@
 #include <stdio.h>
 
 #include "compression/shared.h"
-#include "util/allocator.h"
+#include "util/basics.h"
 #include "hwy/base.h"
 
 #endif  // THIRD_PARTY_GEMMA_CPP_COMPRESSION_NUQ_INL_H_

diff --git a/gemma/gemma-inl.h b/gemma/gemma-inl.h
@@ -1176,14 +1176,15 @@ SampleFunc ChooseSampleFunc(const RuntimeConfig& runtime_config) {
 
   // Fast path for top-1 with no accept_token.
   if (kTopK == 1 && !runtime_config.accept_token) {
-    return [](float* logits, size_t vocab_size) -> TokenAndProb {
+    return [](float* logits, size_t vocab_size) -> TokenAndProb HWY_ATTR {
       PROFILER_ZONE("Gen.Sample Top1");
       return Top1OfSoftmax(logits, vocab_size);
     };
   }
 
   // General case: Softmax with top-k sampling.
-  return [&runtime_config](float* logits, size_t vocab_size) -> TokenAndProb {
+  return [&runtime_config](float* logits,
+                           size_t vocab_size) -> TokenAndProb HWY_ATTR {
     PROFILER_ZONE("Gen.Sample general");
     Softmax(logits, vocab_size);
     const int token = SampleTopK<kTopK>(logits, vocab_size, *runtime_config.gen,

diff --git a/gemma/gemma.h b/gemma/gemma.h
@@ -28,7 +28,8 @@
 #include "gemma/kv_cache.h"
 #include "gemma/tokenizer.h"
 #include "paligemma/image.h"
-#include "util/allocator.h"
+#include "util/allocator.h"  // RowVectorBatch
+#include "util/basics.h"     // TokenAndProb
 #include "util/threading.h"
 #include "hwy/contrib/thread_pool/thread_pool.h"
 #include "hwy/timer.h"

diff --git a/ops/dot_test.cc b/ops/dot_test.cc
@@ -840,7 +840,7 @@ class DotStats {
 
     ASSERT_INSIDE(kPairwise, 4.5E-4, s_rels[kPairwise].GeometricMean(), 1.5E-2);
     // Extremely high error on aarch64.
-    ASSERT_INSIDE(kPairwise, 1.1E-3f, s_rels[kPairwise].Max(), 1250.f);
+    ASSERT_INSIDE(kPairwise, 1.1E-3f, s_rels[kPairwise].Max(), 2E3f);
   }
 
   // Backward relative error, lower is better.

diff --git a/ops/matmul.h b/ops/matmul.h
@@ -19,9 +19,10 @@
 #include <stddef.h>
 
 #include "util/allocator.h"  // RowVectorBatch
-#include "util/threading.h"  // PerClusterPools
+#include "util/threading.h"
+#include "hwy/aligned_allocator.h"  // IWYU pragma: export
 #include "hwy/base.h"
-#include "hwy/contrib/thread_pool/thread_pool.h"
+#include "hwy/contrib/thread_pool/thread_pool.h"  // IWYU pragma: export
 #include "hwy/per_target.h"
 
 namespace gcpp {