Restructure internal initialization.

PiperOrigin-RevId: 769507096

BUILD.bazel

@@ -501,7 +501,6 @@ cc_library(
         ":threading",
         ":threading_context",
         ":weights",
-        # Placeholder for internal dep, do not remove.,
         "//io:blob_store",
         "//io",
         "//paligemma:image",
@@ -588,6 +587,7 @@ cc_test(
         ":configs",
         ":gemma_lib",
         "@googletest//:gtest_main",  # buildcleaner: keep
+        "//io",
         "@highway//:hwy",
         "@highway//:hwy_test_util",
     ],

evals/gemma_test.cc

@@ -22,6 +22,7 @@
 
 #include "evals/benchmark_helper.h"
 #include "gemma/configs.h"
+#include "io/io.h"
 #include "hwy/base.h"
 #include "hwy/tests/hwy_gtest.h"
 
@@ -175,6 +176,7 @@ TEST_F(GemmaTest, CrossEntropySmall) {
 
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
+  gcpp::InternalInit();
   gcpp::GemmaTest::InitEnv(argc, argv);
   int ret = RUN_ALL_TESTS();
   gcpp::GemmaTest::DeleteEnv();

gemma/gemma.cc

@@ -43,7 +43,6 @@
 
 #include <vector>
 
-// Placeholder for internal header, do not modify.
 #include "gemma/configs.h"
 #include "gemma/model_store.h"
 #include "gemma/tokenizer.h"
@@ -636,11 +635,7 @@ HWY_EXPORT(GenerateSingleT);
 HWY_EXPORT(GenerateBatchT);
 HWY_EXPORT(GenerateImageTokensT);
 
-// Internal init must run before I/O. This helper function takes care of that,
-// plus calling `SetArgs`.
 MatMulEnv MakeMatMulEnv(const ThreadingArgs& threading_args) {
-  // Placeholder for internal init, do not modify.
-
   ThreadingContext::SetArgs(threading_args);
   return MatMulEnv(ThreadingContext::Get());
 }

gemma/gemma.h

@@ -97,6 +97,7 @@ struct TimingInfo {
   size_t tokens_generated = 0;
 };
 
+// Returns the `MatMulEnv` after calling `SetArgs`.
 MatMulEnv MakeMatMulEnv(const ThreadingArgs& threading_args);
 
 using KVCaches = hwy::Span<KVCache>;

gemma/run.cc

@@ -294,6 +294,7 @@ void Run(const LoaderArgs& loader, const ThreadingArgs& threading,
 }  // namespace gcpp
 
 int main(int argc, char** argv) {
+  gcpp::InternalInit();
   {
     PROFILER_ZONE("Startup.misc");
 

io/BUILD.bazel

@@ -40,6 +40,7 @@ cc_library(
         "//conditions:default": [],
     }),
     deps = [
+        # Placeholder for internal dep, do not remove.,
         "//:allocator",
         "@highway//:hwy",
     ] + FILE_DEPS,

io/io.cc

@@ -72,6 +72,7 @@
 #include <sys/stat.h>  // O_RDONLY
 #include <unistd.h>    // read, write, close
 
+// Placeholder for internal header, do not modify.
 #include "util/allocator.h"
 
 namespace gcpp {
@@ -218,6 +219,9 @@ bool IOBatch::Add(void* mem, size_t bytes) {
   return true;
 }
 
+void InternalInit() {
+}
+
 uint64_t IOBatch::Read(const File& file) const {
 #if GEMMA_IO_PREADV
   HWY_ASSERT(!spans_.empty());

io/io.h

@@ -144,6 +144,10 @@ struct Path {
 // Aborts on error.
 std::string ReadFileToString(const Path& path);
 
+// No-op in open-source. Must be called at the beginning of a binary, before
+// any I/O or flag usage.
+void InternalInit();
+
 }  // namespace gcpp
 
 #endif  // THIRD_PARTY_GEMMA_CPP_COMPRESSION_IO_H_

paligemma/BUILD.bazel

@@ -45,6 +45,7 @@ cc_test(
         "//:configs",
         "//:gemma_lib",
         "//compression:types",
+        "//io",
         "@highway//:hwy",
         "@highway//:hwy_test_util",
     ],

paligemma/paligemma_test.cc

@@ -22,6 +22,7 @@
 #include "evals/benchmark_helper.h"
 #include "gemma/configs.h"
 #include "gemma/gemma.h"
+#include "io/io.h"
 #include "util/allocator.h"
 #include "hwy/base.h"
 #include "hwy/tests/hwy_gtest.h"
@@ -120,10 +121,11 @@ TEST_F(PaliGemmaTest, QueryObjects) {
 }  // namespace gcpp
 
 int main(int argc, char** argv) {
+  testing::InitGoogleTest(&argc, argv);
+  gcpp::InternalInit();
+
   gcpp::GemmaEnv env(argc, argv);
   gcpp::s_env = &env;
 
-  testing::InitGoogleTest(&argc, argv);
-
   return RUN_ALL_TESTS();
 }