flash: Add DMA config logic

tests/regression/flash_attention/kernel.cpp

@@ -617,6 +617,26 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
     threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
   }
 
+  if constexpr (GEMMINI_DMA) {
+    if (tid_in_threadblock == 0) {
+      gemmini_extended_config_ex(WEIGHT_STATIONARY, 0, 0, 1, 0, 0);
+      // gemmini_extended_config_ex(dataflow, act & 3, 0, 1, a_transpose,
+      // b_transpose);
+
+      // configure DMA for Q tile
+      gemmini_extended3_config_ld(HEADDIM * sizeof(elem_t), MVIN_SCALE_IDENTITY,
+                                  false, 0);
+      // configure DMA for K tile
+      gemmini_extended3_config_ld(B_COL * sizeof(elem_t), MVIN_SCALE_IDENTITY,
+                                  false, 1);
+      // configure DMA for Q*K store
+      gemmini_extended_config_st(B_COL * sizeof(elem_t), 0,
+                                 MVIN_SCALE_IDENTITY);
+
+      gemmini_fence();
+    }
+  }
+
   // read Q and K into SMEM before the loop starts
   //
   static_assert(B_ROW == B_COL, "currently only supports square tiles");