[TFUtils] Expose untyped accessor to evaluation result tensors

[mtrofin]

Summary

These were implementation detail, but become necessary for generic data
copying.

Also added const variations to them, and move assignment, since we had a
move ctor (and the move assignment helps in a subsequent patch).

[yundiqian]

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

120–122

why making them public?

[mtrofin]

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

120–122

See motivation in patch description - they are necessary for generic data copying (for upcoming patch)

[MaskRay]

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

107

You can write = default

[MaskRay]

The tag [llvm] can be replaced by a more specific tag, e.g. TFUtils or some other abbreviation for the framework you are working on.

[mtrofin]

updated description

[mtrofin]

In D85262#2196818, @MaskRay wrote:

The tag [llvm] can be replaced by a more specific tag, e.g. TFUtils or some other abbreviation for the framework you are working on.

done

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

107

That doesn't work, the "impl" in the pImpl idiom is not available here. Just like the move ctor, the implementation needs to be in the .cpp file.

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

 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 &operator=(const EvaluationResult &Other) = delete;
+
     EvaluationResult(EvaluationResult &&Other);
+    EvaluationResult &operator=(EvaluationResult &&Other);
+
     ~EvaluationResult();
 
-    /// Get a pointer to the first element of the tensor at Index.
+    /// Get a (const) pointer to the first element of the tensor at Index.
     template <typename T> T *getTensorValue(size_t Index) {
       return static_cast<T *>(getUntypedTensorValue(Index));
     }
 
+    template <typename T> const T *getTensorValue(size_t Index) const {
+      return static_cast<T *>(getUntypedTensorValue(Index));
+    }
+
+    /// Get a (const) pointer to the untyped data of the tensor.
+    void *getUntypedTensorValue(size_t Index);
+    const void *getUntypedTensorValue(size_t Index) const;
+
   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<TensorSpec> &InputSpecs,
                    const std::vector<TensorSpec> &OutputSpecs,
                    const char *Tags = "serve");
   ~TFModelEvaluator();

llvm/lib/Analysis/TFUtils.cpp

 
 TFModelEvaluator::EvaluationResult::EvaluationResult(
     std::unique_ptr<EvaluationResultImpl> Impl)
     : Impl(std::move(Impl)) {}
 
 TFModelEvaluator::EvaluationResult::EvaluationResult(EvaluationResult &&Other)
     : Impl(std::move(Other.Impl)) {}
 
+TFModelEvaluator::EvaluationResult &
+TFModelEvaluator::EvaluationResult::operator=(EvaluationResult &&Other) {
+  Impl = std::move(Other.Impl);
+  return *this;
+}
+
 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; }
 
 TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_IMPL)
 
 #undef TFUTILS_GETDATATYPE_IMPL
 
 TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
 TFModelEvaluator::~TFModelEvaluator() {}
 #endif // defined(LLVM_HAVE_TF_API)

llvm/unittests/Analysis/TFUtilsTest.cpp

 static std::string getModelPath() {
   SmallString<128> InputsDir = unittest::getInputFileDirectory(TestMainArgv0);
   llvm::sys::path::append(InputsDir, "ir2native_x86_64_model");
   return std::string(InputsDir);
 }
 
 // Test observable behavior when no model is provided.
 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)
   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(), InputSpecs, OutputSpecs);
   EXPECT_TRUE(Evaluator.isValid());
 
   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_TRUE(ER.hasValue());
     float Ret = *ER->getTensorValue<float>(0);
     EXPECT_EQ(static_cast<size_t>(Ret), 80);
+    EXPECT_EQ(ER->getUntypedTensorValue(0),
+              reinterpret_cast<const void *>(ER->getTensorValue<float>(0)));
   }
   // The input vector should be unchanged
   for (auto I = 0; I < KnownSize; ++I) {
     EXPECT_EQ(V[I], 1);
   }
   // Zero-out the unused position '0' of the instruction histogram, which is
   // after the first 9 calculated values. Should the the same result.
   V[9] = 0;
   {
     auto ER = Evaluator.evaluate();
     EXPECT_TRUE(ER.hasValue());
     float Ret = *ER->getTensorValue<float>(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)
   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(), InputSpecs, OutputSpecs);
   EXPECT_TRUE(Evaluator.isValid());
 
   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());
 }
 
 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", 
TEST(TFUtilsTest, TensorSpecSizesAndTypes) {
   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));
-}
+}