[ML] Add final reward logging facility.

[mtrofin]

Summary

Allow logging final rewards. A final reward is logged only once, and is
serialized as all-zero values, except for the last one.

[yundiqian]

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

142–145

It seems over-complicated to pass the flag FinalReward to writeRawTensorsAsFeatureLists function and treat the case separately.

How about making the RawLogData ready-to-print (reward vector is <0, 0, ..., 0, reward>) so that we don't need to change writeRawTensorsAsFeatureLists function? basically the user is supposed to make sure the data in RawLogData is ready-to-print and writeRawTensorsAsFeatureLists only takes care of printing format.

We can either:

1. logReward(0) in each step
2. have a function in Logger called logFinalReward(T value) or overwriteFinalReward(T value), which overwrites the value in RawLogData.back()->back()

or:

1. don't logReward(0) in each step
2. have a function in Logger called logFinalReward(T value) that fills in 0 in RawLogData.back() except putting reward in the last, we can tell the length by looking at feature length already logged in RawLogData

[mtrofin]

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

142–145

That would mean keeping around a potentially large array containing 0 (except the last value).

I would rather not add memory overhead if it can be avoided, and the extra consideration in code isn't that hard to grok.

[yundiqian]

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

142–145

Since it's for development mode, I guess it is debatable whether it worth adding code complexity to trade for better efficiency. imo it does not worth it, but I'm open to either options.

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

Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
         RawLogData(FeatureSpecs.size() + IncludeReward),
         IncludeReward(IncludeReward) {}
 
   template <typename T> void logReward(T Value) {
     assert(IncludeReward);
     logTensorValue(RawLogData.size() - 1, &Value);
   }
 
+  template <typename T> void logFinalReward(T Value) {
+    assert(RawLogData.back().empty());
+    logReward(Value);
+  }
+
   template <typename T>
   void logTensorValue(size_t FeatureID, const T *Value, size_t Size = 1) {
     const char *Start = reinterpret_cast<const char *>(Value);
     const char *End = Start + sizeof(T) * Size;
     RawLogData[FeatureID].insert(RawLogData[FeatureID].end(), Start, End);
   }
 
   void print(raw_ostream &OS);

llvm/lib/Analysis/TFUtils.cpp

void writeTensorValues(raw_ostream &OutFile, const char *TensorData,
   for (size_t I = 0; I < ElemCount; ++I) {
     if (I > 0)
       OutFile << ", ";
     OutFile << TypedData[I];
   }
   OutFile << "]";
 }
 
