[llvm][NFC] TensorSpec abstraction for ML evaluator

[mtrofin]

Summary

Further abstracting the specification of a tensor, to more easily
support different types and shapes of tensor, and also to perform
initialization up-front, at TFModelEvaluator construction time.

[mtrofin]

Removed a spurious empty line and an unnecessary friend relationship.

[davidxl]

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

39

document the class briefly?

[mtrofin]

Doc for TensorSpec

[ebrevdo]

llvm/lib/Analysis/TFUtils.cpp

136

OutputSpec

163

InputSpec

170

OutputSpec

202–203

OutputSpec

[mtrofin]

var renames

[mtrofin]

ptal

[ebrevdo]

LGTM

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

 ///   setting internal scalars, for all dimensions (tensors are row-major:
 ///   https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/c/c_api.h#L205)
 /// - call evaluate. The input tensors' values are not consumed after this, and
 ///   may still be read.
 /// - use the outputs in the output vector
 class TFModelEvaluatorImpl;
 class EvaluationResultImpl;
 
+/// TensorSpec encapsulates the specification of a tensor: its dimensions, or
+/// "shape" (row-major), its type (see TensorSpec::getDataType specializations
+/// for supported types), its name and port (see "TensorFlow: Large-Scale
+/// Machine Learning on Heterogeneous Distributed Systems", section 4.2, para 2:
+/// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/45166.pdf)
+///
+/// TensorSpec is used to set up a TFModelEvaluator by describing the expected
+/// inputs and outputs.
+class TensorSpec final {
+public:
+  template <typename T>
+  static TensorSpec createSpec(const std::string &Name,
+                               const std::vector<int64_t> &Shape,
+                               int Port = 0) {
+    return TensorSpec(Name, Port, getDataType<T>(), Shape);
+  }
+
+  const std::string &name() const { return Name; }
+  int port() const { return Port; }
+  int typeIndex() const { return TypeIndex; }
+  const std::vector<int64_t> &shape() const { return Shape; }
+
+private:
+  TensorSpec(const std::string &Name, int Port, int TypeIndex,
+             const std::vector<int64_t> &Shape)
+      : Name(Name), Port(Port), TypeIndex(TypeIndex), Shape(Shape) {}
+
+  template <typename T> static int getDataType() {
+    llvm_unreachable("Undefined tensor type");
+  }
+
+  std::string Name;
+  int Port = 0;
+  int TypeIndex = 0;
+  std::vector<int64_t> Shape;
+};
+
 class TFModelEvaluator final {
 public:
   /// The result of a model evaluation. Handles the lifetime of the output
   /// tensors, which means that their values need to be used before
   /// the EvaluationResult's dtor is called.
   class EvaluationResult {
   public:
     EvaluationResult(const EvaluationResult &) = delete;
     EvaluationResult(EvaluationResult &&Other);
     ~EvaluationResult();
 
     /// Get a pointer to the first element of the tensor at Index.
     template <typename T> T *getTensorValue(size_t Index) {
       return static_cast<T *>(getUntypedTensorValue(Index));
     }
 
   private:
     friend class TFModelEvaluator;
     EvaluationResult(std::unique_ptr<EvaluationResultImpl> Impl);
     void *getUntypedTensorValue(size_t Index);
     std::unique_ptr<EvaluationResultImpl> Impl;
   };
 
   TFModelEvaluator(StringRef SavedModelPath,
-                   const std::vector<std::string> &InputNames,
-                   const std::vector<std::string> &OutputNames,
+                   const std::vector<TensorSpec> &InputSpecs,
+                   const std::vector<TensorSpec> &OutputSpecs,
                    const char *Tags = "serve");
   ~TFModelEvaluator();
   TFModelEvaluator(const TFModelEvaluator &) = delete;
   TFModelEvaluator(TFModelEvaluator &&) = delete;
 
   /// Evaluate the model, assuming it is valid. Returns None if the evaluation
   /// fails or the model is invalid, or an EvaluationResult otherwise. The
   /// inputs are assumed to have been already provided via getInput(). When
   /// returning None, it also invalidates this object.
   Optional<EvaluationResult> evaluate();
 
   /// Provides access to the input vector.
   template <typename T> T *getInput(size_t Index) {
     return static_cast<T *>(getUntypedInput(Index));
   }
 
   /// Returns true if the tensorflow model was loaded successfully, false
   /// otherwise.
   bool isValid() const { return !!Impl; }
 
-  /// Initialize the input at Index as a tensor of the given type and
-  /// dimensions.
-  template <typename T>
-  void initInput(size_t Index, const std::vector<int64_t> &Dimensions) {
-    return initInput(Index, getModelTypeIndex<T>(), Dimensions);
-  }
-
 private:
   void *getUntypedInput(size_t Index);
-  template <typename T> int getModelTypeIndex();
-  void initInput(size_t Index, int TypeIndex,
-                 const std::vector<int64_t> &Dimensions);
-
   std::unique_ptr<TFModelEvaluatorImpl> Impl;
 };
 
