[MLGO] Add support for multiple training traces per module

[mtrofin]

Summary

This happens in e.g. regalloc, where we trace decisions per function,
but wouldn't want to spew N log files (i.e. one per function). So we
output a key-value association, where the key is an ID for the
sub-module object, and the value is the tensorflow::SequenceExample.

The current relation with protobuf is tenuous, so we're avoiding a
custom message type in favor of using the Struct message, but that
requires the values be wire-able strings, hence base64 encoding.

We plan on resolving the protobuf situation shortly, and improve the
encoding of such logs, but this is sufficient for now for setting up
regalloc training.

[yundiqian]

llvm/include/llvm/Analysis/Utils/TFUtils.h

184

from the same module in format of google::protobuf::Struct, e.g., ...

[snehasish]

llvm/unittests/Analysis/TFUtilsTest.cpp

311

This should be ASSERT_TRUE since if parsing from Result fails the following EXPECT checks should not be executed.

https://testing.googleblog.com/2008/07/tott-expect-vs-assert.html

[mtrofin]

llvm/unittests/Analysis/TFUtilsTest.cpp

311

Good catch, thanks - that goes for the other tests too, I'll fix separately first.

[mtrofin]

feedback

[mtrofin]

llvm/unittests/Analysis/TFUtilsTest.cpp

311

rGb7f298f17416 takes care of the others.

llvm/include/llvm/Analysis/Utils/TFUtils.h

public:
   // to piecemeal cast their int32_t values.
   // FIXME: let's drop int32_t support. While it's supported by evaluator, it's
   // not supported by the tensorflow::SequenceExample proto. For small values,
   // we can consider using bytes.
   char *addEntryAndGetFloatOrInt64Buffer(size_t FeatureID);
 
   // Flush the content of the log to the stream, clearing the stored data in the
   // process.
+  void flush(std::string *Str);
   void flush(raw_ostream &OS);
 
