Implement dynamic fusion softmax operator

 - Return aux tensorInfo by get_aux_tensors() at runtime to init the aux
   tensor with the right size.
 - Keep softmax unfusable for this commit
 - Hence, added Tensor3D to template writer arguments declaration, for sake of
   keeping dynamic fusion softmax componenets' kernels matching their cl
   counterparts.

Resolves: COMPMID-5523
Change-Id: I667f39545db925f667036ef448302c79a0330373
Signed-off-by: Ramy Elgammal <ramy.elgammal@arm.com>
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/c/VisualCompute/ComputeLibrary/+/483924
Tested-by: bsgcomp <bsgcomp@arm.com>
Reviewed-by: Gunes Bayir <gunes.bayir@arm.com>
Comments-Addressed: bsgcomp <bsgcomp@arm.com>
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/8986
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Jakub Sujak <jakub.sujak@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Benchmark: Arm Jenkins <bsgcomp@arm.com>

11 files changed, 775 insertions, 55 deletions

diff --git a/src/dynamic_fusion/sketch/gpu/GpuKernelComponentGraph.cpp b/src/dynamic_fusion/sketch/gpu/GpuKernelComponentGraph.cpp
index 4cf7a7fece..b70a192775 100644
--- a/src/dynamic_fusion/sketch/gpu/GpuKernelComponentGraph.cpp
+++ b/src/dynamic_fusion/sketch/gpu/GpuKernelComponentGraph.cpp
@@ -59,8 +59,13 @@ GpuKernelComponentStream GpuKernelComponentGraph::fuse(const MemoryDescriptorMap
     {
         const auto component = _components.at(op.op).get();
         const auto success   = stream.add_component(component);
-        ARM_COMPUTE_ERROR_ON(!success);
-        ARM_COMPUTE_UNUSED(success);
+        if(!success) // Assume first failure was because the root component is unfusable
+        {
+            stream.new_component_group();
+            const auto success = stream.add_component(component);
+            ARM_COMPUTE_ERROR_ON(!success);
+            ARM_COMPUTE_UNUSED(success);
+        }
     }
 
     return stream;
diff --git a/src/dynamic_fusion/sketch/gpu/GpuWorkloadSourceCode.h b/src/dynamic_fusion/sketch/gpu/GpuWorkloadSourceCode.h
index 2375f5c6c6..d1d0bdf77f 100644
--- a/src/dynamic_fusion/sketch/gpu/GpuWorkloadSourceCode.h
+++ b/src/dynamic_fusion/sketch/gpu/GpuWorkloadSourceCode.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022 Arm Limited.
+ * Copyright (c) 2022-2023 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
diff --git a/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DMaxShiftExpSum.cpp b/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DMaxShiftExpSum.cpp
index 8ab1853e84..279c77e227 100644
--- a/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DMaxShiftExpSum.cpp
+++ b/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DMaxShiftExpSum.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022 Arm Limited.
+ * Copyright (c) 2022-2023 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -27,9 +27,8 @@
 #include "arm_compute/core/CL/CLHelpers.h"
 #include "arm_compute/core/Validate.h"
 #include "arm_compute/core/utils/misc/ShapeCalculator.h"
-#include "arm_compute/dynamic_fusion/sketch/attributes/DepthwiseConv2dAttributes.h"
+#include "arm_compute/dynamic_fusion/sketch/attributes/SoftmaxAttributes.h"
 #include "src/core/CL/CLValidate.h"
-#include "src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateDepthwiseConv2d.h"
 #include "src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DMaxShiftExpSum.h"
 
 namespace arm_compute
@@ -49,7 +48,9 @@ Status ClComponentLogits1DMaxShiftExpSum::validate(
     const ITensorInfo *sum = tensors.get_const_tensor(TensorType::ACL_DST_0);
     const ITensorInfo *dst = tensors.get_const_tensor(TensorType::ACL_DST_1);
 
-    ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(src, sum, dst);
+    ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(src);
+    ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(sum);
+    ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(dst);
 
     // 1. Check validity
     // All tensor infos are initialized
@@ -61,11 +62,6 @@ Status ClComponentLogits1DMaxShiftExpSum::validate(
     ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(src, dst, sum);
     ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(src, dst);
 
-    const uint32_t src_dimensions = src->tensor_shape().num_dimensions();
-    const uint32_t sum_dimensions = sum->tensor_shape().num_dimensions();
-
-    ARM_COMPUTE_RETURN_ERROR_ON(src_dimensions != sum_dimensions + 1);
-
     // Device requirements are met
     ARM_COMPUTE_RETURN_ERROR_ON_F16_UNSUPPORTED(src);
 
@@ -94,4 +90,4 @@ const IGpuTemplateComponentWriter *ClComponentLogits1DMaxShiftExpSum::template_w
 }
 } // namespace dynamic_fusion
 } // namespace experimental