-template <> int TFModelEvaluator::getModelTypeIndex<float>();
-template <> int TFModelEvaluator::getModelTypeIndex<double>();
-template <> int TFModelEvaluator::getModelTypeIndex<int8_t>();
-template <> int TFModelEvaluator::getModelTypeIndex<uint8_t>();
-template <> int TFModelEvaluator::getModelTypeIndex<int16_t>();
-template <> int TFModelEvaluator::getModelTypeIndex<uint16_t>();
-template <> int TFModelEvaluator::getModelTypeIndex<int32_t>();
-template <> int TFModelEvaluator::getModelTypeIndex<uint32_t>();
-template <> int TFModelEvaluator::getModelTypeIndex<int64_t>();
-template <> int TFModelEvaluator::getModelTypeIndex<uint64_t>();
+template <> int TensorSpec::getDataType<float>();
+template <> int TensorSpec::getDataType<double>();
+template <> int TensorSpec::getDataType<int8_t>();
+template <> int TensorSpec::getDataType<uint8_t>();
+template <> int TensorSpec::getDataType<int16_t>();
+template <> int TensorSpec::getDataType<uint16_t>();
+template <> int TensorSpec::getDataType<int32_t>();
+template <> int TensorSpec::getDataType<uint32_t>();
+template <> int TensorSpec::getDataType<int64_t>();
+template <> int TensorSpec::getDataType<uint64_t>();
 
 } // namespace llvm
 
 #endif // LLVM_HAVE_TF_API
 #endif // LLVM_ANALYSIS_UTILS_TFUTILS_H

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

public:
 
   void setFeature(FeatureIndex Index, int64_t Value) override;
   int64_t getFeature(int Index) const override;
   bool isValid() const { return !!Evaluator; }
 
 private:
   std::unique_ptr<TFModelEvaluator> Evaluator;
 
-  // The training framework needs some additional features, that just need to
-  // be set to 0.
-  struct TensorSpec {
-    std::string Name;
-    std::function<void(TFModelEvaluator *, size_t Index,
-                       const std::vector<int64_t> &Dim)>
-        Initializer;
-  };
-
+  // The training framework needs some additional features.
   const std::vector<TensorSpec> TrainingOnlyFeatures{
-      {"inlining_default",
-       [](TFModelEvaluator *Evaluator, size_t Index,
-          const std::vector<int64_t> &Dim) {
-         Evaluator->initInput<int64_t>(Index, Dim);
-       }},
-      {"discount",
-       [](TFModelEvaluator *Evaluator, size_t Index,
-          const std::vector<int64_t> &Dim) {
-         Evaluator->initInput<float>(Index, Dim);
-       }},
-      {"reward",
-       [](TFModelEvaluator *Evaluator, size_t Index,
-          const std::vector<int64_t> &Dim) {
-         Evaluator->initInput<float>(Index, Dim);
-       }},
-      {"step_type", [](TFModelEvaluator *Evaluator, size_t Index,
-                       const std::vector<int64_t> &Dim) {
-         Evaluator->initInput<int32_t>(Index, Dim);
-       }}};
+      TensorSpec::createSpec<int64_t>(TFFeedPrefix + "inlining_default", {1}),
+      TensorSpec::createSpec<float>(TFFeedPrefix + "discount", {1}),
+      TensorSpec::createSpec<float>(TFFeedPrefix + "reward", {1}),
+      TensorSpec::createSpec<int32_t>(TFFeedPrefix + "step_type", {1})};
 };
 } // namespace
 
 DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(
     Module &M, ModuleAnalysisManager &MAM,
     std::unique_ptr<MLModelRunner> ModelRunner,
     std::function<bool(CallBase &)> GetDefaultAdvice, bool IsDoingInference)
     : MLInlineAdvisor(M, MAM, std::move(ModelRunner)),
