diff options
author | Dimitry Andric <dim@FreeBSD.org> | 2021-07-29 20:15:26 +0000 |
---|---|---|
committer | Dimitry Andric <dim@FreeBSD.org> | 2021-07-29 20:15:26 +0000 |
commit | 344a3780b2e33f6ca763666c380202b18aab72a3 (patch) | |
tree | f0b203ee6eb71d7fdd792373e3c81eb18d6934dd /llvm/lib/Analysis/TFUtils.cpp | |
parent | b60736ec1405bb0a8dd40989f67ef4c93da068ab (diff) | |
download | src-344a3780b2e33f6ca763666c380202b18aab72a3.tar.gz src-344a3780b2e33f6ca763666c380202b18aab72a3.zip |
Vendor import of llvm-project main 88e66fa60ae5, the last commit beforevendor/llvm-project/llvmorg-13-init-16847-g88e66fa60ae5vendor/llvm-project/llvmorg-12.0.1-rc2-0-ge7dac564cd0evendor/llvm-project/llvmorg-12.0.1-0-gfed41342a82f
the upstream release/13.x branch was created.
Diffstat (limited to 'llvm/lib/Analysis/TFUtils.cpp')
-rw-r--r-- | llvm/lib/Analysis/TFUtils.cpp | 249 |
1 files changed, 143 insertions, 106 deletions
diff --git a/llvm/lib/Analysis/TFUtils.cpp b/llvm/lib/Analysis/TFUtils.cpp index 1377cac217ab..e93dc303ae63 100644 --- a/llvm/lib/Analysis/TFUtils.cpp +++ b/llvm/lib/Analysis/TFUtils.cpp @@ -15,6 +15,7 @@ #include "llvm/ADT/Twine.h" #include "llvm/Analysis/Utils/TFUtils.h" +#include "llvm/Support/CommandLine.h" #include "llvm/Support/Debug.h" #include "llvm/Support/JSON.h" #include "llvm/Support/ManagedStatic.h" @@ -22,14 +23,22 @@ #include "llvm/Support/Path.h" #include "llvm/Support/raw_ostream.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; +using google::protobuf::Message; +using google::protobuf::TextFormat; + +static cl::opt<bool> + ProtobufTextMode("tfutils-text-log", cl::init(false), cl::Hidden, + cl::desc("Output textual (human-readable) protobuf.")); + namespace { using TFGraphPtr = std::unique_ptr<TF_Graph, decltype(&TF_DeleteGraph)>; @@ -64,89 +73,6 @@ TFStatusPtr createTFStatus() { TFSessionOptionsPtr createTFSessionOptions() { return TFSessionOptionsPtr(TF_NewSessionOptions(), &TF_DeleteSessionOptions); } - -/// Write the values of one tensor as a list. -template <typename T> -void writeTensorValues(raw_ostream &OutFile, const char *TensorData, - size_t ElemCount) { - OutFile << "["; - const T *TypedData = reinterpret_cast<const T *>(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). -void writeRawTensorsAsFeatureLists(raw_ostream &OutFile, - const LoggedFeatureSpec &LoggedSpec, - 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"; - ValueWriter = [&](const char *Data) { - writeTensorValues<int64_t>(OutFile, Data, Spec.getElementCount()); - }; - } else if (Spec.isElementType<int32_t>()) { - FieldName = "int64_list"; - ValueWriter = [&](const char *Data) { - writeTensorValues<int32_t>(OutFile, Data, Spec.getElementCount()); - }; - - } else if (Spec.isElementType<float>()) { - FieldName = "float_list"; - ValueWriter = [&](const char *Data) { - writeTensorValues<float>(OutFile, Data, Spec.getElementCount()); - }; - - } else { - llvm_unreachable("Unsupported tensor type."); - } - - OutFile << " feature_list: {\n"; - OutFile << " key: " - << "\"" - << (LoggedSpec.LoggingName ? *LoggedSpec.LoggingName : Spec.name()) - << "\" "; - OutFile << "value: {\n"; - size_t TensorByteSize = Spec.getElementCount() * Spec.getElementByteSize(); - - 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; - } - } - - WriteFeatureProto(CurrentTensor); - - OutFile << " }\n"; - OutFile << " }\n"; -} } // namespace namespace llvm { @@ -332,6 +258,76 @@ private: bool checkReportAndInvalidate(const TF_Output &Output, const TensorSpec &OutputSpec); }; + +class LoggerDataImpl { + const std::vector<LoggedFeatureSpec> LoggedFeatureSpecs; + const TensorSpec RewardSpec; + + tensorflow::SequenceExample SE; + std::vector<tensorflow::FeatureList *> FeatureLists; + tensorflow::FeatureList *Reward = nullptr; + +public: + LoggerDataImpl(const std::vector<LoggedFeatureSpec> &LoggedSpecs, + const TensorSpec &RewardSpec, bool IncludeReward) + : LoggedFeatureSpecs(LoggedSpecs), RewardSpec(RewardSpec) { + auto *FL = SE.mutable_feature_lists()->mutable_feature_list(); + if (IncludeReward) + Reward = &(*FL)[RewardSpec.name()]; + // Allocate first the map entries, then capture their address. We will not + // mutate the set of features after this (i.e. the pointers won't dangle). + for (const auto &LFS : LoggedSpecs) { + (*FL)[LFS.LoggingName ? *LFS.LoggingName : LFS.Spec.name()] = {}; + } + for (const auto &LFS : LoggedSpecs) + FeatureLists.push_back( + &(*FL)[LFS.LoggingName ? *LFS.LoggingName : LFS.Spec.name()]); + } + + void print(raw_ostream &OS) { + std::string OutStr; + if (ProtobufTextMode) + google::protobuf::TextFormat::PrintToString(SE, &OutStr); + else + OutStr = SE.SerializeAsString(); + + OS << OutStr; + } + + char *addNewTensor(size_t FeatureID) { + const auto &Spec = LoggedFeatureSpecs[FeatureID].Spec; + if (Spec.isElementType<float>()) { + auto *RF = FeatureLists[FeatureID] + ->add_feature() + ->mutable_float_list() + ->mutable_value(); + RF->Resize(Spec.getElementCount(), 0.0); + return reinterpret_cast<char *>(RF->mutable_data()); + } else if (Spec.isElementType<int32_t>() || Spec.isElementType<int64_t>()) { + auto *RF = FeatureLists[FeatureID] + ->add_feature() + ->mutable_int64_list() + ->mutable_value(); + RF->Resize(Spec.getElementCount(), 0); + return reinterpret_cast<char *>(RF->mutable_data()); + } + llvm_unreachable("Unsupported tensor type."); + } + + template <typename T> void logReward(T Value) { + if (RewardSpec.isElementType<float>()) + Reward->add_feature()->mutable_float_list()->add_value(Value); + else if (RewardSpec.isElementType<int32_t>() || + RewardSpec.isElementType<int64_t>()) + Reward->add_feature()->mutable_int64_list()->add_value(Value); + else + llvm_unreachable("Unsupported tensor type."); + } + + size_t getNrRecords() const { + return FeatureLists.empty() ? 0 : FeatureLists[0]->feature().size(); + } +}; } // namespace llvm TFModelEvaluatorImpl::TFModelEvaluatorImpl( @@ -476,30 +472,71 @@ TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_IMPL) TFModelEvaluator::EvaluationResult::~EvaluationResult() {} TFModelEvaluator::~TFModelEvaluator() {} -void Logger::print(raw_ostream &OS) { - 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; +Logger::Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs, + const TensorSpec &RewardSpec, bool IncludeReward) + : FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec), + IncludeReward(IncludeReward), + LoggerData(std::make_unique<LoggerDataImpl>(FeatureSpecs, RewardSpec, + IncludeReward)) {} - OS << "feature_lists: {\n"; - for (size_t I = 0; I < FeatureSpecs.size(); ++I) - writeRawTensorsAsFeatureLists(OS, FeatureSpecs[I], RawLogData[I].data(), - NumberOfRecords); +Logger::~Logger() {} - if (IncludeReward) - writeRawTensorsAsFeatureLists(OS, {RewardSpec, None}, - RawLogData.back().data(), NumberOfRecords, - NumberOfRewards == 1); +#define LOG_REWARD(NAME, TYPE) \ + void Logger::log##NAME##Reward(TYPE Value) { \ + assert(IncludeReward); \ + LoggerData->logReward(Value); \ + } - OS << "}\n"; +LOG_REWARD(Float, float) +LOG_REWARD(Int32, int32_t) +LOG_REWARD(Int64, int64_t) +#undef LOG_REWARD + +#define LOG_FINAL_REWARD(NAME, TYPE) \ + void Logger::log##NAME##FinalReward(TYPE Value) { \ + assert(RewardSpec.isElementType<TYPE>()); \ + for (size_t I = 1; I < LoggerData->getNrRecords(); ++I) \ + log##NAME##Reward(0); \ + log##NAME##Reward(Value); \ + } + +LOG_FINAL_REWARD(Float, float) +LOG_FINAL_REWARD(Int32, int32_t) +LOG_FINAL_REWARD(Int64, int64_t) +#undef LOG_FINAL_REWARD + +void Logger::logFloatValue(size_t FeatureID, const float *Value) { + assert(FeatureSpecs[FeatureID].Spec.isElementType<float>()); + logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value)); +} + +void Logger::logInt64Value(size_t FeatureID, const int64_t *Value) { + assert(FeatureSpecs[FeatureID].Spec.isElementType<int64_t>()); + logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value)); +} + +void Logger::logInt32Value(size_t FeatureID, const int32_t *Value) { + assert(FeatureSpecs[FeatureID].Spec.isElementType<int32_t>()); + logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value)); +} + +void Logger::logSpecifiedTensorValue(size_t FeatureID, const char *RawData) { + const auto &Spec = FeatureSpecs[FeatureID].Spec; + char *Buff = addEntryAndGetFloatOrInt64Buffer(FeatureID); + if (Spec.isElementType<int32_t>()) + for (size_t I = 0; I < Spec.getElementCount(); ++I) + (reinterpret_cast<int64_t *>(Buff))[I] = + static_cast<int64_t>((reinterpret_cast<const int32_t *>(RawData))[I]); + else if (Spec.isElementType<int64_t>() || Spec.isElementType<float>()) + std::memcpy(Buff, RawData, + Spec.getElementCount() * Spec.getElementByteSize()); + else + llvm_unreachable("Unsupported tensor type"); } + +char *Logger::addEntryAndGetFloatOrInt64Buffer(size_t FeatureID) { + return reinterpret_cast<char *>(LoggerData->addNewTensor(FeatureID)); +} + +void Logger::print(raw_ostream &OS) { LoggerData->print(OS); } #endif // defined(LLVM_HAVE_TF_API) |