COMPMID-2342: Add layer normalization support in CLLSTMLayer

Change-Id: I25d974aa94e69c5f79a0bd99d5869a351d6d954d
Signed-off-by: Michele Di Giorgio <michele.digiorgio@arm.com>
Reviewed-on: https://review.mlplatform.org/c/1324
Reviewed-by: Manuel Bottini <manuel.bottini@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Michalis Spyrou <michalis.spyrou@arm.com>

2 files changed, 125 insertions, 55 deletions

diff --git a/arm_compute/runtime/CL/functions/CLLSTMLayer.h b/arm_compute/runtime/CL/functions/CLLSTMLayer.h
index 3add152878..44b3baf81b 100644
--- a/arm_compute/runtime/CL/functions/CLLSTMLayer.h
+++ b/arm_compute/runtime/CL/functions/CLLSTMLayer.h
@@ -39,6 +39,7 @@
 #include "arm_compute/runtime/CL/functions/CLElementwiseOperations.h"
 #include "arm_compute/runtime/CL/functions/CLFullyConnectedLayer.h"
 #include "arm_compute/runtime/CL/functions/CLGEMM.h"
+#include "arm_compute/runtime/CL/functions/CLMeanStdDevNormalizationLayer.h"
 #include "arm_compute/runtime/IMemoryManager.h"
 #include "arm_compute/runtime/common/LSTMParams.h"
 
@@ -76,17 +77,22 @@ public:
      * @param[out] output                      Destination tensor. Output is a 2D tensor with dimensions [output_size, batch_size].
      *                                         Data types supported: Same as @p input.
      * @param[in]  lstm_params                 (Optional) Weights tensors used in peephole optimization:
-     *                                         input_to_input_weights       2D weights tensor with dimensions [input_size, num_units]. Data type supported: Same as @p input.
-     *                                         recurrent_to_input_weights   2D weights tensor with dimensions [output_size, num_units]. Data type supported: Same as @p input.
-     *                                         cell_to_input_weights        1D weights tensor with dimensions [num_units]. Can be nullptr. Data type supported: Same as @p input.
-     *                                         cell_to_forget_weights       1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
-     *                                         cell_to_output_weights       1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
-     *                                         input_gate_bias              1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input
-     *                                         projection_weights           2D weights tensor with dimensions [output_size, num_units]. Data type supported: Same as @p input.
-     *                                         projection_bias              1D weights tensor with dimensions [output_size]. Data type supported: Same as @p input.
+     *                                         input_to_input_weights         2D weights tensor with dimensions [input_size, num_units]. Data type supported: Same as @p input.
+     *                                         recurrent_to_input_weights     2D weights tensor with dimensions [output_size, num_units]. Data type supported: Same as @p input.
+     *                                         cell_to_input_weights          1D weights tensor with dimensions [num_units]. Can be nullptr. Data type supported: Same as @p input.
+     *                                         cell_to_forget_weights         1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
+     *                                         cell_to_output_weights         1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
+     *                                         input_gate_bias                1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input
+     *                                         projection_weights             2D weights tensor with dimensions [output_size, num_units]. Data type supported: Same as @p input.
+     *                                         projection_bias                1D weights tensor with dimensions [output_size]. Data type supported: Same as @p input.
+     *                                         input_layer_norm_coefficients  1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
+     *                                         forget_layer_norm_coefficients 1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
+     *                                         cell_layer_norm_coefficients   1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
+     *                                         output_layer_norm_coefficients 1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
      * @param[in]  activation_info             Contains activation information described in @ref ActivationLayerInfo.
