[mlgo][nfc] Decouple TensorSpec from tensorflow.

[mtrofin]

Summary

The motivation is twofold:

1. Allow plugging in a different training-time evaluator, e.g. TFLite-based, etc.

2. Allow using TensorSpec for AOT, too, to support evolution: we start by extracting a superset of the features currently supported by a model. For the tensors the model does not support, we just return a valid, but useless, buffer. This makes using a 'smaller' model (less supported tensors) transparent to the compiler. The key is to dimension the buffer appropriately, and we already have TensorSpec modeling that info.

The only coupling was due to the reliance of a TF internal API for
getting the element size, but for the types we are interested in,
sizeof is sufficient.

A subsequent change will yank out TensorSpec in its own module.

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

 /// 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.
+
+/// Known tensor types.
+#define SUPPORTED_TENSOR_TYPES(M)                                              \
+  M(float, Float)                                                              \
+  M(double, Double)                                                            \
+  M(int8_t, Int8)                                                              \
+  M(uint8_t, UInt8)                                                            \
+  M(int16_t, Int16)                                                            \
+  M(uint16_t, UInt16)                                                          \
+  M(int32_t, Int32)                                                            \
+  M(uint32_t, UInt32)                                                          \
+  M(int64_t, Int64)                                                            \
+  M(uint64_t, UInt64)
+
+enum class TensorType {
+  Invalid,
+#define _TENSOR_TYPE_ENUM_MEMBERS(_, Name) Name,
+  SUPPORTED_TENSOR_TYPES(_TENSOR_TYPE_ENUM_MEMBERS)
+#undef _TENSOR_TYPE_ENUM_MEMBERS
+};
+
 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);
+    return TensorSpec(Name, Port, getDataType<T>(), sizeof(T), Shape);
   }
 
   const std::string &name() const { return Name; }
   int port() const { return Port; }
-  int typeIndex() const { return TypeIndex; }
+  TensorType type() const { return Type; }
   const std::vector<int64_t> &shape() const { return Shape; }
 
   bool operator==(const TensorSpec &Other) const {
-    return Name == Other.Name && Port == Other.Port &&
-           TypeIndex == Other.TypeIndex && Shape == Other.Shape;
+    return Name == Other.Name && Port == Other.Port && Type == Other.Type &&
+           Shape == Other.Shape;
   }
 
   bool operator!=(const TensorSpec &Other) const { return !(*this == Other); }
 
   /// Get the number of elements in a tensor with this shape.
   size_t getElementCount() const { return ElementCount; }
   /// Get the size, in bytes, of one element.
-  size_t getElementByteSize() const;
+  size_t getElementByteSize() const { return ElementSize; }
 
   template <typename T> bool isElementType() const {
-    return getDataType<T>() == TypeIndex;
+    return getDataType<T>() == Type;
   }
 
 private:
-  TensorSpec(const std::string &Name, int Port, int TypeIndex,
-             const std::vector<int64_t> &Shape);
+  TensorSpec(const std::string &Name, int Port, TensorType Type,
+             size_t ElementSize, const std::vector<int64_t> &Shape);
 
-  template <typename T> static int getDataType() {
-    llvm_unreachable("Undefined tensor type");
-  }
+  template <typename T> static TensorType getDataType();
 
   std::string Name;
   int Port = 0;
-  int TypeIndex = 0;
+  TensorType Type = TensorType::Invalid;
   std::vector<int64_t> Shape;
   size_t ElementCount = 0;
+  size_t ElementSize = 0;
 };
 
 /// Construct a TensorSpec from a JSON dictionary of the form:
 /// { "name": <string>,
 ///   "port": <int>,
 ///   "type": <string. Use LLVM's types, e.g. float, double, int64_t>,
 ///   "shape": <array of ints> }
 /// For the "type" field, see the C++ primitive types used in
public:
 
   /// Untyped access to input.
   void *getUntypedInput(size_t Index);
 
 private:
   std::unique_ptr<TFModelEvaluatorImpl> Impl;
 };
 
-/// List of supported types, as a pair:
-/// - C++ type
-/// - enum name (implementation-specific)
-#define TFUTILS_SUPPORTED_TYPES(M)                                             \
-  M(float, TF_FLOAT)                                                           \
-  M(double, TF_DOUBLE)                                                         \
-  M(int8_t, TF_INT8)                                                           \
-  M(uint8_t, TF_UINT8)                                                         \
-  M(int16_t, TF_INT16)                                                         \
-  M(uint16_t, TF_UINT16)                                                       \
-  M(int32_t, TF_INT32)                                                         \
-  M(uint32_t, TF_UINT32)                                                       \
-  M(int64_t, TF_INT64)                                                         \
-  M(uint64_t, TF_UINT64)
-
-#define TFUTILS_GETDATATYPE_DEF(T, E)                                          \
-  template <> int TensorSpec::getDataType<T>();
-
-TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_DEF)
+#define TFUTILS_GETDATATYPE_DEF(T, Name)                                       \
+  template <> TensorType TensorSpec::getDataType<T>();
+SUPPORTED_TENSOR_TYPES(TFUTILS_GETDATATYPE_DEF)
 
 #undef TFUTILS_GETDATATYPE_DEF
 } // namespace llvm
 
 #endif // LLVM_HAVE_TF_API
 #endif // LLVM_ANALYSIS_UTILS_TFUTILS_H

llvm/lib/Analysis/TFUtils.cpp

 //===- TFUtils.cpp - tensorflow evaluation utilities ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utilities for interfacing with tensorflow C APIs.
 //
 //===----------------------------------------------------------------------===//
