Refactor the Tensor3dMapper object

- Return TileVariable instead of std::string
- Accept TensorSamplerFormat instead of TensorSamples in the list of
constructor input arguments

Resolves COMPMID-6496

Signed-off-by: Gian Marco Iodice <gianmarco.iodice@arm.com>
Change-Id: Idb555cbed45d741b049dfbbc1d4bbddf139deae4
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/10131
Benchmark: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Gunes Bayir <gunes.bayir@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>

2 files changed, 83 insertions, 67 deletions

diff --git a/compute_kernel_writer/src/Tensor3dMapper.cpp b/compute_kernel_writer/src/Tensor3dMapper.cpp
index 60f4f57466..7384b924da 100644
--- a/compute_kernel_writer/src/Tensor3dMapper.cpp
+++ b/compute_kernel_writer/src/Tensor3dMapper.cpp
@@ -27,114 +27,129 @@
 #include "ckw/Error.h"
 #include "ckw/types/TensorSamplerTypes.h"
 #include "src/ITensor.h"
-
+#include "src/ITile.h"
 
 namespace ckw
 {
-Tensor3dMapper::Tensor3dMapper(ITensor *tensor, TensorSampler sampler)
-        : _tensor(tensor), _sampler(sampler)
+Tensor3dMapper::Tensor3dMapper(ITensor *tensor, TensorSamplerFormat format)
+        : _tensor(tensor), _format(format)
 {
 }
 
-std::string Tensor3dMapper::tensor_component_x() const
+TileVariable Tensor3dMapper::dim_x() const
 {
-    const TensorSamplerFormat format = _sampler.format();
-    switch(format)
+    switch(_format)
     {
         case TensorSamplerFormat::Dim0_Dim1xDim2_1:
         case TensorSamplerFormat::Dim0_Dim1_Dim2:
-            return _tensor->component(TensorComponentType::Dim0).scalar(0,0).str;
+            return _tensor->component(TensorComponentType::Dim0).scalar(0, 0);
         default:
             CKW_THROW_MSG("Unsupported tensor format");
-            return "";
+            return _tensor->component(TensorComponentType::Unknown).scalar(0, 0);
     }
 }
 
-std::string Tensor3dMapper::tensor_component_y() const
+TileVariable Tensor3dMapper::dim_y() const
 {
-    const TensorSamplerFormat format = _sampler.format();
-    switch(format)
+    switch(_format)
     {
         case TensorSamplerFormat::Dim0_Dim1xDim2_1:
-            return _tensor->component(TensorComponentType::Dim1xDim2).scalar(0,0).str;
+            return _tensor->component(TensorComponentType::Dim1xDim2).scalar(0, 0);
         case TensorSamplerFormat::Dim0_Dim1_Dim2:
-            return _tensor->component(TensorComponentType::Dim1).scalar(0,0).str;
+            return _tensor->component(TensorComponentType::Dim1).scalar(0, 0);
         default:
             CKW_THROW_MSG("Unsupported tensor format");
-            return "";
+            return _tensor->component(TensorComponentType::Unknown).scalar(0, 0);
     }
 }
 
-std::string Tensor3dMapper::tensor_component_z() const
+TileVariable Tensor3dMapper::dim_z() const
 {
-    const TensorSamplerFormat format = _sampler.format();
-    switch(format)
+    TileVariable dim_one;
+
+    switch(_format)
     {
         case TensorSamplerFormat::Dim0_Dim1xDim2_1:
-            return "1";
+            dim_one = _tensor->component(TensorComponentType::Dim3).scalar(0, 0);
+            dim_one.str = "1";
+            return dim_one;
         case TensorSamplerFormat::Dim0_Dim1_Dim2:
-            return _tensor->component(TensorComponentType::Dim2).scalar(0,0).str;
+            return _tensor->component(TensorComponentType::Dim2).scalar(0, 0);
         default:
             CKW_THROW_MSG("Unsupported tensor format");
-            return "";
+            return _tensor->component(TensorComponentType::Unknown).scalar(0, 0);
     }
 }
 
-std::string Tensor3dMapper::tensor_component_stride_x() const
+TileVariable Tensor3dMapper::dim_batch() const
 {
-    const TensorSamplerFormat format = _sampler.format();
-    switch(format)
+    TileVariable dim_one;
+
+    switch(_format)
     {
         case TensorSamplerFormat::Dim0_Dim1xDim2_1:
         case TensorSamplerFormat::Dim0_Dim1_Dim2:
-            return _tensor->component(TensorComponentType::Stride0).scalar(0,0).str;
+            return _tensor->component(TensorComponentType::Dim3).scalar(0, 0);
         default:
             CKW_THROW_MSG("Unsupported tensor format");
-            return "";
+            return _tensor->component(TensorComponentType::Unknown).scalar(0, 0);
     }
 }
 
-std::string Tensor3dMapper::tensor_component_stride_y() const
+TileVariable Tensor3dMapper::stride_x() const
 {
-    const TensorSamplerFormat format = _sampler.format();
-    switch(format)
+    switch(_format)
     {
         case TensorSamplerFormat::Dim0_Dim1xDim2_1:
         case TensorSamplerFormat::Dim0_Dim1_Dim2:
-            return _tensor->component(TensorComponentType::Stride1).scalar(0,0).str;
+            return _tensor->component(TensorComponentType::Stride0).scalar(0, 0);
         default:
             CKW_THROW_MSG("Unsupported tensor format");
-            return "";
+            return _tensor->component(TensorComponentType::Unknown).scalar(0, 0);
     }
 }
 
-std::string Tensor3dMapper::tensor_component_stride_z() const
+TileVariable Tensor3dMapper::stride_y() const
 {
-    const TensorSamplerFormat format = _sampler.format();
-    switch(format)
+    switch(_format)
     {
         case TensorSamplerFormat::Dim0_Dim1xDim2_1:
-            return "0";
         case TensorSamplerFormat::Dim0_Dim1_Dim2:
-            return _tensor->component(TensorComponentType::Stride2).scalar(0,0).str;
+            return _tensor->component(TensorComponentType::Stride1).scalar(0, 0);
         default:
             CKW_THROW_MSG("Unsupported tensor format");
-            return "";
+            return _tensor->component(TensorComponentType::Unknown).scalar(0, 0);
     }
 }
 
-std::string Tensor3dMapper::tensor_component_stride_batch() const
+TileVariable Tensor3dMapper::stride_z() const
 {
-    return _tensor->component(TensorComponentType::Stride3).scalar(0,0).str;
-}
+    TileVariable stride_zero;
 
-TensorSampler Tensor3dMapper::sampler() const
-{
-    return _sampler;
+    switch(_format)
+    {
+        case TensorSamplerFormat::Dim0_Dim1xDim2_1:
+            stride_zero = _tensor->component(TensorComponentType::Stride3).scalar(0, 0);
+            stride_zero.str = "0";
+            return stride_zero;
+        case TensorSamplerFormat::Dim0_Dim1_Dim2:
+            return _tensor->component(TensorComponentType::Stride2).scalar(0, 0);
+        default:
+            CKW_THROW_MSG("Unsupported tensor format");
+            return _tensor->component(TensorComponentType::Unknown).scalar(0, 0);
+    }
 }
 
-ITensor *Tensor3dMapper::tensor() const
+TileVariable Tensor3dMapper::stride_batch() const
 {
-    return _tensor;
+    switch(_format)
+    {
+        case TensorSamplerFormat::Dim0_Dim1xDim2_1:
+        case TensorSamplerFormat::Dim0_Dim1_Dim2:
+            return _tensor->component(TensorComponentType::Stride3).scalar(0, 0);
+        default:
+            CKW_THROW_MSG("Unsupported tensor format");
+            return _tensor->component(TensorComponentType::Unknown).scalar(0, 0);
+    }
 }
 } // namespace ckw