-     * @param[in]  cell_threshold              The clipping threshold for the cell state, such that values are bound within [-cell_clip, cell_clip]. If set to 0.0 then clipping is disabled.
-     * @param[in]  projection_threshold        The clipping threshold for the output from the projection layer, such that values are bound within [-proj_clip, proj_clip]. If set to 0.0 then clipping is disabled.
+     * @param[in]  cell_threshold              The clipping threshold for the cell state, such that values are bound within [-cell_clip, cell_clip]. If set to 0.0f then clipping is disabled.
+     * @param[in]  projection_threshold        The clipping threshold for the output from the projection layer, such that values are bound within [-proj_clip, proj_clip].
+     *                                         If set to 0.0f then clipping is disabled.
      */
     void configure(const ICLTensor *input,
                    const ICLTensor *input_to_forget_weights, const ICLTensor *input_to_cell_weights, const ICLTensor *input_to_output_weights,
@@ -98,35 +104,40 @@ public:
 
     /** Static function to check if given info will lead to a valid configuration of @ref CLLSTMLayer
      *
-     * @param[in] input                       Source tensor. Input is a 2D tensor with dimensions [input_size, batch_size]. Data types supported: F16/F32.
-     * @param[in] input_to_forget_weights     2D weights tensor with dimensions [input_size, num_units]. Data type supported: Same as @p input.
-     * @param[in] input_to_cell_weights       2D weights tensor with dimensions [input_size, num_units]. Data type supported: Same as @p input.
-     * @param[in] input_to_output_weights     2D weights tensor with dimensions [input_size, num_units]. Data type supported: Same as @p input.
-     * @param[in] recurrent_to_forget_weights 2D weights tensor with dimensions [output_size, num_units]. Data type supported: Same as @p input.
-     * @param[in] recurrent_to_cell_weights   2D weights tensor with dimensions [output_size, num_units]. Data type supported: Same as @p input.
-     * @param[in] recurrent_to_output_weights 2D weights tensor with dimensions [output_size, num_units]. Data type supported: Same as @p input.
-     * @param[in] forget_gate_bias            1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
-     * @param[in] cell_bias                   1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
-     * @param[in] output_gate_bias            1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
-     * @param[in] output_state_in             2D weights tensor with dimensions [output_size, batch_size]. Data type supported: Same as @p input.
-     * @param[in] cell_state_in               2D tensor with dimensions [num_units, batch_size]. Data type supported: Same as @p input.
-     * @param[in] scratch_buffer              2D tensor with dimensions [num_units * 4, batch_size] with CIFG or [num_units * 3, batch_size] without CIGF. Data type supported: Same as @p input.
-     * @param[in] output_state_out            2D weights tensor with dimensions [output_size, batch_size]. Data type supported: Same as @p input.
-     * @param[in] cell_state_out              2D tensor with dimensions [num_units, batch_size]. Data type supported: Same as @p input.
-     * @param[in] output                      Destination tensor. Output is a 2D tensor with dimensions [output_size, batch_size].
-     *                                        Data types supported: Same as @p input.
-     * @param[in] lstm_params                 (Optional) Weights tensors used in peephole optimization:
-     *                                        input_to_input_weights       2D weights tensor with dimensions [input_size, num_units]. Data type supported: Same as @p input.
-     *                                        recurrent_to_input_weights   2D weights tensor with dimensions [output_size, num_units]. Data type supported: Same as @p input.
-     *                                        cell_to_input_weights        1D weights tensor with dimensions [num_units]. Can be nullptr. Data type supported: Same as @p input.
-     *                                        cell_to_forget_weights       1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
-     *                                        cell_to_output_weights       1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input.
-     *                                        input_gate_bias              1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input
-     *                                        projection_weights           2D weights tensor with dimensions [output_size, num_units]. Data type supported: Same as @p input.
-     *                                        projection_bias              1D weights tensor with dimensions [output_size]. Data type supported: Same as @p input.
+     * @param[in] input                       Source tensor info. Input is a 2D tensor with dimensions [input_size, batch_size]. Data types supported: F16/F32.
+     * @param[in] input_to_forget_weights     2D weights tensor info with dimensions [input_size, num_units]. Data type supported: Same as @p input.
+     * @param[in] input_to_cell_weights       2D weights tensor info with dimensions [input_size, num_units]. Data type supported: Same as @p input.
+     * @param[in] input_to_output_weights     2D weights tensor info with dimensions [input_size, num_units]. Data type supported: Same as @p input.
+     * @param[in] recurrent_to_forget_weights 2D weights tensor info with dimensions [output_size, num_units]. Data type supported: Same as @p input.
+     * @param[in] recurrent_to_cell_weights   2D weights tensor info with dimensions [output_size, num_units]. Data type supported: Same as @p input.
+     * @param[in] recurrent_to_output_weights 2D weights tensor info with dimensions [output_size, num_units]. Data type supported: Same as @p input.
+     * @param[in] forget_gate_bias            1D weights tensor info with dimensions [num_units]. Data type supported: Same as @p input.
+     * @param[in] cell_bias                   1D weights tensor info with dimensions [num_units]. Data type supported: Same as @p input.
+     * @param[in] output_gate_bias            1D weights tensor info with dimensions [num_units]. Data type supported: Same as @p input.
+     * @param[in] output_state_in             2D weights tensor info with dimensions [output_size, batch_size]. Data type supported: Same as @p input.
+     * @param[in] cell_state_in               2D tensor info with dimensions [num_units, batch_size]. Data type supported: Same as @p input.
+     * @param[in] scratch_buffer              2D tensor info with dimensions [num_units * 4, batch_size] with CIFG or [num_units * 3, batch_size] without CIGF.
+     *                                        Data type supported: Same as @p input.
+     * @param[in] output_state_out            2D weights tensor info with dimensions [output_size, batch_size]. Data type supported: Same as @p input.
+     * @param[in] cell_state_out              2D tensor info with dimensions [num_units, batch_size]. Data type supported: Same as @p input.
+     * @param[in] output                      Destination tensor info. Output is a 2D tensor with dimensions [output_size, batch_size]. Data types supported: Same as @p input.
+     * @param[in] lstm_params                 (Optional) Weights tensors info used in peephole optimization:
+     *                                        input_to_input_weights         2D weights tensor info with dimensions [input_size, num_units]. Data type supported: Same as @p input.
+     *                                        recurrent_to_input_weights     2D weights tensor info with dimensions [output_size, num_units]. Data type supported: Same as @p input.
+     *                                        cell_to_input_weights          1D weights tensor info with dimensions [num_units]. Can be nullptr. Data type supported: Same as @p input.
+     *                                        cell_to_forget_weights         1D weights tensor info with dimensions [num_units]. Data type supported: Same as @p input.
+     *                                        cell_to_output_weights         1D weights tensor info with dimensions [num_units]. Data type supported: Same as @p input.
+     *                                        input_gate_bias                1D weights tensor info with dimensions [num_units]. Data type supported: Same as @p input
+     *                                        projection_weights             2D weights tensor info with dimensions [output_size, num_units]. Data type supported: Same as @p input.
+     *                                        projection_bias                1D weights tensor info with dimensions [output_size]. Data type supported: Same as @p input.
+     *                                        input_layer_norm_coefficients  1D weights tensor info with dimensions [num_units]. Data type supported: Same as @p input.
+     *                                        forget_layer_norm_coefficients 1D weights tensor info with dimensions [num_units]. Data type supported: Same as @p input.
+     *                                        cell_layer_norm_coefficients   1D weights tensor info with dimensions [num_units]. Data type supported: Same as @p input.
+     *                                        output_layer_norm_coefficients 1D weights tensor info with dimensions [num_units]. Data type supported: Same as @p input.
      * @param[in] activation_info             Contains activation information described in @ref ActivationLayerInfo.
-     * @param[in] cell_threshold              The clipping threshold for the cell state, such that values are bound within [-cell_clip, cell_clip]. If set to 0.0 then clipping is disabled.
-     * @param[in] projection_threshold        The clipping threshold for the output from the projection layer, such that values are bound within [-proj_clip, proj_clip]. If set to 0.0 then clipping is disabled.
+     * @param[in] cell_threshold              The clipping threshold for the cell state, such that values are bound within [-cell_clip, cell_clip]. If set to 0.0f then clipping is disabled.
+     * @param[in] projection_threshold        The clipping threshold for the output from the projection layer, such that values are bound within [-proj_clip, proj_clip].
+     *                                        If set to 0.0f then clipping is disabled.
      *
      * @return a status
      */
@@ -145,23 +156,16 @@ public:
 private:
     CLMemoryGroup                        _memory_group;
     CLFullyConnectedLayer                _fully_connected_input_gate;
-    CLGEMM                               _gemm_input_gate;
-    CLTransposeKernel                    _transpose_input_gate;
-    CLSaturatedArithmeticOperationKernel _accum_input_gate1;
-    CLArithmeticAddition                 _accum_input_gate2;
+    CLArithmeticAddition                 _accum_input_gate1;
     CLSaturatedArithmeticOperationKernel _subtract_input_gate;
     CLPixelWiseMultiplicationKernel      _pixelwise_mul_input_gate;
     CLActivationLayerKernel              _activation_input_gate;
     CLFullyConnectedLayer                _fully_connected_forget_gate;
-    CLGEMM                               _gemm_forget_gate;
-    CLTransposeKernel                    _transpose_forget_gate;
-    CLSaturatedArithmeticOperationKernel _accum_forget_gate1;
-    CLArithmeticAddition                 _accum_forget_gate2;
+    CLArithmeticAddition                 _accum_forget_gate1;
     CLPixelWiseMultiplicationKernel      _pixelwise_mul_forget_gate;
     CLActivationLayerKernel              _activation_forget_gate;
     CLFullyConnectedLayer                _fully_connected_cell_state;
     CLGEMM                               _gemm_cell_state1;
-    CLGEMM                               _gemm_cell_state2;
     CLTransposeKernel                    _transpose_cell_state;
     CLSaturatedArithmeticOperationKernel _accum_cell_state1;
     CLSaturatedArithmeticOperationKernel _accum_cell_state2;
@@ -170,17 +174,12 @@ private:
     CLActivationLayerKernel              _cell_clip;
     CLPixelWiseMultiplicationKernel      _pixelwise_mul_cell_state2;
     CLFullyConnectedLayer                _fully_connected_output;
-    CLGEMM                               _gemm_output;
     CLPixelWiseMultiplicationKernel      _pixelwise_mul_output_state1;
-    CLTransposeKernel                    _transpose_output;
-    CLSaturatedArithmeticOperationKernel _accum_output1;
-    CLArithmeticAddition                 _accum_output2;
+    CLArithmeticAddition                 _accum_output1;
     CLActivationLayerKernel              _activation_output;
     CLActivationLayerKernel              _activation_output_state;
     CLPixelWiseMultiplicationKernel      _pixelwise_mul_output_state2;
     CLFullyConnectedLayer                _fully_connected_output_state;
-    CLGEMM                               _gemm_output_state;
-    CLSaturatedArithmeticOperationKernel _accum_output_state;
     CLActivationLayerKernel              _projection_clip;
     CLCopyKernel                         _copy_cell_state;
     CLCopyKernel                         _copy_output;
@@ -190,6 +189,18 @@ private:
     CLWidthConcatenate2TensorsKernel     _concat_weights_input_gate;
     CLWidthConcatenate2TensorsKernel     _concat_weights_output;
     CLMemsetKernel                       _ones_memset_kernel;
+    CLMeanStdDevNormalizationLayer       _mean_std_norm_input_gate;
+    CLPixelWiseMultiplicationKernel      _pixelwise_mul_input_gate_coeff;
+    CLSaturatedArithmeticOperationKernel _accum_input_gate_bias;
+    CLMeanStdDevNormalizationLayer       _mean_std_norm_forget_gate;
+    CLPixelWiseMultiplicationKernel      _pixelwise_mul_forget_gate_coeff;
+    CLSaturatedArithmeticOperationKernel _accum_forget_gate_bias;
+    CLMeanStdDevNormalizationLayer       _mean_std_norm_cell_gate;
+    CLPixelWiseMultiplicationKernel      _pixelwise_mul_cell_gate_coeff;
+    CLSaturatedArithmeticOperationKernel _accum_cell_gate_bias;
+    CLMeanStdDevNormalizationLayer       _mean_std_norm_output_gate;
+    CLPixelWiseMultiplicationKernel      _pixelwise_mul_output_gate_coeff;
+    CLSaturatedArithmeticOperationKernel _accum_output_gate_bias;
     CLTensor                             _input_gate_out1;
     CLTensor                             _input_gate_out2;
     CLTensor                             _input_gate_out3;
@@ -212,12 +223,21 @@ private:
     CLTensor                             _cell_state_activation;
     CLTensor                             _output_state1;
     CLTensor                             _ones;
+    CLTensor                             _input_layer_norm_out1;
+    CLTensor                             _input_layer_norm_out2;
+    CLTensor                             _forget_layer_norm_out1;
+    CLTensor                             _forget_layer_norm_out2;
+    CLTensor                             _cell_layer_norm_out1;
+    CLTensor                             _cell_layer_norm_out2;
+    CLTensor                             _output_layer_norm_out1;
+    CLTensor                             _output_layer_norm_out2;
     bool                                 _run_peephole_opt;
     bool                                 _run_cifg_opt;
     bool                                 _perform_cell_clipping;
     bool                                 _has_projection_weights;
     bool                                 _perform_projection_clipping;
     bool                                 _is_prepared;
+    bool                                 _is_layer_norm_lstm;
 };
 } // namespace arm_compute
 #endif /* __ARM_COMPUTE_CLLSTMLAYER_H__ */