-/// Untyped implementation of the API above.
+/// Write a list of tensors as a sequence of TensorFlow FeatureList protobufs.
+/// The tensors are assumed to be stored contiguously, in row-major format,
+/// in the TensorData buffer. Each tensor has the shape given by Spec. The
+/// feature name in the output is either the provided LoggingName, if
+/// specified, otherwise it's the name of the tensor (as given by Spec).
 void writeRawTensorsAsFeatureLists(raw_ostream &OutFile,
                                    const Logger::LoggedFeatureSpec &LoggedSpec,
-                                   const char *TensorData, size_t TensorCount) {
+                                   const char *TensorData, size_t TensorCount,
+                                   bool FinalReward = false) {
   const char *FieldName = "<invalid>";
   std::function<void(const char *)> ValueWriter;
   const auto &Spec = LoggedSpec.Spec;
   // The 'Feature' protobuf only has 3 possible fields: float_list,
   // int64_list, or bytes_list, so we capture int32 values as int64. We don't
   // support any other types.
   if (Spec.isElementType<int64_t>()) {
     FieldName = "int64_list";
void writeRawTensorsAsFeatureLists(raw_ostream &OutFile,
 
   OutFile << "  feature_list: {\n";
   OutFile << "    key: "
           << "\""
           << (LoggedSpec.LoggingName ? *LoggedSpec.LoggingName : Spec.name())
           << "\" ";
   OutFile << "value: {\n";
   size_t TensorByteSize = Spec.getElementCount() * Spec.getElementByteSize();
-  for (const char *P = TensorData,
-                  *E = TensorData + TensorByteSize * TensorCount;
-       P < E; P += TensorByteSize) {
+
+  auto WriteFeatureProto = [&](const char *P) {
     OutFile << "      feature: { " << FieldName << ": { value: ";
     ValueWriter(P);
     OutFile << " } }\n";
+  };
+
+  const char *CurrentTensor = TensorData;
+  static int64_t Zero = 0;
+  // Write all but the last value. If this is the final reward, don't increment
+  // the CurrentTensor, and just write 0.
+  for (size_t I = 0; I < TensorCount - 1; ++I) {
+    if (FinalReward)
+      WriteFeatureProto(reinterpret_cast<const char *>(&Zero));
+    else {
+      WriteFeatureProto(CurrentTensor);
+      CurrentTensor += TensorByteSize;
     }
-  OutFile << "    }\n";
-  OutFile << "  }\n";
   }
 
-/// Write a list of tensors as a sequence of TensorFlow FeatureList protobufs.
-/// The tensors are assumed to be stored contiguously, in row-major format,
-/// in the TensorData buffer. Each tensor has the shape given by Spec. The
-/// feature name in the output is either the provided LoggingName, if
-/// specified, otherwise it's the name of the tensor (as given by Spec).
-template <typename T>
-void writeTensorsAsFeatureLists(raw_ostream &OutFile,
-                                const Logger::LoggedFeatureSpec &Spec,
-                                const T *TensorData, size_t TensorCount) {
-  writeRawTensorsAsFeatureLists(
-      OutFile, Spec, reinterpret_cast<const char *>(TensorData), TensorCount);
+  WriteFeatureProto(CurrentTensor);
+
+  OutFile << "    }\n";
+  OutFile << "  }\n";
 }
 } // namespace
 
 namespace llvm {
 class EvaluationResultImpl {
 public:
   EvaluationResultImpl(size_t OutputSize)
       : OutputSize(OutputSize), Output(OutputSize){};
if (RawLogData.empty())
     return;
   if (RawLogData[0].empty())
     return;
   size_t Tensor0Size = FeatureSpecs[0].Spec.getElementCount() *
                        FeatureSpecs[0].Spec.getElementByteSize();
   size_t NumberOfRecords = RawLogData[0].size() / Tensor0Size;
   if (NumberOfRecords == 0)
     return;
+  size_t RewardSize =
+      RewardSpec.getElementCount() * RewardSpec.getElementByteSize();
+  size_t NumberOfRewards = RawLogData.back().size() / RewardSize;
 
   OS << "feature_lists: {\n";
   for (size_t I = 0; I < FeatureSpecs.size(); ++I)
-    writeTensorsAsFeatureLists(OS, FeatureSpecs[I], RawLogData[I].data(),
+    writeRawTensorsAsFeatureLists(OS, FeatureSpecs[I], RawLogData[I].data(),
                                   NumberOfRecords);
 
   if (IncludeReward)
-    writeTensorsAsFeatureLists(OS, {RewardSpec, None}, RawLogData.back().data(),
-                               NumberOfRecords);
+    writeRawTensorsAsFeatureLists(OS, {RewardSpec, None},
+                                  RawLogData.back().data(), NumberOfRecords,
+                                  NumberOfRewards == 1);
 
   OS << "}\n";
 }
 #endif // defined(LLVM_HAVE_TF_API)

llvm/unittests/Analysis/TFUtilsTest.cpp

feature_list: {
     }
   }
 }
 )";
   std::string Result;
   raw_string_ostream OS(Result);
   L.print(OS);
   EXPECT_EQ(Result, Expected);
-}
+}
+
+TEST(TFUtilsTest, LoggerFinalReward) {
+  std::vector<Logger::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});
+  Logger L(Features, Rewards, true);
+  for (size_t I = 0; I < 3; ++I) {
+    float F = static_cast<float>(I);
+    L.logTensorValue(0, &F);
+    L.logTensorValue(1, &I);
+  }
+  L.logFinalReward<float>(3.14);
+  const auto *Expected = R"(feature_lists: {
+  feature_list: {
+    key: "the_float" value: {
+      feature: { float_list: { value: [0.000000e+00] } }
+      feature: { float_list: { value: [1.000000e+00] } }
+      feature: { float_list: { value: [2.000000e+00] } }
+    }
+  }
+  feature_list: {
+    key: "the_int" value: {
+      feature: { int64_list: { value: [0] } }
+      feature: { int64_list: { value: [1] } }
+      feature: { int64_list: { value: [2] } }
+    }
+  }
+  feature_list: {
+    key: "reward" value: {
+      feature: { float_list: { value: [0.000000e+00] } }
+      feature: { float_list: { value: [0.000000e+00] } }
+      feature: { float_list: { value: [3.140000e+00] } }
+    }
+  }
+}
+)";
+  std::string Result;
+  raw_string_ostream OS(Result);
+  L.print(OS);
+  EXPECT_EQ(Result, Expected);
+}