[MLInliner] Simplify TFUTILS_SUPPORTED_TYPES

[mtrofin]

Summary

We only need the C++ type and the corresponding TF Enum. The other
parameter was used for the output spec json file, but we can just
standardize on the C++ type name there.

[ebrevdo]

This TypeSpec will no longer match specs generated by stringifying tensorflow types. For example.

In [4]: tf.int64.name 
Out[4]: 'int64'

meaning we will have to do the string conversion in python instead of in C++; either way we can't really get away from it. Your choice where you want us to do it.

[mtrofin]

This is intentional. We mainly wanted to avoid a third type alias, i.e. we have LLVM's aliases (float, doble, int32_t, etc); and tensorflow's enums (TF_FLOAT, TF_DOUBLE, TF_INT32). The uses of these 2 groups is typechecked. The one we eliminated, the string, which was only used in json, was not. Arguably, the resulting code is more robust.

In D86549#2236733, @ebrevdo wrote:

This TypeSpec will no longer match specs generated by stringifying tensorflow types. For example.

In [4]: tf.int64.name 
Out[4]: 'int64'

meaning we will have to do the string conversion in python instead of in C++; either way we can't really get away from it. Your choice where you want us to do it.

[yundiqian]

Yes, Mircea and I discussed about it. Doing conversion in either place is fine. The argument of doing the it here is better isolation of TF and LLVM and easier testing in LLVM

In D86549#2236733, @ebrevdo wrote:

This TypeSpec will no longer match specs generated by stringifying tensorflow types. For example.

In [4]: tf.int64.name 
Out[4]: 'int64'

meaning we will have to do the string conversion in python instead of in C++; either way we can't really get away from it. Your choice where you want us to do it.

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

private:
 
   std::string Name;
   int Port = 0;
   int TypeIndex = 0;
   std::vector<int64_t> Shape;
   size_t ElementCount = 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
+/// TFUTILS_SUPPORTED_TYPES.
 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.
public:
   /// otherwise.
   bool isValid() const { return !!Impl; }
 
 private:
   void *getUntypedInput(size_t Index);
   std::unique_ptr<TFModelEvaluatorImpl> Impl;
 };
 
-/// List of supported types, as a triple:
-/// C++ type
-/// short name (for strings, for instance)
-/// capitalized short name (for enums, for instance)
+/// List of supported types, as a pair:
+/// - C++ type
+/// - enum name (implementation-specific)
 #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)
+  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, S, C)                                       \
+#define TFUTILS_GETDATATYPE_DEF(T, E)                                          \
   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

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)                                                        \
+#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;
 }
 
 class TFModelEvaluatorImpl {
 public:
void *TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) {
   return TF_TensorData(Impl->getOutput()[Index]);
 }
 
 const void *
 TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) const {
   return TF_TensorData(Impl->getOutput()[Index]);
 }
 
-#define TFUTILS_GETDATATYPE_IMPL(T, S, E)                                      \
-  template <> int TensorSpec::getDataType<T>() { return TF_##E; }
+#define TFUTILS_GETDATATYPE_IMPL(T, E)                                         \
+  template <> int TensorSpec::getDataType<T>() { return E; }
 
 TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_IMPL)
 
 #undef TFUTILS_GETDATATYPE_IMPL
 
 TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
 TFModelEvaluator::~TFModelEvaluator() {}
 #endif // defined(LLVM_HAVE_TF_API)

llvm/lib/Analysis/models/inliner/output_spec.json

 [
     {
         "logging_name": "inlining_decision",
         "tensor_spec": {
             "name": "StatefulPartitionedCall",
             "port": 0,
-            "type": "int64",
+            "type": "int64_t",
             "shape": [
                 1
             ]
         }
     }
 ]
  
 No newline at end of file

llvm/test/Transforms/Inline/ML/Inputs/test_output_spec.json

 [
     {
         "logging_name": "inlining_decision",
         "tensor_spec": {
             "name": "StatefulPartitionedCall",
             "port": 0,
-            "type": "int64",
+            "type": "int64_t",
             "shape": [
                 1
             ]
         }
     },
     {
         "logging_name": "fake_extra_output",
         "tensor_spec": {
             "name": "StatefulPartitionedCall",
             "port": 0,
-            "type": "int64",
+            "type": "int64_t",
             "shape": [
                 1
             ]
         }
     }
 ]
  
 No newline at end of file

llvm/unittests/Analysis/TFUtilsTest.cpp

TEST(TFUtilsTest, EvalError) {
   EXPECT_FALSE(ER.hasValue());
   EXPECT_FALSE(Evaluator.isValid());
 }
 
 TEST(TFUtilsTest, JSONParsing) {
   auto Value = json::parse(
       R"({"name": "tensor_name", 
         "port": 2, 
-        "type": "int32", 
+        "type": "int32_t", 
         "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));
 }