From c47c3d6c29ebc942048ae5df4af931cb4650a741 Mon Sep 17 00:00:00 2001
From: Ouadie EL FAROUKI <ouadie.elfarouki@codeplay.com>
Date: Tue, 30 May 2023 22:15:45 +0100
Subject: [PATCH] Fixed non-local transpose-add kernel for edge cases

---
 src/operations/extension/transpose.hpp | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/operations/extension/transpose.hpp b/src/operations/extension/transpose.hpp
index 2974df76d..533d015f3 100644
--- a/src/operations/extension/transpose.hpp
+++ b/src/operations/extension/transpose.hpp
@@ -214,31 +214,35 @@ TransposeAdd<both_trans, Tile_size, local_memory, in1_t, in2_t, out_t,
 
     if constexpr (both_trans) {
       // Compute sum & then transpose
-      auto j = idx / N_;
-      auto i = idx - j * N_;
+      auto j = idx / N_pad_;
+      auto i = idx - j * N_pad_;
 
       auto in_index_a = i + j * lda_;
       auto in_index_b = i + j * ldb_;
 
-      auto temp_sum = alpha_ * A[in_index_a] + beta_ * B[in_index_b];
+      if (i < N_ && j < M_) {
+        auto temp_sum = alpha_ * A[in_index_a] + beta_ * B[in_index_b];
 
-      auto out_index_c = i * ldc_ + j;
+        auto out_index_c = i * ldc_ + j;
 
-      C[out_index_c] = temp_sum;
+        C[out_index_c] = temp_sum;
+      }
 
     } else {
       // Transpose A then compute sum (Applies to B as well)
-      auto j = idx / M_;
-      auto i = idx - j * M_;
+      auto j = idx / M_pad_;
+      auto i = idx - j * M_pad_;
 
       auto in_index_at = j + i * lda_;
       auto in_index_b = i + j * ldb_;
 
-      auto temp_sum = alpha_ * A[in_index_at] + beta_ * B[in_index_b];
+      if (i < M_ && j < N_) {
+        auto temp_sum = alpha_ * A[in_index_at] + beta_ * B[in_index_b];
 
-      auto out_index_c = i + j * ldc_;
+        auto out_index_c = i + j * ldc_;
 
-      C[out_index_c] = temp_sum;
+        C[out_index_c] = temp_sum;
+      }
     }
   }
 }