[llvm] 9d93a98 - [MLGO] Force persistency in tflite buffers.

[Jacob Hegna via llvm-commits]

Author: Jacob Hegna
Date: 2022-10-12T19:50:36Z
New Revision: 9d93a98f853a5c5cf02c0e5e794c8a80fb11fa3a

URL: https://github.com/llvm/llvm-project/commit/9d93a98f853a5c5cf02c0e5e794c8a80fb11fa3a
DIFF: https://github.com/llvm/llvm-project/commit/9d93a98f853a5c5cf02c0e5e794c8a80fb11fa3a.diff

LOG: [MLGO] Force persistency in tflite buffers.

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".

Differential Revision: https://reviews.llvm.org/D135739

Added: 
    

Modified: 
    llvm/lib/Analysis/TFLiteUtils.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/TFLiteUtils.cpp b/llvm/lib/Analysis/TFLiteUtils.cpp
index 41c9847ad64af..9a7a441e0d269 100644
--- a/llvm/lib/Analysis/TFLiteUtils.cpp
+++ b/llvm/lib/Analysis/TFLiteUtils.cpp
@@ -121,6 +121,15 @@ TFModelEvaluatorImpl::TFModelEvaluatorImpl(
   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 
diff erent 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();