[MLGO] Force persistency in tflite buffers.

[jacobhegna]

Summary

When training large models, we encounter use-after-free bugs when
writing to the input tensors for various MLGO models. This patch fixes the
issue by marking the tensors as "persistent".

[mtrofin]

Thanks!

[mtrofin]

llvm/lib/Analysis/TFLiteUtils.cpp

129

Move the !Interpreter check before this

[jacobhegna]

Move interpreter check before modifying the allocation type.

llvm/lib/Analysis/TFLiteUtils.cpp

if (!Model) {
     invalidate();
     return;
   }
 
   tflite::ops::builtin::BuiltinOpResolver Resolver;
   tflite::InterpreterBuilder Builder(*Model, Resolver);
   Builder(&Interpreter);
 
+  // We assume the input buffers are valid for the lifetime of the interpreter.
+  // By default, tflite allocates memory in an arena and will periodically take
+  // away memory and reallocate it in a different location after evaluations in
+  // order to improve utilization of the buffers owned in the arena. So, we
+  // explicitly mark our input buffers as persistent to avoid this behavior.
+  for (size_t I = 0; I < Interpreter->inputs().size(); ++I)
+    Interpreter->tensor(I)->allocation_type =
+        TfLiteAllocationType::kTfLiteArenaRwPersistent;
+
   if (!Interpreter ||
       Interpreter->AllocateTensors() != TfLiteStatus::kTfLiteOk) {
     invalidate();
     return;
   }
   // Known inputs and outputs
   StringMap<int> InputsMap;
   StringMap<int> OutputsMap;