From 15092dfdba4007dc7475cdf56f0dfa72ac280355 Mon Sep 17 00:00:00 2001 From: suleshahid <110432064+suleshahid@users.noreply.github.com> Date: Thu, 21 Mar 2024 21:54:02 +0000 Subject: [PATCH] IF kernel nulltpr fix (#2515) In some cases if you IF op has 0 outputs, it can reference a nullptr in trying to access node->outputs->size(). This is in the case if your model was flatbuffer_aligned which most likely optimizes away the empty array to nullptr. BUG=[328527454](https://buganizer.corp.google.com/issues/328527454) --- tensorflow/lite/micro/kernels/if.cc | 2 +- tensorflow/lite/micro/kernels/kernel_util.cc | 1 + tensorflow/lite/micro/micro_interpreter_graph.cc | 4 +++- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/micro/kernels/if.cc b/tensorflow/lite/micro/kernels/if.cc index 7aa8eb15..029846b2 100644 --- a/tensorflow/lite/micro/kernels/if.cc +++ b/tensorflow/lite/micro/kernels/if.cc @@ -67,7 +67,7 @@ TfLiteStatus IfPrepare(TfLiteContext* context, TfLiteNode* node) { // passed to the branch subgraphs. Therefore, the number of subgraph inputs // will be the number of node inputs - 1. size_t num_inputs = node->inputs->size - 1; - size_t num_outputs = node->outputs->size; + size_t num_outputs = NumOutputs(node); MicroGraph& graph_info = micro_context->graph(); diff --git a/tensorflow/lite/micro/kernels/kernel_util.cc b/tensorflow/lite/micro/kernels/kernel_util.cc index 12adf361..a509f5d6 100644 --- a/tensorflow/lite/micro/kernels/kernel_util.cc +++ b/tensorflow/lite/micro/kernels/kernel_util.cc @@ -252,6 +252,7 @@ TfLiteStatus CopySubgraphOutputsToOpOutputs(TfLiteContext* context, TfLiteNode* node, MicroGraph* graph_info, int subgraph_idx) { + if (graph_info->NumSubgraphOutputs(subgraph_idx) == 0) return kTfLiteOk; TF_LITE_ENSURE(context, static_cast(node->outputs->size) == graph_info->NumSubgraphOutputs(subgraph_idx)); for (int i = 0; i < node->outputs->size; i++) { diff --git a/tensorflow/lite/micro/micro_interpreter_graph.cc b/tensorflow/lite/micro/micro_interpreter_graph.cc index dbd7072a..0d18fe73 100644 --- a/tensorflow/lite/micro/micro_interpreter_graph.cc +++ b/tensorflow/lite/micro/micro_interpreter_graph.cc @@ -259,7 +259,9 @@ TfLiteEvalTensor* MicroInterpreterGraph::GetSubgraphInput(int subgraph_idx, } size_t MicroInterpreterGraph::NumSubgraphOutputs(int subgraph_idx) { - return model_->subgraphs()->Get(subgraph_idx)->outputs()->size(); + return model_->subgraphs()->Get(subgraph_idx)->outputs() == nullptr + ? 0 + : model_->subgraphs()->Get(subgraph_idx)->outputs()->size(); } TfLiteEvalTensor* MicroInterpreterGraph::GetSubgraphOutput(int subgraph_idx,