diff options
author | Won Jeon <won.jeon@arm.com> | 2023-08-10 10:33:01 +0000 |
---|---|---|
committer | Won Jeon <won.jeon@arm.com> | 2023-08-18 15:21:15 -0700 |
commit | a21b2e88d19d8cb11a9120d40bacbb594d600565 (patch) | |
tree | 3bc8a40db72a31c1e552a3bd6339627a1175686e /reference_model/src/ops | |
parent | e0247481eb1f83f6eb7161d3f7ac2690b180952a (diff) | |
download | reference_model-a21b2e88d19d8cb11a9120d40bacbb594d600565.tar.gz |
Add DIM operator to reference model
Signed-off-by: Won Jeon <won.jeon@arm.com>
Change-Id: Iea11ee5d3d98773e9c5e9b827593c05afb41ce3b
Diffstat (limited to 'reference_model/src/ops')
-rw-r--r-- | reference_model/src/ops/data_layout.cc | 62 | ||||
-rw-r--r-- | reference_model/src/ops/data_layout.h | 21 | ||||
-rw-r--r-- | reference_model/src/ops/op_factory.cc | 9 | ||||
-rw-r--r-- | reference_model/src/ops/template_types.h | 5 |
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; |