/*
 * Iris Vulkan Acceleration
 *
 * GPU-accelerated matrix multiplication using Vulkan compute shaders.
 * Targets cross-platform GPUs (initially Linux). Provides a GEMM-offload
 * backend: the transformer/VAE/text-encoder run their CPU orchestration
 * path, but the heavy matrix multiplies are dispatched to the GPU.
 *
 * Only the Vulkan loader (libvulkan) is required at link time; no other
 * external libraries.
 */

#ifndef IRIS_VULKAN_H
#define IRIS_VULKAN_H

#include <stddef.h>
#include <stdint.h>

#ifdef __cplusplus
extern "C" {
#endif

/*
 * Initialize Vulkan acceleration: enumerate a compute-capable device,
 * create the queue, command pool, and GEMM pipelines.
 * Returns 1 on success, 0 if Vulkan is unavailable. Safe to call repeatedly.
 */
int iris_vulkan_init(void);

/*
 * Check whether Vulkan acceleration is available and initialized.
 */
int iris_vulkan_available(void);

/*
 * Human-readable name of the selected GPU (valid after init).
 */
const char *iris_vulkan_device_name(void);

/*
 * Release all Vulkan resources (buffers, caches, pipelines, device).
 */
void iris_vulkan_cleanup(void);

/*
 * Free all GPU-resident cached weight buffers (VRAM) while keeping the device
 * and pipelines alive. Cached weights re-upload lazily on next use. Used to
 * reclaim transformer weights before VAE decode; call only between batches.
 */
void iris_vulkan_release_weight_cache(void);

/*
 * GPU matrix multiply (generic, f32):
 *   C[M,N] = alpha * op(A)[M,K] @ op(B)[K,N] + beta * C[M,N]
 * transpose_a / transpose_b select op(). B is uploaded every call, so this
 * is safe for dynamic matrices (e.g. attention K/V temporaries).
 */
void iris_vulkan_sgemm(int transpose_a, int transpose_b, int M, int N, int K, float alpha, const float *A, int lda, const float *B, int ldb, float beta,
                       float *C, int ldc);

/*
 * Same as iris_vulkan_sgemm() but caches the B matrix on the GPU keyed by
 * its pointer. Use only when B is immutable across calls (model weights).
 */
void iris_vulkan_sgemm_cached(int transpose_a, int transpose_b, int M, int N, int K, float alpha, const float *A, int lda, const float *B, int ldb, float beta,
                              float *C, int ldc);

/*
 * GPU matrix multiply with bf16 weights (B is bfloat16, A and C are f32):
 *   C[M,N] = alpha * op(A)[M,K] @ op(B)[K,N] + beta * C[M,N]
 * B is cached on the GPU by pointer (weights are immutable).
 */
void iris_vulkan_sgemm_bf16(int transpose_a, int transpose_b, int M, int N, int K, float alpha, const float *A, int lda, const uint16_t *B_bf16, int ldb,
                            float beta, float *C, int ldc);

/*
 * GPU matrix multiply with GGML Q8_0 block-quantized weights (B is Q8_0,
 * A and C are f32), dequantized in-shader:
 *   C[M,N] = alpha * op(A)[M,K] @ op(B)[K,N] + beta * C[M,N]
 * B points at the raw 34-byte-per-32-element block stream and is cached on the
 * GPU by pointer (weights are immutable). ldb is in logical elements.
 */
void iris_vulkan_sgemm_q8(int transpose_a, int transpose_b, int M, int N, int K, float alpha, const float *A, int lda, const void *B_q8, int ldb, float beta,
                          float *C, int ldc);

/* ========================================================================
 * Resident VAE-decode tensor API.
 *
 * Mirrors the subset of the Metal iris_gpu_* surface used by vae_decode_gpu
 * in iris_vae.c, so that path compiles and runs GPU-resident under Vulkan.
 * Tensors are device-local (VRAM) on discrete GPUs; data stays on the GPU
 * between ops. Wrap a sequence in iris_gpu_batch_begin/end to fold many
 * dispatches into a single submit.
 * ======================================================================== */

typedef struct iris_gpu_tensor *iris_gpu_tensor_t;

iris_gpu_tensor_t iris_gpu_tensor_create(const float *data, size_t num_elements);
iris_gpu_tensor_t iris_gpu_tensor_alloc(size_t num_elements);
void              iris_gpu_tensor_free(iris_gpu_tensor_t tensor);
void              iris_gpu_tensor_read(iris_gpu_tensor_t tensor, float *out);
void              iris_gpu_tensor_write(iris_gpu_tensor_t tensor, const float *data);

void iris_gpu_copy_f32(iris_gpu_tensor_t dst, iris_gpu_tensor_t src, size_t n);

iris_gpu_tensor_t iris_gpu_conv2d_f32(iris_gpu_tensor_t x, const float *weight, const float *bias, int batch, int in_ch, int out_ch, int H, int W, int kH,
                                      int kW, int stride, int padding);

void iris_gpu_group_norm_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t x, const float *gamma, const float *beta, int batch, int channels, int spatial,
                             int num_groups, float eps);

