[onert] Introduce Pool2D dynamic shape inference (Samsung#12861)

This commit adds Pool2D dynamic shape inference. ONE-DCO-1.0-Signed-off-by: Hyeongseok Oh <[email protected]>
lemmaa · Apr 12, 2024 · d619e49 · d619e49
1 parent 1d49752
commit d619e49
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/runtime/onert/core/include/exec/DynamicShapeInferer.h b/runtime/onert/core/include/exec/DynamicShapeInferer.h
@@ -72,6 +72,7 @@ class DynamicShapeInferer : public ir::OperationVisitor
   void visit(const ir::operation::Pack &op) override;
   void visit(const ir::operation::Pad &op) override;
   void visit(const ir::operation::Permute &op) override;
+  void visit(const ir::operation::Pool2D &op) override;
   void visit(const ir::operation::Pow &op) override;
   // TODO write op starting from Q
   void visit(const ir::operation::Range &op) override;

diff --git a/runtime/onert/core/src/exec/DynamicShapeInferer.cc b/runtime/onert/core/src/exec/DynamicShapeInferer.cc
@@ -705,6 +705,26 @@ void DynamicShapeInferer::visit(const ir::operation::Permute & /* op */)
   // on-the-fly, as it must support inter-backend inference/allocation.
 }
 
+void DynamicShapeInferer::visit(const ir::operation::Pool2D &op)
+{
+  // check if input is not dynamic
+  auto input_ind = op.getInputs().at(ir::operation::Pool2D::INPUT);
+  auto input = _tensor_registry->getITensor(input_ind);
+
+  if (!input->is_dynamic())
+    return;
+
+  ir::Shape input_shape = input->getShape();
+
+  auto output_ind = op.getOutputs().at(0);
+  auto output = _tensor_registry->getITensor(output_ind);
+
+  ir::Shape output_shape = shape_inference::inferPoolShape(input_shape, op.param());
+
+  output->applyShape(output_shape);
+  assert(output->buffer() != nullptr);
+}
+
 void DynamicShapeInferer::visit(const ir::operation::Pow &op)
 {
   handleBinaryArithmeticOp(op, op.getInputs().at(ir::operation::Pow::Input::LHS),