[NFC][TFUtils] also include output specs lookup logic in loadOutputSpecs

[mtrofin]

Summary

The lookup logic is also reusable.

Also refactored the API to return the loaded vector - this makes it more
clear what state it is in in the case of error (as it won't be
returned).

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

 Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
                                            const json::Value &Value);
 
 struct LoggedFeatureSpec {
   TensorSpec Spec;
   Optional<std::string> LoggingName;
 };
 
-bool loadOutputSpecs(LLVMContext &Ctx, StringRef FileName,
-                     StringRef ExpectedDecisionName,
-                     std::vector<LoggedFeatureSpec> &Ret);
+/// Load the output specs. If SpecFileOverride is not empty, that path is used.
+/// Otherwise, the file is assumed to be called 'output_spec.json' and be found
+/// under ModelPath (the model directory).
+/// The first output tensor name must match ExpectedDecisionName.
+/// In case of error, the return is None and the error is logged.
+Optional<std::vector<LoggedFeatureSpec>>
+loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
+                StringRef ModelPath, StringRef SpecFileOverride = StringRef());
 
 /// 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.
 /// The assumption is that, for an event to be logged (i.e. a set of feature
 /// values and a reward), the user calls the log* API for each feature exactly
 /// once, providing the index matching the position in the feature spec list
 /// provided at construction:

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

 
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/InlineSizeEstimatorAnalysis.h"
 #include "llvm/Analysis/MLInlineAdvisor.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ManagedStatic.h"
-#include "llvm/Support/Path.h"
 
 #include <vector>
 
 using namespace llvm;
 
 static cl::opt<std::string> TrainingLog(
     "training-log", cl::Hidden,
     cl::desc("Path where the development - mode inlining log is saved."));
ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx,
                                                    const std::string &ModelPath)
     : MLModelRunner(Ctx) {
   std::vector<TensorSpec> InputSpecs;
   for (size_t I = 0; I < NumberOfFeatures; ++I)
     InputSpecs.push_back(
         TensorSpec::createSpec<int64_t>(TFFeedPrefix + FeatureNameMap[I], {1}));
   InputSpecs.insert(InputSpecs.end(), TrainingOnlyFeatures.begin(),
                     TrainingOnlyFeatures.end());
-  SmallVector<char, 128> OutputSpecsPath;
-  StringRef OutputSpecPath = TFOutputSpecOverride;
-  if (OutputSpecPath.empty()) {
-    llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json");
-    OutputSpecPath = {OutputSpecsPath.data(), OutputSpecsPath.size()};
-  }
-
-  if (!loadOutputSpecs(Ctx, OutputSpecPath, DecisionName, OutputSpecs))
+  if (auto MaybeOutSpecs =
+          loadOutputSpecs(Ctx, DecisionName, ModelPath, TFOutputSpecOverride))
+    OutputSpecs = std::move(*MaybeOutSpecs);
+  else
     return;
 
   Evaluator = std::make_unique<TFModelEvaluator>(
       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();

llvm/lib/Analysis/TFUtils.cpp

 #if defined(LLVM_HAVE_TF_API)
 
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/Utils/TFUtils.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 "tensorflow/c/c_api.h"
 #include "tensorflow/c/c_api_experimental.h"
 
 #include <cassert>
 #include <numeric>
 

 #define PARSE_TYPE(T, E)                                                       \
   if (TensorType == #T)                                                        \
     return TensorSpec::createSpec<T>(TensorName, TensorShape, TensorPort);
   TFUTILS_SUPPORTED_TYPES(PARSE_TYPE)
 #undef PARSE_TYPE
   return None;
 }
 
-bool loadOutputSpecs(LLVMContext &Ctx, StringRef FileName,
-                     StringRef ExpectedDecisionName,
-                     std::vector<LoggedFeatureSpec> &Ret) {
+Optional<std::vector<LoggedFeatureSpec>>
+loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
+                StringRef ModelPath, StringRef SpecFileOverride) {
+  SmallVector<char, 128> OutputSpecsPath;
+  StringRef FileName = SpecFileOverride;
+  if (FileName.empty()) {
+    llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json");
+    FileName = {OutputSpecsPath.data(), OutputSpecsPath.size()};
+  }
+
   auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName);
   if (!BufferOrError) {
     Ctx.emitError("Error opening output specs file: " + FileName + " : " +
                   BufferOrError.getError().message());
-    return false;
+    return None;
   }
   auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer());
   if (!ParsedJSONValues) {
     Ctx.emitError("Could not parse specs file: " + FileName);
-    return false;
+    return None;
   }
   auto ValuesArray = ParsedJSONValues->getAsArray();
   if (!ValuesArray) {
     Ctx.emitError("Expected an array of {tensor_spec:<TensorSpec>, "
                   "logging_name:<name>} dictionaries");
-    return false;
+    return None;
   }
-
+  std::vector<LoggedFeatureSpec> Ret;
   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<int64_t>() &&
                 !TensorSpec->isElementType<int32_t>() &&
                 !TensorSpec->isElementType<float>()) {
               Ctx.emitError(
                   "Only int64, int32, and float tensors are supported. "
                   "Found unsupported type for tensor named " +
                   TensorSpec->name());
-              return false;
+              return None;
             }
             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;
+    return None;
   }
   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 None;
   }
-  return true;
+  return Ret;
 }
 
 class TFModelEvaluatorImpl {
 public:
   TFModelEvaluatorImpl(StringRef SavedModelPath,
                        const std::vector<TensorSpec> &InputSpecs,
                        function_ref<TensorSpec(size_t)> GetOutputSpecs,
                        size_t OutputSpecsSize, const char *Tags);