void iris_gpu_swish_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t x, int n);
void iris_gpu_add_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t a, iris_gpu_tensor_t b, int n);

iris_gpu_tensor_t iris_gpu_upsample_nearest_2x_f32(iris_gpu_tensor_t x, int channels, int H, int W);

/* In-place-safe LeakyReLU: out[i] = x[i] >= 0 ? x[i] : slope*x[i]. */
void iris_gpu_leaky_relu_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t x, int n, float slope);

/* Scaled residual add: out[i] = scale*a[i] + b[i]. */
void iris_gpu_scale_add_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t a, iris_gpu_tensor_t b, float scale, int n);

void iris_gpu_batch_begin(void);
void iris_gpu_batch_end(void);

/* ========================================================================
 * Resident Qwen3 text-encoder ops.
 *
 * Fully GPU-resident forward path: activations stay in f32 device tensors
 * across all layers; weights are bf16 (cached by pointer) and applied
 * in-shader. Mirrors the subset of the Metal iris_gpu_*_bf16 surface that
 * qwen3_forward uses, but keeps activations in f32 for simplicity. Wrap a
 * whole forward pass in iris_gpu_batch_begin/end to fold every layer's
 * dispatches into a single submit (one GPU sync instead of ~150).
 *
 * Returns 1 if the resident ops are available (pipelines built), else 0.
 * ======================================================================== */
int iris_vk_qwen_available(void);

/* out[seq, out_dim] = x[seq, in_dim] @ weight[out_dim, in_dim]^T.
 * weight is bf16, cached on the GPU by pointer (immutable model weight). */
void iris_vk_qwen_linear(iris_gpu_tensor_t out, iris_gpu_tensor_t x, const uint16_t *weight_bf16, int seq, int in_dim, int out_dim);

/* Same as iris_vk_qwen_linear but weight is GGML Q8_0 block-quantized,
 * dequantized in-shader. weight_q8 points at the raw 34-byte-per-32-element
 * block stream, cached on the GPU by pointer (immutable model weight). */
void iris_vk_qwen_linear_q8(iris_gpu_tensor_t out, iris_gpu_tensor_t x, const void *weight_q8, int seq, int in_dim, int out_dim);

/* Per-row RMSNorm over hidden, scaled by a bf16 weight[hidden]. */
void iris_vk_qwen_rms_norm(iris_gpu_tensor_t out, iris_gpu_tensor_t x, const uint16_t *weight_bf16, int seq, int hidden, float eps);

/* Per-head RMSNorm (over head_dim) of t in place, bf16 weight[head_dim]. */
void iris_vk_qwen_head_rms_norm(iris_gpu_tensor_t t, const uint16_t *weight_bf16, int seq, int num_heads, int head_dim, float eps);

