/*
 * Vortex Exact Attention isolated capability probes.
 *
 * The MMA operand signatures are adapted from SageAttention 2.2.0 csrc/mma.cuh
 * (commit d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5), Apache-2.0, which records
 * prior adaptation from FlashInfer 0.1.5. This file does not implement attention.
 */

#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>

#include <cuda.h>
#include <cuda/barrier>
#include <cuda_runtime.h>

#include <array>
#include <cstdint>
#include <vector>

#define CUDA_CHECK(call) do { \
  cudaError_t error_ = (call); \
  TORCH_CHECK(error_ == cudaSuccess, cudaGetErrorString(error_)); \
} while (0)

namespace {

constexpr int kThreads = 320;
constexpr int kProducerThreads = 192;
constexpr int kSlotBytes = 8192;
constexpr int kKvSlotBytes = 8192;
constexpr int kDynamicSharedBytes = 50 * 1024;

__device__ __forceinline__ void ready_arrive(int slot) {
  if (slot == 0) asm volatile("bar.arrive 0, 320;" ::: "memory");
  else asm volatile("bar.arrive 1, 320;" ::: "memory");
}

__device__ __forceinline__ void ready_sync(int slot) {
  if (slot == 0) asm volatile("bar.sync 0, 320;" ::: "memory");
  else asm volatile("bar.sync 1, 320;" ::: "memory");
}

__device__ __forceinline__ void free_arrive(int slot) {
  if (slot == 0) asm volatile("bar.arrive 2, 320;" ::: "memory");
  else asm volatile("bar.arrive 3, 320;" ::: "memory");
}

__device__ __forceinline__ void free_sync(int slot) {
  if (slot == 0) asm volatile("bar.sync 2, 320;" ::: "memory");
  else asm volatile("bar.sync 3, 320;" ::: "memory");
}

template <bool UseMbarrier>
__global__ __launch_bounds__(kThreads, 1)
void vea_b_handoff_probe(int64_t* result, int epochs, bool payload) {
  extern __shared__ uint8_t shared[];
  uint8_t* scores = shared;
  uint8_t* k_slot = shared + 2 * kSlotBytes;
  uint8_t* v_slot = k_slot + 2 * kKvSlotBytes;
  using BlockBarrier = cuda::barrier<cuda::thread_scope_block>;
  BlockBarrier* barriers = reinterpret_cast<BlockBarrier*>(shared + 48 * 1024);
  const int tid = threadIdx.x;
  const bool producer = tid < kProducerThreads;
  int64_t checksum = 0;
  int completed_epoch = -1;

  if constexpr (UseMbarrier) {
    if (tid < 4) init(barriers + tid, kThreads);
    __syncthreads();
  }

  if (producer) {
    for (int epoch = 0; epoch < epochs; ++epoch) {
      const int slot = epoch & 1;
      if (epoch >= 2) {
        if constexpr (UseMbarrier) barriers[2 + slot].arrive_and_wait();
        else free_sync(slot);
      }
      if (payload) {
        if (tid < 64) {
          for (int index = tid; index < kKvSlotBytes; index += 64) {
            k_slot[slot * kKvSlotBytes + index] = static_cast<uint8_t>((epoch + index * 3) & 255);
            v_slot[slot * kKvSlotBytes + index] = static_cast<uint8_t>((epoch + index * 5) & 255);
          }
        } else {
          const int row = tid - 64;
          uint8_t* row_slot = scores + slot * kSlotBytes + row * 64;
          #pragma unroll
          for (int column = 0; column < 64; ++column) {
            row_slot[column] = static_cast<uint8_t>((epoch + row * 17 + column) & 255);
          }
        }
      }
      __threadfence_block();
      if constexpr (UseMbarrier) (void)barriers[slot].arrive();
      else ready_arrive(slot);
      completed_epoch = epoch;
    }
    if constexpr (UseMbarrier) barriers[2 + ((epochs - 1) & 1)].arrive_and_wait();
    else free_sync((epochs - 1) & 1);
    if (epochs > 1) {
      if constexpr (UseMbarrier) barriers[2 + ((epochs - 2) & 1)].arrive_and_wait();
      else free_sync((epochs - 2) & 1);
    }
  } else {
    const int consumer = tid - kProducerThreads;
    for (int epoch = 0; epoch < epochs; ++epoch) {
      const int slot = epoch & 1;
      if constexpr (UseMbarrier) barriers[slot].arrive_and_wait();
      else ready_sync(slot);
      if (payload) {
        const uint8_t* row_slot = scores + slot * kSlotBytes + consumer * 64;
        #pragma unroll
        for (int column = 0; column < 64; ++column) {
          const uint8_t expected = static_cast<uint8_t>((epoch + consumer * 17 + column) & 255);
          const uint8_t actual = row_slot[column];
          checksum += actual;
          if (actual != expected) atomicAdd(reinterpret_cast<unsigned long long*>(result + blockIdx.x * 4 + 2), 1ULL);
        }
        const int kv_index = consumer;
        const uint8_t expected_k = static_cast<uint8_t>((epoch + kv_index * 3) & 255);
        const uint8_t expected_v = static_cast<uint8_t>((epoch + kv_index * 5) & 255);
        const uint8_t actual_k = k_slot[slot * kKvSlotBytes + kv_index];
        const uint8_t actual_v = v_slot[slot * kKvSlotBytes + kv_index];
        if (actual_k != expected_k || actual_v != expected_v) {
          atomicAdd(reinterpret_cast<unsigned long long*>(result + blockIdx.x * 4 + 2), 1ULL);
        }
        checksum += actual_k + actual_v;
      }
      completed_epoch = epoch;
      __threadfence_block();
      if constexpr (UseMbarrier) (void)barriers[2 + slot].arrive();
      else free_arrive(slot);
    }
  }
  __syncthreads();
  if (tid == 0) result[blockIdx.x * 4] = completed_epoch + 1;
  if (tid == kProducerThreads) result[blockIdx.x * 4 + 1] = checksum;
  if (!producer && completed_epoch == epochs - 1) {
    atomicAdd(reinterpret_cast<unsigned long long*>(result + blockIdx.x * 4 + 3), 1ULL);
  }
}

__device__ __forceinline__ void int8_mma(int32_t* c, const uint32_t* a, const uint32_t* b) {
  asm volatile(
      "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
      "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13};"
      : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
        "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
}

__device__ __forceinline__ void fp8_mma(uint32_t* c, const uint32_t* a, const uint32_t* b) {
  asm volatile(
      "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
      "{%0,%1},{%2,%3,%4,%5},{%6,%7},{%8,%9};"
      : "=r"(c[0]), "=r"(c[1])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
        "r"(c[0]), "r"(c[1]));
}

__global__ __launch_bounds__(kThreads, 1)
void vea_b_tensor_issue_probe(int64_t* values, int64_t* clocks, int iterations, int mode) {
  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const bool run_int = warp >= 2 && warp <= 5 && mode != 2;
  const bool run_fp = warp >= 6 && warp <= 9 && mode != 1;
  uint32_t a[4] = {0x01010101u + static_cast<uint32_t>(lane), 0x01010101u, 0x01010101u, 0x01010101u};
  uint32_t b[2] = {0x01010101u, 0x01010101u + static_cast<uint32_t>(lane)};
  int32_t ci[4] = {0, 0, 0, 0};
  uint32_t cf[2] = {0, 0};
  const uint64_t started = clock64();
  if (run_int) {
    for (int index = 0; index < iterations; ++index) int8_mma(ci, a, b);
  } else if (run_fp) {
    a[0] = a[1] = a[2] = a[3] = 0x38383838u;
    b[0] = b[1] = 0x38383838u;
    for (int index = 0; index < iterations; ++index) fp8_mma(cf, a, b);
  }
  const uint64_t finished = clock64();
  const int offset = (blockIdx.x * 10 + warp);
  if (lane == 0) {
    clocks[offset * 2] = static_cast<int64_t>(started);
    clocks[offset * 2 + 1] = static_cast<int64_t>(finished);
    values[offset] = run_int ? static_cast<int64_t>(ci[0]) : (run_fp ? static_cast<int64_t>(cf[0]) : 0);
  }
}

template <int Role>
__global__ __launch_bounds__(kThreads, 1)
void vea_b_register_probe(float* output, int iterations) {
  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  float total = 0.0f;
  if constexpr (Role == 0 || Role == 2) {
    if (warp >= 2 && warp <= 5) {
      int32_t rs[32];
      float m[8], d[8];
      #pragma unroll
      for (int i = 0; i < 32; ++i) rs[i] = i + lane;
      #pragma unroll
      for (int i = 0; i < 8; ++i) { m[i] = i + lane * 0.01f; d[i] = 1.0f + i; }
      for (int step = 0; step < iterations; ++step) {
        #pragma unroll
        for (int i = 0; i < 32; ++i) rs[i] += (step + i) & 7;
        #pragma unroll
        for (int i = 0; i < 8; ++i) { m[i] = fmaxf(m[i], step * 0.001f); d[i] = fmaf(d[i], 0.9999f, m[i]); }
      }
      #pragma unroll
      for (int i = 0; i < 32; ++i) total += static_cast<float>(rs[i]);
      #pragma unroll
      for (int i = 0; i < 8; ++i) total += m[i] + d[i];
    }
  }
  if constexpr (Role == 1 || Role == 2) {
    if (warp >= 6 && warp <= 9) {
      float ro[128];
      #pragma unroll
      for (int i = 0; i < 128; ++i) ro[i] = i * 0.001f + lane;
      for (int step = 0; step < iterations; ++step) {
        #pragma unroll
        for (int i = 0; i < 128; ++i) ro[i] = fmaf(ro[i], 0.99999f, (step + i) * 0.00001f);
      }
      #pragma unroll
      for (int i = 0; i < 128; ++i) total += ro[i];
    }
  }
  output[blockIdx.x * kThreads + threadIdx.x] = total;
}

template <typename Kernel>
std::array<int64_t, 8> attributes(Kernel kernel, int dynamic_shared_bytes) {
  cudaFuncAttributes value{};
  CUDA_CHECK(cudaFuncGetAttributes(&value, kernel));
  int active_blocks = 0;
  CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active_blocks, kernel, kThreads, dynamic_shared_bytes));
  return {value.numRegs, static_cast<int64_t>(value.sharedSizeBytes), static_cast<int64_t>(value.localSizeBytes),
          value.maxThreadsPerBlock, value.binaryVersion, value.ptxVersion,
          static_cast<int64_t>(value.maxDynamicSharedSizeBytes), active_blocks};
}

}  // namespace

