[MLGO] Make TFLiteUtils throw an error if some features haven't been passed to the model

[aidengrossman]

Summary

In the Tensorflow C lib utilities, an error gets thrown if some features
haven't gotten passed into the model (due to differences in ordering
which now don't exist with the transition to TFLite). However, this is
not currently the case when using TFLiteUtils. This patch makes some
minor changes to throw an error when not all inputs of the model have
been passed, which when not handled will result in a seg fault within
TFLite.

[mtrofin]

Ah, good catch. Can you add a test, perhaps TFUtilsTest.cpp?

[aidengrossman]

Add unit test to ensure that the Evaluator is invalid if an evaluation is made
where there is a feature that is missing.

llvm/lib/Analysis/TFLiteUtils.cpp

TFModelEvaluatorImpl::TFModelEvaluatorImpl(
   // Known inputs and outputs
   StringMap<int> InputsMap;
   StringMap<int> OutputsMap;
   for (size_t I = 0; I < Interpreter->inputs().size(); ++I)
     InputsMap[Interpreter->GetInputName(I)] = I;
   for (size_t I = 0; I < Interpreter->outputs().size(); ++I)
     OutputsMap[Interpreter->GetOutputName(I)] = I;
 
+  size_t NumberFeaturesPassed = 0;
   for (size_t I = 0; I < InputSpecs.size(); ++I) {
     auto &InputSpec = InputSpecs[I];
     auto MapI = InputsMap.find(InputSpec.name() + ":" +
                                std::to_string(InputSpec.port()));
     if (MapI == InputsMap.end()) {
       Input[I] = nullptr;
       continue;
     }
     Input[I] = Interpreter->tensor(MapI->second);
     if (!checkReportAndInvalidate(Input[I], InputSpec))
       return;
     std::memset(Input[I]->data.data, 0,
                 InputSpecs[I].getTotalTensorBufferSize());
+    ++NumberFeaturesPassed;
+  }
+
+  if (NumberFeaturesPassed < Interpreter->inputs().size()) {
+    // we haven't passed all the required features to the model, throw an error.
+    errs() << "Required feature(s) have not been passed to the ML model";
+    invalidate();
+    return;
   }
 
   for (size_t I = 0; I < OutputSpecsSize; ++I) {
     auto OutputSpec = GetOutputSpecs(I);
     Output[I] = Interpreter->output_tensor(
         OutputsMap[OutputSpec.name() + ":" +
                    std::to_string(OutputSpec.port())]);
     if (!checkReportAndInvalidate(Output[I], OutputSpec))

llvm/unittests/Analysis/TFUtilsTest.cpp

TEST(TFUtilsTest, UnsupportedFeature) {
   float Ret = Evaluator->evaluate<float>();
   EXPECT_EQ(static_cast<int64_t>(Ret), 80);
   // The input vector should be unchanged
   for (auto I = 0; I < KnownSize; ++I)
     EXPECT_EQ(V[I], 1);
   for (auto I = 0; I < 2 * 5; ++I)
     EXPECT_FLOAT_EQ(F[I], 3.14 + I);
 }
+
+TEST(TFUtilsTest, MissingFeature) {
+  std::vector<TensorSpec> InputSpecs{};
+  std::vector<TensorSpec> OutputSpecs{
+      TensorSpec::createSpec<float>("StatefulPartitionedCall", {1})};
+
+  TFModelEvaluator Evaluator(getModelPath(), InputSpecs, OutputSpecs);
+  EXPECT_FALSE(Evaluator.isValid());
+}