From 9d93a98f853a5c5cf02c0e5e794c8a80fb11fa3a Mon Sep 17 00:00:00 2001 From: Jacob Hegna Date: Wed, 12 Oct 2022 02:05:19 +0000 Subject: [PATCH] [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 --- llvm/lib/Analysis/TFLiteUtils.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/llvm/lib/Analysis/TFLiteUtils.cpp b/llvm/lib/Analysis/TFLiteUtils.cpp index 41c9847..9a7a441 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 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(); -- 2.7.4