diff --git a/arm_compute/runtime/common/LSTMParams.h b/arm_compute/runtime/common/LSTMParams.h
index 5b33e2e937..6979f90721 100644
--- a/arm_compute/runtime/common/LSTMParams.h
+++ b/arm_compute/runtime/common/LSTMParams.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018 ARM Limited.
+ * Copyright (c) 2018-2019 ARM Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -41,7 +41,8 @@ public:
     /** Constructor */
     LSTMParams()
         : _input_to_input_weights(nullptr), _recurrent_to_input_weights(nullptr), _cell_to_input_weights(nullptr), _input_gate_bias(nullptr), _cell_to_forget_weights(nullptr),
-          _cell_to_output_weights(nullptr), _projection_weights(nullptr), _projection_bias(nullptr), _has_peephole_opt(false), _has_projection(false), _has_cifg_opt(true)
+          _cell_to_output_weights(nullptr), _projection_weights(nullptr), _projection_bias(nullptr), _input_layer_norm_weights(nullptr), _forget_layer_norm_weights(nullptr), _cell_layer_norm_weights(nullptr),
+          _output_layer_norm_weights(nullptr), _has_peephole_opt(false), _has_projection(false), _has_cifg_opt(true), _use_layer_norm(false)
     {
     }
     /** Prevent instances of this class from being copied (As this class contains pointers) */
