Support WHILE with growing output tensor in CpuExecutor

1. Adds a WHILE test that produces a tensor of dynamic size. The loop
   body model has inputs and outputs of unknown rank.
2. Adds an IF test where the branch models have unkown input and output
   ranks.
3. Modifies the partitioner to avoid trying to delegate referenced
   models if an IF or WHILE or their associated referenced models have
   unknown dimensions.
4. Adds dynamic loop output shape support to CpuExecutor by allocating
   new temporaries on each iteration.

Bug: 132458982
Bug: 154597673
Test: NNT_static --gtest_filter='*_dynamic*'
Change-Id: I8da1aa92abf090adbdcbf647c113fe27f4c002d8
Merged-In: I8da1aa92abf090adbdcbf647c113fe27f4c002d8
(cherry picked from commit a871fb7ecc8a1ec56107ac32b1f9c92cc8d5f74c)

1 files changed, 24 insertions, 4 deletions

diff --git a/nn/common/CpuExecutor.cpp b/nn/common/CpuExecutor.cpp
index d8582ed58..9373d59e2 100644
--- a/nn/common/CpuExecutor.cpp
+++ b/nn/common/CpuExecutor.cpp
@@ -1800,6 +1800,15 @@ int CpuExecutor::executeWhileOperation(const Operation& operation, RunTimeOperan
     std::vector<uint8_t*> tmp1(bodySubgraph.outputIndexes.size());
     std::vector<uint8_t*> tmp2(bodySubgraph.outputIndexes.size());
 
+    // For body outputs with unknown shape, we skip double buffering and
+    // allocate on each iteration instead. This allows growing output tensors
+    // inside a WHILE loop.
+    std::vector<bool> bodyOutputHasUnknownShape(bodySubgraph.outputIndexes.size());
+    for (uint32_t i = 0, n = bodySubgraph.outputIndexes.size(); i < n; ++i) {
+        const Operand& operand = bodySubgraph.operands[bodySubgraph.outputIndexes[i]];
+        bodyOutputHasUnknownShape[i] = nonExtensionOperandSizeOfData(operand) == 0;
+    }
+
     // Initialize condition inputs from outer operands.
     for (uint32_t i = 0, n = condSubgraph.inputIndexes.size(); i < n; ++i) {
         setInfoExceptLifetime(&condOperands[condSubgraph.inputIndexes[i]],
@@ -1842,16 +1851,27 @@ int CpuExecutor::executeWhileOperation(const Operation& operation, RunTimeOperan
         for (uint32_t i = 0, n = bodySubgraph.inputIndexes.size(); i < n; ++i) {
             bodyOperands[bodySubgraph.inputIndexes[i]] = condOperands[condSubgraph.inputIndexes[i]];
         }
-        // Switch body outputs.
+        // Set body outputs.
         auto& outputBuffer = iteration % 2 == 0 ? tmp1 : tmp2;
-        auto& otherBuffer = iteration % 2 == 0 ? tmp2 : tmp1;
         for (uint32_t i = 0, n = bodySubgraph.outputIndexes.size(); i < n; ++i) {
             RunTimeOperandInfo& info = bodyOperands[bodySubgraph.outputIndexes[i]];
-            otherBuffer[i] = info.buffer;
+            if (bodyOutputHasUnknownShape[i]) {
+                // Reset dimensions and buffer.
+                info.dimensions = bodySubgraph.operands[bodySubgraph.outputIndexes[i]].dimensions;
+                if (outputBuffer[i] != nullptr) {
+                    delete[] outputBuffer[i];
+                    outputBuffer[i] = nullptr;
+                }
+            }
             info.buffer = outputBuffer[i];
         }
 
         NN_RETURN_IF_ERROR(executeSubgraph(bodySubgraph, bodyOperands.data()));
+
+        // Update output buffer information in case we have allocated new buffers.
+        for (uint32_t i = 0, n = bodySubgraph.outputIndexes.size(); i < n; ++i) {
+            outputBuffer[i] = bodyOperands[bodySubgraph.outputIndexes[i]].buffer;
+        }
     }
 
     // Copy body outputs to outer outputs.
@@ -1863,7 +1883,7 @@ int CpuExecutor::executeWhileOperation(const Operation& operation, RunTimeOperan
         }
         CHECK_EQ(outerOperand.length, innerOperand.length);
         // TODO: Use the outer buffer as tmp1 to avoid copies.
-        memcpy(outerOperand.buffer, innerOperand.buffer, innerOperand.length);
+        std::memcpy(outerOperand.buffer, innerOperand.buffer, innerOperand.length);
     }
 
     auto freeLoopOutputs = [](const std::vector<uint8_t*>& tmp) {