-} // namespace arm_compute
\ No newline at end of file
+} // namespace arm_compute
diff --git a/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DMaxShiftExpSum.h b/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DMaxShiftExpSum.h
index af47609a2c..b5db458248 100644
--- a/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DMaxShiftExpSum.h
+++ b/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DMaxShiftExpSum.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022 Arm Limited.
+ * Copyright (c) 2022-2023 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -64,9 +64,9 @@ public:
 
     /** Validate the component
      *
-     * @param[in]     properties Component properties @ref Properties
-     * @param[in,out] tensors    Tensor arguments to the component
-     * @param[in]     attributes Component attributes @ref Attributes
+     * @param[in] properties Component properties @ref Properties
+     * @param[in] tensors    Tensor arguments to the component
+     * @param[in] attributes Component attributes @ref Attributes
      *
      * @return Status        Validation results
      *
diff --git a/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DNorm.cpp b/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DNorm.cpp
new file mode 100644
index 0000000000..7864d56d29
--- /dev/null
+++ b/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DNorm.cpp
@@ -0,0 +1,95 @@
+/*
+ * Copyright (c) 2023 Arm Limited.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+#include "src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DNorm.h"
+
+#include "arm_compute/core/CL/CLHelpers.h"
+#include "arm_compute/core/Validate.h"
+#include "arm_compute/core/utils/misc/ShapeCalculator.h"
+#include "arm_compute/dynamic_fusion/sketch/attributes/SoftmaxAttributes.h"
+#include "src/core/CL/CLValidate.h"
+#include "src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DNorm.h"
+
+namespace arm_compute
+{
+namespace experimental
+{
+namespace dynamic_fusion
+{
+Status ClComponentLogits1DNorm::validate(
+    const Properties                &properties,
+    const ArgumentPack<ITensorInfo> &tensors,
+    const Attributes                &attributes)
+{
+    ARM_COMPUTE_UNUSED(properties, attributes);
+
+    const ITensorInfo *src = tensors.get_const_tensor(TensorType::ACL_SRC_0);
+    const ITensorInfo *sum = tensors.get_const_tensor(TensorType::ACL_SRC_1);
+    const ITensorInfo *dst = tensors.get_const_tensor(TensorType::ACL_DST_0);
+
+    ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(src);
+    ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(sum);
+    ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(dst);
+
+    // 1. Check validity
+    // All tensor infos are initialized
+    ARM_COMPUTE_RETURN_ERROR_ON(src->tensor_shape().total_size() == 0);
+    ARM_COMPUTE_RETURN_ERROR_ON(sum->tensor_shape().total_size() == 0);
+    ARM_COMPUTE_RETURN_ERROR_ON(dst->tensor_shape().total_size() == 0);
+
+    // Check for mismatches in shapes and data types
+    ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(src, dst, sum);
+    ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(src, dst);
+
+    ARM_COMPUTE_RETURN_ERROR_ON(attributes.is_log_softmax() && !is_data_type_float(src->data_type()));
+
+    // Device requirements are met
+    ARM_COMPUTE_RETURN_ERROR_ON_F16_UNSUPPORTED(src);
+
+    // 2. Check support level
+    ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(src, 1, DataType::F16, DataType::F32);
+
+    return Status{};
+}
+
+ClComponentLogits1DNorm::ClComponentLogits1DNorm(ComponentId                      id,
+                                                 const Properties                &properties,
+                                                 const ArgumentPack<ITensorInfo> &tensors,
+                                                 const Attributes                &attributes)
+    : IGpuKernelComponent{ id, properties, tensors },
+      _component_writer{ std::make_unique<ClTemplateLogits1DNorm>(id, tensors, attributes) }
+{
+}
+
+ClComponentLogits1DNorm::~ClComponentLogits1DNorm()
+{
+}
+
+const IGpuTemplateComponentWriter *ClComponentLogits1DNorm::template_writer() const
+{
+    return _component_writer.get();
+}
+} // namespace dynamic_fusion
+} // namespace experimental
+} // namespace arm_compute
diff --git a/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DNorm.h b/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DNorm.h
new file mode 100644
index 0000000000..5bd350b9bd
--- /dev/null
+++ b/src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DNorm.h
@@ -0,0 +1,128 @@
+/*
+ * Copyright (c) 2023 Arm Limited.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+#ifndef SRC_DYNAMIC_FUSION_SKETCH_GPU_COMPONENTS_CL_CLCOMPONENTLOGITS1DNORM
+#define SRC_DYNAMIC_FUSION_SKETCH_GPU_COMPONENTS_CL_CLCOMPONENTLOGITS1DNORM
+
+#include "arm_compute/dynamic_fusion/sketch/attributes/SoftmaxAttributes.h"
+#include "src/dynamic_fusion/sketch/gpu/components/IGpuKernelComponent.h"
+
+namespace arm_compute
+{
+/** Forward declaration */
+class ITensorInfo;
+namespace experimental
+{
+namespace dynamic_fusion
+{
+/** Forward declaration */
+template <typename T>
+class ArgumentPack;
+
+/** Forward declaration */
+class ClTemplateLogits1DNorm;
+
+/** Component to calculate the final step of the Softmax Layer
+ * where each logit value is multiplied by the inverse of the sum of the logits.
+ *
+ *  1D example:
+ *
+ *      (input)  src: [x1 x2 ... xn], shape: (1 x d)
+ *      (input)  sum: [x1 + x2 + ... + xn], shape: (1 x 1)
+ *      (output) dst: [x1/sum x2/sum ... xn/sum], shape: (1 x d)
+ *
+ *  This component is used by the softmax operator to get the final result.
+*/
+class ClComponentLogits1DNorm final : public IGpuKernelComponent
+{
+public:
+    /** Attributes are a set of backend-agnostic parameters that define what a component does */
+    using Attributes = SoftmaxAttributes;
+
+    /** Validate the component
+     *
+     * @param[in] properties Component properties @ref Properties
+     * @param[in] tensors    Tensor arguments to the component
+     * @param[in] attributes Component attributes @ref Attributes
+     *
+     * @return Status        Validation results
+     *
+     * Tensor argument names:
+     * - ACL_SRC_0: Input
+     * - ACL_SRC_1: Input
+     * - ACL_DST_0: Output
+     *
+     * Tensor argument constness:
+     * - ACL_SRC_0: Const
+     * - ACL_SRC_1: Const
+     * - ACL_DST_0: Const
+     *
+     * Valid data layouts:
+     * - All
+     *
+     ** Valid data type configurations:
+     * |ACL_SRC_0  |ACL_SRC_1  |ACL_DST_0  |
+     * |:----------|:----------|:----------|
+     * |F16        | F16       | F16       |
+     * |F32        | F32       | F32       |
+     */
+    static Status validate(
+        const Properties                &properties,
+        const ArgumentPack<ITensorInfo> &tensors,
+        const Attributes                &attributes);
+
+    /** Constructor
+     *
+     * Similar to @ref ClComponentLogits1DNorm::validate()
+     */
+    ClComponentLogits1DNorm(ComponentId                      id,
+                            const Properties                &properties,
+                            const ArgumentPack<ITensorInfo> &tensors,
+                            const Attributes                &attributes);
+
+    /** Destructor */
+    ~ClComponentLogits1DNorm() override;
+    /** Prevent instances of this class from being copy constructed */
+    ClComponentLogits1DNorm(const ClComponentLogits1DNorm &component) = delete;
+    /** Prevent instances of this class from being copied */
+    ClComponentLogits1DNorm &operator=(const ClComponentLogits1DNorm &component) = delete;
+    /** Allow instances of this class to be move constructed */
+    ClComponentLogits1DNorm(ClComponentLogits1DNorm &&component) = default;
+    /** Allow instances of this class to be moved */
+    ClComponentLogits1DNorm &operator=(ClComponentLogits1DNorm &&component) = default;
+    /** Get template writer for the component */
+    const IGpuTemplateComponentWriter *template_writer() const override;
+    /** Get component type */
+    GpuComponentType type() const override
+    {
+        return GpuComponentType::Unfusable;
+    }
+
+private:
+    std::unique_ptr<ClTemplateLogits1DNorm> _component_writer;
+};
+} // namespace dynamic_fusion
+} // namespace experimental
+} // namespace arm_compute
+
+#endif /* SRC_DYNAMIC_FUSION_SKETCH_GPU_COMPONENTS_CL_CLCOMPONENTLOGITS1DNORM */
diff --git a/src/dynamic_fusion/sketch/gpu/operators/GpuSoftmax.cpp b/src/dynamic_fusion/sketch/gpu/operators/GpuSoftmax.cpp
new file mode 100644
index 0000000000..db74538f91
--- /dev/null
+++ b/src/dynamic_fusion/sketch/gpu/operators/GpuSoftmax.cpp
@@ -0,0 +1,198 @@
+/*
+ * Copyright (c) 2023 Arm Limited.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+#include "arm_compute/dynamic_fusion/sketch/gpu/operators/GpuSoftmax.h"
+#include "arm_compute/core/Error.h"
+#include "src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DMaxShiftExpSum.h"
+#include "src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DNorm.h"
+
+#include "src/common/utils/Log.h"
+#include "src/core/helpers/AutoConfiguration.h"
+#include "src/dynamic_fusion/sketch/ArgumentPack.h"
+#include "src/dynamic_fusion/sketch/gpu/GpuOperatorProperties.h"
+#include "src/dynamic_fusion/sketch/gpu/GpuWorkloadSketchImpl.h"
+
+namespace arm_compute
+{
+namespace experimental
+{
+namespace dynamic_fusion
+{
+namespace
+{
+GpuOperatorType operator_type = GpuOperatorType::Unfusable;
+} // namespace
+
+Status GpuSoftmax::is_supported_op(const GpuWorkloadContext &context,
+                                   const ITensorInfo        *src,
+                                   const ITensorInfo        *dst,
+                                   const Attributes         &attributes)
+{
+    ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(src, dst);
+    TensorInfo dst_info_to_validate;
+
+    // Auto initialize dst tensor info
+    if(dst != nullptr)
+    {
+        dst_info_to_validate = *dst;
+    }
+    else
+    {
+        auto_init_if_empty(dst_info_to_validate, *src->clone());
+    }
+    // Check components
+    if(context.gpu_language() == GpuLanguage::OpenCL)
+    {
+        const auto cl_compile_ctx = context.cl_compile_context();
+        ARM_COMPUTE_RETURN_ERROR_ON(cl_compile_ctx == nullptr);
+        const KernelProperties properties = IGpuKernelComponent::Properties().stage(UnitWorkloadStage{ UnitWorkloadStage::Stage::Run });
+
+        TensorShape logits_sum_shape = src->tensor_shape();
+        TensorInfo  logits(src->clone()->set_tensor_shape(logits_sum_shape));
+
+        // The sum tensor dim0 only need one element
+        logits_sum_shape.set(0, 1);
+        TensorInfo sum(src->clone()->set_tensor_shape(logits_sum_shape));
+
+        // Validate Component
+        ArgumentPack<ITensorInfo> arguments_exp_sum;
+        ArgumentPack<ITensorInfo> arguments_norm;
+
+        arguments_exp_sum.add_const_tensor(ACL_SRC_0, src);
+        arguments_exp_sum.add_const_tensor(ACL_DST_0, &sum);
+        arguments_exp_sum.add_const_tensor(ACL_DST_1, &logits);
+
+        arguments_norm.add_const_tensor(ACL_SRC_0, &logits);
+        arguments_norm.add_const_tensor(ACL_SRC_1, &sum);
+        arguments_norm.add_const_tensor(ACL_DST_0, &dst_info_to_validate);
+
+        ARM_COMPUTE_RETURN_ON_ERROR(ClComponentLogits1DMaxShiftExpSum::validate(properties, arguments_exp_sum, attributes));
+        ARM_COMPUTE_RETURN_ON_ERROR(ClComponentLogits1DNorm::validate(properties, arguments_norm, attributes));
+    }
+    else
+    {
+        ARM_COMPUTE_RETURN_ERROR_MSG("Unimplemented Gpu language");
+    }
+
+    return Status{};
+}
+
+Status GpuSoftmax::validate_op(const GpuWorkloadSketch &sketch,
+                               const ITensorInfo       *src,
+                               const ITensorInfo       *dst,
+                               const Attributes        &attributes)
+{
+    ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(src, dst);
+    ARM_COMPUTE_RETURN_ERROR_ON(!src->has_valid_id() || !dst->has_valid_id());
+    ARM_COMPUTE_RETURN_ERROR_ON_MSG(src->num_dimensions() > 4, "Only up to 4 dimensions are supported");
+    ARM_COMPUTE_RETURN_ERROR_ON(attributes.axis() < static_cast<int32_t>(-src->num_dimensions()) || static_cast<int32_t>(src->num_dimensions()) <= attributes.axis());
+
+    // Auto initialize dst tensor info
+    TensorInfo dst_info_to_validate = *dst;
+    auto_init_if_empty(dst_info_to_validate, *src->clone());
+
+    const size_t actual_axis   = static_cast<size_t>(wrap_around(attributes.axis(), static_cast<int32_t>(src->num_dimensions())));
+    const bool   needs_permute = actual_axis != 0;
+    ARM_COMPUTE_RETURN_ERROR_ON_MSG(needs_permute, "Dynamic fusion softmax on axis!=0 not supported yet.");
+
+    // Perform fusion test and check if the operator meets the fusion constraints
+    ArgumentPack<ITensorInfo> tensors;
+    tensors.add_const_tensor(ACL_SRC_0, src);
+    tensors.add_const_tensor(ACL_DST_0, &dst_info_to_validate);
+
+    const auto op = sketch.implementation().operator_group().new_operator(operator_type, tensors);
+    ARM_COMPUTE_RETURN_ERROR_ON_MSG(!sketch.implementation().operator_group().try_add_operator(op),
+                                    "Operator fusion test failed. This operator cannot be fused into the workload");
+
+    // Check if configuration is supported
+    return is_supported_op(*sketch.gpu_context(), src, &dst_info_to_validate, attributes);
+}
+
+void GpuSoftmax::create_op(GpuWorkloadSketch &sketch,
+                           ITensorInfo       *src,
+                           ITensorInfo       *dst,
+                           const Attributes &attributes)
+{
+    ARM_COMPUTE_ERROR_ON_NULLPTR(src, dst);
+    ARM_COMPUTE_LOG_PARAMS(src, dst, attributes);
+    TensorShape  logits_sum_shape = src->tensor_shape();
+    ITensorInfo *logits           = sketch.implementation().create_auxiliary_tensor(src->clone()->set_tensor_shape(logits_sum_shape));
+    logits_sum_shape.set(0, 1);
+    ITensorInfo *sum = sketch.implementation().create_auxiliary_tensor(src->clone()->set_tensor_shape(logits_sum_shape));
+
+    // Auto initialize dst tensor info and the auxiliary tensor infos as well
+    auto_init_if_empty(*dst, *src->clone());
+
+    // Assert validation
+    ARM_COMPUTE_ERROR_THROW_ON(GpuSoftmax::validate_op(sketch, src, dst, attributes));
+    ARM_COMPUTE_ERROR_ON_NULLPTR(logits, sum);
+
+    // Translate into components and add to component graph
+    auto      &comp_graph = sketch.implementation().component_graph();
+    const auto sketch_ctx = sketch.implementation().context();
+
+    if(sketch_ctx->gpu_language() == GpuLanguage::OpenCL)
+    {
+        const auto cl_compile_ctx = sketch_ctx->cl_compile_context();
+        ARM_COMPUTE_ERROR_ON(cl_compile_ctx == nullptr);
+
+        // Add Direct Conv2d Component
+        {
+            auto properties = IGpuKernelComponent::Properties();
+            properties.stage(UnitWorkloadStage{ UnitWorkloadStage::Stage::Run });
+
+            ArgumentPack<ITensorInfo> arguments_exp_sum;
+            ArgumentPack<ITensorInfo> arguments_norm;
+
+            arguments_exp_sum.add_const_tensor(ACL_SRC_0, src);
+            arguments_exp_sum.add_const_tensor(ACL_DST_0, sum);
+            arguments_exp_sum.add_const_tensor(ACL_DST_1, logits);
+
+            arguments_norm.add_const_tensor(ACL_SRC_0, logits);
+            arguments_norm.add_const_tensor(ACL_SRC_1, sum);
+            arguments_norm.add_const_tensor(ACL_DST_0, dst);
+
+            comp_graph.add_new_component<ClComponentLogits1DMaxShiftExpSum>(properties, arguments_exp_sum, attributes);
+            comp_graph.add_new_component<ClComponentLogits1DNorm>(properties, arguments_norm, attributes);
+        }
+    }
+    else
+    {
+        ARM_COMPUTE_ERROR("Unimplemented Gpu language");
+    }
+
+    // Set up fusion test by adding to the Operator Group
+    // Note this has to be performed after all the components have been successfully added to the component graph
+
+    // Pack tensor infos
+    ArgumentPack<ITensorInfo> tensors;
+    tensors.add_const_tensor(ACL_SRC_0, src);
+    tensors.add_const_tensor(ACL_DST_0, dst);
+
+    const auto op = sketch.implementation().operator_group().new_operator(operator_type, tensors);
+    sketch.implementation().operator_group().add_operator(op);
+}
+
+} // namespace dynamic_fusion
+} // namespace experimental
+} // namespace arm_compute
diff --git a/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DMaxShiftExpSum.cpp b/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DMaxShiftExpSum.cpp
index 8f1ed95351..6de0ba7617 100644
--- a/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DMaxShiftExpSum.cpp
+++ b/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DMaxShiftExpSum.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022 Arm Limited.
+ * Copyright (c) 2022-2023 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -34,6 +34,10 @@ namespace experimental
 {
 namespace dynamic_fusion
 {
+namespace
+{
+    constexpr unsigned int serial_vector_size = 8;
+} // namespace
 ClTemplateLogits1DMaxShiftExpSum::ClTemplateLogits1DMaxShiftExpSum(ComponentId                      id,
                                                                    const ArgumentPack<ITensorInfo> &tensors,
                                                                    const Attributes                &attributes)
@@ -46,7 +50,9 @@ ClTemplateLogits1DMaxShiftExpSum::ClTemplateLogits1DMaxShiftExpSum(ComponentId
     _src = this->tensors().get_const_tensor(TensorType::ACL_SRC_0);
     _sum = this->tensors().get_const_tensor(TensorType::ACL_DST_0);
     _dst = this->tensors().get_const_tensor(TensorType::ACL_DST_1);
-    ARM_COMPUTE_ERROR_ON_NULLPTR(_src, _sum, _dst);
+    ARM_COMPUTE_ERROR_ON_NULLPTR(_src);
+    ARM_COMPUTE_ERROR_ON_NULLPTR(_sum);
+    ARM_COMPUTE_ERROR_ON_NULLPTR(_dst);
 }
 
 std::string ClTemplateLogits1DMaxShiftExpSum::get_name() const
@@ -65,7 +71,6 @@ std::string ClTemplateLogits1DMaxShiftExpSum::get_component_code(const Component
 {
     __global uchar *src_addr = {{src}}_ptr + {{src}}_offset_first_element_in_bytes + g_ind_1 * {{src}}_stride_y + g_ind_2 * {{src}}_stride_z;
     __global uchar *dst_addr = {{dst}}_ptr + {{dst}}_offset_first_element_in_bytes + g_ind_1 * {{dst}}_stride_y + g_ind_2 * {{dst}}_stride_z;
-
     Image sum = CONVERT_TENSOR3D_TO_IMAGE_STRUCT({{sum}});
     VEC_TYPE max_val_vec = (VEC_TYPE)({{MINVAL}});
 )_";
@@ -139,44 +144,41 @@ std::string ClTemplateLogits1DMaxShiftExpSum::get_component_code(const Component
     sum1D += data;
 )_";
     }
-    else
-    {
-        code += R"_(
+    code += R"_(
     for(uint i = PARTIAL_N0; i < {{SRC_WIDTH}}; i += N0)
     {
         VEC_TYPE data = VLOAD(N0)(0, (__global {{DATA_TYPE}} *)(src_addr + i * sizeof({{DATA_TYPE}})));
         data -= max_val;
 )_";
 
-        if(beta_defined)
-        {
-            code += R"_(
-        data *= beta;
+    if(beta_defined)
+    {
+        code += R"_(
+    data *= beta;
 )_";
-        }
+    }
 
-        if(_attributes.is_log_softmax())
-        {
-            code += R"_(
-        VSTORE(N0)
-        (data, 0, (__global {{DATA_TYPE}} *)(dst_addr + i * sizeof({{DATA_TYPE}})));
-        data = exp(data);
+    if(_attributes.is_log_softmax())
+    {
+        code += R"_(
+    VSTORE(N0)
+    (data, 0, (__global {{DATA_TYPE}} *)(dst_addr + i * sizeof({{DATA_TYPE}})));
+    data = exp(data);
 )_";
-        }
-        else
-        {
-            code += R"_(
-        data = exp(data);
-        VSTORE(N0)
-        (data, 0, (__global {{DATA_TYPE}} *)(dst_addr + i * sizeof({{DATA_TYPE}})));
+    }
+    else
+    {
+        code += R"_(
+    data = exp(data);
+    VSTORE(N0)
+    (data, 0, (__global {{DATA_TYPE}} *)(dst_addr + i * sizeof({{DATA_TYPE}})));
 )_";
-        }
+    }
 
-        code += R"_(
-        sum1D += data;
+    code += R"_(
+    sum1D += data;
     }
 )_";
-    }
 
     code += R"_(
     *((__global {{DATA_TYPE}} *)sum.ptr) = SUM_REDUCE(sum1D, N0);
@@ -192,19 +194,19 @@ void ClTemplateLogits1DMaxShiftExpSum::declare_variables(GpuKernelVariableTable
     vtable.declare_variable(
         comp_group,
         _src,
-        GpuKernelArgumentInfo(GpuKernelArgumentInfo::Type::Tensor_4D_t_Buffer),
+        GpuKernelArgumentInfo(GpuKernelArgumentInfo::Type::Tensor_3D),
         "src");
 
     vtable.declare_variable(
         comp_group,
         _sum,
-        GpuKernelArgumentInfo(GpuKernelArgumentInfo::Type::Tensor_4D_t_Buffer),
+        GpuKernelArgumentInfo(GpuKernelArgumentInfo::Type::Tensor_3D),
         "sum");
 
     vtable.declare_variable(
         comp_group,
         _dst,
-        GpuKernelArgumentInfo(GpuKernelArgumentInfo::Type::Tensor_4D_t_Buffer),
+        GpuKernelArgumentInfo(GpuKernelArgumentInfo::Type::Tensor_3D),
         "dst");
 }
 
@@ -237,7 +239,6 @@ CLBuildOptions ClTemplateLogits1DMaxShiftExpSum::get_build_options(const Compone
     ARM_COMPUTE_UNUSED(comp_group);
     CLBuildOptions build_opts{};
 
-    constexpr unsigned int serial_vector_size = 8;
     const unsigned int     reduction_dim_size = _src->dimension(0);
     const unsigned int     vector_size        = adjust_vec_size(serial_vector_size, reduction_dim_size);
 
@@ -261,7 +262,7 @@ std::string ClTemplateLogits1DMaxShiftExpSum::get_config_id() const
 
 std::set<std::string> ClTemplateLogits1DMaxShiftExpSum::get_headers_list() const
 {
-    return std::set<std::string>{ "helpers.h" };
+    return std::set<std::string>{ "helpers.h", "tile_helpers.h" };
 }
 
 Window ClTemplateLogits1DMaxShiftExpSum::get_window() const
diff --git a/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DNorm.cpp b/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DNorm.cpp
new file mode 100644
index 0000000000..0e1c9ef28f
--- /dev/null
+++ b/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DNorm.cpp
@@ -0,0 +1,185 @@
+/*
+ * Copyright (c) 2023 Arm Limited.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+#include "src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DNorm.h"
+
+#include "src/core/helpers/WindowHelpers.h"
+#include "src/dynamic_fusion/sketch/gpu/GpuKernelComponentGroup.h"
+#include "support/StringSupport.h"
+
+namespace arm_compute
+{
+namespace experimental
+{
+namespace dynamic_fusion
+{
+ClTemplateLogits1DNorm::ClTemplateLogits1DNorm(ComponentId                      id,
+                                               const ArgumentPack<ITensorInfo> &tensors,
+                                               const Attributes                &attributes)
+    : IGpuTemplateComponentWriter{ id, tensors },
+      _src{},
+      _sum{},
+      _dst{},
+      _attributes{ attributes }
+{
+    _src = this->tensors().get_const_tensor(TensorType::ACL_SRC_0);
+    _sum = this->tensors().get_const_tensor(TensorType::ACL_SRC_1);
+    _dst = this->tensors().get_const_tensor(TensorType::ACL_DST_0);
+    ARM_COMPUTE_ERROR_ON_NULLPTR(_src);
+    ARM_COMPUTE_ERROR_ON_NULLPTR(_sum);
+    ARM_COMPUTE_ERROR_ON_NULLPTR(_dst);
+}
+
+std::string ClTemplateLogits1DNorm::get_name() const
+{
+    return "logits_1d_norm";
+}
+
+std::string ClTemplateLogits1DNorm::get_component_code(const ComponentGroup &comp_group) const
+{
+    ARM_COMPUTE_UNUSED(comp_group);
+
+    std::string code = R"_(
+//------------------ START KERNEL {{meta_kernel_id}} ---------------------
+{
+    const int x_offs = g_ind_0 * sizeof({{DATA_TYPE}});
+    __global uchar *src_addr = {{src}}_ptr + {{src}}_offset_first_element_in_bytes + x_offs + g_ind_1 * {{src}}_stride_y + g_ind_2 * {{src}}_stride_z;
+    __global uchar *dst_addr = {{dst}}_ptr + {{dst}}_offset_first_element_in_bytes + x_offs + g_ind_1 * {{dst}}_stride_y + g_ind_2 * {{dst}}_stride_z;
+    Image sum = CONVERT_TENSOR3D_TO_IMAGE_STRUCT_NO_STEP({{sum}});
+)_";
+    // Load max value of 1D logits vector (row)
+    code += R"_(
+    {{DATA_TYPE}} sum_val = *((__global {{DATA_TYPE}} *)offset(&sum, 0, g_ind_1));
+    VEC_DATA_TYPE({{DATA_TYPE}}, N0)
+    data0 = VLOAD(N0)(0, (__global {{DATA_TYPE}} *)src_addr);
+)_";
+
+    if(_attributes.is_log_softmax())
+    {
+        code += R"_(
+    sum_val = log(sum_val);
+    data0 -= sum_val;
+)_";
+    }
+    else
+    {
+        code += R"_(
+    data0 /= sum_val;
+)_";
+    }
+
+    code += R"_(
+    STORE_VECTOR_SELECT(data, {{DATA_TYPE}}, dst_addr, N0, PARTIAL_N0, PARTIAL_N0 != 0 && g_ind_0 == 0);
+}
+//------------------ END KERNEL {{meta_kernel_id}} ---------------------
+)_";
+
+    return code;
+}
+
+void ClTemplateLogits1DNorm::declare_variables(GpuKernelVariableTable &vtable, const ComponentGroup &comp_group) const
+{
+    vtable.declare_variable(
+        comp_group,
+        _src,
+        GpuKernelArgumentInfo(GpuKernelArgumentInfo::Type::Tensor_3D),
+        "src");
+
+    vtable.declare_variable(
+        comp_group,
+        _sum,
+        GpuKernelArgumentInfo(GpuKernelArgumentInfo::Type::Tensor_3D),
+        "sum");
+
+    vtable.declare_variable(
+        comp_group,
+        _dst,
+        GpuKernelArgumentInfo(GpuKernelArgumentInfo::Type::Tensor_3D),
+        "dst");
+}
+
+TagLUT ClTemplateLogits1DNorm::get_tag_lut(const GpuKernelVariableTable &vtable, const ComponentGroup &comp_group) const
+{
+    ARM_COMPUTE_UNUSED(comp_group);
+
+    TagLUT lut{};
+
+    // Arguments and global shared variables
+    lut["src"] = vtable.get_variable(_src);
+    lut["sum"] = vtable.get_variable(_sum);
+    lut["dst"] = vtable.get_variable(_dst);
+
+    // Local build options
+    lut["meta_kernel_id"] = id();
+
+    const DataType data_type = _src->data_type();
+
+    lut["DATA_TYPE"] = get_cl_type_from_data_type(data_type);
+
+    return lut;
+}
+
+CLBuildOptions ClTemplateLogits1DNorm::get_build_options(const ComponentGroup &comp_group) const
+{
+    ARM_COMPUTE_UNUSED(comp_group);
+    CLBuildOptions build_opts{};
+
+    const auto         root_window = comp_group.get_root_component()->template_writer()->get_window();
+    const unsigned int n0          = root_window.x().step();
+    build_opts.add_option("-DN0=" + support::cpp11::to_string(n0));
+    build_opts.add_option("-DPARTIAL_N0=" + support::cpp11::to_string((_src->dimension(0) % n0)));
+
+    return build_opts;
+}
+
+std::string ClTemplateLogits1DNorm::get_config_id() const
+{
+    std::string config_id = get_name();
+
+    config_id += "_";
+    config_id += support::cpp11::to_string(_src->dimension(0));
+    config_id += "_";
+    config_id += string_from_data_type(_src->data_type());
+
+    return config_id;
+}
+
+std::set<std::string> ClTemplateLogits1DNorm::get_headers_list() const
+{
+    return std::set<std::string>{ "helpers.h", "tile_helpers.h" };
+}
+
+Window ClTemplateLogits1DNorm::get_window() const
+{
+    ARM_COMPUTE_ERROR_ON_MSG(_dst->tensor_shape().total_size() == 0U, "Destination tensor is not initialized");
+    constexpr unsigned int serial_vector_size = 16;
+    const unsigned int vector_size = adjust_vec_size(serial_vector_size, _src->dimension(0));
+
+    Window win = calculate_max_window(*_src, Steps(vector_size));
+    return win.collapse(win, Window::DimZ);
+}
+
+} // namespace dynamic_fusion
+} // namespace experimental
+} // namespace arm_compute
diff --git a/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DNorm.h b/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DNorm.h
new file mode 100644
index 0000000000..5a74be5842
--- /dev/null
+++ b/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateLogits1DNorm.h
@@ -0,0 +1,106 @@
+/*
+ * Copyright (c) 2023 Arm Limited.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+#ifndef SRC_DYNAMIC_FUSION_SKETCH_GPU_TEMPLATE_WRITER_CL_CLTEMPLATELOGITS1DNORM
+#define SRC_DYNAMIC_FUSION_SKETCH_GPU_TEMPLATE_WRITER_CL_CLTEMPLATELOGITS1DNORM
+
+#include "src/dynamic_fusion/sketch/gpu/components/cl/ClComponentLogits1DNorm.h"
+#include "src/dynamic_fusion/sketch/gpu/template_writer/GpuKernelVariableTable.h"
+#include "src/dynamic_fusion/sketch/gpu/template_writer/IGpuTemplateComponentWriter.h"
+
+namespace arm_compute
+{
+namespace experimental
+{
+namespace dynamic_fusion
+{
+class ClTemplateLogits1DNorm final : public IGpuTemplateComponentWriter
+{
+public:
+    using Attributes = ClComponentLogits1DNorm::Attributes;
+
+    /** Constructor
+     *
+     * @param[in] id         Component id
+     * @param[in] tensors    Tensor arguments to the components
+     * @param[in] attributes Component attributes
+     */
+    ClTemplateLogits1DNorm(ComponentId id, const ArgumentPack<ITensorInfo> &tensors, const Attributes &attributes);
+    /** Prevent instances of this class from being copy constructed */
+    ClTemplateLogits1DNorm(const ClTemplateLogits1DNorm &) = delete;
+    /** Prevent instances of this class from being copied */
+    ClTemplateLogits1DNorm &operator=(const ClTemplateLogits1DNorm &) = delete;
+    /** Allow instances of this class to be move constructed */
+    ClTemplateLogits1DNorm(ClTemplateLogits1DNorm &&) = default;
+    /** Allow instances of this class to be moved */
+    ClTemplateLogits1DNorm &operator=(ClTemplateLogits1DNorm &&) = default;
+    /** Generate kernel component name */
+    std::string get_name() const override;
+    /** Generate kernel component code template
+     *
+     * @param[in] comp_group Component group of which the component is a part of
+     *
+     * @return std::string Component code
+     */
+    std::string get_component_code(const ComponentGroup &comp_group) const override;
+    /** Declare all variables used by the component in the @p vtable
+     *
+     * @param[out] vtable     Variable table
+     * @param[in]  comp_group Component group of which the component is a part of
+     */
+    void declare_variables(GpuKernelVariableTable &vtable, const ComponentGroup &comp_group) const override;
+    /** Generate the tag look-up table used to instantiate the component code.
+     *
+     * @param[in] vtable     Variable table
+     * @param[in] comp_group Component group of which the component is a part of
+     *
+     * @return TagLUT  Tag lookup table
+     */
+    TagLUT get_tag_lut(const GpuKernelVariableTable &vtable, const ComponentGroup &comp_group) const override;
+    /** Generate the build options used in the component
+     *
+     * @param[in] comp_group Component group of which the component is a part of
+     *
+     * @return CLBuildOptions Build options
+     */
+    CLBuildOptions get_build_options(const ComponentGroup &comp_group) const override;
+    /** Generate the component config id string used for tuning */
+    std::string get_config_id() const override;
+    /** Generate the header list used in the component */
+    std::set<std::string> get_headers_list() const override;
+    /** Generate the execution window for the component */
+    Window get_window() const override;
+
+private:
+    const ITensorInfo *_src; // exponentiated input
+    const ITensorInfo *_sum; // exponentiated and summed input
+    const ITensorInfo *_dst; // normalization of input with _sum
+
+    Attributes _attributes;
+};
+} // namespace dynamic_fusion
+} // namespace experimental
+} // namespace arm_compute
+
+#endif /* SRC_DYNAMIC_FUSION_SKETCH_GPU_TEMPLATE_WRITER_CL_CLTEMPLATELOGITS1DNORM */
diff --git a/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateWriter.cpp b/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateWriter.cpp
index 2ab6316947..eda15f1d95 100644
--- a/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateWriter.cpp
+++ b/src/dynamic_fusion/sketch/gpu/template_writer/cl/ClTemplateWriter.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022 Arm Limited.
+ * Copyright (c) 2022-2023 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -180,7 +180,8 @@ std::string ClTemplateWriter::write_code()
     }
 
     auto arguments = _components.get_argument_tensors();
