Add DIM operator to reference model

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

8 files changed, 141 insertions, 1 deletions

diff --git a/reference_model/src/operators.cc b/reference_model/src/operators.cc
index 1ae0683..ae5963d 100644
--- a/reference_model/src/operators.cc
+++ b/reference_model/src/operators.cc
@@ -68,6 +68,8 @@ tosa::DType translate_client_datatype(tosa_datatype_t type)
             return tosa::DType::DType_UINT16;
         case tosa_datatype_uint8_t:
             return tosa::DType::DType_UINT8;
+        case tosa_datatype_shape_t:
+            return tosa::DType::DType_SHAPE;
         default:
             return tosa::DType::DType_UNKNOWN;
     }
@@ -1978,6 +1980,34 @@ extern "C"
         return tosa_status_valid;
     }
 
+    tosa_status_t tosa_run_dim(tosa_tensor_t client_input1, const int32_t client_axis, tosa_tensor_t client_output)
+    {
+        // Create operator attributes
+        TosaAxisAttribute attr(client_axis);
+
+        // Create tensors
+        tosa::TosaSerializationTensor* input1 = translate_client_tensor(client_input1, "input1");
+        tosa::TosaSerializationTensor* output = translate_client_tensor(client_output, "output");
+
+        // Create operator
+        auto op = new tosa::TosaSerializationOperator(tosa::Op::Op_DIM, tosa::Attribute::Attribute_AxisAttribute, &attr,
+                                                      { input1->GetName() }, { output->GetName() });
+
+        // Create a tosa single-op basic block
+        tosa::TosaSerializationBasicBlock block("dim", "main", { op }, { input1, output }, { input1->GetName() },
+                                                { output->GetName() });
+
+        // Setup model
+        TosaReference::ModelRunnerImpl runner;
+        TOSA_RETURN_ON_GRAPH_STATUS_ERROR(runner.initialize(block));
+        TOSA_RETURN_ON_ERROR(runner.setInput(input1->GetName(), client_input1.data, client_input1.size));
+
+        // Execute
+        TOSA_RETURN_ON_ERROR(runner.getOutput(output->GetName(), client_output.data, client_output.size));
+
+        return tosa_status_valid;
+    }
+
     tosa_status_t tosa_run_reshape(tosa_tensor_t client_input1,
                                    const int32_t client_new_shape_len,
                                    const int32_t client_new_shape[],
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;
diff --git a/reference_model/src/subgraph_traverser.cc b/reference_model/src/subgraph_traverser.cc
index 5675be9..186cb8b 100644
--- a/reference_model/src/subgraph_traverser.cc
+++ b/reference_model/src/subgraph_traverser.cc
@@ -478,7 +478,8 @@ int SubgraphTraverser::allocateTensor(std::string name)
                 tensor->setTensorValueInt32(i32_data.size(), i32_data.data());
             }
             break;
-            case DType_INT48: {
+            case DType_INT48:
+            case DType_SHAPE: {
                 std::vector<int64_t> i64_data;
                 TosaSerializationHandler::ConvertU8toI48(ts->GetData(), tensor->getElementCount(), i64_data);
                 tensor->setTensorValueInt64(i64_data.size(), i64_data.data());
diff --git a/reference_model/src/tensor.cc b/reference_model/src/tensor.cc
index 4982c99..1aabe5b 100644
--- a/reference_model/src/tensor.cc
+++ b/reference_model/src/tensor.cc
@@ -137,6 +137,7 @@ int TosaReference::Tensor::readFromNpyFile(const char* filename)
             nperror = NumpyUtilities::readFromNpyFile(filename, elements, i32databuf);
             break;
         case DType_INT48:
+        case DType_SHAPE:
             i64databuf = (int64_t*)calloc(sizeof(int64_t), elements);
             ASSERT_MEM(i64databuf);
 
@@ -220,6 +221,7 @@ int TosaReference::Tensor::readFromNpyFile(const char* filename)
             }
             break;
         case TOSA_REF_TYPE_INT48:
+        case TOSA_REF_TYPE_SHAPE:
             if (setTensorValueInt64(elements, i64databuf))
             {
                 free(i64databuf);
@@ -379,6 +381,7 @@ int TosaReference::Tensor::writeToNpyFile(const char* filename) const
             free(i32databuf);
             break;
         case TOSA_REF_TYPE_INT48:
+        case TOSA_REF_TYPE_SHAPE:
             i64databuf = (int64_t*)calloc(sizeof(int64_t), elements);
             ASSERT_MEM(i64databuf);
 
@@ -672,6 +675,7 @@ int TosaReference::Tensor::readfromVector(const ArrayProxy<int64_t> vals)
     switch (getDtype())
     {
         case TOSA_REF_TYPE_INT48:
+        case TOSA_REF_TYPE_SHAPE:
             if (vals.size() != elements)
             {
                 WARNING("The input size (%ld) doesn't match the number of elements (%d) assigned to the tensor.",
@@ -847,6 +851,7 @@ int TosaReference::Tensor::writeToVector(ArrayProxy<int64_t> vals)
     switch (getDtype())
     {
         case TOSA_REF_TYPE_INT48:
+        case TOSA_REF_TYPE_SHAPE:
             if (vals.size() != elements)
             {
                 WARNING("The output size (%ld) doesn't match the number of elements (%d) assigned to the tensor.",
diff --git a/reference_model/src/tensor.h b/reference_model/src/tensor.h
index f59a5e1..74f57ed 100644
--- a/reference_model/src/tensor.h
+++ b/reference_model/src/tensor.h
@@ -811,6 +811,13 @@ public:
                         return new Tensor6<int64_t>(tensorName_, dtype_, shape_);
                 }
                 break;
+            case TOSA_REF_TYPE_SHAPE:
+                // if shape information is not already set, set it here.
+                if (shape_.size() == 0)
+                {
+                    shape_ = { 1 };
+                }
+                return new Tensor0<int64_t>(tensorName_, dtype_, shape_);
             case TOSA_REF_TYPE_BOOL:
                 switch (rank)
                 {