Fix FlashAttention as GEMM kernel requires 64b-aligned matrices

Author: colluca

sw/kernels/blas/gemm/src/gemm_fp16.h

@@ -306,7 +306,7 @@ void gemm_fp16_opt_ex(uint32_t setup_ssr, uint32_t partition_banks,
 #ifdef SNRT_SUPPORTS_FREP
     __fp16* A = (__fp16*)A_p;
     __fp16* B = (__fp16*)B_p;
-    __fp16* C = (__fp16*)C_p;
+    __fp16* C = (__fp16*)C_p;  // Should be double-aligned (see fsd below)
 
     // Unrolling factor of most inner loop.
     // Should be at least as high as the FMA delay

sw/kernels/blas/gemm/src/gemm_fp32.h

@@ -214,7 +214,7 @@ void gemm_fp32_opt(uint32_t setup_ssr, uint32_t partition_banks,
     // cast void pointers to float pointers
     float* A = (float*)A_p;
     float* B = (float*)B_p;
-    float* C = (float*)C_p;
+    float* C = (float*)C_p;  // Should be double-aligned (see fsd below)
     // Unrolling factor of most inner loop.
     // Should be at least as high as the FMA delay
     // for maximum utilization

sw/kernels/dnn/flashattention_2/src/flashattention_2_fp16.h

@@ -56,27 +56,28 @@ static inline void flashattention_2_fp16(flashattention_2_layer_t layer) {
     uint32_t shifted_exp_size = B_r * sizeof(float);
 
     // allocate memory in TCDM
+    // align to size of double since this is required for some GEMM arrays
     __fp16 *Q_fa =
-        (__fp16 *)snrt_l1_alloc_cluster_local(q_fa_size, alignof(__fp16));
+        (__fp16 *)snrt_l1_alloc_cluster_local(q_fa_size, alignof(double));
     __fp16 *K_fa =
-        (__fp16 *)snrt_l1_alloc_cluster_local(k_fa_size, alignof(__fp16));
+        (__fp16 *)snrt_l1_alloc_cluster_local(k_fa_size, alignof(double));
     __fp16 *V_fa =
-        (__fp16 *)snrt_l1_alloc_cluster_local(v_fa_size, alignof(__fp16));
+        (__fp16 *)snrt_l1_alloc_cluster_local(v_fa_size, alignof(double));
     __fp16 *S_fa =
-        (__fp16 *)snrt_l1_alloc_cluster_local(s_fa_size, alignof(__fp16));
+        (__fp16 *)snrt_l1_alloc_cluster_local(s_fa_size, alignof(double));
     __fp16 *P_fa =
-        (__fp16 *)snrt_l1_alloc_cluster_local(p_fa_size, alignof(__fp16));
+        (__fp16 *)snrt_l1_alloc_cluster_local(p_fa_size, alignof(double));
     __fp16 *O_fa =
-        (__fp16 *)snrt_l1_alloc_cluster_local(o_fa_size, alignof(__fp16));
-    float *m_i = (float *)snrt_l1_alloc_cluster_local(m_i_size, alignof(float));
+        (__fp16 *)snrt_l1_alloc_cluster_local(o_fa_size, alignof(double));
+    float *m_i = (float *)snrt_l1_alloc_cluster_local(m_i_size, alignof(double));
     float *m_i_prev =
-        (float *)snrt_l1_alloc_cluster_local(m_i_size, alignof(float));
-    float *l_i = (float *)snrt_l1_alloc_cluster_local(l_i_size, alignof(float));
+        (float *)snrt_l1_alloc_cluster_local(m_i_size, alignof(double));
+    float *l_i = (float *)snrt_l1_alloc_cluster_local(l_i_size, alignof(double));
 
     // Allocate space for V^t
     __fp16 *V_t;
     if (!baseline) {
-        V_t = (__fp16 *)snrt_l1_alloc_cluster_local(v_fa_size, alignof(__fp16));
+        V_t = (__fp16 *)snrt_l1_alloc_cluster_local(v_fa_size, alignof(double));
     }
 
     float shifted_exp;

sw/kernels/dnn/flashattention_2/src/flashattention_2_fp32.h

@@ -56,27 +56,28 @@ static inline void flashattention_2_fp32(flashattention_2_layer_t layer) {
     uint32_t shifted_exp_size = B_r * sizeof(float);
 
     // allocate memory in TCDM
+    // align to size of double since this is required for some GEMM arrays
     float *Q_fa =
-        (float *)snrt_l1_alloc_cluster_local(q_fa_size, alignof(float));
+        (float *)snrt_l1_alloc_cluster_local(q_fa_size, alignof(double));
     float *K_fa =
-        (float *)snrt_l1_alloc_cluster_local(k_fa_size, alignof(float));
+        (float *)snrt_l1_alloc_cluster_local(k_fa_size, alignof(double));
     float *V_fa =
-        (float *)snrt_l1_alloc_cluster_local(v_fa_size, alignof(float));
+        (float *)snrt_l1_alloc_cluster_local(v_fa_size, alignof(double));
     float *S_fa =
-        (float *)snrt_l1_alloc_cluster_local(s_fa_size, alignof(float));
+        (float *)snrt_l1_alloc_cluster_local(s_fa_size, alignof(double));
     float *P_fa =
-        (float *)snrt_l1_alloc_cluster_local(p_fa_size, alignof(float));
+        (float *)snrt_l1_alloc_cluster_local(p_fa_size, alignof(double));
     float *O_fa =
-        (float *)snrt_l1_alloc_cluster_local(o_fa_size, alignof(float));
-    float *m_i = (float *)snrt_l1_alloc_cluster_local(m_i_size, alignof(float));
+        (float *)snrt_l1_alloc_cluster_local(o_fa_size, alignof(double));
+    float *m_i = (float *)snrt_l1_alloc_cluster_local(m_i_size, alignof(double));
     float *m_i_prev =
-        (float *)snrt_l1_alloc_cluster_local(m_i_size, alignof(float));
-    float *l_i = (float *)snrt_l1_alloc_cluster_local(l_i_size, alignof(float));
+        (float *)snrt_l1_alloc_cluster_local(m_i_size, alignof(double));
+    float *l_i = (float *)snrt_l1_alloc_cluster_local(l_i_size, alignof(double));
 
     // allocate space for V^t when using optimized kernels
     float *V_t;
     if (!baseline) {
-        V_t = (float *)snrt_l1_alloc_cluster_local(v_fa_size, alignof(float));
+        V_t = (float *)snrt_l1_alloc_cluster_local(v_fa_size, alignof(double));
     }
 
     float shifted_exp;