Merge branch 'ae' into ae-flash-ampere

kernels/flash_attention/Makefile

@@ -1,8 +1,8 @@
 PROJECT = flash_attention
 
-# VX_SRCS = kernel.cpp
+VX_SRCS = kernel.cpp
 # VX_SRCS = kernel.gemmini.warpspec.cpp
-VX_SRCS = kernel.gemmini.cpp
+# VX_SRCS = kernel.gemmini.cpp
 VX_INCLUDES = flash_impl.hpp ../sgemm_tcore/sgemm_impl.hpp
 
 OPTS ?= -n16

kernels/flash_attention/flash_impl.hpp

@@ -17,9 +17,9 @@ constexpr uint32_t ROWMAX_SETS = 3;
 // constexpr bool WARP_SPECIALIZED = true;
 // constexpr bool GEMMINI_WARP_SPECIALIZED = false;
 // constexpr bool TENSOR_CORE = true;
-constexpr bool WARP_SPECIALIZED = false;
+constexpr bool WARP_SPECIALIZED = true;
 constexpr bool GEMMINI_WARP_SPECIALIZED = false;
-constexpr bool TENSOR_CORE = false;
+constexpr bool TENSOR_CORE = true;
 
 // temporary safety stop for wrong configs
 static_assert(NUM_CORES == 4);

kernels/sgemm_tcore/sgemm_impl.hpp

@@ -6,7 +6,7 @@
 #include "include/gemmini.h"
 #include "gemmini_mmio.h"
 
-#define FP_SIZE 16
+#define FP_SIZE 32
 
 // "fake" fp16 type that only has the correct data width.
 using float16_t = uint16_t;

lib/gemmini/include/gemmini_params.h

@@ -1 +1 @@
-gemmini_params.dim16fp16.h
+gemmini_params.dim8fp32.h

lib/include/VX_config.h

@@ -84,7 +84,7 @@
 #endif
 
 #ifndef NUM_CORES
-#define NUM_CORES 8
+#define NUM_CORES 4
 #endif
 
 #ifndef NUM_WARPS

lib/include/gemmini_mmio.h

@@ -12,9 +12,9 @@
 // 64KB
 // #define SMEM_SIZE 0x10000
 // 128KB (FP16 GEMM)
-#define SMEM_SIZE 0x20000
+// #define SMEM_SIZE 0x20000
 // 256KB (FlashAttention)
-// #define SMEM_SIZE 0x40000
+#define SMEM_SIZE 0x40000
 
 #define SMEM_MASK (SMEM_SIZE - 1)
 #define SMEM_ADDR_END (SMEM_BASE + SMEM_SIZE)

lib/include/vx_spawn.h

@@ -18,7 +18,7 @@
 #include <stdio.h>
 
 #ifndef CORES_PER_CLUSTER
-#define CORES_PER_CLUSTER 8
+#define CORES_PER_CLUSTER 4
 #endif
 
 #ifdef __cplusplus