[llvm] Expose type and element count-related APIs on TensorSpec

[mtrofin]

Summary

Added a mechanism to check the element type, get the total element
count, and the size of an element.

[davidxl]

What is the use case of the new APIs?

[mtrofin]

In D85250#2194715, @davidxl wrote:

What is the use case of the new APIs?

They are a small but necessary part of a change that allows us to import a specification of additional outputs that the model can produce and then need to be captured in the training log. This enables "final reward" training algorithms. Since they are relatively "general goodness", I didn't add that context here.

[davidxl]

lgtm (as Eugene pointed out -- use std::accumulate instead ).

[mtrofin]

using accumulate

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

public:
 
   bool operator==(const TensorSpec &Other) const {
     return Name == Other.Name && Port == Other.Port &&
            TypeIndex == Other.TypeIndex && 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;
+
+  template <typename T> bool isElementType() const {
+    return getDataType<T>() == TypeIndex;
+  }
+
 private:
   TensorSpec(const std::string &Name, int Port, int TypeIndex,
-             const std::vector<int64_t> &Shape)
-      : Name(Name), Port(Port), TypeIndex(TypeIndex), Shape(Shape) {}
+             const std::vector<int64_t> &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;
+  size_t ElementCount = 0;
 };
 
 Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
                                            const json::Value &Value);
 
 class TFModelEvaluator final {
 public:
   /// The result of a model evaluation. Handles the lifetime of the output

llvm/lib/Analysis/TFUtils.cpp

 #include "llvm/Support/JSON.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/c/c_api_experimental.h"
 
 #include <cassert>
+#include <numeric>
 
 using namespace llvm;
 
 namespace {
 
 using TFGraphPtr = std::unique_ptr<TF_Graph, decltype(&TF_DeleteGraph)>;
 using TFSessionOptionsPtr =
     std::unique_ptr<TF_SessionOptions, decltype(&TF_DeleteSessionOptions)>;
public:
   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),
+      ElementCount(std::accumulate(Shape.begin(), Shape.end(), 1,
+                                   std::multiplies<int64_t>())) {}
+
 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);
     return None;

llvm/unittests/Analysis/TFUtilsTest.cpp

auto Value = json::parse(
         "shape":[1,4]
         }
       )");
   EXPECT_TRUE(!!Value);
   LLVMContext Ctx;
   auto Spec = getTensorSpecFromJSON(Ctx, *Value);
   EXPECT_FALSE(Spec.hasValue());
 }
+
+TEST(TFUtilsTest, TensorSpecSizesAndTypes) {
+  auto Spec1D = TensorSpec::createSpec<int16_t>("Hi1", {1});
+  auto Spec2D = TensorSpec::createSpec<int16_t>("Hi2", {1, 1});
+  auto Spec1DLarge = TensorSpec::createSpec<float>("Hi3", {10});
+  auto Spec3DLarge = TensorSpec::createSpec<float>("Hi3", {2, 4, 10});
+  EXPECT_TRUE(Spec1D.isElementType<int16_t>());
+  EXPECT_FALSE(Spec3DLarge.isElementType<double>());
+  EXPECT_EQ(Spec1D.getElementCount(), 1);
+  EXPECT_EQ(Spec2D.getElementCount(), 1);
+  EXPECT_EQ(Spec1DLarge.getElementCount(), 10);
+  EXPECT_EQ(Spec3DLarge.getElementCount(), 80);
+  EXPECT_EQ(Spec3DLarge.getElementByteSize(), sizeof(float));
+  EXPECT_EQ(Spec1D.getElementByteSize(), sizeof(int16_t));
+}
 No newline at end of file