From 87a74effff65f6fa1b0e565818e02c3b414ae1cf Mon Sep 17 00:00:00 2001
From: Georgios Pinitas <georgios.pinitas@arm.com>
Date: Fri, 20 Aug 2021 17:26:45 +0100
Subject: Re-use auxiliary memory withing CpuWinogradConv2d operators

Input/Output transformation operations are independent and done in
different time-steps of the algorithm, this memory can be re-used
between this transformation stages.

Moreover, reduce the allocation when extracting workspace sizes for
Winograd trasformations. There is a mix return of sizes in bytes and
elements, thus ensure the correct is in place. storage_size() member
functions return elements while working_space() function bytes.

Resolves: COMPMID-4781

Signed-off-by: Georgios Pinitas <georgios.pinitas@arm.com>
Change-Id: I705445ba7ca818cead48369db3cacd49684c7192
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/6145
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Michele Di Giorgio <michele.digiorgio@arm.com>
---
 src/core/cpu/kernels/CpuWinogradConv2dKernel.cpp | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

(limited to 'src/core/cpu/kernels')

diff --git a/src/core/cpu/kernels/CpuWinogradConv2dKernel.cpp b/src/core/cpu/kernels/CpuWinogradConv2dKernel.cpp
index 5620d36e2c..9456f96354 100644
--- a/src/core/cpu/kernels/CpuWinogradConv2dKernel.cpp
+++ b/src/core/cpu/kernels/CpuWinogradConv2dKernel.cpp
@@ -194,8 +194,8 @@ template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, in
 unsigned int CpuWinogradConv2dTransformWeightsKernel<T, OutputTileRows, OutputTileCols, KernelRows, KernelCols>::get_weight_storage_size(int num_output_channels, int num_input_channels) const
 {
     const KernelShape shape(num_output_channels, KernelRows, KernelCols, num_input_channels);
-    return static_cast<unsigned int>(
-               WinogradConv::get_kernel_storage_size(num_input_channels, num_output_channels));
+    // WinogradConv returns the size in bytes, we divide by `sizeof(T)` to express that in units of T
+    return static_cast<unsigned int>(WinogradConv::get_kernel_storage_size(num_input_channels, num_output_channels) / sizeof(T));
 }
 
 template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
@@ -297,7 +297,8 @@ unsigned int CpuWinogradConv2dTransformInputKernel<T, OutputTileRows, OutputTile
     // Construct shapes for the input and kernel tensors.
     const Tensor4DShape input_shape(num_batches, num_rows, num_cols, num_channels);
     const KernelShape   kern_shape(1, KernelRows, KernelCols, num_channels);
-    return static_cast<unsigned int>(WinogradConv::get_input_storage_size(num_batches, num_rows, num_cols, num_channels, same_padding));
+    // Return the size, converted into units of TIn
+    return static_cast<unsigned int>(WinogradConv::get_input_storage_size(num_batches, num_rows, num_cols, num_channels, same_padding) / sizeof(T));
 }
 
 template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
@@ -432,8 +433,9 @@ unsigned int CpuWinogradConv2dTransformOutputKernel<T, OutputTileRows, OutputTil
     // Construct shapes for the input and kernel tensors.
     const Tensor4DShape input_shape(num_batches, num_rows, num_cols, 1);
     const KernelShape   kern_shape(num_output_channels, KernelRows, KernelCols, 1);
+    // Return the size, converted into units of TOut
     return static_cast<unsigned int>(
-               WinogradConv::get_output_storage_size(num_batches, num_rows, num_cols, num_output_channels));
+               WinogradConv::get_output_storage_size(num_batches, num_rows, num_cols, num_output_channels) / sizeof(T));
 }
 
 template <typename T, int OutputTileRows, int OutputTileCols, int KernelRows, int KernelCols>
-- 
cgit v1.2.1