Merge pull request #1810 from IntelPython/tweak-to-gemm

Author: oleksandr-pavlyk

dpctl/tensor/libtensor/include/kernels/linalg_functions/gemm.hpp

@@ -970,6 +970,7 @@ class GemmBatchFunctorThreadNM_vecm
         size_t i = block_i * wg_delta_n * wi_delta_n;
         size_t j = block_j * wg_delta_m * wi_total_delta_m;
 
+        using slmA_t = typename LocAccT1::value_type;
         using slmB_t = typename LocAccT2::value_type;
 
         const size_t a_st0 = k;
@@ -1057,16 +1058,29 @@ class GemmBatchFunctorThreadNM_vecm
             const std::uint32_t lo_lhs_st_k = (wg_delta_n * wi_delta_n);
             const std::uint32_t lo_rhs_rk_k = (wg_delta_m * wi_delta_m_vecs);
             for (std::uint32_t pr_k = 0; pr_k < wi_delta_k; ++pr_k) {
+                std::array<slmA_t, wi_delta_n> pr_lhs{};
+#pragma unroll
+                for (std::uint32_t pr_i = 0; pr_i < wi_delta_n; ++pr_i) {
+                    pr_lhs[pr_i] =
+                        local_lhs_block[pr_k * lo_lhs_st_k +
+                                        (local_i + pr_i * wg_delta_n)];
+                }
+
+                std::array<slmB_t, wi_delta_m_vecs> pr_rhs{};
+#pragma unroll
+                for (std::uint32_t pr_j = 0; pr_j < wi_delta_m_vecs; ++pr_j) {
+                    pr_rhs[pr_j] =
+                        local_rhs_block[pr_k * lo_rhs_rk_k +
+                                        (local_j + pr_j * wg_delta_m)];
+                }
+
 #pragma unroll
                 for (std::uint32_t pr_i = 0; pr_i < wi_delta_n; ++pr_i) {
 #pragma unroll
                     for (std::uint32_t pr_j = 0; pr_j < wi_delta_m_vecs; ++pr_j)
                     {
                         private_C[pr_i * wi_delta_m_vecs + pr_j] +=
-                            local_lhs_block[pr_k * lo_lhs_st_k +
-                                            (local_i + pr_i * wg_delta_n)] *
-                            local_rhs_block[pr_k * lo_rhs_rk_k +
-                                            (local_j + pr_j * wg_delta_m)];
+                            pr_lhs[pr_i] * pr_rhs[pr_j];
                     }
                 }
             }