# Vulkan Device Capabilities Inspector
# Run: Rscript inst/examples/vulkan_caps.R

library(ggmlR)

cat("=== Vulkan Device Capabilities ===\n\n")

# --- Availability ---
if (!ggml_vulkan_available()) {
  cat("Vulkan: NOT COMPILED\n")
  cat("  Reinstall: R CMD INSTALL . (requires libvulkan-dev + glslc)\n")
  quit(status = 1)
}

cat("Vulkan: compiled OK\n")

n <- ggml_vulkan_device_count()
cat("Devices found:", n, "\n\n")

if (n == 0) {
  cat("No Vulkan devices. Check driver installation.\n")
  quit(status = 1)
}

# --- Per-device info ---
for (i in seq_len(n)) {
  idx <- i - 1L

  desc <- ggml_vulkan_device_description(idx)
  mem  <- ggml_vulkan_device_memory(idx)
  caps <- ggml_vulkan_device_caps(idx)

  cat(sprintf("Device [%d]: %s\n", idx, desc))
  cat(sprintf("  Memory : %.2f GB free / %.2f GB total\n",
              mem$free / 1e9, mem$total / 1e9))
  cat("\n  --- Capabilities ---\n")

  cat(sprintf("  arch               : %s\n", caps$arch))

  cat(sprintf("  fp16               : %s",  if (caps$fp16) "YES" else "NO"))
  cat("   (FP16 arithmetic, required for fast inference)\n")

  cat(sprintf("  bf16               : %s",  if (caps$bf16) "YES" else "NO"))
  cat("   (BF16 native arithmetic; Flux/SD3 weights natively BF16)\n")

  cat(sprintf("  integer_dot_product: %s",  if (caps$integer_dot_product) "YES" else "NO"))
  cat("   (4x8-bit packed dot product; accelerates Q4/Q8 GEMM)\n")

  cat(sprintf("  coopmat_support    : %s",  if (caps$coopmat_support) "YES" else "NO"))
  cat("   (VK_KHR_cooperative_matrix, enables fast GEMM kernels)\n")

  cat(sprintf("  coopmat1_fa_support: %s",  if (caps$coopmat1_fa_support) "YES" else "NO"))
  cat("   (coopmat v1 flash-attention path, RDNA4/Ampere+)\n")

  cat(sprintf("  subgroup_size      : %d", caps$subgroup_size))
  cat("   (warp size used by ggml)\n")

  cat(sprintf("  subgroup_min_size  : %d", caps$subgroup_min_size))
  cat("   (min from VK_EXT_subgroup_size_control)\n")

  cat(sprintf("  subgroup_max_size  : %d", caps$subgroup_max_size))
  cat("   (max from VK_EXT_subgroup_size_control)\n")

  cat(sprintf("  subgroup_no_shmem  : %s",  if (caps$subgroup_no_shmem) "YES" else "NO"))
  cat("   (subgroup ops without shared memory, affects shader selection)\n")

  cat(sprintf("  wavefronts_per_simd: %d", caps$wavefronts_per_simd))
  cat("   (AMD only: 16=RDNA4, 20=RDNA1, 8=RDNA2/3; 0=non-AMD)\n")

  if (caps$coopmat_support && caps$coopmat_m > 0) {
    cat(sprintf("  coopmat tile       : M=%d N=%d K=%d", caps$coopmat_m, caps$coopmat_n, caps$coopmat_k))
    cat("   (KHR coopmat subgroup tile sizes)\n")
  }

  # --- BF16 coopmat path (Flux/SD3 native weights) ---
  # caps$bf16 reports the VK_KHR_shader_bfloat16 extension only. The actual
  # BF16 coopmat matmul pipeline is additionally gated at BUILD time by
  # -DGGML_VULKAN_BFLOAT16_GLSLC_SUPPORT, set in configure ONLY if the glslc
  # used at install time can compile GL_EXT_bfloat16. An old system glslc
  # silently disables the whole BF16 coopmat path even though the hardware
  # and the C++ pipeline both support it. This block reproduces that exact
  # configure feature test so the gate is visible without a rebuild.
  if (i == 1L) {  # build gate is global; report once
    cat("\n  --- BF16 coopmat build gate (Flux/SD3) ---\n")

    bf16_test <- system.file("ggml-vulkan/vulkan-shaders/feature-tests/bfloat16.comp",
                             package = "ggmlR")
    if (!nzchar(bf16_test)) {
      # installed package strips src/; fall back to repo-relative path
      cand <- "src/ggml-vulkan/vulkan-shaders/feature-tests/bfloat16.comp"
      if (file.exists(cand)) bf16_test <- cand
    }

    glslc <- Sys.which("glslc")
    if (!nzchar(glslc)) {
      cat("  glslc                : NOT FOUND in PATH\n")
      cat("    Cannot verify BF16 build gate. Source the Vulkan SDK\n")
      cat("    setup-env.sh before installing ggmlR.\n")
    } else {
      cat(sprintf("  glslc                : %s\n", glslc))
      ver <- tryCatch(
        system2(glslc, "--version", stdout = TRUE, stderr = TRUE)[1],
        error = function(e) NA_character_)
      if (!is.na(ver)) cat(sprintf("    (%s)\n", trimws(ver)))

      if (nzchar(bf16_test) && file.exists(bf16_test)) {
        rc <- suppressWarnings(system2(
          glslc, c("--target-env=vulkan1.2", "-o", "/dev/null",
                   shQuote(bf16_test)),
          stdout = FALSE, stderr = FALSE))
        gate_ok <- identical(rc, 0L)
        cat(sprintf("  GL_EXT_bfloat16      : %s\n",
                    if (gate_ok) "COMPILES (gate would PASS)"
                    else         "REJECTED (gate FAILS)"))
        if (gate_ok && caps$bf16) {
          cat("  => BF16 coopmat ACTIVE: Flux/SD3 BF16 weights use\n")
          cat("     native BF16xBF16 coopmat (no BF16->F16 fallback)\n")
        } else if (gate_ok && !caps$bf16) {
          cat("  => glslc OK but device lacks VK_KHR_shader_bfloat16\n")
          cat("     (BF16 coopmat unavailable on this GPU)\n")
        } else {
          cat("  => BF16 coopmat DISABLED at build time.\n")
          cat("     Hardware/pipeline support it, but this glslc is too\n")
          cat("     old. Rebuild with Vulkan SDK glslc first in PATH:\n")
          cat("       source <vulkansdk>/setup-env.sh\n")
          cat("       which glslc   # must NOT be /usr/bin/glslc\n")
          cat("       ./cleanup && R CMD INSTALL . --configure-args=\"--with-simd\"\n")
        }
      } else {
        cat("  bfloat16.comp        : test shader not found\n")
        cat("    (installed package; run this from the ggmlR repo root\n")
        cat("     to verify the build gate)\n")
      }
    }
  }

  # --- Summary verdict ---
  cat("\n  --- Verdict ---\n")

  if (caps$fp16 && caps$coopmat1_fa_support) {
    cat("  BEST: coopmat flash-attention path active (fastest)\n")
  } else if (caps$fp16 && caps$coopmat_support) {
    cat("  GOOD: coopmat GEMM path active, no flash-attention\n")
  } else if (caps$fp16) {
    cat("  OK:   FP16 active, no coopmat (scalar/subgroup shaders)\n")
  } else {
    cat("  WARN: FP32 only — slow, check driver/device support\n")
  }

  # --- Subgroup-shuffle mmq pipeline ---
  cat("\n  --- Subgroup-shuffle mmq (Q4_K/Q5_K/Q6_K) ---\n")
  shuffle_ok <- caps$integer_dot_product && caps$subgroup_size >= 64L
  if (shuffle_ok) {
    cat(sprintf("  subgroup_no_shmem mmq : ACTIVE (subgroup_size=%d, Q4_K/Q5_K/Q6_K)\n",
                caps$subgroup_size))
  } else if (caps$integer_dot_product) {
    cat(sprintf("  USE_SUBGROUP_NO_SHMEM: INACTIVE (subgroup_size=%d < 64)\n",
                caps$subgroup_size))
    cat("    Standard mmq path (shmem staging) used instead\n")
  } else {
    cat("  USE_SUBGROUP_NO_SHMEM: INACTIVE (integer_dot_product not supported)\n")
    cat("    mmq path unavailable, using dequant+F32 matmul\n")
  }

  # --- Push constants ---
  cat("\n  --- Push constants ---\n")
  cat(sprintf("  maxPushConstantsSize   : %d bytes", caps$max_push_constants_size))
  cat("   (Vulkan spec minimum: 128; ggmlR requires 256 for 5D ops)\n")
  cat(sprintf("  supports_256_push_const: %s",
              if (caps$supports_256_push_constants) "YES" else "NO"))
  if (caps$supports_256_push_constants) {
    cat("   (5D tensor ops fully supported)\n")
  } else {
    cat("   (WARNING: 5D ops will error at runtime — update driver)\n")
  }

  cat("\n")
}