+#include "tensorflow/c/tf_datatype.h"
 #include "llvm/Config/config.h"
+#include "llvm/Support/ErrorHandling.h"
 #if defined(LLVM_HAVE_TF_API)
 
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
 #include "llvm/Support/Base64.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/JSON.h"

 
 void serialize(const Message &SE, std::string *OutStr) {
   if (ProtobufTextMode) {
     TextFormat::PrintToString(SE, OutStr);
   } else {
     *OutStr = SE.SerializeAsString();
   }
 }
+
+int getTFTypeIndex(TensorType TType) {
+  switch (TType) {
+  case TensorType::Double:
+    return TF_DOUBLE;
+  case TensorType::Float:
+    return TF_FLOAT;
+  case TensorType::Int8:
+    return TF_INT8;
+  case TensorType::UInt8:
+    return TF_UINT8;
+  case TensorType::Int16:
+    return TF_INT16;
+  case TensorType::UInt16:
+    return TF_UINT16;
+  case TensorType::Int32:
+    return TF_INT32;
+  case TensorType::UInt32:
+    return TF_UINT32;
+  case TensorType::Int64:
+    return TF_INT64;
+  case TensorType::UInt64:
+    return TF_UINT64;
+  case TensorType::Invalid:
+    llvm_unreachable("Unknown tensor type");
+  }
+}
 } // namespace
 
 namespace llvm {
 class EvaluationResultImpl {
 public:
   EvaluationResultImpl(size_t OutputSize)
       : OutputSize(OutputSize), Output(OutputSize){};
 
   ~EvaluationResultImpl() {
     for (auto *P : Output)
       if (P)
         TF_DeleteTensor(P);
   }
 
   EvaluationResultImpl(const EvaluationResultImpl &) = delete;
   EvaluationResultImpl(EvaluationResultImpl &&Other) = delete;
   std::vector<TF_Tensor *> &getOutput() { return Output; }
 
 private:
   const size_t OutputSize;
   std::vector<TF_Tensor *> Output;
 };
 
-size_t TensorSpec::getElementByteSize() const {
-  return TF_DataTypeSize(static_cast<TF_DataType>(TypeIndex));
-}
-
-TensorSpec::TensorSpec(const std::string &Name, int Port, int TypeIndex,
-                       const std::vector<int64_t> &Shape)
-    : Name(Name), Port(Port), TypeIndex(TypeIndex), Shape(Shape),
+TensorSpec::TensorSpec(const std::string &Name, int Port, TensorType Type,
+                       size_t ElementSize, const std::vector<int64_t> &Shape)
+    : Name(Name), Port(Port), Type(Type), Shape(Shape),
       ElementCount(std::accumulate(Shape.begin(), Shape.end(), 1,
-                                   std::multiplies<int64_t>())) {}
+                                   std::multiplies<int64_t>())),
+      ElementSize(ElementSize) {}
 
 Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
                                            const json::Value &Value) {
   auto EmitError = [&](const llvm::Twine &Message) -> Optional<TensorSpec> {
     std::string S;
     llvm::raw_string_ostream OS(S);
     OS << Value;
     Ctx.emitError("Unable to parse JSON Value as spec (" + Message + "): " + S);
Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
   if (!Mapper.map<int>("port", TensorPort))
     return EmitError("'port' property not present or not an int");
   if (!Mapper.map<std::vector<int64_t>>("shape", TensorShape))
     return EmitError("'shape' property not present or not an int array");
 
 #define PARSE_TYPE(T, E)                                                       \
   if (TensorType == #T)                                                        \
     return TensorSpec::createSpec<T>(TensorName, TensorShape, TensorPort);
-  TFUTILS_SUPPORTED_TYPES(PARSE_TYPE)
+  SUPPORTED_TENSOR_TYPES(PARSE_TYPE)
 #undef PARSE_TYPE
   return None;
 }
 
 Optional<std::vector<LoggedFeatureSpec>>
 loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
                 StringRef ModelPath, StringRef SpecFileOverride) {
   SmallVector<char, 128> OutputSpecsPath;
TFModelEvaluatorImpl::TFModelEvaluatorImpl(
   }
   for (size_t I = 0; I < InputSpecs.size(); ++I) {
     auto &InputSpec = InputSpecs[I];
     InputFeed[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()),
+    initInput(I, static_cast<TF_DataType>(getTFTypeIndex(InputSpec.type())),
               InputSpec.shape());
   }
   for (size_t I = 0; I < OutputSpecsSize; ++I) {
     auto OutputSpec = GetOutputSpecs(I);
     OutputFeed[I] = {
         TF_GraphOperationByName(Graph.get(), (OutputSpec.name()).c_str()),
         OutputSpec.port()};
     if (!checkReportAndInvalidate(OutputFeed[I], OutputSpec))

 }
 
 const void *
 TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) const {
   return TF_TensorData(Impl->getOutput()[Index]);
 }
 
 #define TFUTILS_GETDATATYPE_IMPL(T, E)                                         \
-  template <> int TensorSpec::getDataType<T>() { return E; }
+  template <> TensorType TensorSpec::getDataType<T>() { return TensorType::E; }
 
-TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_IMPL)
+SUPPORTED_TENSOR_TYPES(TFUTILS_GETDATATYPE_IMPL)
 
 #undef TFUTILS_GETDATATYPE_IMPL
 
 TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
 TFModelEvaluator::~TFModelEvaluator() {}
 
 Logger::Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
                const TensorSpec &RewardSpec, bool IncludeReward)