Add flash attention kernel skeleton

tests/regression/flash_attention/Makefile

@@ -0,0 +1,12 @@
+PROJECT = flash_attention
+
+SRCS = main.cpp common.h
+
+VX_SRCS = kernel.cpp
+VX_INCLUDES = ../sgemm_tcore/sgemm_impl.hpp
+
+OPTS ?= -n16
+
+VX_CFLAGS += -I../sgemm_tcore
+
+include ../common.mk

tests/regression/flash_attention/common.h

@@ -0,0 +1,18 @@
+#ifndef _COMMON_H_
+#define _COMMON_H_
+
+#include <cstdint>
+
+#define KERNEL_ARG_DEV_MEM_ADDR 0x9fff0000
+#define DEV_SMEM_START_ADDR 0xff000000
+
+typedef struct {
+  uint32_t dim_m;
+  uint32_t dim_n;
+  uint32_t dim_k;
+  uint64_t addr_a;
+  uint64_t addr_b;
+  uint64_t addr_c;
+} kernel_arg_t;
+
+#endif

tests/regression/flash_attention/kernel.cpp

@@ -0,0 +1,158 @@
+#include <stdint.h>
+#include <vx_intrinsics.h>
+#include <vx_print.h>
+#include <vx_spawn.h>
+#include "common.h"
+#include "sgemm_impl.hpp"
+#include "include/gemmini.h"
+#include "gemmini_mmio.h"
+
+// using float_type = float;
+using float_type = float16_t;
+
+inline void thread_block_flashattn(float *S, float *gmem,
+                                   const uint32_t tid_in_threadblock,
+                                   const uint32_t threads_per_threadblock,
+                                   const uint32_t threadblock_id_in_cluster,
+                                   uint8_t *sharedmem_per_threadblock) {
+  asm volatile("thread_block_flashattn_start_%=:" ::);
+
+  constexpr uint32_t Brow = BM; // FIXME
+  constexpr uint32_t Bcol = BN; // FIXME
+  const uint32_t tid_in_warp = tid_in_threadblock % NUM_THREADS;
+  const uint32_t warp_id = tid_in_threadblock / NUM_THREADS;
+  const uint32_t warps_in_threadblock = threads_per_threadblock / NUM_THREADS;
+  const uint32_t warps_per_threadblock_per_core =
+      NUM_WARPS / threads_per_threadblock;
+
+  // float ft[8];
+  // asm volatile("fmv.s %0, f16" : "=f"(ft[0]));
+  // asm volatile("fmv.s %0, f17" : "=f"(ft[1]));
+  // asm volatile("fmv.s %0, f18" : "=f"(ft[2]));
+  // asm volatile("fmv.s %0, f19" : "=f"(ft[3]));
+  // asm volatile("fmv.s %0, f20" : "=f"(ft[4]));
+  // asm volatile("fmv.s %0, f21" : "=f"(ft[5]));
+  // asm volatile("fmv.s %0, f22" : "=f"(ft[6]));
+  // asm volatile("fmv.s %0, f23" : "=f"(ft[7]));
+  //
+  // one warp handles one row in tile; iterate enough times to cover all the
+  // rows
+  for (int warp_offset = 0; warp_offset < Brow;
+       warp_offset += warps_in_threadblock) {
+    const uint32_t row = warp_offset + warp_id;
+    const uint32_t first_thread_offset = Bcol * row;
+    uint32_t thread_offset = first_thread_offset + tid_in_warp;
+
+    constexpr uint32_t load_iter = Bcol / NUM_THREADS;
+    float curr_max = S[first_thread_offset];
+#pragma GCC unroll
+    for (int iter = 0; iter < load_iter; iter++) {
+      asm volatile("fmax.s %0, %1, %2"
+                   : "=f"(curr_max)
+                   : "f"(curr_max), "f"(S[thread_offset]));
+      thread_offset += NUM_THREADS;
+    }
+    // get max value across the same-warp threads using smem
+    float *warp_smem = S + (row * NUM_THREADS);
+    warp_smem[tid_in_warp] = curr_max;
+
+    // sync writes to warp_smem
+    threadblock_barrier(threadblock_id_in_cluster,
+                        warps_per_threadblock_per_core);
+
+    // 0-th thread collects all other thread's values in the warp
+    if (tid_in_warp == 0) {
+      for (int iter = 1; iter < NUM_THREADS; iter++) {
+        float other = warp_smem[iter];
+        asm volatile("fmax.s %0, %1, %2"
+                     : "=f"(curr_max)
+                     : "f"(curr_max), "f"(other));
+      }
+      gmem[row] = curr_max;
+    }
+  }
+
+  asm volatile("thread_block_flashattn_finish_%=:" ::);
+}
+
+void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
+  // @perf: All threads are running these compute whose result is mostly same
+  // across the threadblock
+
+#ifdef RADIANCE
+  constexpr uint32_t cores_per_cluster = CORES_PER_CLUSTER;
+#else
+  constexpr uint32_t cores_per_cluster = 1;
+#endif
+
+  uint32_t threads_per_threadblock = (BM * BN) / (ELEM_PER_THREAD);
+  const uint32_t hw_threads_per_cluster =
+      cores_per_cluster * vx_num_threads() * vx_num_warps();
+  // cap maximum threadblock size to # of HW threads in cluster, to prevent
+  // multiple "wave" invocations which slows down the kernel
+  if (threads_per_threadblock > hw_threads_per_cluster) {
+    threads_per_threadblock = hw_threads_per_cluster;
+  }
+  const uint32_t threadblocks_per_cluster =
+      hw_threads_per_cluster / threads_per_threadblock;
+
+  const int threadblock_id = task_id / threads_per_threadblock;
+  const int threadblock_id_in_cluster =
+      threadblock_id % threadblocks_per_cluster;
+  const int tid_in_threadblock = task_id % threads_per_threadblock;
+
+  const uint32_t dim_m = arg->dim_m;
+  const uint32_t dim_n = arg->dim_n;
+  const uint32_t dim_n_in_blocks = dim_n / BN;
+  const int threadblock_id_x = threadblock_id % dim_n_in_blocks;
+  const int threadblock_id_y = threadblock_id / dim_n_in_blocks;
+  const uint32_t problem_size = (dim_m * dim_n) / (ELEM_PER_THREAD);
+  const uint32_t num_threadblocks = problem_size / threads_per_threadblock;
+
+  // "static" shared memory allocation.  This would determine threadblock
+  // occupancy of a single cluster
+  uint8_t *sharedmem_per_threadblock = reinterpret_cast<uint8_t *>(
+      DEV_SMEM_START_ADDR + sizeof(float_type) * 2 /*overkill for non-dma*/ *
+                                (2 * BM * BK) * threadblock_id_in_cluster);
+
+  uint8_t *smem_S = sharedmem_per_threadblock;
+
+  thread_block_gemm<float_type, /*write_to_gmem=*/true>(
+      (const float_type *)arg->addr_a, (const float_type *)arg->addr_b,
+      (float *)smem_S /*write result to SMEM */, arg->dim_m, arg->dim_n,
+      arg->dim_k, tid_in_threadblock, threads_per_threadblock,
+      threadblocks_per_cluster, threadblock_id_in_cluster,
+      sharedmem_per_threadblock);
+
+  // sync writes of GEMM results before softmax
+  const uint32_t warps_per_threadblock_per_core =
+      NUM_WARPS / threads_per_threadblock;
+  threadblock_barrier(threadblock_id_in_cluster,
+                      warps_per_threadblock_per_core);
+
+  thread_block_flashattn((float *)smem_S, (float *)arg->addr_c,
+                         tid_in_threadblock, threads_per_threadblock,
+                         threadblock_id_in_cluster, sharedmem_per_threadblock);
+}
+
+int main() {
+  kernel_arg_t *arg = (kernel_arg_t *)KERNEL_ARG_DEV_MEM_ADDR;
+
+  const uint32_t problem_size = (arg->dim_m * arg->dim_n) / (ELEM_PER_THREAD);
+  const uint32_t hw_threads_per_cluster =
+      CORES_PER_CLUSTER * vx_num_threads() * vx_num_warps();
+  // prevent launching more threads than the necessary problem size
+  // TODO: this does not take into account multiple clusters
+  const uint32_t grid_size = (problem_size > hw_threads_per_cluster)
+                                 ? hw_threads_per_cluster
+                                 : problem_size;
+
+#ifdef RADIANCE
+  vx_spawn_tasks_cluster(grid_size, (vx_spawn_tasks_cb)kernel_body, arg);
+#else
+  // NOTE: This kernel assumes contiguous thread scheduling for efficient shared
+  // memory allocation, and therefore does not work with original vx_spawn_tasks
+  vx_spawn_tasks_contiguous(grid_size, (vx_spawn_tasks_cb)kernel_body, arg);
+#endif
+  return 0;
+}