+  // Flush a set of logs that are produced from the same module, e.g.
+  // per-function regalloc traces, as a google::protobuf::Struct message.
+  static void flushLogs(raw_ostream &OS,
+                        const StringMap<std::unique_ptr<Logger>> &Loggers);
+
 private:
   std::vector<LoggedFeatureSpec> FeatureSpecs;
   TensorSpec RewardSpec;
   const bool IncludeReward;
   std::unique_ptr<LoggerDataImpl> LoggerData;
 };
 
 class TFModelEvaluator final {

llvm/lib/Analysis/TFUtils.cpp

 //===- TFUtils.cpp - tensorflow evaluation utilities ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utilities for interfacing with tensorflow C APIs.
 //
 //===----------------------------------------------------------------------===//
 #include "llvm/Config/config.h"
 #if defined(LLVM_HAVE_TF_API)
 
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
+#include "llvm/Support/Base64.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include "google/protobuf/struct.pb.h"
 #include "google/protobuf/text_format.h"
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/c/c_api_experimental.h"
 #include "tensorflow/core/example/example.pb.h"
 #include <cassert>
 #include <numeric>
 
 using namespace llvm;

 
 TFStatusPtr createTFStatus() {
   return TFStatusPtr(TF_NewStatus(), &TF_DeleteStatus);
 }
 
 TFSessionOptionsPtr createTFSessionOptions() {
   return TFSessionOptionsPtr(TF_NewSessionOptions(), &TF_DeleteSessionOptions);
 }
+
+void serialize(const Message &SE, std::string *OutStr) {
+  if (ProtobufTextMode) {
+    TextFormat::PrintToString(SE, OutStr);
+  } else {
+    *OutStr = SE.SerializeAsString();
+  }
+}
 } // namespace
 
 namespace llvm {
 class EvaluationResultImpl {
 public:
   EvaluationResultImpl(size_t OutputSize)
       : OutputSize(OutputSize), Output(OutputSize){};
 

 
 public:
   LoggerDataImpl(const std::vector<LoggedFeatureSpec> &LoggedSpecs,
                  const TensorSpec &RewardSpec, bool IncludeReward)
       : LoggedFeatureSpecs(LoggedSpecs), RewardSpec(RewardSpec),
         IncludeReward(IncludeReward), FeatureLists(LoggedFeatureSpecs.size()) {}
 
   // flush the logged info to a stream and clear the log contents.
-  void flush(raw_ostream &OS) {
+  void flush(std::string *Str) {
     size_t NrRecords = getNrRecords();
     (void)NrRecords;
     tensorflow::SequenceExample SE;
     transferLog(SE);
     assert(isSelfConsistent(SE, NrRecords));
-    std::string OutStr;
-    if (ProtobufTextMode)
-      google::protobuf::TextFormat::PrintToString(SE, &OutStr);
-    else
-      OutStr = SE.SerializeAsString();
-
-    OS << OutStr;
+    serialize(SE, Str);
   }
 
   char *addNewTensor(size_t FeatureID) {
     const auto &Spec = LoggedFeatureSpecs[FeatureID].Spec;
     if (Spec.isElementType<float>()) {
       auto *RF = FeatureLists[FeatureID]
                      .add_feature()
                      ->mutable_float_list()
void Logger::logSpecifiedTensorValue(size_t FeatureID, const char *RawData) {
   else
     llvm_unreachable("Unsupported tensor type");
 }
 
 char *Logger::addEntryAndGetFloatOrInt64Buffer(size_t FeatureID) {
   return reinterpret_cast<char *>(LoggerData->addNewTensor(FeatureID));
 }
 
-void Logger::flush(raw_ostream &OS) { LoggerData->flush(OS); }
+void Logger::flush(std::string *Str) { LoggerData->flush(Str); }
+
+void Logger::flush(raw_ostream &OS) {
+  std::string Buff;
+  LoggerData->flush(&Buff);
+  OS << Buff;
+}
+
+void Logger::flushLogs(raw_ostream &OS,
+                       const StringMap<std::unique_ptr<Logger>> &Loggers) {
+  google::protobuf::Struct Msg;
+  for (const auto &NamedLogger : Loggers) {
+    tensorflow::SequenceExample SE;
+    const auto &Logger = NamedLogger.second;
+    std::string Unencoded;
+    if (Logger->LoggerData->getNrRecords() > 0)
+      Logger->flush(&Unencoded);
+
+    (*Msg.mutable_fields())[NamedLogger.first().str()]
+        .mutable_string_value()
+        ->append(ProtobufTextMode ? Unencoded : encodeBase64(Unencoded));
+  }
+
+  std::string OutStr;
+  serialize(Msg, &OutStr);
+  OS << OutStr;
+}
 #endif // defined(LLVM_HAVE_TF_API)

llvm/unittests/Analysis/TFUtilsTest.cpp

 //===- TFUtilsTest.cpp - test for TFUtils ---------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/Utils/TFUtils.h"
+#include "google/protobuf/struct.pb.h"
 #include "tensorflow/core/example/example.pb.h"
 #include "tensorflow/core/example/feature.pb.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Path.h"
TEST(TFUtilsTest, LoggerFinalReward) {
   const float Zero[]{0.0};
   const float R[]{3.14};
   tensorflow::SequenceExample Expected;
   ASSERT_TRUE(Expected.ParseFromString(Result));
   PROTO_CHECKER("reward", float_list, 0, Zero);
   PROTO_CHECKER("reward", float_list, 1, Zero);
   PROTO_CHECKER("reward", float_list, 2, R);
 }
+
+TEST(TFUtilsTest, LoggerGroup) {
+  std::vector<LoggedFeatureSpec> Features;
+  Features.push_back({TensorSpec::createSpec<float>("the_float", {1}), None});
+  Features.push_back({TensorSpec::createSpec<int64_t>("the_int", {1}), None});
+
+  auto Rewards = TensorSpec::createSpec<float>("reward", {1});
+  StringMap<std::unique_ptr<Logger>> Loggers;
+  std::vector<std::string> Names{"a", "b"};
+  size_t Bump = 0;
+  for (auto Name : Names) {
+    auto L = std::make_unique<Logger>(Features, Rewards, true);
+    for (int64_t I = 0; I < 3; ++I) {
+      float F = static_cast<float>(I) + Bump;
+      L->logFloatValue(0, &F);
+      L->logInt64Value(1, &I);
+    }
+    L->logFloatFinalReward(3.14 + Bump);
+    Loggers.insert(std::make_pair(Name, std::move(L)));
+  }
+  std::string Result;
+  raw_string_ostream OS(Result);
+  Logger::flushLogs(OS, Loggers);
+  google::protobuf::Struct Expected;
+  ASSERT_TRUE(Expected.ParseFromString(Result));
+  EXPECT_EQ(Expected.fields_size(), 2);
+  EXPECT_TRUE(Expected.fields().contains("a"));
+  EXPECT_TRUE(Expected.fields().contains("b"));
+}