torch::Tensor launch_handoff_probe(int64_t blocks, int64_t epochs, bool payload) {
  TORCH_CHECK(blocks > 0 && epochs > 1, "blocks and epochs must be positive");
  c10::cuda::CUDAGuard guard(0);
  auto result = torch::zeros({blocks, 4}, torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA));
  CUDA_CHECK(cudaFuncSetAttribute(vea_b_handoff_probe<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicSharedBytes));
  vea_b_handoff_probe<false><<<blocks, kThreads, kDynamicSharedBytes, at::cuda::getDefaultCUDAStream()>>>(
      result.data_ptr<int64_t>(), static_cast<int>(epochs), payload);
  CUDA_CHECK(cudaGetLastError());
  return result;
}

torch::Tensor launch_mbarrier_handoff_probe(int64_t blocks, int64_t epochs, bool payload) {
  TORCH_CHECK(blocks > 0 && epochs > 1, "blocks and epochs must be positive");
  c10::cuda::CUDAGuard guard(0);
  auto result = torch::zeros({blocks, 4}, torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA));
  CUDA_CHECK(cudaFuncSetAttribute(vea_b_handoff_probe<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicSharedBytes));
  vea_b_handoff_probe<true><<<blocks, kThreads, kDynamicSharedBytes, at::cuda::getDefaultCUDAStream()>>>(
      result.data_ptr<int64_t>(), static_cast<int>(epochs), payload);
  CUDA_CHECK(cudaGetLastError());
  return result;
}