/* Apply RoPE to q and k in place using f32 cos/sin tables (cached by pointer). */
void iris_vk_qwen_rope(iris_gpu_tensor_t q, iris_gpu_tensor_t k, const float *cos_table, const float *sin_table, int seq, int num_q_heads, int num_kv_heads,
                       int head_dim);

/* GQA causal attention with padding mask (f32 tensor, 1.0 valid / 0.0 pad). */
void iris_vk_qwen_attention(iris_gpu_tensor_t out, iris_gpu_tensor_t q, iris_gpu_tensor_t k, iris_gpu_tensor_t v, iris_gpu_tensor_t mask, int seq,
                            int num_heads, int num_kv_heads, int head_dim, float scale);

/* SwiGLU: gate = silu(gate) * up, in place on gate. */
void iris_vk_qwen_silu_mul(iris_gpu_tensor_t gate, iris_gpu_tensor_t up, int n);

/* ========================================================================
 * Resident transformer (denoising) ops.
 *
 * Fully GPU-resident forward path: activations stay in f32 device tensors
 * across all blocks; immutable weights are bf16 (cached by pointer), and
 * dynamic per-block f32 vectors (modulation-fused norm weights, gates, AdaLN
 * shift/scale) are uploaded fresh per call. Mirrors the subset of the Metal
 * iris_gpu_* surface the resident transformer path uses, so it compiles and
 * runs under Vulkan.
 *
 * Returns 1 if the resident pipelines are available, else 0.
 * ======================================================================== */
int iris_vk_res_available(void);

/* Allocate an uninitialized bf16 (2 bytes/element) resident tensor. */
iris_gpu_tensor_t iris_gpu_tensor_alloc_f16(size_t num_elements);

/* Host pointer to a tensor's f32 contents (valid until the next op touches it
 * or it is freed). On discrete GPUs this downloads into a per-tensor shadow. */
float *iris_gpu_tensor_data(iris_gpu_tensor_t tensor);

/* No-op under Vulkan (tensors are freed explicitly; there is no pool). */
void iris_gpu_tensor_set_persistent(iris_gpu_tensor_t tensor, int persistent);

/* out[seq, out_dim] = x @ W^T + b, with f32 weight W[out_dim, in_dim] and
 * optional f32 bias[out_dim] (NULL for none). Returns a new tensor. */
iris_gpu_tensor_t iris_gpu_linear(iris_gpu_tensor_t x, const float *W, const float *b, int seq_len, int in_dim, int out_dim);

/* out[seq, out_dim] = x @ W^T with bf16 weight W[out_dim, in_dim] (cached). */
iris_gpu_tensor_t iris_gpu_linear_bf16(iris_gpu_tensor_t x, const uint16_t *W_bf16, int seq_len, int in_dim, int out_dim);
int               iris_gpu_linear_bf16_into(iris_gpu_tensor_t out, iris_gpu_tensor_t x, const uint16_t *W_bf16, int seq_len, int in_dim, int out_dim);

/* Same as iris_gpu_linear_bf16_into but W is a GGML Q8_0 block stream,
 * dequantized in-shader (cached on the GPU by pointer). */
int iris_gpu_linear_q8_into(iris_gpu_tensor_t out, iris_gpu_tensor_t x, const void *W_q8, int seq_len, int in_dim, int out_dim);

/* Per-row RMSNorm scaled by an f32 weight[hidden]. */
void iris_gpu_rms_norm_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t x, const float *weight, int seq, int hidden, float eps);

/* Per-head RMSNorm of q and k in place, f32 weights[head_dim]. */
void iris_gpu_qk_rms_norm(iris_gpu_tensor_t q, iris_gpu_tensor_t k, const float *q_weight, const float *k_weight, int seq, int heads, int head_dim, float eps);

/* Consecutive-pair RoPE applied to q and k in place using f32 cos/sin tables
 * [seq, head_dim] (cached by pointer). */
void iris_gpu_rope_single_pair_f32(iris_gpu_tensor_t q, iris_gpu_tensor_t k, const float *cos_freq, const float *sin_freq, int seq, int heads, int head_dim);

