#pragma once

#include <cstddef>
#include <cstdint>

namespace cism {
using DotKernel = float (*)(const float*, const float*, std::size_t);
using Int8DotKernel = float (*)(const std::int8_t*, const float*, std::size_t);
// IEEE-754 binary16 conversion (round-to-nearest-even), portable software
// implementation for packing and scalar fallback (no F16C requirement).
std::uint16_t fp32_to_fp16(float value);
float fp16_to_fp32(std::uint16_t bits);
using Fp16DotKernel = float (*)(const std::uint16_t*, const float*, std::size_t);
// int4/fp4 dots consume a permuted activation stream (even/odd 32-blocks,
// built by the caller) for full blocks plus the original row for the scalar
// tail; the 5th/4th scale args are per-32 FP32 / per-16 E4M3 bytes.
using Int4DotKernel = float (*)(const std::uint8_t*, const float*, const float*, const float*, std::size_t);
// FP4 (E2M1 elements, E4M3 scale per 16 elements): nibbles, permuted + original
// activations, and scale bytes.
using Fp4DotKernel = float (*)(const std::uint8_t*, const float*, const float*, const std::uint8_t*, std::size_t);
// Quantized-activation kernels: the caller quantizes the shared activation
// to int16 per 32-element block (32767/absmax, block-dequant scales passed
// alongside); weights decode to int16 through pshufb LUTs and multiply
// accumulate in integer pmaddwd lanes with independent accumulators.
using Int8Q8DotKernel = float (*)(const std::int8_t*, const std::int16_t*, const float*, std::size_t);
using Int4Q8DotKernel = float (*)(const std::uint8_t*, const std::int16_t*, const float*, const float*, std::size_t);
using Fp4Q8DotKernel = float (*)(const std::uint8_t*, const std::int16_t*, const std::uint8_t*, const float*, std::size_t);
DotKernel fp32_kernel();
Int8DotKernel int8_kernel();
Fp16DotKernel fp16_kernel();
Int4DotKernel int4_kernel();
Fp4DotKernel fp4_kernel();
// Quantized-activation selectors: null when AVX2 is unavailable (callers
// fall back to the fp32-activation kernels above).
Int8Q8DotKernel int8_q8_kernel();
Int4Q8DotKernel int4_q8_kernel();
Fp4Q8DotKernel fp4_q8_kernel();
const char* kernel_name();
// Portable-ISA dispatch (Intel/AMD/x86 + ARM phones).
// Highest SIMD tier this binary can actually execute on this CPU:
// "scalar" (any CPU) < "avx2" (x86-64, FMA) < "avx-vnni" (Zen4+/ADL+,
// integer MAC fast path) < "avx512-vnni" (server AVX512) / "neon" (ARM64
// phones, Apple Silicon, Raspberry Pi). kernel_name() above keeps reporting
// the fp32-kernel family ("scalar"/"avx2"/"neon"); kernel_variant() reports
// this overall tier for telemetry and tests.
const char* kernel_variant();
bool has_avx2_cpu();
bool has_f16c_cpu();
bool has_sse41_cpu();
bool has_avx_cpu();
bool has_avx_vnni_cpu();
bool has_avx512_vnni_cpu();
bool has_neon_cpu();
// int8-activation VNNI path: quantize_row_i8 shares quantize_row's blocking
// (int8 127/absmax per 32, one fp32 dequant scale per block) so the VNNI
// kernel sees the same blocking the int16 path uses. Selector is null unless
// the binary was built with VNNI support AND the CPU has it (otherwise the
// int16 Q8 path above is used, so pre-VNNI machines are untouched).
void quantize_row_i8(const float* input, std::size_t n, std::int8_t* values, float* scales);
void quantize_row_i8_avx2(const float* input, std::size_t n, std::int8_t* values, float* scales);
using Int8Q8VnniKernel = float (*)(const std::int8_t*, const std::int8_t*, const float*, std::size_t);
Int8Q8VnniKernel int8_q8_vnni_kernel();
float dot_scalar(const float* a, const float* b, std::size_t n);
float dot_fp16_scalar(const std::uint16_t* weights, const float* input, std::size_t n);
float dot_int8_scalar(const std::int8_t* weights, const float* input, std::size_t n);
float dot_int4_scalar(const std::uint8_t* weights, const float* input, const float* scales, std::size_t n);
float dot_fp4_scalar(const std::uint8_t* weights, const float* input, const std::uint8_t* scales, std::size_t n);
// 5-arg adapters over the 4-arg scalars (ignore the permuted buffer) so the
// permuted kernel type keeps a working non-AVX2 fallback.
float dot_int4_scalar5(const std::uint8_t* weights, const float* act_perm,
                       const float* input, const float* scales, std::size_t n);
float dot_fp4_scalar5(const std::uint8_t* weights, const float* act_perm,
                      const float* input, const std::uint8_t* scales, std::size_t n);
// E2M1 quantization support shared by the packer and the kernels.
// element_lut: raw nibble -> E2M1 value times two as int8 (exact small ints).
// scale_lut: E4M3 bit pattern -> float scale, already halved so a decoded
// element (half value) times the table entry equals E2M1 * E4M3 exactly.
// encode_scale: positive float -> E4M3 bit pattern, round-half-even like
// torch.float8_e4m3fn, saturating at 448.
const std::int8_t* fp4_element_lut();
// Kernel tuning notes (Zen3 AVX2+FMA, generic /arch:AVX2, no VNNI):
// - int8 is bandwidth-bound; 4-bit is kernel-bound (~26 uops/32 structural).
// - Inner blocking below (64-wide unroll) and T0 prefetches are the only
//   kernel-internal levers that keep scalar-fallback bitwise semantics.
// - Row-slice / thread-threshold / GEMM-blocking / SiLU-fusion live in
//   runtime.cpp (not owned here); this header documents the shared constants
//   so a future runtime wiring stays consistent. SIMD order may differ from
//   scalar; numeric changes stay PPL-gated (see VALIDATION.md noise note).
constexpr std::size_t kParallelThreshold = 1u << 16;  // rows*cols, runtime-owned
constexpr std::size_t kDotUnroll = 64;                // kernel inner blocking
constexpr std::size_t kPrefetchDistance = 512;        // bytes, T0, never faults
// Permuted-activation helper for int4/fp4 (even/odd 32-blocks). Scalar
// semantics; the AVX2 build offers an unrolled variant. Not wired to the
// runtime path here (session-persistent scratch already reuses the buffer);
// provided so future runtime wiring has one canonical implementation.
void permute_act32_blocked(const float* input, std::size_t n, float* out);
// Gated SiLU*up epilogue (stable sigmoid, matches Session::forward_tokens).
// Test-gated only: silu_fusion_enabled() is false unless CISM_FUSE_SILU=1,
// and the runtime path does not call these (prior A/B showed fusion ~7%
// slower on Zen3 streaming hardware). Scalar is the bitwise reference.
void silu_mul_scalar(const float* gate, const float* up, float* out, std::size_t n);
bool silu_fusion_enabled();
// Vector-exp activation blocks for the decode hot path (Surjo GDN/MLP/XSA).
// AVX2 TU evaluates exp with a ≤1-ULP polynomial (deterministic: same input
// bits always give same output bits); the scalar fallback uses libm and is
// the bitwise reference on non-AVX2 machines. In-place forms allow x==out.
// Unified stable sigmoid 1/(1+exp(-x)) (exact for x>=0, ≤1 ULP elsewhere).
// Numeric changes vs libm stay PPL-gated (VALIDATION.md noise note).
void act_exp(float* x, std::size_t n);
void act_silu(float* x, std::size_t n);
void act_silu_mul(const float* gate, const float* up, float* out, std::size_t n);
void act_sigmoid(float* x, std::size_t n);
void act_sigmoid_mul(float* o, const float* g, std::size_t n);
// Clamp-free SiLU*up for the dense path (no +-15 clamp there).
void act_silu_mul_plain(const float* gate, const float* up, float* out, std::size_t n);
// Exact GELU for the FWKV path (0.5*x*(1+erf(x/sqrt(2))), PyTorch exact).
void act_gelu(float* x, std::size_t n);
#ifdef CISM_HAVE_AVX2
void permute_act32_avx2(const float* input, std::size_t n, float* out);
void silu_mul_avx2(const float* gate, const float* up, float* out, std::size_t n);
#endif
const float* fp4_scale_lut();
float fp4_decode_scale(std::uint8_t bits);
std::uint8_t fp4_encode_scale(float value);
#ifdef CISM_HAVE_AVX2
float dot_avx2(const float* a, const float* b, std::size_t n);
float dot_fp16_avx2(const std::uint16_t* weights, const float* input, std::size_t n);
// Vector-exp act blocks (defined in kernels_avx2.cpp; scalar twins below).
void act_exp_avx2(float* x, std::size_t n);
void act_silu_avx2(float* x, std::size_t n);
void act_silu_mul_avx2(const float* gate, const float* up, float* out, std::size_t n);
void act_silu_mul_plain_avx2(const float* gate, const float* up, float* out, std::size_t n);
void act_gelu_avx2(float* x, std::size_t n);
void act_sigmoid_avx2(float* x, std::size_t n);
void act_sigmoid_mul_avx2(float* o, const float* g, std::size_t n);
float dot_int8_avx2(const std::int8_t* weights, const float* input, std::size_t n);
float dot_int4_avx2(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const float* scales, std::size_t n);
float dot_fp4_avx2(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const std::uint8_t* scales, std::size_t n);
float dot_int8_q8_avx2(const std::int8_t* weights, const std::int16_t* act, const float* act_scales, std::size_t n);
float dot_int4_q8_avx2(const std::uint8_t* weights, const std::int16_t* act, const float* scales,
                       const float* act_scales, std::size_t n);
float dot_fp4_q8_avx2(const std::uint8_t* weights, const std::int16_t* act,
                      const std::uint8_t* scales, const float* act_scales, std::size_t n);
#endif
// AVX-VNNI / AVX512-VNNI integer-MAC path (kernels_vnni.cpp; scalar-correct
// fallback when the compiler lacks VNNI — dispatch never selects it then).
float dot_int8_q8_vnni(const std::int8_t* weights, const std::int8_t* act,
                       const float* act_scales, std::size_t n);
// AVX512 zmm fp32 dot (kernels_vnni.cpp; needs AVX512F, gated on VNNI
// presence like the rest of that TU). Different chunking/reduce from AVX2:
// matches within fp32 rounding (PPL-gated), not bitwise-identical.
#ifdef CISM_HAVE_AVX512_VNNI
float dot_fp32_avx512(const float* a, const float* b, std::size_t n);
float dot_int8_avx512(const std::int8_t* weights, const float* input, std::size_t n);
#endif
#if defined(CISM_HAVE_NEON) && defined(__aarch64__)
float dot_neon(const float* a, const float* b, std::size_t n);
float dot_int8_neon(const std::int8_t* weights, const float* input, std::size_t n);
float dot_int4_neon(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const float* scales, std::size_t n);
float dot_fp4_neon(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const std::uint8_t* scales, std::size_t n);
void permute_act32_neon(const float* input, std::size_t n, float* out);
#endif
// SSE4.1 + AVX fallback TU (kernels_sse.cpp) for pre-AVX2 x86
// (Nehalem/Sandy/Bulldozer era). Same signatures as the AVX2 kernels;
// dispatch prefers AVX2 whenever has_avx2_cpu() is true.
#ifdef CISM_HAVE_SSE41
float dot_sse41(const float* a, const float* b, std::size_t n);
float dot_int8_sse41(const std::int8_t* weights, const float* input, std::size_t n);
float dot_int4_sse41(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const float* scales, std::size_t n);
float dot_fp4_sse41(const std::uint8_t* weights, const float* act_perm, const float* act_orig, const std::uint8_t* scales, std::size_t n);
#endif
#ifdef CISM_HAVE_AVX
float dot_avx_fp32(const float* a, const float* b, std::size_t n);
#endif
}  // namespace cism