std::vector<torch::Tensor> launch_tensor_issue_probe(int64_t blocks, int64_t iterations, int64_t mode) {
  TORCH_CHECK(blocks > 0 && iterations > 0 && mode >= 0 && mode <= 2, "invalid tensor issue arguments");
  c10::cuda::CUDAGuard guard(0);
  auto options = torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA);
  auto values = torch::zeros({blocks, 10}, options);
  auto clocks = torch::zeros({blocks, 10, 2}, options);
  vea_b_tensor_issue_probe<<<blocks, kThreads, 0, at::cuda::getDefaultCUDAStream()>>>(
      values.data_ptr<int64_t>(), clocks.data_ptr<int64_t>(), static_cast<int>(iterations), static_cast<int>(mode));
  CUDA_CHECK(cudaGetLastError());
  return {values, clocks};
}

torch::Tensor launch_register_probe(int64_t blocks, int64_t iterations, int64_t role) {
  TORCH_CHECK(blocks > 0 && iterations > 0 && role >= 0 && role <= 2, "invalid register probe arguments");
  c10::cuda::CUDAGuard guard(0);
  auto output = torch::zeros({blocks, kThreads}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA));
  if (role == 0) vea_b_register_probe<0><<<blocks, kThreads, 0, at::cuda::getDefaultCUDAStream()>>>(output.data_ptr<float>(), iterations);
  else if (role == 1) vea_b_register_probe<1><<<blocks, kThreads, 0, at::cuda::getDefaultCUDAStream()>>>(output.data_ptr<float>(), iterations);
  else vea_b_register_probe<2><<<blocks, kThreads, 0, at::cuda::getDefaultCUDAStream()>>>(output.data_ptr<float>(), iterations);
  CUDA_CHECK(cudaGetLastError());
  return output;
}

torch::Tensor probe_kernel_attributes(int64_t dynamic_shared_bytes) {
  CUDA_CHECK(cudaFuncSetAttribute(
      vea_b_handoff_probe<false>, cudaFuncAttributeMaxDynamicSharedMemorySize,
      static_cast<int>(dynamic_shared_bytes)));
  CUDA_CHECK(cudaFuncSetAttribute(
      vea_b_handoff_probe<true>, cudaFuncAttributeMaxDynamicSharedMemorySize,
      static_cast<int>(dynamic_shared_bytes)));
  std::array<std::array<int64_t, 8>, 6> rows = {
      attributes(vea_b_handoff_probe<false>, static_cast<int>(dynamic_shared_bytes)),
      attributes(vea_b_handoff_probe<true>, static_cast<int>(dynamic_shared_bytes)),
      attributes(vea_b_tensor_issue_probe, 0),
      attributes(vea_b_register_probe<0>, 0),
      attributes(vea_b_register_probe<1>, 0),
      attributes(vea_b_register_probe<2>, 0),
  };
  auto output = torch::empty({6, 8}, torch::TensorOptions().dtype(torch::kInt64));
  auto accessor = output.accessor<int64_t, 2>();
  for (int row = 0; row < 6; ++row) for (int column = 0; column < 8; ++column) accessor[row][column] = rows[row][column];
  return output;
}