From be898a9c8dbf8979523cf569c0b6fbfacffa668d Mon Sep 17 00:00:00 2001
From: Ouadie EL FAROUKI <ouadie.elfarouki@codeplay.com>
Date: Thu, 29 Feb 2024 12:55:18 +0000
Subject: [PATCH] updated gemm template for joint matrix case

---
 src/operations/blas3/gemm_local.hpp              | 4 ++--
 src/operations/blas3/gemm_local_joint_matrix.hpp | 7 ++++---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/operations/blas3/gemm_local.hpp b/src/operations/blas3/gemm_local.hpp
index 6bb0c6629..d39681346 100644
--- a/src/operations/blas3/gemm_local.hpp
+++ b/src/operations/blas3/gemm_local.hpp
@@ -113,7 +113,7 @@ class Gemm<input_t, output_t, DoubleBuffer, NbcA, NbcB, ClSize, TileType,
   static constexpr bool symm_a = SymmA;
   static constexpr bool symm_b = SymmB;
 
-  //! @brief Number of elements which fit within a cache line.
+  //! @brief Number of input elements which fit within a cache line.
   static constexpr index_t cl_elems = ClSize / sizeof(element_in_t);
   //! @brief Number of work items within a work group
   static constexpr index_t wg_size = wg_rows * wg_cols;
@@ -128,7 +128,7 @@ class Gemm<input_t, output_t, DoubleBuffer, NbcA, NbcB, ClSize, TileType,
 
   static_assert(wg_size % cl_elems == 0,
                 "Work group size should be a multiple "
-                "of elements in a cache line\n"
+                "of input elements in a cache line\n"
                 " --- this is ensured iff:"
                 " cl_size | sizeof(element_in_t) * wg_rows * wg_cols");
 
diff --git a/src/operations/blas3/gemm_local_joint_matrix.hpp b/src/operations/blas3/gemm_local_joint_matrix.hpp
index 440229fef..d6e37c1fd 100644
--- a/src/operations/blas3/gemm_local_joint_matrix.hpp
+++ b/src/operations/blas3/gemm_local_joint_matrix.hpp
@@ -62,6 +62,7 @@ namespace blas {
  * @tparam TransA  iff true, matrix A will be transposed on the fly
  * @tparam TransB  iff true, matrix B will be transposed on the fly
  * @tparam element_t  type of matrix elements
+ * @tparam element_out_t  type of output matrix elements (C) and scalars
  * @tparam is_beta_zero True if beta == 0.
  * @tparam VectorSize The packet size to be used for vectorization.
  * @tparam batch_type the type of batch strideded /interleaved
@@ -70,10 +71,10 @@ namespace blas {
  */
 template <typename input_t, typename output_t, bool DoubleBuffer, bool NbcA,
           bool NbcB, int ClSize, typename TileType, bool TransA, bool TransB,
-          bool SymmA, bool SymmB, typename element_t, bool is_beta_zero,
-          int VectorSize>
+          bool SymmA, bool SymmB, typename element_t, typename element_out_t,
+          bool is_beta_zero, int VectorSize>
 class Gemm<input_t, output_t, DoubleBuffer, NbcA, NbcB, ClSize, TileType,
-           TransA, TransB, SymmA, SymmB, element_t, is_beta_zero,
+           TransA, TransB, SymmA, SymmB, element_t, element_out_t, is_beta_zero,
            static_cast<int>(gemm_memory_t::local),
            static_cast<int>(gemm_algorithm_t::standard),
            static_cast<int>(gemm_vectorization_t::none), VectorSize,