[llvm] Add a parser from JSON to TensorSpec

[mtrofin]

Summary

A JSON->TensorSpec utility we will use subsequently to specify
additional outputs needed for certain training scenarios.

[mtrofin]

One more test

[MaskRay]

llvm/unittests/Analysis/TFUtilsTest.cpp

116

How about:

R"({

})");

You can remove two lines.

125

Expand auto if the type is not obvious.

https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable

126

No newline at end of file

[mtrofin]

feedback

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

 //===- TFUtils.h - utilities for tensorflow C API ---------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 #ifndef LLVM_ANALYSIS_UTILS_TFUTILS_H
 #define LLVM_ANALYSIS_UTILS_TFUTILS_H
 
 #include "llvm/Config/llvm-config.h"
 
 #ifdef LLVM_HAVE_TF_API
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/JSON.h"
 
 #include <memory>
 #include <vector>
 
 namespace llvm {
 
 /// Load a SavedModel, find the given inputs and outputs, and setup storage
 /// for input tensors. The user is responsible for correctly dimensioning the
static TensorSpec createSpec(const std::string &Name,
     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; }
 
+  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); }
+
 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;
 };
 
+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
   /// tensors, which means that their values need to be used before
   /// the EvaluationResult's dtor is called.
   class EvaluationResult {
   public:
     EvaluationResult(const EvaluationResult &) = delete;
public:
   /// otherwise.
   bool isValid() const { return !!Impl; }
 
 private:
   void *getUntypedInput(size_t Index);
   std::unique_ptr<TFModelEvaluatorImpl> Impl;
 };
 
-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>();
+/// List of supported types, as a triple:
+/// C++ type
+/// short name (for strings, for instance)
+/// capitalized short name (for enums, for instance)
+#define TFUTILS_SUPPORTED_TYPES(M)                                             \
+  M(float, float, FLOAT)                                                       \
+  M(double, double, DOUBLE)                                                    \
+  M(int8_t, int8, INT8)                                                        \
+  M(uint8_t, uint8, UINT8)                                                     \
+  M(int16_t, int16, INT16)                                                     \
+  M(uint16_t, uint16, UINT16)                                                  \
+  M(int32_t, int32, INT32)                                                     \
+  M(uint32_t, uint32, UINT32)                                                  \
+  M(int64_t, int64, INT64)                                                     \
+  M(uint64_t, uint64, UINT64)
+
+#define TFUTILS_GETDATATYPE_DEF(T, S, C)                                       \
+  template <> int TensorSpec::getDataType<T>();
+
+TFUTILS_SUPPORTED_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 ----------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utilities for interfacing with tensorflow C APIs.
 //
 //===----------------------------------------------------------------------===//
 #include "llvm/Config/config.h"
 #if defined(LLVM_HAVE_TF_API)
 
-#include "llvm/Analysis/Utils/TFUtils.h"
 #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/raw_ostream.h"
 
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/c/c_api_experimental.h"
 
 #include <cassert>
 
public:
   EvaluationResultImpl(EvaluationResultImpl &&Other) = delete;
   std::vector<TF_Tensor *> &getOutput() { return Output; }
 
 private:
   const size_t OutputSize;
   std::vector<TF_Tensor *> Output;
 };
 
+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;
+  };
+  json::ObjectMapper Mapper(Value);
+  if (!Mapper)
+    return EmitError("Value is not a dict");
+
+  std::string TensorName;
+  int TensorPort = -1;
+  std::string TensorType;
+  std::vector<int64_t> TensorShape;
+
+  if (!Mapper.map<std::string>("name", TensorName))
+    return EmitError("'name' property not present or not a string");
+  if (!Mapper.map<std::string>("type", TensorType))
+    return EmitError("'type' property not present or not a string");
+  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, S, E)                                                    \
+  if (TensorType == #S)                                                        \
+    return TensorSpec::createSpec<T>(TensorName, TensorShape, TensorPort);
+  TFUTILS_SUPPORTED_TYPES(PARSE_TYPE)
+#undef PARSE_TYPE
+  return None;
+}
+
 class TFModelEvaluatorImpl {
 public:
   TFModelEvaluatorImpl(StringRef SavedModelPath,
                        const std::vector<TensorSpec> &InputSpecs,
                        const std::vector<TensorSpec> &OutputSpecs,
                        const char *Tags);
 
   bool isValid() const { return IsValid; }

 
 TFModelEvaluator::EvaluationResult::EvaluationResult(EvaluationResult &&Other)
     : Impl(std::move(Other.Impl)) {}
 
 void *TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) {
   return TF_TensorData(Impl->getOutput()[Index]);
 }
 
-template <> int TensorSpec::getDataType<float>() { return TF_FLOAT; }
-
-template <> int TensorSpec::getDataType<double>() { return TF_DOUBLE; }
-
-template <> int TensorSpec::getDataType<int8_t>() { return TF_INT8; }
-
-template <> int TensorSpec::getDataType<uint8_t>() { return TF_UINT8; }
-
-template <> int TensorSpec::getDataType<int16_t>() { return TF_INT16; }
-
-template <> int TensorSpec::getDataType<uint16_t>() { return TF_UINT16; }
-
-template <> int TensorSpec::getDataType<int32_t>() { return TF_INT32; }
-
-template <> int TensorSpec::getDataType<uint32_t>() { return TF_UINT32; }
+#define TFUTILS_GETDATATYPE_IMPL(T, S, E)                                      \
+  template <> int TensorSpec::getDataType<T>() { return TF_##E; }
 
-template <> int TensorSpec::getDataType<int64_t>() { return TF_INT64; }
+TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_IMPL)
 
-template <> int TensorSpec::getDataType<uint64_t>() { return TF_UINT64; }
+#undef TFUTILS_GETDATATYPE_IMPL
 
 TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
 TFModelEvaluator::~TFModelEvaluator() {}
 #endif // defined(LLVM_HAVE_TF_API)

llvm/unittests/Analysis/TFUtilsTest.cpp

TEST(TFUtilsTest, EvalError) {
   // 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());
 }
+
+TEST(TFUtilsTest, JSONParsing) {
+  auto Value = json::parse(
+      R"({"name": "tensor_name", 
+        "port": 2, 
+        "type": "int32", 
+        "shape":[1,4]
+        })");
+  EXPECT_TRUE(!!Value);
+  LLVMContext Ctx;
+  Optional<TensorSpec> Spec = getTensorSpecFromJSON(Ctx, *Value);
+  EXPECT_TRUE(Spec.hasValue());
+  EXPECT_EQ(*Spec, TensorSpec::createSpec<int32_t>("tensor_name", {1, 4}, 2));
+}
+
+TEST(TFUtilsTest, JSONParsingInvalidTensorType) {
+  auto Value = json::parse(
+      R"(
+        {"name": "tensor_name", 
+        "port": 2, 
+        "type": "no such type", 
+        "shape":[1,4]
+        }
+      )");
+  EXPECT_TRUE(!!Value);
+  LLVMContext Ctx;
+  auto Spec = getTensorSpecFromJSON(Ctx, *Value);
+  EXPECT_FALSE(Spec.hasValue());
+}