Handle CONSTANT_REFERENCE memory in ExecutionPlan::getBuffer()

Fix: 153432185
Bug: 137836124
Test: NNT_static
Change-Id: Icc63966088cd92c4466ec9d94e068309c01342e9

1 files changed, 11 insertions, 0 deletions

diff --git a/nn/runtime/ExecutionPlan.cpp b/nn/runtime/ExecutionPlan.cpp
index 568d19acb..4a6a8e9ab 100644
--- a/nn/runtime/ExecutionPlan.cpp
+++ b/nn/runtime/ExecutionPlan.cpp
@@ -1034,6 +1034,7 @@ std::optional<ExecutionPlan::Buffer> ExecutionPlan::getBuffer(
     const auto& sourceOperandToOffsetOfTemporary = controller->mSourceOperandToOffsetOfTemporary;
     const auto& sourceOperandToInputIndex = controller->mSourceOperandToInputIndex;
     const auto& sourceOperandToOutputIndex = controller->mSourceOperandToOutputIndex;
+    const auto& sourceOperandToConstantReference = controller->mSourceOperandToConstantReference;
     if (auto it = sourceOperandToOffsetOfTemporary.find(operandIndex);
         it != sourceOperandToOffsetOfTemporary.end()) {
         const uint32_t offset = it->second;
@@ -1047,6 +1048,14 @@ std::optional<ExecutionPlan::Buffer> ExecutionPlan::getBuffer(
                it != sourceOperandToOutputIndex.end()) {
         const ModelArgumentInfo& info = controller->mExecutionBuilder->getOutputInfo(it->second);
         return getBufferFromModelArgumentInfo(info, controller->mExecutionBuilder);
+    } else if (auto it = sourceOperandToConstantReference.find(operandIndex);
+               it != sourceOperandToConstantReference.end()) {
+        const ConstantReferenceLocation& location = it->second;
+        const std::optional<RunTimePoolInfo> info = location.memory->getRunTimePoolInfo();
+        if (info == std::nullopt) {
+            return std::nullopt;
+        }
+        return Buffer(info->getBuffer() + location.offset, location.length);
     }
     return std::nullopt;
 }
@@ -1366,6 +1375,7 @@ int ExecutionPlan::nextCompound(const WhileStep* step, std::shared_ptr<Controlle
             const SourceOperandIndex& outerOperand = step->outerOutputOperands[i];
             std::optional<Buffer> outerBuffer = getBuffer(controller, outerOperand);
             if (outerBuffer == std::nullopt) {
+                LOG(ERROR) << "Unable to get outerBuffer for operand " << toString(outerOperand);
                 return ANEURALNETWORKS_OP_FAILED;
             }
             const Operand& sourceOperand =
@@ -1374,6 +1384,7 @@ int ExecutionPlan::nextCompound(const WhileStep* step, std::shared_ptr<Controlle
             CHECK_NE(size, 0u);
             std::optional<Buffer> innerBuffer = getBuffer(controller, innerOperand);
             if (innerBuffer == std::nullopt) {
+                LOG(ERROR) << "Unable to get innerBuffer for operand " << toString(innerOperand);
                 return ANEURALNETWORKS_OP_FAILED;
             }
             CHECK_LE(size, innerBuffer->getSize());