[onert] Check Gather operand data type (#14054)

This commit adds operation validation for Gather to check operand data type. ONE-DCO-1.0-Signed-off-by: Hyeongseok Oh <[email protected]>
Samsung · Sep 24, 2024 · 646c7f1 · 646c7f1
1 parent 102163e
commit 646c7f1
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/runtime/onert/core/src/ir/OperationValidator.cc b/runtime/onert/core/src/ir/OperationValidator.cc
@@ -363,6 +363,21 @@ void OperationValidator::visit(const operation::Fill &node)
                           {DataType::FLOAT32, DataType::INT32, DataType::INT64, DataType::BOOL8}));
 }
 
+void OperationValidator::visit(const operation::Gather &node)
+{
+  const auto output_index{node.getOutputs().at(0)};
+  const auto input_index{node.getInputs().at(operation::Gather::INPUT)};
+  const auto indices_index{node.getInputs().at(operation::Gather::INDICES)};
+
+  const auto input_type = operandType(input_index);
+  if (input_type == DataType::QUANT_GGML_Q4_0 || input_type == DataType::QUANT_GGML_Q8_0)
+    OP_REQUIRES(isValidType(output_index, {DataType::FLOAT32}));
+  else
+    OP_REQUIRES(isSameType(output_index, input_index));
+
+  OP_REQUIRES(isValidType(indices_index, {DataType::INT32, DataType::INT64}));
+}
+
 void OperationValidator::visit(const operation::HashtableLookup &node)
 {
   const auto hits_index{node.getOutputs().at(operation::HashtableLookup::Output::HITS)};

diff --git a/runtime/onert/core/src/ir/OperationValidator.h b/runtime/onert/core/src/ir/OperationValidator.h
@@ -62,6 +62,7 @@ class OperationValidator : public OperationVisitor
   void visit(const operation::EmbeddingLookup &node) override;
   void visit(const operation::ExpandDims &node) override;
   void visit(const operation::Fill &node) override;
+  void visit(const operation::Gather &node) override;
   void visit(const operation::HashtableLookup &node) override;
   void visit(const operation::Pack &node) override;
   void visit(const operation::Pad &node) override;