tests/regression/flash_attention/main.cpp

@@ -0,0 +1,308 @@
+#include <iostream>
+#include <fstream>
+#include <unistd.h>
+#include <string.h>
+#include <vortex.h>
+#include <vector>
+#include <cassert>
+#include "common.h"
+#include "half.hpp"
+
+using half_float::half;
+using half_float::half_cast;
+
+#define RT_CHECK(_expr)                                         \
+   do {                                                         \
+     int _ret = _expr;                                          \
+     if (0 == _ret)                                             \
+       break;                                                   \
+     printf("Error: '%s' returned %d!\n", #_expr, (int)_ret);   \
+	 cleanup();			                                              \
+     exit(-1);                                                  \
+   } while (false)
+
+///////////////////////////////////////////////////////////////////////////////
+
+const char* kernel_file = "kernel.bin";
+uint32_t count = 0;
+
+template <typename T> std::vector<T> src_a_data;
+template <typename T> std::vector<T> src_b_data;
+std::vector<float> ref_data;
+
+vx_device_h device = nullptr;
+std::vector<uint8_t> staging_buf;
+kernel_arg_t kernel_arg = {};
+
+static void show_usage() {
+   std::cout << "Vortex Test." << std::endl;
+   std::cout << "Usage: [-k: kernel] [-n words] [-h: help]" << std::endl;
+}
+
+static void parse_args(int argc, char **argv) {
+  int c;
+  while ((c = getopt(argc, argv, "n:k:h?")) != -1) {
+    switch (c) {
+    case 'n':
+      count = atoi(optarg);
+      break;
+    case 'k':
+      kernel_file = optarg;
+      break;
+    case 'h':
+    case '?': {
+      show_usage();
+      exit(0);
+    } break;
+    default:
+      show_usage();
+      exit(-1);
+    }
+  }
+}
+
+void cleanup() {
+  if (device) {
+    // vx_mem_free(device, kernel_arg.addr_a);
+    // vx_mem_free(device, kernel_arg.addr_b);
+    // vx_mem_free(device, kernel_arg.addr_c);
+    vx_dev_close(device);
+  }
+}
+
+template <typename T>
+void generate_source_matrix(uint32_t dim_m, uint32_t dim_n, uint32_t dim_k) {
+  static_assert(std::is_same_v<half, T> || std::is_same_v<float, T>,
+                "unsupported floating point datatype");
+
+  src_a_data<T>.resize(dim_m * dim_k);
+  src_b_data<T>.resize(dim_k * dim_n);
+
+  for (uint32_t i = 0; i < src_a_data<T>.size(); ++i) {
+    if constexpr (std::is_same_v<half, T>) {
+      src_a_data<T>[i] = half_cast<half>(static_cast<float>(i));
+    } else if (std::is_same_v<float, T>) {
+      src_a_data<T>[i] = static_cast<float>(i);
+    }
+    std::cout << "A: " << i << ": value=" << src_a_data<T>[i] << std::endl;
+  }
+  for (uint32_t i = 0; i < src_b_data<T>.size(); ++i) {
+    if constexpr (std::is_same_v<half, T>) {
+      src_b_data<T>[i] = half_cast<half>(static_cast<float>(i));
+    } else if (std::is_same_v<float, T>) {
+      src_b_data<T>[i] = static_cast<float>(i);
+    }
+    std::cout << "B: " << i << ": value=" << src_b_data<T>[i] << std::endl;
+  }
+}
+
+template <typename T>
+void generate_reference_matmul(uint32_t dim_m, uint32_t dim_n, uint32_t dim_k) {
+  static_assert(std::is_same_v<half, T> || std::is_same_v<float, T>,
+                "unsupported floating point datatype");
+
+  ref_data.resize(dim_m * dim_n);
+
+  for (uint32_t i = 0; i < dim_m; ++i) {
+    for (uint32_t j = 0; j < dim_n; ++j) {
+      float ref = 0.0f;
+      for (uint32_t k = 0; k < dim_k; ++k) {
+        ref += static_cast<float>(src_a_data<T>[dim_k * i + k]) *
+               static_cast<float>(src_b_data<T>[dim_n * k + j]);
+      }
+      ref_data.at(dim_n * i + j) = ref;
+    }
+  }
+}
+
+int run_test(const kernel_arg_t& kernel_arg,
+             uint32_t buf_size,
+             uint32_t dim_m, uint32_t dim_n) {
+  // start device
+  std::cout << "start device" << std::endl;
+  RT_CHECK(vx_start(device));
+
+  // wait for completion
+  std::cout << "wait for completion" << std::endl;
+  RT_CHECK(vx_ready_wait(device, VX_MAX_TIMEOUT));
+
+  // download destination buffer
+  std::cout << "download destination buffer" << std::endl;
+  RT_CHECK(vx_copy_from_dev(device, staging_buf.data(), kernel_arg.addr_c, buf_size));
+
+  // verify result
+  std::cout << "verify result" << std::endl;
+  {
+    int errors = 0;
+    auto buf_ptr = (float*)staging_buf.data();
+    for (uint32_t i = 0; i < dim_m * dim_n; ++i) {
+      float ref = ref_data.at(i);
+      float cur = buf_ptr[i];
+      if (std::abs((cur - ref) / ref) > 1e-6) {
+        std::cout << "error at result #" << std::dec << i
+                  << std::hex << ": actual=" << cur << ", expected=" << ref << std::endl;
+        ++errors;
+      }
+    }
+    if (errors != 0) {
+      std::cout << "Found " << std::dec << errors << " errors!" << std::endl;
+      std::cout << "FAILED!" << std::endl;
+      return 1;
+    }
+  }
+
+  return 0;
+}
+
+int main(int argc, char *argv[]) {
+  // parse command arguments
+  parse_args(argc, argv);
+
+  if (count == 0) {
+    count = 1;
+  }
+
+  std::srand(50);
+
+  // open device connection
+  std::cout << "open device connection" << std::endl;
+  RT_CHECK(vx_dev_open(&device));
+
+  // FIXME: hardcoded
+  uint32_t dim_m = 128;
+  uint32_t dim_n = 128;
+  uint32_t dim_k = 128;
+
+  using float_type = half;
+
+  generate_source_matrix<float_type>(dim_m, dim_n, dim_k);
+  generate_reference_matmul<float_type>(dim_m, dim_n, dim_k);
+
+  std::cout << "write reference output" << std::endl;
+  std::ofstream ref_file("reference.c.bin", std::ios::binary | std::ios::out);
+  if (!ref_file) {
+    std::cerr << "error: failed to open reference.c.bin for writing\n";
+    exit(EXIT_FAILURE);
+  }
+  ref_file.write(reinterpret_cast<char *>(ref_data.data()), ref_data.size() * sizeof(ref_data[0]));
+  ref_file.close();
+
+  uint32_t src_a_buf_size = src_a_data<float_type>.size() * sizeof(src_a_data<float_type>[0]);
+  uint32_t src_b_buf_size = src_b_data<float_type>.size() * sizeof(src_b_data<float_type>[0]);
+  uint32_t dst_buf_size = ref_data.size() * sizeof(src_a_data<float_type>[0]);
+
+  std::cout << "buffer size: " << dst_buf_size << " bytes" << std::endl;
+
+  // upload program
+  std::cout << "upload program" << std::endl;
+  RT_CHECK(vx_upload_kernel_file(device, kernel_file));
+
+  // allocate device memory
+  std::cout << "allocate device memory" << std::endl;
+  // RT_CHECK(vx_mem_alloc(device, src_a_buf_size, VX_MEM_TYPE_GLOBAL, &kernel_arg.addr_a));
+  // RT_CHECK(vx_mem_alloc(device, src_b_buf_size, VX_MEM_TYPE_GLOBAL, &kernel_arg.addr_b));
+  // RT_CHECK(vx_mem_alloc(device, dst_buf_size, VX_MEM_TYPE_GLOBAL, &kernel_arg.addr_c));
+  kernel_arg.addr_a = 0xa0000000;
+  kernel_arg.addr_b = 0xa1000000;
+  kernel_arg.addr_c = 0xc0000000;
+
+  kernel_arg.dim_m = dim_m;
+  kernel_arg.dim_n = dim_n;
+  kernel_arg.dim_k = dim_k;
+
+  std::cout << "dev_addr_a=0x" << std::hex << kernel_arg.addr_a << std::endl;
+  std::cout << "dev_addr_b=0x" << std::hex << kernel_arg.addr_b << std::endl;
+  std::cout << "dev_addr_c=0x" << std::hex << kernel_arg.addr_c << std::endl;
+
+  // allocate staging buffer
+  {
+    std::cout << "allocate staging buffer" << std::endl;
+    uint32_t staging_buf_size = std::max<uint32_t>(
+        src_a_buf_size,
+        std::max<uint32_t>(
+            src_b_buf_size,
+            std::max<uint32_t>(dst_buf_size, sizeof(kernel_arg_t))));
+    staging_buf.resize(staging_buf_size);
+  }
+
+  // upload kernel argument
+  {
+    std::cout << "upload kernel argument" << std::endl;
+    auto buf_ptr = staging_buf.data();
+    memcpy(buf_ptr, &kernel_arg, sizeof(kernel_arg_t));
+    RT_CHECK(vx_copy_to_dev(device, KERNEL_ARG_DEV_MEM_ADDR, staging_buf.data(), sizeof(kernel_arg_t)));
+
+    std::cout << "uploading argument buffer to device, device mem address="
+              << std::hex << KERNEL_ARG_DEV_MEM_ADDR << ", size=" << std::dec
+              << sizeof(kernel_arg_t) << " bytes\n";
+    std::ofstream file("args.bin", std::ios::binary | std::ios::out);
+    if (!file) {
+        std::cerr << "error: failed to open args.bin for writing\n";
+        exit(EXIT_FAILURE);
+    }
+    file.write(reinterpret_cast<char *>(staging_buf.data()),
+               sizeof(kernel_arg_t));
+    file.close();
+  }
+
+  // upload source buffer
+  {
+    {
+        auto buf_ptr = staging_buf.data();
+        memcpy(buf_ptr, src_a_data<float_type>.data(),
+               src_a_data<float_type>.size() * sizeof(float_type));
+        RT_CHECK(vx_copy_to_dev(device, kernel_arg.addr_a, staging_buf.data(),
+                                src_a_buf_size));
+
+        std::cout << "uploading source A matrix to device, device mem address="
+                  << std::hex << kernel_arg.addr_a << ", size=" << std::dec
+                  << src_a_buf_size << " bytes\n";
+        std::ofstream file("input.a.bin", std::ios::binary | std::ios::out);
+        if (!file) {
+        std::cerr << "error: failed to open args.bin for writing\n";
+        exit(EXIT_FAILURE);
+        }
+        file.write(reinterpret_cast<char *>(buf_ptr), src_a_buf_size);
+        file.close();
+    }
+    {
+        auto buf_ptr = staging_buf.data();
+        memcpy(buf_ptr, src_b_data<float_type>.data(),
+               src_b_data<float_type>.size() * sizeof(float_type));
+        RT_CHECK(vx_copy_to_dev(device, kernel_arg.addr_b, staging_buf.data(),
+                                src_b_buf_size));
+
+        std::cout << "uploading source B matrix to device, device mem address="
+                  << std::hex << kernel_arg.addr_b << ", size=" << std::dec
+                  << src_b_buf_size << " bytes\n";
+        std::ofstream file("input.b.bin", std::ios::binary | std::ios::out);
+        if (!file) {
+        std::cerr << "error: failed to open args.bin for writing\n";
+        exit(EXIT_FAILURE);
+        }
+        file.write(reinterpret_cast<char *>(buf_ptr), src_b_buf_size);
+        file.close();
+    }
+  }
+
+  // clear destination buffer
+  {
+    std::cout << "clear destination buffer" << std::endl;
+    auto buf_ptr = (int32_t*)staging_buf.data();
+    for (uint32_t i = 0; i < ref_data.size(); ++i) {
+      buf_ptr[i] = 0xdeadbeef;
+    }
+    RT_CHECK(vx_copy_to_dev(device, kernel_arg.addr_c, staging_buf.data(), dst_buf_size));
+  }
+
+  // run tests
+  std::cout << "run tests" << std::endl;
+  RT_CHECK(run_test(kernel_arg, dst_buf_size, kernel_arg.dim_m, kernel_arg.dim_n));
+  std::cout << "PASSED!" << std::endl;
+
+  // cleanup
+  std::cout << "cleanup" << std::endl;
+  cleanup();
+
+  return 0;
+}