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/kernel.cpp

@@ -95,8 +95,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
   constexpr uint32_t quartile = (128 << 10) >> 2; // 128KB / 4
   static_assert((quartile * 4) == SMEM_SIZE, "wrong quartile constant");
 
-  MARK_BEG();
-
   constexpr uint32_t smem_a_offset = 0;
   constexpr uint32_t smem_a_dbuf_offset = 1 * quartile;
   constexpr uint32_t smem_b_offset =
@@ -121,8 +119,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
                       threadblocks_per_cluster, threadblock_id_in_cluster,
                       sharedmem_per_threadblock);
 
-  MARK_END();
-
   float *gmem_tmp_d0 = reinterpret_cast<float *>(0xd0000000UL);
   float *gmem_tmp_d1 = reinterpret_cast<float *>(0xd1000000UL);
   float *gmem_tmp_d2 = reinterpret_cast<float *>(0xd2000000UL);