COMPMID-911: Allow GEMM to work with 3D tensors

Change-Id: I8c4823a0d909e19e9ef548f00b9ae98c66de61dd
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/123569
Tested-by: Jenkins <bsgcomp@arm.com>
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>

1 files changed, 19 insertions, 4 deletions

diff --git a/arm_compute/core/utils/misc/ShapeCalculator.h b/arm_compute/core/utils/misc/ShapeCalculator.h
index deab181aee..9666702749 100644
--- a/arm_compute/core/utils/misc/ShapeCalculator.h
+++ b/arm_compute/core/utils/misc/ShapeCalculator.h
@@ -360,11 +360,26 @@ inline TensorShape compute_rnn_shape(const ITensorInfo *input, const unsigned in
 }
 inline TensorShape compute_mm_shape(const ITensorInfo &input0, const ITensorInfo &input1, bool is_interleaved_transposed, const GEMMReshapeInfo &reshape_info)
 {
-    TensorShape tensor_shape{ input0.tensor_shape() };
-    tensor_shape.set(0, is_interleaved_transposed ? reshape_info.n() : input1.dimension(0));
-    tensor_shape.set(1, is_interleaved_transposed ? reshape_info.m() : input0.dimension(1));
+    ARM_COMPUTE_ERROR_ON_MSG(input0.num_dimensions() > 4, "The number of dimensions for the matrix A must be <= 4");
 
-    return tensor_shape;
+    const bool is_gemm3d = reshape_info.depth_output_gemm3d() != 1;
+
+    // If the output of GEMM has to be reinterpreted as 3D, the number of input0 rows (M) is obtained collapsing the second and third
+    // dimension of the output tensor
+    const int dim0 = is_interleaved_transposed ? reshape_info.n() : input1.dimension(0);
+    const int dim1 = is_interleaved_transposed ? reshape_info.m() / reshape_info.depth_output_gemm3d() : input0.dimension(1) / reshape_info.depth_output_gemm3d();
+    const int dim2 = input0.tensor_shape()[2];
+    const int dim3 = input0.tensor_shape()[3];
+
+    TensorShape output_shape{ input0.tensor_shape() };
+
+    output_shape.set(0, dim0);
+    output_shape.set(1, dim1);
+    output_shape.set(2, is_gemm3d ? reshape_info.depth_output_gemm3d() : dim2);
+    output_shape.set(3, is_gemm3d ? dim2 : dim3);
+    output_shape.set(4, is_gemm3d ? dim3 : 1);
+
+    return output_shape;
 }
 
 template <typename T>