Add DIM operator to reference model

Signed-off-by: Won Jeon <won.jeon@arm.com>
Change-Id: Iea11ee5d3d98773e9c5e9b827593c05afb41ce3b

4 files changed, 97 insertions, 0 deletions

diff --git a/reference_model/src/ops/data_layout.cc b/reference_model/src/ops/data_layout.cc
index bc97c89..2d1fdb0 100644
--- a/reference_model/src/ops/data_layout.cc
+++ b/reference_model/src/ops/data_layout.cc
@@ -209,6 +209,60 @@ int OpPad<Rank, Dtype>::eval()
     return GraphNode::eval();
 }
 
+template <int Rank, TOSA_REF_TYPE Dtype>
+OpDim<Rank, Dtype>::OpDim(SubgraphTraverser* sgt_, TosaAttributeBase* attribute_, uint64_t id_)
+    : GraphNode(sgt_, Op_DIM, id_)
+{
+    setRequiredOperands(1, 1);
+
+    INIT_ATTRIBUTE(Axis);
+}
+
+template <int Rank, TOSA_REF_TYPE Dtype>
+OpDim<Rank, Dtype>::~OpDim()
+{
+    if (attribute)
+        delete attribute;
+}
+
+template <int Rank, TOSA_REF_TYPE Dtype>
+int OpDim<Rank, Dtype>::checkTensorAttributes()
+{
+    // Check Tosa Level
+    auto tosa_level = g_func_config.tosa_level;
+    LEVEL_CHECK(Rank <= tosa_level.MAX_RANK, "Rank should be smaller than or equal to MAX_RANK");
+
+    if (validateRequiredOperands())
+        return 1;
+
+    if (validateRequiredRank(inputs[0]))
+        return 1;
+
+    if (attribute->axis() < 0 || (size_t)attribute->axis() >= Rank)
+    {
+        printNodeValidationError("OpDim: axis must between [0, input_rank - 1]");
+        return 1;
+    }
+
+    in  = dynamic_cast<TosaReference::TensorTemplate<TIn>*>(inputs[0]);
+    out = dynamic_cast<TosaReference::TensorTemplate<TOut>*>(outputs[0]);
+
+    ASSERT_MEM(in && out);
+
+    return 0;
+}
+
+template <int Rank, TOSA_REF_TYPE Dtype>
+int OpDim<Rank, Dtype>::eval()
+{
+    int32_t axis    = attribute->axis();
+    int64_t out_val = in->getShape()[axis];
+
+    this->out->getTensor().setConstant(out_val);
+
+    return GraphNode::eval();
+}
+
 template <int InRank, int OutRank, TOSA_REF_TYPE Dtype>
 OpReshape<InRank, OutRank, Dtype>::OpReshape(SubgraphTraverser* sgt_, TosaAttributeBase* attribute_, uint64_t id_)
     : GraphNode(sgt_, Op_RESHAPE, id_)
@@ -780,6 +834,14 @@ DEF_INSTANTIATE_RANK1_6_ONE_RANK_ONE_TYPE(OpPad, INT32);
 DEF_INSTANTIATE_RANK1_6_ONE_RANK_ONE_TYPE(OpPad, BOOL);
 DEF_INSTANTIATE_RANK1_6_ONE_RANK_ONE_TYPE(OpPad, FP64);
 
+DEF_INSTANTIATE_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, FP16);
+DEF_INSTANTIATE_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, BF16);
+DEF_INSTANTIATE_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, FP32);
+DEF_INSTANTIATE_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, INT8);
+DEF_INSTANTIATE_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, INT16);
+DEF_INSTANTIATE_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, INT32);
+DEF_INSTANTIATE_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, BOOL);
+
 DEF_INSTANTIATE_RESHAPE(OpReshape, FP16);
 DEF_INSTANTIATE_RESHAPE(OpReshape, BF16);
 DEF_INSTANTIATE_RESHAPE(OpReshape, FP32);
diff --git a/reference_model/src/ops/data_layout.h b/reference_model/src/ops/data_layout.h
index 94ce248..024f9a2 100644
--- a/reference_model/src/ops/data_layout.h
+++ b/reference_model/src/ops/data_layout.h
@@ -66,6 +66,27 @@ protected:
     TosaPadAttribute* attribute;
 };
 
+template <int Rank, TOSA_REF_TYPE Dtype>
+class OpDim : public GraphNode
+{
+public:
+    OpDim(SubgraphTraverser* sgt_, TosaAttributeBase* attribute_, uint64_t id_);
+    virtual ~OpDim();
+
+    virtual int checkTensorAttributes();
+    virtual int eval();
+
+    using InEigenType  = typename GetEigenType<Dtype>::type;
+    using OutEigenType = typename GetEigenType<TOSA_REF_TYPE_SHAPE>::type;
+    using TIn          = Eigen::Tensor<InEigenType, Rank>;
+    using TOut         = Eigen::Tensor<OutEigenType, 0>;
+
+protected:
+    TosaReference::TensorTemplate<TIn>* in;
+    TosaReference::TensorTemplate<TOut>* out;
+    TosaAxisAttribute* attribute;
+};
+
 template <int InRank, int OutRank, TOSA_REF_TYPE Dtype>
 class OpReshape : public GraphNode
 {
diff --git a/reference_model/src/ops/op_factory.cc b/reference_model/src/ops/op_factory.cc
index a3069dc..d834b74 100644
--- a/reference_model/src/ops/op_factory.cc
+++ b/reference_model/src/ops/op_factory.cc
@@ -419,6 +419,15 @@ GraphNode* OpFactory::newOp(SubgraphTraverser* sgt,
             DEF_FACTORY_RANK1_6_ONE_RANK_ONE_TYPE(OpPad, BOOL);
             DEF_FACTORY_RANK1_6_ONE_RANK_ONE_TYPE(OpPad, FP64);
             break;
+        case Op_DIM:
+            DEF_FACTORY_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, FP16);
+            DEF_FACTORY_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, BF16);
+            DEF_FACTORY_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, FP32);
+            DEF_FACTORY_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, INT32);
+            DEF_FACTORY_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, INT8);
+            DEF_FACTORY_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, INT16);
+            DEF_FACTORY_RANK1_6_ONE_RANK_ONE_TYPE(OpDim, BOOL);
+            break;
         case Op_RESHAPE:
             DEF_FACTORY_RESHAPE(OpReshape, FP16);
             DEF_FACTORY_RESHAPE(OpReshape, BF16);
diff --git a/reference_model/src/ops/template_types.h b/reference_model/src/ops/template_types.h
index 6dd6e76..342d5c2 100644
--- a/reference_model/src/ops/template_types.h
+++ b/reference_model/src/ops/template_types.h
@@ -98,6 +98,11 @@ struct GetEigenType<TOSA_REF_TYPE_INT48>
     using type = int64_t;
 };
 template <>
+struct GetEigenType<TOSA_REF_TYPE_SHAPE>
+{
+    using type = int64_t;
+};
+template <>
 struct GetEigenType<TOSA_REF_TYPE_BOOL>
 {
     using type = bool;