diff --git a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp --- a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp +++ b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp @@ -81,13 +81,84 @@ void print(raw_fd_ostream &OutFile); private: + /// Write the values of one tensor as a list. template - void writeTensor(raw_fd_ostream &OutFile, StringRef TensorName, - const std::vector &Tensor); + void writeTensorValues(raw_fd_ostream &OutFile, const char *TensorData, + size_t ElemCount) const { + OutFile << "["; + const T *TypedData = reinterpret_cast(TensorData); + for (size_t I = 0; I < ElemCount; ++I) { + if (I > 0) + OutFile << ", "; + OutFile << TypedData[I]; + } + OutFile << "]"; + } + + /// 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 + void + writeTensorsAsFeatureLists(raw_fd_ostream &OutFile, const TensorSpec &Spec, + const T *TensorData, size_t TensorCount, + Optional LoggingName = None) const { + writeRawTensorsAsFeatureLists(OutFile, Spec, + reinterpret_cast(TensorData), + TensorCount, LoggingName); + } + + /// Untyped implementation of the API above. + void + writeRawTensorsAsFeatureLists(raw_fd_ostream &OutFile, const TensorSpec &Spec, + const char *TensorData, size_t TensorCount, + Optional LoggingName = None) const { + const char *FieldName = ""; + std::function ValueWriter; + // 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()) { + FieldName = "int64_list"; + ValueWriter = [&](const char *Data) { + writeTensorValues(OutFile, Data, Spec.getElementCount()); + }; + } else if (Spec.isElementType()) { + FieldName = "int64_list"; + ValueWriter = [&](const char *Data) { + writeTensorValues(OutFile, Data, Spec.getElementCount()); + }; + + } else if (Spec.isElementType()) { + FieldName = "float_list"; + ValueWriter = [&](const char *Data) { + writeTensorValues(OutFile, Data, Spec.getElementCount()); + }; + + } else + llvm_unreachable("Unsupported tensor type."); + + OutFile << " feature_list: {\n"; + OutFile << " key: " + << "\"" << (LoggingName ? *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) { + OutFile << " feature: { " << FieldName << ": { value: "; + ValueWriter(P); + OutFile << " } }\n"; + } + OutFile << " }\n"; + OutFile << " }\n"; + } std::vector Features; - std::vector DefaultDecisions; - std::vector Decisions; + std::vector DefaultDecisions; + std::vector Decisions; std::vector Effects; std::vector Rewards; }; @@ -294,35 +365,30 @@ } void TrainingLogger::print(raw_fd_ostream &OutFile) { - if (DefaultDecisions.empty()) + size_t NumberOfRecords = Decisions.size(); + if (NumberOfRecords == 0) return; - OutFile << "feature_lists: {\n"; - for (size_t I = 0; I < Features.size(); I++) { - writeTensor(OutFile, FeatureNameMap.at(I), Features[I]); - } - writeTensor(OutFile, DefaultDecisionName, DefaultDecisions); - writeTensor(OutFile, DecisionName, Decisions); - writeTensor(OutFile, RewardName, Rewards); + OutFile << "feature_lists: {\n"; + for (size_t I = 0; I < Features.size(); ++I) + writeTensorsAsFeatureLists( + OutFile, TensorSpec::createSpec(FeatureNameMap.at(I), {1}), + Features[I].data(), NumberOfRecords); + + writeTensorsAsFeatureLists( + OutFile, TensorSpec::createSpec(DefaultDecisionName, {1}), + DefaultDecisions.data(), NumberOfRecords); + + writeTensorsAsFeatureLists(OutFile, + TensorSpec::createSpec(DecisionName, {1}), + Decisions.data(), NumberOfRecords); + writeTensorsAsFeatureLists(OutFile, + TensorSpec::createSpec(RewardName, {1}), + Rewards.data(), NumberOfRecords); OutFile << "}\n"; } -template -void TrainingLogger::writeTensor(raw_fd_ostream &OutFile, StringRef TensorName, - const std::vector &Tensor) { - OutFile << " feature_list: {\n"; - OutFile << " key: " - << "\"" << TensorName << "\" "; - OutFile << "value: {\n"; - for (const auto &Feature : Tensor) { - OutFile << " feature: { int64_list: { value: [" << Feature - << "] } }\n"; - } - OutFile << " }\n"; - OutFile << " }\n"; -} - DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor( Module &M, ModuleAnalysisManager &MAM, std::unique_ptr ModelRunner,