/* Split a fused projection into separate q/k/v and/or gate/up tensors.
 * hidden>0 splits the leading 3*hidden into q,k,v; mlp_hidden>0 splits the
 * trailing 2*mlp_hidden into gate,up. */
void iris_gpu_split_qkv_mlp(iris_gpu_tensor_t fused, iris_gpu_tensor_t q, iris_gpu_tensor_t k, iris_gpu_tensor_t v, iris_gpu_tensor_t gate,
                            iris_gpu_tensor_t up, int seq, int hidden, int mlp_hidden);

/* SwiGLU: gate = silu(gate) * up, in place on gate. */
void iris_gpu_silu_mul(iris_gpu_tensor_t gate, iris_gpu_tensor_t up, int n);

/* AdaLN: out = (1 + scale[i]) * layernorm(x)[s,i] + shift[i]. */
void iris_gpu_adaln_norm(iris_gpu_tensor_t out, iris_gpu_tensor_t x, const float *shift, const float *scale, int seq, int hidden, float eps);

/* Gated residual: out[s,i] += gate[i] * proj[s,i]. */
void iris_gpu_gated_add(iris_gpu_tensor_t out, const float *gate, iris_gpu_tensor_t proj, int seq, int hidden);

/* Full (bidirectional) f32 self-attention; returns 1 on success. */
int iris_gpu_attention_fused(iris_gpu_tensor_t out, iris_gpu_tensor_t Q, iris_gpu_tensor_t K, iris_gpu_tensor_t V, int seq_q, int seq_k, int num_heads,
                             int head_dim, float scale);

/* bf16 attention paths are not implemented under Vulkan; they return 0 so the
 * caller falls back to the f32 path above. */
int iris_gpu_attention_bf16(iris_gpu_tensor_t out, iris_gpu_tensor_t Q, iris_gpu_tensor_t K, iris_gpu_tensor_t V, int seq_q, int seq_k, int num_heads,
                            int head_dim, float scale);
int iris_gpu_attention_fused_bf16(iris_gpu_tensor_t out, iris_gpu_tensor_t Q, iris_gpu_tensor_t K, iris_gpu_tensor_t V, int seq_q, int seq_k, int num_heads,
                                  int head_dim, float scale);
int iris_gpu_convert_f32_to_bf16_into(iris_gpu_tensor_t bf16_out, iris_gpu_tensor_t f32_in);
int iris_gpu_convert_bf16_to_f32_into(iris_gpu_tensor_t f32_out, iris_gpu_tensor_t bf16_in);

/* GPU blit copy for f32 tensors with element offsets. */
void iris_gpu_copy_region_f32(iris_gpu_tensor_t dst, size_t dst_offset, iris_gpu_tensor_t src, size_t src_offset, size_t n);

/* ========================================================================
 * Resident Flux transformer (denoising) ops.
 *
 * The Flux double/single blocks reuse the resident op surface above
 * (adaln/qk-rmsnorm/rope/attention/linear/gated-add/silu) for f32 activations.
 * The one extra primitive they need is a strided row-block copy, used to split
 * the fused [Q,K,V,gate,up] single-block projection into separate streams and
 * to assemble the [attn|mlp] concat feeding the output projection.
 *
 * Returns 1 if the Flux resident pipeline is available (rowcopy + the resident
 * ops it depends on), else 0.
 * ======================================================================== */
int iris_vk_flux_available(void);

/* dst[s*dst_stride + dst_off + e] = src[s*src_stride + src_off + e],
 * for e in [0, w), s in [0, seq). Strides/offsets are in elements (f32). */
void iris_gpu_row_copy_f32(iris_gpu_tensor_t dst, int dst_stride, int dst_off, iris_gpu_tensor_t src, int src_stride, int src_off, int seq, int w);

#ifdef __cplusplus
}
#endif

#endif /* IRIS_VULKAN_H */