for (auto &F : M) {
     Ret += getNativeSizeEstimate(F);
   }
   return Ret;
 }
 
 ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx,
                                                    const std::string &ModelPath)
     : MLModelRunner(Ctx) {
-  std::vector<std::string> InputNames;
-  std::vector<std::string> OutputNames;
+  std::vector<TensorSpec> InputSpecs;
+  std::vector<TensorSpec> OutputSpecs;
   for (size_t I = 0; I < NumberOfFeatures; ++I)
-    InputNames.push_back(TFFeedPrefix + FeatureNameMap[I]);
-  for (size_t I = 0; I < TrainingOnlyFeatures.size(); ++I)
-    InputNames.push_back(TFFeedPrefix + TrainingOnlyFeatures[I].Name);
-  OutputNames.push_back(TFDecisionName);
+    InputSpecs.push_back(
+        TensorSpec::createSpec<int64_t>(TFFeedPrefix + FeatureNameMap[I], {1}));
+  InputSpecs.insert(InputSpecs.end(), TrainingOnlyFeatures.begin(),
+                    TrainingOnlyFeatures.end());
+  OutputSpecs.push_back(TensorSpec::createSpec<int64_t>(TFDecisionName, {1}));
 
   Evaluator =
-      std::make_unique<TFModelEvaluator>(ModelPath, InputNames, OutputNames);
+      std::make_unique<TFModelEvaluator>(ModelPath, InputSpecs, OutputSpecs);
   if (!Evaluator || !Evaluator->isValid()) {
     Ctx.emitError("Failed to create inliner saved model evaluator");
     Evaluator.reset();
     return;
   }
-
-  static const std::vector<int64_t> Dim{1};
-
-  size_t InputIndex = 0;
-  for (; InputIndex < NumberOfFeatures; ++InputIndex) {
-    Evaluator->initInput<int64_t>(InputIndex, Dim);
-  }
-
-  for (; InputIndex < InputNames.size(); ++InputIndex) {
-    TrainingOnlyFeatures[InputIndex - NumberOfFeatures].Initializer(
-        Evaluator.get(), InputIndex, Dim);
-  }
 }
 
 bool ModelUnderTrainingRunner::run() {
   auto ER = Evaluator->evaluate();
   if (!ER.hasValue()) {
     Ctx.emitError("Error evaluating model.");
     return false;
   }

llvm/lib/Analysis/InlineSizeEstimatorAnalysis.cpp

 bool InlineSizeEstimatorAnalysis::isEvaluatorRequested() {
   return !TFIR2NativeModelPath.empty();
 }
 
 InlineSizeEstimatorAnalysis::InlineSizeEstimatorAnalysis() {
   if (!isEvaluatorRequested()) {
     return;
   }
-  std::vector<std::string> InputNames{"serving_default_input_1"};
-  std::vector<std::string> OutputName{"StatefulPartitionedCall"};
+  std::vector<TensorSpec> InputSpecs{TensorSpec::createSpec<int32_t>(
+      "serving_default_input_1",
+      {1, static_cast<int64_t>(
+              IRToNativeSizeLearning::FunctionFeatures::FeatureCount)})};
+  std::vector<TensorSpec> OutputSpecs{
+      TensorSpec::createSpec<float>("StatefulPartitionedCall", {1})};
   Evaluator = std::make_unique<TFModelEvaluator>(
-      TFIR2NativeModelPath.getValue().c_str(), InputNames, OutputName);
+      TFIR2NativeModelPath.getValue().c_str(), InputSpecs, OutputSpecs);
   if (!Evaluator || !Evaluator->isValid()) {
     Evaluator.reset();
     return;
   }
-  static const std::vector<int64_t> Dim{
-      1, static_cast<int64_t>(
-             IRToNativeSizeLearning::FunctionFeatures::FeatureCount)};
-
-  Evaluator->initInput<int32_t>(0, Dim);
 }
 
 InlineSizeEstimatorAnalysis::Result
 InlineSizeEstimatorAnalysis::run(const Function &F,
                                  FunctionAnalysisManager &FAM) {
   if (!Evaluator)
     return None;
   auto Features = IRToNativeSizeLearning::getFunctionFeatures(

llvm/lib/Analysis/TFUtils.cpp

 private:
   const size_t OutputSize;
   std::vector<TF_Tensor *> Output;
 };
 
 class TFModelEvaluatorImpl {
 public:
   TFModelEvaluatorImpl(StringRef SavedModelPath,
-                       const std::vector<std::string> &InputNames,
-                       const std::vector<std::string> &OutputNames,
+                       const std::vector<TensorSpec> &InputSpecs,
+                       const std::vector<TensorSpec> &OutputSpecs,
                        const char *Tags);
 
   bool isValid() const { return IsValid; }
   size_t OutputSize() const { return OutputFeed.size(); }
 
   void evaluate(TF_Tensor **Output, TF_Status *Status) {
     TF_SessionRun(Session, nullptr, InputFeed.data(), Input.data(),
                   Input.size(), OutputFeed.data(), Output, OutputFeed.size(),
private:
   std::vector<TF_Output> OutputFeed;
 
   void invalidate() { IsValid = false; }
 
   bool IsValid = true;
 
   /// Reusable utility for ensuring we can bind the requested Name to a node in
   /// the SavedModel Graph.
-  bool checkReportAndInvalidate(const TF_Output &Output, StringRef Name);
+  bool checkReportAndInvalidate(const TF_Output &Output,
+                                const TensorSpec &OutputSpec);
 };
 } // namespace llvm
 
 TFModelEvaluatorImpl::TFModelEvaluatorImpl(
-    StringRef SavedModelPath, const std::vector<std::string> &InputNames,
-    const std::vector<std::string> &OutputNames, const char *Tags)
+    StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
+    const std::vector<TensorSpec> &OutputSpecs, const char *Tags)
     : Graph(createTFGraph()), Options(createTFSessionOptions()),
-      InputFeed(InputNames.size()), Input(InputNames.size()),
-      OutputFeed(OutputNames.size()) {
+      InputFeed(InputSpecs.size()), Input(InputSpecs.size()),
+      OutputFeed(OutputSpecs.size()) {
   if (!ensureInitTF()) {
     errs() << "Tensorflow should have been initialized";
     return;
   }
   auto Status = createTFStatus();
 
   Session = TF_LoadSessionFromSavedModel(Options.get(), nullptr,
                                          SavedModelPath.str().c_str(), &Tags, 1,
                                          Graph.get(), nullptr, Status.get());
   if (TF_GetCode(Status.get()) != TF_Code::TF_OK) {
     errs() << TF_Message(Status.get());
     invalidate();
   }
-  for (size_t I = 0; I < InputNames.size(); ++I) {
+  for (size_t I = 0; I < InputSpecs.size(); ++I) {
+    auto &InputSpec = InputSpecs[I];
     InputFeed[I] = {
-        TF_GraphOperationByName(Graph.get(), (InputNames[I]).c_str()), 0};
-    if (!checkReportAndInvalidate(InputFeed[I], InputNames[I]))
+        TF_GraphOperationByName(Graph.get(), (InputSpec.name()).c_str()),
+        InputSpec.port()};
+    if (!checkReportAndInvalidate(InputFeed[I], InputSpec))
       return;
+    initInput(I, static_cast<TF_DataType>(InputSpec.typeIndex()),
+              InputSpec.shape());
   }
-  for (size_t I = 0; I < OutputNames.size(); ++I) {
+  for (size_t I = 0; I < OutputSpecs.size(); ++I) {
+    auto &OutputSpec = OutputSpecs[I];
     OutputFeed[I] = {
-        TF_GraphOperationByName(Graph.get(), (OutputNames[I]).c_str()), 0};
-    if (!checkReportAndInvalidate(OutputFeed[I], OutputNames[I]))
+        TF_GraphOperationByName(Graph.get(), (OutputSpec.name()).c_str()),
+        OutputSpec.port()};
+    if (!checkReportAndInvalidate(OutputFeed[I], OutputSpec))
       return;
   }
 }
 
 TFModelEvaluator::TFModelEvaluator(StringRef SavedModelPath,
-                                   const std::vector<std::string> &InputNames,
-                                   const std::vector<std::string> &OutputNames,
+                                   const std::vector<TensorSpec> &InputSpecs,
+                                   const std::vector<TensorSpec> &OutputSpecs,
                                    const char *Tags)
-    : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputNames, OutputNames,
+    : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, OutputSpecs,
                                     Tags)) {
   if (!Impl->isValid())
     Impl.reset();
 }
 
 TFModelEvaluatorImpl::~TFModelEvaluatorImpl() {
   for (auto *T : Input) {
     TF_DeleteTensor(T);
   }
   if (Session == nullptr)
     return;
   auto Status = createTFStatus();
   TF_DeleteSession(Session, Status.get());
   Session = nullptr;
   if (TF_GetCode(Status.get()) != TF_Code::TF_OK)
     errs() << "Could not delete TF session";
 }
 