@@ -96,6 +97,25 @@ public:
         _has_peephole_opt       = true;
         return *this;
     }
+    /** Set layer normalization tensor parameters.
+     *
+     * @param[in] input_layer_norm_weights  1D weights tensor with dimensions [num_units]. Data type supported: Data types supported: F16/F32.
+     * @param[in] forget_layer_norm_weights 1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input_layer_norm_weights.
+     * @param[in] cell_layer_norm_weights   1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input_layer_norm_weights.
+     * @param[in] output_layer_norm_weights 1D weights tensor with dimensions [num_units]. Data type supported: Same as @p input_layer_norm_weights.
+     *
+     * @return Reference to this LSTMParams object
+     */
+    LSTMParams &set_layer_normalization_params(const T *input_layer_norm_weights, const T *forget_layer_norm_weights,
+                                               const T *cell_layer_norm_weights, const T *output_layer_norm_weights)
+    {
+        _input_layer_norm_weights  = input_layer_norm_weights;
+        _forget_layer_norm_weights = forget_layer_norm_weights;
+        _cell_layer_norm_weights   = cell_layer_norm_weights;
+        _output_layer_norm_weights = output_layer_norm_weights;
+        _use_layer_norm            = true;
+        return *this;
+    }
 
     const T *input_to_input_weights() const
     {
@@ -137,6 +157,26 @@ public:
         return _projection_bias;
     }
 
