diff --git a/llvm/include/llvm/Analysis/Utils/TFUtils.h b/llvm/include/llvm/Analysis/Utils/TFUtils.h --- a/llvm/include/llvm/Analysis/Utils/TFUtils.h +++ b/llvm/include/llvm/Analysis/Utils/TFUtils.h @@ -100,6 +100,15 @@ Optional getTensorSpecFromJSON(LLVMContext &Ctx, const json::Value &Value); +struct LoggedFeatureSpec { + TensorSpec Spec; + Optional LoggingName; +}; + +bool loadOutputSpecs(LLVMContext &Ctx, StringRef FileName, + StringRef ExpectedDecisionName, + std::vector &Ret); + /// Logging utility - given an ordered specification of features, and assuming /// a scalar reward, allow logging feature values and rewards, and then print /// as tf.train.SequenceExample text protobuf. @@ -121,11 +130,6 @@ /// At the end, call print to generate the protobuf. class Logger final { public: - struct LoggedFeatureSpec { - TensorSpec Spec; - Optional LoggingName; - }; - /// Construct a Logger. If IncludeReward is false, then logReward shouldn't /// be called, and the reward feature won't be printed out. Logger(const std::vector &FeatureSpecs, @@ -201,6 +205,11 @@ const std::vector &InputSpecs, const std::vector &OutputSpecs, const char *Tags = "serve"); + TFModelEvaluator(StringRef SavedModelPath, + const std::vector &InputSpecs, + function_ref GetOutputSpecs, + size_t OutputSpecsSize, const char *Tags = "serve"); + ~TFModelEvaluator(); TFModelEvaluator(const TFModelEvaluator &) = delete; TFModelEvaluator(TFModelEvaluator &&) = delete; 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 @@ -298,9 +298,9 @@ int64_t getFeature(int Index) const override; bool isValid() const { return !!Evaluator; } - const std::vector &outputNames() const { return OutputNames; } - - const std::vector &outputSpecs() const { return OutputSpecs; } + const std::vector &outputLoggedFeatureSpecs() const { + return OutputSpecs; + } const Optional & lastEvaluationResult() const { @@ -309,12 +309,9 @@ private: std::unique_ptr Evaluator; - std::vector OutputNames; - std::vector OutputSpecs; + std::vector OutputSpecs; Optional LastEvaluationResult; - bool loadOutputSpecs(LLVMContext &Ctx, StringRef FileName); - // The training framework needs some additional features. const std::vector TrainingOnlyFeatures{ TensorSpec::createSpec(TFFeedPrefix + "inlining_default", {1}), @@ -329,14 +326,15 @@ : LogFileName(LogFileName), MUTR(MUTR) { // The first output is the inlining decision. if (MUTR) - OutputCount = MUTR->outputSpecs().size(); - std::vector FT; + OutputCount = MUTR->outputLoggedFeatureSpecs().size(); + std::vector FT; for (size_t I = 0; I < NumberOfFeatures; ++I) FT.push_back( {TensorSpec::createSpec(FeatureNameMap.at(I), {1}), None}); - for (size_t I = 1; I < OutputCount; ++I) - FT.push_back({MUTR->outputSpecs()[I], MUTR->outputNames()[I]}); + if (MUTR && MUTR->outputLoggedFeatureSpecs().size() > 1) + FT.insert(FT.end(), MUTR->outputLoggedFeatureSpecs().begin() + 1, + MUTR->outputLoggedFeatureSpecs().end()); DefaultDecisionPos = FT.size(); FT.push_back( @@ -361,7 +359,7 @@ for (size_t I = 1; I < OutputCount; ++I) { const auto &Result = *MUTR->lastEvaluationResult(); - auto &Spec = MUTR->outputSpecs()[I]; + auto &Spec = MUTR->outputLoggedFeatureSpecs()[I].Spec; const char *RawData = reinterpret_cast(Result.getUntypedTensorValue(I)); L->logTensorValue(CurrentFeature, RawData, @@ -480,11 +478,13 @@ llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json"); OutputSpecPath = {OutputSpecsPath.data(), OutputSpecsPath.size()}; } - if (!loadOutputSpecs(Ctx, OutputSpecPath)) + + if (!loadOutputSpecs(Ctx, OutputSpecPath, DecisionName, OutputSpecs)) return; - Evaluator = - std::make_unique(ModelPath, InputSpecs, OutputSpecs); + Evaluator = std::make_unique( + ModelPath, InputSpecs, [&](size_t I) { return OutputSpecs[I].Spec; }, + OutputSpecs.size()); if (!Evaluator || !Evaluator->isValid()) { Ctx.emitError("Failed to create inliner saved model evaluator"); Evaluator.reset(); @@ -492,63 +492,6 @@ } } -bool ModelUnderTrainingRunner::loadOutputSpecs(LLVMContext &Ctx, - StringRef FileName) { - auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName); - if (!BufferOrError) { - Ctx.emitError("Error opening output specs file: " + FileName + " : " + - BufferOrError.getError().message()); - return false; - } - auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer()); - if (!ParsedJSONValues) { - Ctx.emitError("Could not parse specs file: " + FileName); - return false; - } - auto ValuesArray = ParsedJSONValues->getAsArray(); - if (!ValuesArray) { - Ctx.emitError("Expected an array of {tensor_spec:, " - "logging_name:} dictionaries"); - return false; - } - - for (const auto &Value : *ValuesArray) - if (const auto *Obj = Value.getAsObject()) - if (const auto *SpecPart = Obj->get("tensor_spec")) - if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart)) - if (auto LoggingName = Obj->getString("logging_name")) { - if (!TensorSpec->isElementType() && - !TensorSpec->isElementType() && - !TensorSpec->isElementType()) { - Ctx.emitError( - "Only int64, int32, and float tensors are supported. " - "Found unsupported type for tensor named " + - TensorSpec->name()); - return false; - } - OutputNames.push_back(LoggingName->str()); - OutputSpecs.push_back(*TensorSpec); - } - - if (ValuesArray->size() != OutputNames.size()) { - Ctx.emitError( - "Unable to parse output spec. It should be a json file containing an " - "array of dictionaries. Each dictionary must have a 'tensor_spec' key, " - "with a json object describing a TensorSpec; and a 'logging_name' key, " - "which is a string to use as name when logging this tensor in the " - "training log."); - return false; - } - assert(OutputNames.size() == OutputSpecs.size()); - if (OutputNames.empty() || OutputNames[0] != DecisionName) { - Ctx.emitError("The first output spec must describe the decision tensor, " - "and must have the logging_name " + - StringRef(DecisionName)); - return false; - } - return true; -} - bool ModelUnderTrainingRunner::run() { LastEvaluationResult = Evaluator->evaluate(); if (!LastEvaluationResult.hasValue()) { diff --git a/llvm/lib/Analysis/TFUtils.cpp b/llvm/lib/Analysis/TFUtils.cpp --- a/llvm/lib/Analysis/TFUtils.cpp +++ b/llvm/lib/Analysis/TFUtils.cpp @@ -18,6 +18,7 @@ #include "llvm/Support/Debug.h" #include "llvm/Support/JSON.h" #include "llvm/Support/ManagedStatic.h" +#include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/raw_ostream.h" #include "tensorflow/c/c_api.h" @@ -83,7 +84,7 @@ /// 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 LoggedFeatureSpec &LoggedSpec, const char *TensorData, size_t TensorCount, bool FinalReward = false) { const char *FieldName = ""; @@ -215,12 +216,68 @@ return None; } +bool loadOutputSpecs(LLVMContext &Ctx, StringRef FileName, + StringRef ExpectedDecisionName, + std::vector &Ret) { + auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName); + if (!BufferOrError) { + Ctx.emitError("Error opening output specs file: " + FileName + " : " + + BufferOrError.getError().message()); + return false; + } + auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer()); + if (!ParsedJSONValues) { + Ctx.emitError("Could not parse specs file: " + FileName); + return false; + } + auto ValuesArray = ParsedJSONValues->getAsArray(); + if (!ValuesArray) { + Ctx.emitError("Expected an array of {tensor_spec:, " + "logging_name:} dictionaries"); + return false; + } + + for (const auto &Value : *ValuesArray) + if (const auto *Obj = Value.getAsObject()) + if (const auto *SpecPart = Obj->get("tensor_spec")) + if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart)) + if (auto LoggingName = Obj->getString("logging_name")) { + if (!TensorSpec->isElementType() && + !TensorSpec->isElementType() && + !TensorSpec->isElementType()) { + Ctx.emitError( + "Only int64, int32, and float tensors are supported. " + "Found unsupported type for tensor named " + + TensorSpec->name()); + return false; + } + Ret.push_back({*TensorSpec, LoggingName->str()}); + } + + if (ValuesArray->size() != Ret.size()) { + Ctx.emitError( + "Unable to parse output spec. It should be a json file containing an " + "array of dictionaries. Each dictionary must have a 'tensor_spec' key, " + "with a json object describing a TensorSpec; and a 'logging_name' key, " + "which is a string to use as name when logging this tensor in the " + "training log."); + return false; + } + if (Ret.empty() || *Ret[0].LoggingName != ExpectedDecisionName) { + Ctx.emitError("The first output spec must describe the decision tensor, " + "and must have the logging_name " + + StringRef(ExpectedDecisionName)); + return false; + } + return true; +} + class TFModelEvaluatorImpl { public: TFModelEvaluatorImpl(StringRef SavedModelPath, const std::vector &InputSpecs, - const std::vector &OutputSpecs, - const char *Tags); + function_ref GetOutputSpecs, + size_t OutputSpecsSize, const char *Tags); bool isValid() const { return IsValid; } size_t OutputSize() const { return OutputFeed.size(); } @@ -271,10 +328,11 @@ TFModelEvaluatorImpl::TFModelEvaluatorImpl( StringRef SavedModelPath, const std::vector &InputSpecs, - const std::vector &OutputSpecs, const char *Tags) + function_ref GetOutputSpecs, size_t OutputSpecsSize, + const char *Tags = "serve") : Graph(createTFGraph()), Options(createTFSessionOptions()), InputFeed(InputSpecs.size()), Input(InputSpecs.size()), - OutputFeed(OutputSpecs.size()) { + OutputFeed(OutputSpecsSize) { if (!ensureInitTF()) { errs() << "Tensorflow should have been initialized"; return; @@ -298,8 +356,8 @@ initInput(I, static_cast(InputSpec.typeIndex()), InputSpec.shape()); } - for (size_t I = 0; I < OutputSpecs.size(); ++I) { - auto &OutputSpec = OutputSpecs[I]; + for (size_t I = 0; I < OutputSpecsSize; ++I) { + auto OutputSpec = GetOutputSpecs(I); OutputFeed[I] = { TF_GraphOperationByName(Graph.get(), (OutputSpec.name()).c_str()), OutputSpec.port()}; @@ -308,15 +366,23 @@ } } +TFModelEvaluator::TFModelEvaluator( + StringRef SavedModelPath, const std::vector &InputSpecs, + function_ref GetOutputSpecs, size_t OutputSpecsSize, + const char *Tags) + : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, GetOutputSpecs, + OutputSpecsSize, Tags)) { + if (!Impl->isValid()) + Impl.reset(); +} + TFModelEvaluator::TFModelEvaluator(StringRef SavedModelPath, const std::vector &InputSpecs, const std::vector &OutputSpecs, const char *Tags) - : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, OutputSpecs, - Tags)) { - if (!Impl->isValid()) - Impl.reset(); -} + : TFModelEvaluator( + SavedModelPath, InputSpecs, [&](size_t I) { return OutputSpecs[I]; }, + OutputSpecs.size(), Tags) {} TFModelEvaluatorImpl::~TFModelEvaluatorImpl() { for (auto *T : Input) { diff --git a/llvm/unittests/Analysis/TFUtilsTest.cpp b/llvm/unittests/Analysis/TFUtilsTest.cpp --- a/llvm/unittests/Analysis/TFUtilsTest.cpp +++ b/llvm/unittests/Analysis/TFUtilsTest.cpp @@ -144,7 +144,7 @@ } TEST(TFUtilsTest, Logger) { - std::vector Features; + std::vector Features; Features.push_back( {TensorSpec::createSpec("the_float", {2, 3}), None}); Features.push_back({TensorSpec::createSpec("the_int", {2}), @@ -191,7 +191,7 @@ } TEST(TFUtilsTest, LoggerNoReward) { - std::vector Features; + std::vector Features; Features.push_back( {TensorSpec::createSpec("the_float", {2, 3}), None}); Features.push_back({TensorSpec::createSpec("the_int", {2}), @@ -230,7 +230,7 @@ } TEST(TFUtilsTest, LoggerFinalReward) { - std::vector Features; + std::vector Features; Features.push_back({TensorSpec::createSpec("the_float", {1}), None}); Features.push_back({TensorSpec::createSpec("the_int", {1}), None});