\ No newline at end of file
diff --git a/compute_kernel_writer/src/Tensor3dMapper.h b/compute_kernel_writer/src/Tensor3dMapper.h
index 1c3856addd..fa68ac2d15 100644
--- a/compute_kernel_writer/src/Tensor3dMapper.h
+++ b/compute_kernel_writer/src/Tensor3dMapper.h
@@ -25,56 +25,57 @@
 #ifndef CKW_SRC_TENSOR3DMAPPER_H
 #define CKW_SRC_TENSOR3DMAPPER_H
 
-#include "ckw/TensorSampler.h"
-
 #include <string>
+
 namespace ckw
 {
-
+// Forward declarations
 class ITensor;
-class TensorSampler;
+enum class TensorSamplerFormat;
+struct TileVariable;
 
-/** This internal-only class is responsible to map an Nd tensor to a 3d tensor with the
- *  help of TensorSampler object. The aim of the dimensionality reduction is to reduce
+/** This internal-only class is responsible to map an Nd tensor spatial dimensions to a 3d tensor spatial dimensions with the
+ *  help of TensorSamplerFormat.
+ *  Attention: The batch is not considered as a spatial dimension and it is treated as an offset
+ *
+ *  The aim of the dimensionality reduction is primarily to reduce
  *  the address calculation to:
- *      x + y * stride_y + z * stride_z + offset, where offset is determined by the batch.
+ *      x + y * stride_y + z * stride_z + offset, where offset is determined by the batch (for example, b * stride_batch).
+ *
  */
 class Tensor3dMapper
 {
 public:
     /** Constructor */
-    Tensor3dMapper(ITensor *tensor, TensorSampler sampler);
+    Tensor3dMapper(ITensor *tensor, TensorSamplerFormat format);
 
     /** Get dimension x as string */
-    std::string tensor_component_x() const;
+    TileVariable dim_x() const;
 
     /** Get dimension y as string */
-    std::string tensor_component_y() const;
+    TileVariable dim_y() const;
 
     /** Get dimension z as string */
-    std::string tensor_component_z() const;
+    TileVariable dim_z() const;
+
+    /** Get batch dimension as string */
+    TileVariable dim_batch() const;
 
     /** Get stride for dimension x as string */
-    std::string tensor_component_stride_x() const;
+    TileVariable stride_x() const;
 
     /** Get stride for dimension y as string */
-    std::string tensor_component_stride_y() const;
+    TileVariable stride_y() const;
 
     /** Get stride for dimension z as string */
-    std::string tensor_component_stride_z() const;
+    TileVariable stride_z() const;
 
     /** Get stride for batch dimension as string */
-    std::string tensor_component_stride_batch() const;
-
-    /** Get the tensor sampler */
-    TensorSampler sampler() const;
-
-    /** Get the associated tensor */
-    ITensor *tensor() const;
+    TileVariable stride_batch() const;
 
 private:
-    ITensor         *_tensor;
-    TensorSampler    _sampler;
+    ITensor             *_tensor;
+    TensorSamplerFormat  _format;
 };
 } // namespace ckw