-bool TFModelEvaluatorImpl::checkReportAndInvalidate(const TF_Output &Output,
-                                                    StringRef Name) {
+bool TFModelEvaluatorImpl::checkReportAndInvalidate(
+    const TF_Output &Output, const TensorSpec &OutputSpec) {
   if (Output.oper)
     return true;
-  errs() << "Could not find TF_Output named: " + Name;
+  errs() << "Could not find TF_Output named: " + OutputSpec.name();
   IsValid = false;
   return IsValid;
 }
 
 Optional<TFModelEvaluator::EvaluationResult> TFModelEvaluator::evaluate() {
   if (!isValid())
     return None;
   std::unique_ptr<EvaluationResultImpl> Ret =

 
 TFModelEvaluator::EvaluationResult::EvaluationResult(EvaluationResult &&Other)
     : Impl(std::move(Other.Impl)) {}
 
 void *TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) {
   return TF_TensorData(Impl->getOutput()[Index]);
 }
 
-void TFModelEvaluator::initInput(size_t Index, int TypeIndex,
-                                 const std::vector<int64_t> &Dimensions) {
-  Impl->initInput(Index, static_cast<TF_DataType>(TypeIndex), Dimensions);
-}
+template <> int TensorSpec::getDataType<float>() { return TF_FLOAT; }
 
-template <> int TFModelEvaluator::getModelTypeIndex<float>() {
-  return TF_FLOAT;
-}
+template <> int TensorSpec::getDataType<double>() { return TF_DOUBLE; }
 
-template <> int TFModelEvaluator::getModelTypeIndex<double>() {
-  return TF_DOUBLE;
-}
+template <> int TensorSpec::getDataType<int8_t>() { return TF_INT8; }
 
-template <> int TFModelEvaluator::getModelTypeIndex<int8_t>() {
-  return TF_INT8;
-}
+template <> int TensorSpec::getDataType<uint8_t>() { return TF_UINT8; }
 
-template <> int TFModelEvaluator::getModelTypeIndex<uint8_t>() {
-  return TF_UINT8;
-}
+template <> int TensorSpec::getDataType<int16_t>() { return TF_INT16; }
 
-template <> int TFModelEvaluator::getModelTypeIndex<int16_t>() {
-  return TF_INT16;
-}
+template <> int TensorSpec::getDataType<uint16_t>() { return TF_UINT16; }
 
-template <> int TFModelEvaluator::getModelTypeIndex<uint16_t>() {
-  return TF_UINT16;
-}
-
-template <> int TFModelEvaluator::getModelTypeIndex<int32_t>() {
-  return TF_INT32;
-}
+template <> int TensorSpec::getDataType<int32_t>() { return TF_INT32; }
 
-template <> int TFModelEvaluator::getModelTypeIndex<uint32_t>() {
-  return TF_UINT32;
-}
+template <> int TensorSpec::getDataType<uint32_t>() { return TF_UINT32; }
 
-template <> int TFModelEvaluator::getModelTypeIndex<int64_t>() {
-  return TF_INT64;
-}
+template <> int TensorSpec::getDataType<int64_t>() { return TF_INT64; }
 
-template <> int TFModelEvaluator::getModelTypeIndex<uint64_t>() {
-  return TF_UINT64;
-}
+template <> int TensorSpec::getDataType<uint64_t>() { return TF_UINT64; }
 
 TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
 TFModelEvaluator::~TFModelEvaluator() {}
 #endif // defined(LLVM_HAVE_TF_API)

llvm/unittests/Analysis/TFUtilsTest.cpp

TEST(TFUtilsTest, NoModel) {
   TFModelEvaluator Evaluator("", {}, {});
   EXPECT_FALSE(Evaluator.isValid());
 }
 
 // Test we can correctly load a savedmodel and evaluate it.
 TEST(TFUtilsTest, LoadAndExecuteTest) {
   // We use the ir2native model for test. We know it has one feature of
   // dimension (1, 214)
-  std::vector<std::string> InputNames{"serving_default_input_1"};
-  std::vector<std::string> OutputName{"StatefulPartitionedCall"};
   const static int64_t KnownSize = 214;
+  std::vector<TensorSpec> InputSpecs{TensorSpec::createSpec<int32_t>(
+      "serving_default_input_1", {1, KnownSize})};
+  std::vector<TensorSpec> OutputSpecs{
+      TensorSpec::createSpec<float>("StatefulPartitionedCall", {1})};
 
-  TFModelEvaluator Evaluator(getModelPath(), InputNames, OutputName);
-  static const std::vector<int64_t> Dim{1, KnownSize};
-
+  TFModelEvaluator Evaluator(getModelPath(), InputSpecs, OutputSpecs);
   EXPECT_TRUE(Evaluator.isValid());
-  Evaluator.initInput<int32_t>(0, Dim);
 
   int32_t *V = Evaluator.getInput<int32_t>(0);
   // Fill it up with 1's, we know the output.
   for (auto I = 0; I < KnownSize; ++I) {
     V[I] = 1;
   }
   {
     auto ER = Evaluator.evaluate();
V[9] = 0;
     EXPECT_EQ(static_cast<size_t>(Ret), 80);
   }
 }
 
 // Test incorrect input setup
 TEST(TFUtilsTest, EvalError) {
   // We use the ir2native model for test. We know it has one feature of
   // dimension (1, 214)
-  std::vector<std::string> InputNames{"serving_default_input_1"};
-  std::vector<std::string> OutputName{"StatefulPartitionedCall"};
   const static int64_t KnownSize = 213;
+  std::vector<TensorSpec> InputSpecs{TensorSpec::createSpec<int32_t>(
+      "serving_default_input_1", {1, KnownSize})};
+  std::vector<TensorSpec> OutputSpecs{
+      TensorSpec::createSpec<float>("StatefulPartitionedCall", {1})};
 
-  TFModelEvaluator Evaluator(getModelPath(), InputNames, OutputName);
-  static const std::vector<int64_t> Dim{1, KnownSize};
-
+  TFModelEvaluator Evaluator(getModelPath(), InputSpecs, OutputSpecs);
   EXPECT_TRUE(Evaluator.isValid());
-  Evaluator.initInput<int32_t>(0, Dim);
 
   int32_t *V = Evaluator.getInput<int32_t>(0);
   // Fill it up with 1's, we know the output.
   for (auto I = 0; I < KnownSize; ++I) {
     V[I] = 1;
   }
   auto ER = Evaluator.evaluate();
   EXPECT_FALSE(ER.hasValue());
   EXPECT_FALSE(Evaluator.isValid());
 }