adding batch dimension to MatMul.
Signed-off-by: Kevin Cheng <kevin.cheng@arm.com>
Change-Id: I83f75dd5beb60fe7ca2d573ea0f81bac4cd62a07
diff --git a/reference_model/src/ops/tensor_ops.h b/reference_model/src/ops/tensor_ops.h
index 26ce84b..9aaa140 100644
--- a/reference_model/src/ops/tensor_ops.h
+++ b/reference_model/src/ops/tensor_ops.h
@@ -183,15 +183,21 @@
static constexpr DType AccDtype = GetAccDType<Dtype, Dtype>::value;
using InEigenType = typename GetEigenType<Dtype>::type;
using AccEigenType = typename GetEigenType<AccDtype>::type;
- using TIn = Eigen::Tensor<InEigenType, 2>;
- using TAcc = Eigen::Tensor<AccEigenType, 2>;
+ using TIn = Eigen::Tensor<InEigenType, 3>;
+ using TAcc = Eigen::Tensor<AccEigenType, 3>;
+ using TInRank2 = Eigen::Tensor<InEigenType, 2>;
+ using TAccRank2 = Eigen::Tensor<AccEigenType, 2>;
static constexpr int64_t AccQMin = GetQMin<AccDtype>::value;
static constexpr int64_t AccQMax = GetQMax<AccDtype>::value;
protected:
TosaReference::TensorTemplate<TIn>* a;
TosaReference::TensorTemplate<TIn>* b;
- TosaReference::TensorTemplate<TAcc>* c;
+ TosaReference::TensorTemplate<TAcc>* output;
+ int64_t N;
+ int64_t H;
+ int64_t W;
+ int64_t C;
tosa::TosaMatMulQuantInfo* qinfo;
};