+    const T *input_layer_norm_weights() const
+    {
+        return _input_layer_norm_weights;
+    }
+
+    const T *forget_layer_norm_weights() const
+    {
+        return _forget_layer_norm_weights;
+    }
+
+    const T *cell_layer_norm_weights() const
+    {
+        return _cell_layer_norm_weights;
+    }
+
+    const T *output_layer_norm_weights() const
+    {
+        return _output_layer_norm_weights;
+    }
+
     bool has_peephole_opt() const
     {
         return _has_peephole_opt;
@@ -152,6 +192,11 @@ public:
         return _has_cifg_opt;
     }
 
+    bool use_layer_norm() const
+    {
+        return _use_layer_norm;
+    }
+
 private:
     const T *_input_to_input_weights;
     const T *_recurrent_to_input_weights;
@@ -161,9 +206,14 @@ private:
     const T *_cell_to_output_weights;
     const T *_projection_weights;
     const T *_projection_bias;
+    const T *_input_layer_norm_weights;
+    const T *_forget_layer_norm_weights;
+    const T *_cell_layer_norm_weights;
+    const T *_output_layer_norm_weights;
     bool     _has_peephole_opt;
     bool     _has_projection;
     bool     _has_cifg_opt;
+    bool     _use_layer_norm;
 };
 }
 #endif /*__ARM_COMPUTE_LSTMPARAMS_H__ */