| 207 | } |
| 208 | |
| 209 | TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { |
| 210 | const OpData* op_data = reinterpret_cast<OpData*>(node->user_data); |
| 211 | Subgraph* this_subgraph = reinterpret_cast<Subgraph*>(context->impl_); |
| 212 | auto* subgraphs = this_subgraph->GetSubgraphs(); |
| 213 | Subgraph* cond_subgraph = (*subgraphs)[op_data->cond_subgraph_index].get(); |
| 214 | Subgraph* body_subgraph = (*subgraphs)[op_data->body_subgraph_index].get(); |
| 215 | |
| 216 | // The follow graph illustrates the current implementation. |
| 217 | // |
| 218 | // This Subgraph Cond Subgraph Body Subgraph |
| 219 | // +-----------+ (1) +------------+ (3) +------------+ |
| 220 | // | WHILE |-------->| SUBGRAPH |-------->| SUBGRAPH | |
| 221 | // | INPUT | /| INPUT |<----- | INPUT | |
| 222 | // +-----------+ / +------------+ \ +------------+ |
| 223 | // / | \ | |
| 224 | // (6) / | (2) (5) \ | (4) |
| 225 | // / v \ v |
| 226 | // +-----------+ / +------------+ +------------+ |
| 227 | // | WHILE |<-- | SUBGRAPH | | SUBGRAPH | |
| 228 | // | OUTPUT | | OUTPUT | | OUTPUT | |
| 229 | // +-----------+ +------------+ +------------+ |
| 230 | // |
| 231 | // (1) Copy the inputs of WHILE op to the inputs of condition subgraph. |
| 232 | // (2) Invoke condition subgraph. |
| 233 | // Jump to step 5 if result is false. |
| 234 | // (3) Copy the inputs of condition subgraph to the inputs of body subgraph. |
| 235 | // (4) Invoke body subgraph. |
| 236 | // (5) Copy the outputs of body subgraph to the inputs condition subgraph. |
| 237 | // Jump back to step 2! |
| 238 | // (6) Copy the inputs of condition subgraph to the outputs of WHILE op. |
| 239 | // |
| 240 | // If the body subgraph has dynamic sized outputs, it's required to resize the |
| 241 | // tensor before copying in step 1, 3, 4 and 6. |
| 242 | // |
| 243 | // Note the flow is carefully designed to handle the dynamic sized output |
| 244 | // case. The loop invariant is: The newest value is in the inputs of condition |
| 245 | // subgraph. This is always true before step 2. |
| 246 | // |
| 247 | // This is the best we can do without sharing tensor buffer across subgraph |
| 248 | // boundary. Currently we copy the input / output between the subgraphs. This |
| 249 | // isn't optimized yet and a lot of redundant copies are made. |
| 250 | // TODO(b/120234921): Optimize and avoid copying tensors between subgraphs. |
| 251 | |
| 252 | if (op_data->body_has_dynamic_output_tensors) { |
| 253 | // If body subgraph has dynamic outputs, the input of condition subgraph may |
| 254 | // be changed in the last invocation and may need resizing. |
| 255 | TF_LITE_ENSURE_OK( |
| 256 | context, CopyTensorsShapeAndType( |
| 257 | context, this_subgraph, TfLiteIntArrayView(node->inputs), |
| 258 | cond_subgraph, cond_subgraph->inputs(), true)); |
| 259 | TF_LITE_ENSURE_OK(context, cond_subgraph->AllocateTensors()); |
| 260 | } |
| 261 | TF_LITE_ENSURE_OK( |
| 262 | context, |
| 263 | CopyTensorsData(context, this_subgraph, TfLiteIntArrayView(node->inputs), |
| 264 | cond_subgraph, cond_subgraph->inputs())); |
| 265 | |
| 266 | while (true) { |
nothing calls this directly
no test coverage detected