-    std::sort(arguments.begin(), arguments.end(), [](const ITensorInfo *l, const ITensorInfo *r) {
+    std::sort(arguments.begin(), arguments.end(), [](const ITensorInfo * l, const ITensorInfo * r)
+    {
         return l->id() < r->id();
     });
     code += write_kernel_signature(_vtable.get_variable_list(arguments));
@@ -192,16 +193,16 @@ std::string ClTemplateWriter::write_code()
     code += "    //------------------ END KERNEL_BUILDER_COORDINATE ---------------------\n";
 
     {
-        const auto tiles = _components.get_tiles();
+        const auto        tiles = _components.get_tiles();
         std::stringstream tiles_ss;
 
         tiles_ss << "    //------------------ START TILE DECLARATION ---------------------\n";
 
         for(auto tile : tiles)
         {
-            const auto var = _vtable.get_variable(tile);
+            const auto var       = _vtable.get_variable(tile);
             const auto data_type = get_cl_type_from_data_type(tile->data_type());
-            const auto var_name = var.uniq_name;
+            const auto var_name  = var.uniq_name;
 
             tiles_ss << "    TILE(" << data_type << ", M0, N0, " << var_name << ");\n";
         }
@@ -276,6 +277,11 @@ std::string ClTemplateWriter::write_argument_declaration(const GpuKernelVariable
             code += "\n    TENSOR4D_T(" + var.uniq_name + ", IMAGE)";
             break;
         }
+        case GpuKernelArgumentInfo::Type::Tensor_3D:
+        {
+            code += "\n    TENSOR3D_DECLARATION(" + var.uniq_name + ")";
+            break;
+        }
         default:
         {
             ARM_COMPUTE_ERROR("Unsupported declaration generation for GpuKernelArgumentInfo::Type");