/**
 *  @brief SIMD-accelerated Batched Dot Products for Granite Rapids.
 *  @file include/numkong/dots/graniteamx.h
 *  @author Ash Vardanian
 *  @date April 9, 2026
 *
 *  @sa include/numkong/dots.h
 *
 *  Native FP16×FP16→FP32 GEMM kernels using Intel AMX-FP16 (TDPFP16PS) on Granite Rapids CPUs.
 *  Same tile geometry as BF16 (16 rows × 32 FP16 = 1KB per tile), same 2×2 output blocking,
 *  same packing format — only the tile multiply instruction differs.
 *
 *  Tile register allocation:
 *
 *  - TMM0, TMM1: A matrix tiles (row blocks i and i+16)
 *  - TMM2, TMM3: B matrix tiles (column blocks j and j+16)
 *  - TMM4-7: C accumulator tiles (2 × 2 output grid = 32×32 F32 results)
 *
 *  @section amx_fp16_instructions Intel AMX-FP16 Instructions (Granite Rapids+)
 *
 *  FP16 matrix multiply (AMX-FP16):
 *
 *      Intrinsic                   Instruction                     Operation
 *      _tile_dpfp16ps              TDPFP16PS (TMM, TMM, TMM)       C += A × B (fp16 → f32)
 *
 *  TDPFP16PS: 16 × 16 × 32 = 8192 FP16 MACs per instruction (same throughput as TDPBF16PS).
 *
 *  @section ozaki_limitations F32→F64 via Ozaki Scheme — Attempted and Abandoned
 *
 *  We explored using AMX-FP16 tiles to compute F32→F64 GEMMs via the Ozaki decomposition scheme,
 *  splitting each F32 scalar into 2 or 3 FP16 terms and performing cross-product TDPFP16PS operations.
 *
 *  Results on Intel Xeon 6776P (Granite Rapids), single-threaded:
 *
 *  | Variant              | Speed (gso/s) | Precision | Notes                                      |
 *  |----------------------|:-------------:|:---------:|:------------------------------------------:|
 *  | 2-term, 2×1 blocking |    ~150       |  ~22 bits | Split accumulators, N=16 F64 flush         |
 *  | 3-term, 1×1 blocking |    ~110       |  ~22 bits | 3 accumulators by magnitude band           |
 *  | Pipelined 2-term     |    ~156       |  ~22 bits | Double-buffered A split, AMX/AVX-512 overlap|
 *  | MKL SGEMM            |    ~170       |  ~20 bits | Pure F32, no decomposition                 |
 *  | Skylake F64 accum    |     ~50       |  ~48 bits | F32×F32 multiply, F64 accumulation         |
 *
 *  The fundamental bottleneck is TDPFP16PS's internal F32 accumulation: each instruction sums
 *  32 FP16×FP16 products into an F32 register (23-bit mantissa). Even with Ozaki cross-term
 *  separation into distinct TMM accumulators (preventing magnitude mixing) and periodic extraction
 *  to F64 running sums, the per-instruction accumulation of 32 products loses ~5 bits
 *  (log2(32) = 5), capping effective precision at ~28 - 5 = ~23 bits — barely exceeding F32 BLAS.
 *
 *  Approaches attempted:
 *
 *  - 2-term decomposition (a = a_high + a_low): 4 TDPFP16PS per depth tile, ~20-bit products.
 *    With split accumulators (main + correction) merged in F64: ~22-bit effective precision.
 *    Faster than MKL at small depths (≤512) but precision plateaus at ~22 bits.
 *
 *  - 3-term decomposition (a = a_high + a_mid + a_low): 6 TDPFP16PS per depth tile, ~30-bit
 *    products. No precision improvement over 2-term because the F32 TMM accumulation is the
 *    bottleneck, not the decomposition quality. Strictly slower and no more precise.
 *
 *  - Periodic F64 flush (extract TMM accumulators to F64 every N depth tiles): prevents precision
 *    degradation at large depths. With N=16, ~15% overhead. Effective precision still ~24 bits
 *    (limited by per-TDPFP16PS accumulation of 32 products, not by inter-tile accumulation).
 *
 *  - AMX/AVX-512 pipelining (double-buffered A splitting overlapped with AMX compute): ~7%
 *    speedup at large sizes. Does not affect precision.
 *
 *  Conclusion: AMX-FP16's F32 tile accumulation fundamentally limits Ozaki to ~22-24 bits —
 *  comparable to F32 BLAS, far short of the ~48-bit F64 precision needed to justify the complexity.
 *  For F32→F64 GEMM, pure AVX-512 with F64 FMA remains the correct approach.
 */
#ifndef NK_DOTS_GRANITEAMX_H
#define NK_DOTS_GRANITEAMX_H

#if NK_TARGET_X8664_
#if NK_TARGET_GRANITEAMX

#include "numkong/dots/serial.h"
#include "numkong/dots/sapphireamx.h"

#if defined(__cplusplus)
extern "C" {
#endif

#if defined(__clang__)
#pragma clang attribute push(                                                                                                     \
    __attribute__((target(                                                                                                        \
        "avx2,avx512f,avx512vl,avx512bw,avx512dq,avx512fp16,avx512vbmi,f16c,fma,bmi,bmi2,amx-tile,amx-bf16,amx-int8,amx-fp16"))), \
    apply_to = function)
#elif defined(__GNUC__)
#pragma GCC push_options
#pragma GCC target("avx2", "avx512f", "avx512vl", "avx512bw", "avx512dq", "avx512fp16", "avx512vbmi", "f16c", "fma", \
                   "bmi", "bmi2", "amx-tile", "amx-bf16", "amx-int8", "amx-fp16")
#endif

#pragma region Tile Types

typedef struct {
    NK_ALIGN64 nk_f16_t data[16][32]; // 16 rows × 32 columns = 1KB
} nk_dots_f16_a16x32_graniteamx_t;

typedef struct {
    NK_ALIGN64 nk_f16_t data[16][16][2]; // 16 depth-groups × 16 columns × 2 = 1KB (pair-interleaved)
} nk_dots_f16_b32x16_graniteamx_t;

typedef struct {
    NK_ALIGN64 nk_f32_t data[16][16]; // 16 × 16 = 1KB accumulator
} nk_dots_f16_state_graniteamx_t;

typedef struct {
    nk_dots_f16_state_graniteamx_t c[2][2]; // 4KB total (2×2 output blocking)
} nk_dots_f16_state2x2_graniteamx_t;

#pragma endregion Tile Types

#pragma region Helpers

/* Initialize FP16 output state to zero */
NK_INTERNAL void nk_dots_f16_init_graniteamx_(nk_dots_f16_state_graniteamx_t *state) {
    __m512 zero_f32x16 = _mm512_setzero_ps();
    for (nk_size_t row_idx = 0; row_idx < 16; row_idx++) { _mm512_store_ps(state->data[row_idx], zero_f32x16); }
}

/* Load A tile from FP16 row-major source with masking for edge tiles */
NK_INTERNAL void nk_dots_f16_load_a_graniteamx_(        //
    nk_dots_f16_a16x32_graniteamx_t *a_tile,            //
    nk_f16_t const *src, nk_size_t src_stride_elements, //
    nk_size_t valid_rows, nk_size_t valid_cols) {

    __mmask32 column_mask = (valid_cols >= 32) ? 0xFFFFFFFF : ((__mmask32)1 << valid_cols) - 1;
    __m512i zero_i16x32 = _mm512_setzero_si512();

    for (nk_size_t row_idx = 0; row_idx < 16; row_idx++) {
        if (row_idx < valid_rows) {
            __m512i row_i16x32 = _mm512_maskz_loadu_epi16(column_mask, src + row_idx * src_stride_elements);
            _mm512_store_si512((__m512i *)a_tile->data[row_idx], row_i16x32);
        }
        else { _mm512_store_si512((__m512i *)a_tile->data[row_idx], zero_i16x32); }
    }
    nk_compiler_barrier_sapphireamx_();
}

/* Store F32 state to output matrix with masking for edge tiles */
NK_INTERNAL void nk_dots_f16_store_graniteamx_(   //
    nk_dots_f16_state_graniteamx_t const *state,  //
    nk_f32_t *dst, nk_size_t dst_stride_elements, //
    nk_size_t valid_rows, nk_size_t valid_cols) {

    __mmask16 column_mask = (valid_cols >= 16) ? 0xFFFF : ((__mmask16)1 << valid_cols) - 1;

    for (nk_size_t row_idx = 0; row_idx < valid_rows; row_idx++) {
        __m512 row_f32x16 = _mm512_load_ps(state->data[row_idx]);
        _mm512_mask_storeu_ps(dst + row_idx * dst_stride_elements, column_mask, row_f32x16);
    }
}

NK_INTERNAL void nk_dots_f16_output2x2_graniteamx_( //
    nk_dots_f16_state2x2_graniteamx_t const *state, //
    nk_f32_t *dst, nk_size_t dst_stride_elements,   //
    nk_size_t valid_rows, nk_size_t valid_cols) {

    nk_size_t const rows_high = (valid_rows > 16) ? 16 : valid_rows;
    nk_size_t const cols_left = (valid_cols > 16) ? 16 : valid_cols;
    nk_size_t const cols_right = (valid_cols > 16) ? valid_cols - 16 : 0;

    if (rows_high > 0 && cols_left > 0)
        nk_dots_f16_store_graniteamx_(&state->c[0][0], dst, dst_stride_elements, rows_high, cols_left);
    if (rows_high > 0 && cols_right > 0)
        nk_dots_f16_store_graniteamx_(&state->c[0][1], dst + 16, dst_stride_elements, rows_high, cols_right);

    if (valid_rows > 16) {
        nk_size_t const rows_low = valid_rows - 16;
        nk_f32_t *dst_low = dst + 16 * dst_stride_elements;
        if (cols_left > 0)
            nk_dots_f16_store_graniteamx_(&state->c[1][0], dst_low, dst_stride_elements, rows_low, cols_left);
        if (cols_right > 0)
            nk_dots_f16_store_graniteamx_(&state->c[1][1], dst_low + 16, dst_stride_elements, rows_low, cols_right);
    }
}

NK_INTERNAL void nk_dots_f16_update_graniteamx_(     //
    nk_dots_f16_state_graniteamx_t *state,           //
    nk_dots_f16_a16x32_graniteamx_t const *a_tile_0, //
    nk_dots_f16_a16x32_graniteamx_t const *a_tile_1, //
    nk_dots_f16_a16x32_graniteamx_t const *a_tile_2, //
    nk_dots_f16_b32x16_graniteamx_t const *b_tile_0, //
    nk_dots_f16_b32x16_graniteamx_t const *b_tile_1, //
    nk_dots_f16_b32x16_graniteamx_t const *b_tile_2) {

    _tile_loadd(0, state->data, 64);
    _tile_loadd(1, a_tile_0->data, 64);
    _tile_loadd(2, a_tile_1->data, 64);
    _tile_loadd(3, a_tile_2->data, 64);
    _tile_loadd(4, b_tile_0->data, 64);
    _tile_loadd(5, b_tile_1->data, 64);
    _tile_loadd(6, b_tile_2->data, 64);

    _tile_dpfp16ps(0, 1, 4);
    _tile_dpfp16ps(0, 2, 5);
    _tile_dpfp16ps(0, 3, 6);

    _tile_stored(0, state->data, 64);
}

#pragma endregion Helpers

#pragma region F16 Native

NK_PUBLIC nk_size_t nk_dots_packed_size_f16_graniteamx(nk_size_t column_count, nk_size_t depth) {
    nk_size_t const tmm_rows = 16;
    nk_size_t const tmm_cols = 32;
    nk_size_t const tile_bytes = 512 * sizeof(nk_f16_t); // 16 × 32 × 2 = 1KB

    nk_size_t const full_column_tiles = column_count / tmm_rows;
    nk_size_t const tiles_along_depth = nk_size_divide_round_up_(depth, tmm_cols);
    nk_size_t const column_remainder_count = column_count - full_column_tiles * tmm_rows;

    // Header (64 bytes aligned)
    nk_size_t size = sizeof(nk_dots_amx_packed_header_t);

    // All tiles for full column rows (pair-interleaved, depth remainder zero-padded)
    size += full_column_tiles * tiles_along_depth * tile_bytes;

    // Column edge: remaining rows for ALL depth columns, stored row-major
    if (column_remainder_count > 0) size += column_remainder_count * depth * sizeof(nk_f16_t);

    // Per-column norms for angular/euclidean distance (4 bytes each: f32)
    size += column_count * sizeof(nk_f32_t);

    return size;
}

NK_PUBLIC void nk_dots_pack_f16_graniteamx(                     //
    nk_f16_t const *b, nk_size_t column_count, nk_size_t depth, //
    nk_size_t b_stride_in_bytes, void *b_packed) {

    // AMX FP16 tile dimensions: 16 rows × 32 columns (512 FP16 elements = 1KB)
    nk_size_t const tmm_rows = 16;
    nk_size_t const tmm_cols = 32;
    nk_size_t const tile_elements = 512;
    nk_size_t const tile_bytes = tile_elements * sizeof(nk_f16_t);
    nk_size_t const b_stride_elements = b_stride_in_bytes / sizeof(nk_f16_t);

    // Compute layout dimensions
    nk_size_t const column_tiles_count = column_count / tmm_rows;
    nk_size_t const depth_tiles_count = nk_size_divide_round_up_(depth, tmm_cols);
    nk_size_t const column_remainder_count = column_count - column_tiles_count * tmm_rows;
    nk_size_t const total_tiles = column_tiles_count * depth_tiles_count;

    // Write header with layout metadata
    nk_dots_amx_packed_header_t *header = (nk_dots_amx_packed_header_t *)b_packed;
    header->full_column_tiles = (nk_u32_t)column_tiles_count;
    header->full_depth_tiles = (nk_u32_t)depth_tiles_count;
    header->column_remainder_count = (nk_u32_t)column_remainder_count;

    // Compute memory region offsets
    nk_size_t const tiles_offset = sizeof(nk_dots_amx_packed_header_t);
    nk_size_t const column_edge_offset = tiles_offset + total_tiles * tile_bytes;
    header->column_edge_offset = (nk_u32_t)column_edge_offset;

    // Pointers to packed data regions
    nk_f16_t *tiles_ptr = (nk_f16_t *)((char *)b_packed + tiles_offset);
    nk_f16_t *column_edge_ptr = (nk_f16_t *)((char *)b_packed + column_edge_offset);

    // Pack tiles: gather 16 strided rows into aligned temporary, transpose via SIMD, copy to packed buffer.
    // FP16 has the same 16-bit pair-interleaved layout as BF16, so reuse the BF16 transpose.
    for (nk_size_t column_tile_idx = 0; column_tile_idx < column_tiles_count; column_tile_idx++) {
        for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {

            nk_size_t const tile_index = column_tile_idx * depth_tiles_count + depth_tile_idx;
            nk_f16_t *tile_output = tiles_ptr + tile_index * tile_elements;

            nk_size_t const src_row_start = column_tile_idx * tmm_rows;
            nk_size_t const src_column_start = depth_tile_idx * tmm_cols;
            nk_size_t const columns_to_pack = (src_column_start + tmm_cols <= depth) ? tmm_cols
                                                                                     : (depth - src_column_start);

            // Gather 16 strided source rows into a contiguous aligned tile
            nk_dots_bf16_a16x32_sapphireamx_t source_tile;
            if (columns_to_pack == tmm_cols) {
                for (nk_size_t row_idx = 0; row_idx < tmm_rows; row_idx++) {
                    nk_f16_t const *source_row = b + (src_row_start + row_idx) * b_stride_elements + src_column_start;
                    _mm512_store_si512(&source_tile.data[row_idx][0], _mm512_loadu_si512(source_row));
                }
            }
            else {
                __mmask32 depth_mask = (__mmask32)((columns_to_pack < 32) ? ((1U << columns_to_pack) - 1) : ~0U);
                for (nk_size_t row_idx = 0; row_idx < tmm_rows; row_idx++) {
                    nk_f16_t const *source_row = b + (src_row_start + row_idx) * b_stride_elements + src_column_start;
                    _mm512_store_si512(&source_tile.data[row_idx][0], _mm512_maskz_loadu_epi16(depth_mask, source_row));
                }
            }

            // Transpose into aligned local, then copy to (potentially unaligned) packed buffer.
            // BF16 and FP16 share identical 16-bit pair-interleaved layout for TDPBF16PS/TDPFP16PS.
            nk_dots_bf16_b32x16_sapphireamx_t transposed_tile;
            nk_dots_pack_bf16_transposed_sapphireamx_(&source_tile, &transposed_tile);
            for (nk_size_t i = 0; i < tile_bytes; i += 64)
                _mm512_storeu_si512((char *)tile_output + i, _mm512_load_si512((char const *)&transposed_tile + i));
        }
    }

    // Pack column-remainder rows using vectorized masked copies
    if (column_remainder_count > 0) {
        nk_size_t const remainder_start_row = column_tiles_count * tmm_rows;
        for (nk_size_t row_idx = 0; row_idx < column_remainder_count; row_idx++) {
            nk_f16_t const *src_row = b + (remainder_start_row + row_idx) * b_stride_elements;
            nk_f16_t *dst_row = column_edge_ptr + row_idx * depth;
            nk_size_t column_idx = 0;
            for (; column_idx + 32 <= depth; column_idx += 32) {
                _mm512_storeu_si512(dst_row + column_idx, _mm512_loadu_si512(src_row + column_idx));
            }
            if (column_idx < depth) {
                __mmask32 tail_mask = (__mmask32)((1U << (depth - column_idx)) - 1);
                _mm512_mask_storeu_epi16(dst_row + column_idx, tail_mask,
                                         _mm512_maskz_loadu_epi16(tail_mask, src_row + column_idx));
            }
        }
    }

    // Compute and store per-column norms for angular/euclidean distance
    nk_size_t norms_offset = column_edge_offset +
                             (column_remainder_count > 0 ? column_remainder_count * depth * sizeof(nk_f16_t) : 0);
    header->norms_byte_offset = (nk_u32_t)norms_offset;
    nk_f32_t *norms = (nk_f32_t *)((char *)b_packed + norms_offset);
    for (nk_size_t col = 0; col < column_count; col++)
        norms[col] = nk_dots_reduce_sumsq_f16_(b + col * b_stride_elements, depth);
}

NK_PUBLIC void nk_dots_packed_f16_graniteamx(             //
    nk_f16_t const *a, void const *b_packed, nk_f32_t *c, //
    nk_size_t rows_count, nk_size_t cols_count, nk_size_t depth, nk_size_t a_stride_bytes, nk_size_t c_stride_bytes) {
    nk_unused_(cols_count);

    // Parse packed B header
    nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
    nk_size_t const column_tiles_count = header->full_column_tiles;
    nk_size_t const depth_tiles_count = header->full_depth_tiles;
    nk_size_t const column_remainder_count = header->column_remainder_count;

    // Packed B data regions
    nk_f16_t const *b_tiles_base = (nk_f16_t const *)((char const *)b_packed + sizeof(nk_dots_amx_packed_header_t));
    nk_f16_t const *col_edge_ptr = (nk_f16_t const *)((char const *)b_packed + header->column_edge_offset);

    // Stride conversions
    nk_size_t const a_stride_elements = a_stride_bytes / sizeof(nk_f16_t);
    nk_size_t const c_stride_elements = c_stride_bytes / sizeof(nk_f32_t);

    // Tile dimensions
    nk_size_t const tile_depth = 32;
    nk_size_t const tile_size = 512;
    nk_size_t const full_cols = column_tiles_count * 16;

    // Block counts (32 × 32 output blocks = 2 × 2 tiles)
    nk_size_t const row_blocks_count = nk_size_divide_round_up_(rows_count, 32);
    nk_size_t const col_blocks_count = column_tiles_count / 2;

    if (depth_tiles_count == 0) return;

    // Tile buffers for A (only used for edge tiles)
    nk_dots_f16_a16x32_graniteamx_t a_tile_top, a_tile_bottom;
    nk_dots_f16_state2x2_graniteamx_t c_accum_buffer;

    // Precompute: number of full depth-tiles (no masking needed)
    nk_size_t const full_depth_tiles_count = depth / tile_depth;
    nk_size_t const depth_remainder = depth % tile_depth;

    nk_amx_tile_configure_sapphireamx_();

    // Loop order: row_blocks outer, col_blocks inner - maximizes A tile L2 cache reuse
    for (nk_size_t row_block_idx = 0; row_block_idx < row_blocks_count; row_block_idx++) {
        nk_size_t const row_block_start = row_block_idx * 32;
        nk_size_t const valid_rows_count = (row_block_start + 32 <= rows_count) ? 32 : (rows_count - row_block_start);
        nk_size_t const is_full_row_block = (valid_rows_count == 32);

        for (nk_size_t column_block_idx = 0; column_block_idx < col_blocks_count; column_block_idx++) {
            nk_size_t const col_block_start = column_block_idx * 32;
            nk_size_t const b_column_left_base = (column_block_idx * 2) * depth_tiles_count;
            nk_size_t const b_column_right_base = (column_block_idx * 2 + 1) * depth_tiles_count;

            // Zero accumulators (TMM4-7 stay resident across entire depth loop)
            _tile_zero(4);
            _tile_zero(5);
            _tile_zero(6);
            _tile_zero(7);

            // Fast path: full row-block with full depth-tiles → direct A load with 2-deep pipelining
            if (is_full_row_block && full_depth_tiles_count > 0) {
                nk_f16_t const *a_top_base = a + row_block_start * a_stride_elements;
                nk_f16_t const *a_bottom_base = a + (row_block_start + 16) * a_stride_elements;

                nk_dots_f16_b32x16_graniteamx_t const *b_tile_left =
                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base + b_column_left_base * tile_size);
                nk_dots_f16_b32x16_graniteamx_t const *b_tile_right =
                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base + b_column_right_base * tile_size);

                // Prologue: load first depth tile
                _tile_loadd(0, a_top_base, a_stride_bytes);
                _tile_loadd(1, a_bottom_base, a_stride_bytes);
                _tile_loadd(2, b_tile_left->data, 64);
                _tile_loadd(3, b_tile_right->data, 64);

                // Main loop: 2-deep software pipelining
                for (nk_size_t depth_tile_idx = 0; depth_tile_idx < full_depth_tiles_count - 1; depth_tile_idx++) {
                    nk_size_t const next_depth_offset = (depth_tile_idx + 1) * tile_depth;

                    _tile_dpfp16ps(4, 0, 2);
                    _tile_dpfp16ps(5, 0, 3);
                    _tile_dpfp16ps(6, 1, 2);
                    _tile_dpfp16ps(7, 1, 3);

                    _tile_loadd(0, a_top_base + next_depth_offset, a_stride_bytes);
                    _tile_loadd(1, a_bottom_base + next_depth_offset, a_stride_bytes);
                    b_tile_left = (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base +
                                                                            (b_column_left_base + depth_tile_idx + 1) *
                                                                                tile_size);
                    b_tile_right = (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base + (b_column_right_base +
                                                                                             depth_tile_idx + 1) *
                                                                                                tile_size);
                    _tile_loadd(2, b_tile_left->data, 64);
                    _tile_loadd(3, b_tile_right->data, 64);
                }

                // Epilogue: final depth tile
                _tile_dpfp16ps(4, 0, 2);
                _tile_dpfp16ps(5, 0, 3);
                _tile_dpfp16ps(6, 1, 2);
                _tile_dpfp16ps(7, 1, 3);

                // Handle partial depth-tile (if any)
                if (depth_remainder > 0) {
                    nk_size_t const depth_offset = full_depth_tiles_count * tile_depth;

                    nk_dots_f16_load_a_graniteamx_(&a_tile_top, a_top_base + depth_offset, a_stride_elements, 16,
                                                   depth_remainder);
                    nk_dots_f16_load_a_graniteamx_(&a_tile_bottom, a_bottom_base + depth_offset, a_stride_elements, 16,
                                                   depth_remainder);

                    b_tile_left = (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base + (b_column_left_base +
                                                                                            full_depth_tiles_count) *
                                                                                               tile_size);
                    b_tile_right = (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base + (b_column_right_base +
                                                                                             full_depth_tiles_count) *
                                                                                                tile_size);

                    _tile_loadd(0, a_tile_top.data, 64);
                    _tile_loadd(1, a_tile_bottom.data, 64);
                    _tile_loadd(2, b_tile_left->data, 64);
                    _tile_loadd(3, b_tile_right->data, 64);

                    _tile_dpfp16ps(4, 0, 2);
                    _tile_dpfp16ps(5, 0, 3);
                    _tile_dpfp16ps(6, 1, 2);
                    _tile_dpfp16ps(7, 1, 3);
                }
            }
            // Full row-block but only partial depth tile (depth < tile_depth)
            else if (is_full_row_block) {
                nk_f16_t const *a_top_base = a + row_block_start * a_stride_elements;
                nk_f16_t const *a_bottom_base = a + (row_block_start + 16) * a_stride_elements;

                nk_dots_f16_load_a_graniteamx_(&a_tile_top, a_top_base, a_stride_elements, 16, depth_remainder);
                nk_dots_f16_load_a_graniteamx_(&a_tile_bottom, a_bottom_base, a_stride_elements, 16, depth_remainder);

                nk_dots_f16_b32x16_graniteamx_t const *b_tile_left =
                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base + b_column_left_base * tile_size);
                nk_dots_f16_b32x16_graniteamx_t const *b_tile_right =
                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base + b_column_right_base * tile_size);

                _tile_loadd(0, a_tile_top.data, 64);
                _tile_loadd(1, a_tile_bottom.data, 64);
                _tile_loadd(2, b_tile_left->data, 64);
                _tile_loadd(3, b_tile_right->data, 64);

                _tile_dpfp16ps(4, 0, 2);
                _tile_dpfp16ps(5, 0, 3);
                _tile_dpfp16ps(6, 1, 2);
                _tile_dpfp16ps(7, 1, 3);
            }
            // Slow path: edge row-block → buffered load with masking
            else {
                nk_size_t const rows_in_high_tile = (valid_rows_count > 16) ? 16 : valid_rows_count;
                nk_size_t const rows_in_low_tile = (valid_rows_count > 16) ? valid_rows_count - 16 : 0;

                for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {
                    nk_size_t const depth_offset = depth_tile_idx * tile_depth;
                    nk_size_t const valid_depth = (depth_tile_idx < full_depth_tiles_count) ? tile_depth
                                                                                            : depth_remainder;

                    nk_dots_f16_load_a_graniteamx_(&a_tile_top, a + row_block_start * a_stride_elements + depth_offset,
                                                   a_stride_elements, rows_in_high_tile, valid_depth);
                    if (rows_in_low_tile > 0) {
                        nk_dots_f16_load_a_graniteamx_(&a_tile_bottom,
                                                       a + (row_block_start + 16) * a_stride_elements + depth_offset,
                                                       a_stride_elements, rows_in_low_tile, valid_depth);
                    }

                    nk_dots_f16_b32x16_graniteamx_t const *b_tile_left =
                        (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base +
                                                                  (b_column_left_base + depth_tile_idx) * tile_size);
                    nk_dots_f16_b32x16_graniteamx_t const *b_tile_right =
                        (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base +
                                                                  (b_column_right_base + depth_tile_idx) * tile_size);

                    _tile_loadd(0, a_tile_top.data, 64);
                    _tile_loadd(1, a_tile_bottom.data, 64);
                    _tile_loadd(2, b_tile_left->data, 64);
                    _tile_loadd(3, b_tile_right->data, 64);

                    _tile_dpfp16ps(4, 0, 2);
                    _tile_dpfp16ps(5, 0, 3);
                    _tile_dpfp16ps(6, 1, 2);
                    _tile_dpfp16ps(7, 1, 3);
                }
            }

            // Store accumulators to output (once per output block)
            if (is_full_row_block) {
                nk_f32_t *c_block = c + row_block_start * c_stride_elements + col_block_start;
                _tile_stored(4, c_block, c_stride_bytes);
                _tile_stored(5, c_block + 16, c_stride_bytes);
                _tile_stored(6, (nk_f32_t *)((char *)c_block + 16 * c_stride_bytes), c_stride_bytes);
                _tile_stored(7, (nk_f32_t *)((char *)c_block + 16 * c_stride_bytes) + 16, c_stride_bytes);
            }
            else {
                _tile_stored(4, c_accum_buffer.c[0][0].data, 64);
                _tile_stored(5, c_accum_buffer.c[0][1].data, 64);
                _tile_stored(6, c_accum_buffer.c[1][0].data, 64);
                _tile_stored(7, c_accum_buffer.c[1][1].data, 64);
                nk_dots_f16_output2x2_graniteamx_(&c_accum_buffer,
                                                  c + row_block_start * c_stride_elements + col_block_start,
                                                  c_stride_elements, valid_rows_count, 32);
            }
        }
    }

    // Handle odd column-tile (single 16-column tile if column_tiles_count is odd)
    if (column_tiles_count % 2 == 1) {
        nk_size_t const column_tile_idx = column_tiles_count - 1;
        nk_size_t const col_start = column_tile_idx * 16;
        nk_size_t const b_column_base = column_tile_idx * depth_tiles_count;

        for (nk_size_t row_block_idx = 0; row_block_idx < row_blocks_count; row_block_idx++) {
            nk_size_t const row_block_start = row_block_idx * 32;
            nk_size_t const valid_rows_count = (row_block_start + 32 <= rows_count) ? 32
                                                                                    : (rows_count - row_block_start);
            nk_size_t const rows_in_high_tile = (valid_rows_count > 16) ? 16 : valid_rows_count;
            nk_size_t const rows_in_low_tile = (valid_rows_count > 16) ? valid_rows_count - 16 : 0;

            nk_dots_f16_state_graniteamx_t c_high_state, c_low_state;

            _tile_zero(4);
            _tile_zero(6);

            for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {
                nk_size_t const depth_offset = depth_tile_idx * tile_depth;
                nk_size_t const valid_depth = (depth_tile_idx < full_depth_tiles_count) ? tile_depth : depth_remainder;

                nk_dots_f16_load_a_graniteamx_(&a_tile_top, a + row_block_start * a_stride_elements + depth_offset,
                                               a_stride_elements, rows_in_high_tile, valid_depth);
                if (rows_in_low_tile > 0) {
                    nk_dots_f16_load_a_graniteamx_(&a_tile_bottom,
                                                   a + (row_block_start + 16) * a_stride_elements + depth_offset,
                                                   a_stride_elements, rows_in_low_tile, valid_depth);
                }

                nk_dots_f16_b32x16_graniteamx_t const *b_tile =
                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base +
                                                              (b_column_base + depth_tile_idx) * tile_size);

                _tile_loadd(0, a_tile_top.data, 64);
                _tile_loadd(1, a_tile_bottom.data, 64);
                _tile_loadd(2, b_tile->data, 64);

                _tile_dpfp16ps(4, 0, 2);
                _tile_dpfp16ps(6, 1, 2);
            }

            _tile_stored(4, c_high_state.data, 64);
            _tile_stored(6, c_low_state.data, 64);

            nk_dots_f16_store_graniteamx_(&c_high_state, c + row_block_start * c_stride_elements + col_start,
                                          c_stride_elements, rows_in_high_tile, 16);
            if (rows_in_low_tile > 0) {
                nk_dots_f16_store_graniteamx_(&c_low_state, c + (row_block_start + 16) * c_stride_elements + col_start,
                                              c_stride_elements, rows_in_low_tile, 16);
            }
        }
    }

    // Handle column-edge (remaining columns < 16) using AMX with partial tiles
    if (column_remainder_count > 0) {
        for (nk_size_t row_block_idx = 0; row_block_idx < row_blocks_count; row_block_idx++) {
            nk_size_t const row_block_start = row_block_idx * 32;
            nk_size_t const valid_rows_count = (row_block_start + 32 <= rows_count) ? 32
                                                                                    : (rows_count - row_block_start);
            nk_size_t const rows_in_high_tile = (valid_rows_count > 16) ? 16 : valid_rows_count;
            nk_size_t const rows_in_low_tile = (valid_rows_count > 16) ? valid_rows_count - 16 : 0;

            nk_dots_f16_state_graniteamx_t c_high_state, c_low_state;
            nk_dots_bf16_a16x32_sapphireamx_t b_as_a;
            nk_dots_bf16_b32x16_sapphireamx_t b_tile;

            _tile_zero(4);
            _tile_zero(6);

            for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {
                nk_size_t const depth_offset = depth_tile_idx * tile_depth;
                nk_size_t const valid_depth = (depth_tile_idx < full_depth_tiles_count) ? tile_depth : depth_remainder;

                nk_dots_f16_load_a_graniteamx_(&a_tile_top, a + row_block_start * a_stride_elements + depth_offset,
                                               a_stride_elements, rows_in_high_tile, valid_depth);
                if (rows_in_low_tile > 0) {
                    nk_dots_f16_load_a_graniteamx_(&a_tile_bottom,
                                                   a + (row_block_start + 16) * a_stride_elements + depth_offset,
                                                   a_stride_elements, rows_in_low_tile, valid_depth);
                }

                // Load edge columns as BF16-shaped tile (same 16-bit layout) and transpose on-the-fly
                nk_dots_bf16_load_a_sapphireamx_(&b_as_a, (nk_bf16_t const *)(col_edge_ptr + depth_offset), depth,
                                                 column_remainder_count, valid_depth);
                nk_dots_pack_bf16_transposed_sapphireamx_(&b_as_a, &b_tile);

                _tile_loadd(0, a_tile_top.data, 64);
                _tile_loadd(1, a_tile_bottom.data, 64);
                _tile_loadd(2, b_tile.data, 64);

                _tile_dpfp16ps(4, 0, 2);
                _tile_dpfp16ps(6, 1, 2);
            }

            _tile_stored(4, c_high_state.data, 64);
            _tile_stored(6, c_low_state.data, 64);

            nk_dots_f16_store_graniteamx_(&c_high_state, c + row_block_start * c_stride_elements + full_cols,
                                          c_stride_elements, rows_in_high_tile, column_remainder_count);
            if (rows_in_low_tile > 0) {
                nk_dots_f16_store_graniteamx_(&c_low_state, c + (row_block_start + 16) * c_stride_elements + full_cols,
                                              c_stride_elements, rows_in_low_tile, column_remainder_count);
            }
        }
    }

    _tile_release();
}

NK_PUBLIC void nk_dots_symmetric_f16_graniteamx(                                   //
    nk_f16_t const *vectors, nk_size_t vectors_count, nk_size_t depth,             //
    nk_size_t stride_in_bytes, nk_f32_t *result, nk_size_t result_stride_in_bytes, //
    nk_size_t row_start, nk_size_t row_count) {

    nk_size_t const stride_elements = stride_in_bytes / sizeof(nk_f16_t);
    nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);

    nk_size_t const row_end = (row_count == 0)
                                  ? vectors_count
                                  : (row_start + row_count < vectors_count ? row_start + row_count : vectors_count);

    // Round depth up to multiple of 96 (3 tiles × 32 elements)
    nk_size_t const depth_tiles = nk_size_divide_round_up_(depth, 32);
    nk_size_t const depth_tile_groups = nk_size_divide_round_up_(depth_tiles, 3);

    nk_dots_f16_a16x32_graniteamx_t a_tiles[3];
    nk_dots_f16_a16x32_graniteamx_t b_src_tiles[3];
    nk_dots_f16_b32x16_graniteamx_t b_tiles[3];
    nk_dots_f16_state_graniteamx_t state;

    nk_amx_tile_configure_sapphireamx_();

    for (nk_size_t row_tile = row_start; row_tile < row_end; row_tile += 16) {
        nk_size_t const valid_rows = (row_tile + 16 <= row_end) ? 16 : (row_end - row_tile);

        for (nk_size_t col_tile = 0; col_tile < vectors_count; col_tile += 16) {
            nk_size_t const valid_cols = (col_tile + 16 <= vectors_count) ? 16 : (vectors_count - col_tile);

            nk_dots_f16_init_graniteamx_(&state);

            for (nk_size_t depth_group_idx = 0; depth_group_idx < depth_tile_groups; depth_group_idx++) {
                nk_size_t const depth_base = depth_group_idx * 96;

                for (int tile_idx = 0; tile_idx < 3; tile_idx++) {
                    nk_size_t const depth_start = depth_base + tile_idx * 32;
                    nk_size_t const valid_depth = (depth_start + 32 <= depth)
                                                      ? 32
                                                      : (depth > depth_start ? depth - depth_start : 0);

                    nk_dots_f16_load_a_graniteamx_(                         //
                        &a_tiles[tile_idx],                                 //
                        vectors + row_tile * stride_elements + depth_start, //
                        stride_elements, valid_rows, valid_depth);

                    if (row_tile == col_tile) {
                        // Reuse A data as B (self-correlation on diagonal)
                        nk_dots_pack_bf16_transposed_sapphireamx_(
                            (nk_dots_bf16_a16x32_sapphireamx_t const *)&a_tiles[tile_idx],
                            (nk_dots_bf16_b32x16_sapphireamx_t *)&b_tiles[tile_idx]);
                    }
                    else {
                        nk_dots_f16_load_a_graniteamx_(                         //
                            &b_src_tiles[tile_idx],                             //
                            vectors + col_tile * stride_elements + depth_start, //
                            stride_elements, valid_cols, valid_depth);
                        nk_dots_pack_bf16_transposed_sapphireamx_(
                            (nk_dots_bf16_a16x32_sapphireamx_t const *)&b_src_tiles[tile_idx],
                            (nk_dots_bf16_b32x16_sapphireamx_t *)&b_tiles[tile_idx]);
                    }
                }

                nk_dots_f16_update_graniteamx_(                    //
                    &state, &a_tiles[0], &a_tiles[1], &a_tiles[2], //
                    &b_tiles[0], &b_tiles[1], &b_tiles[2]);
            }

            nk_dots_f16_store_graniteamx_(                                     //
                &state, result + row_tile * result_stride_elements + col_tile, //
                result_stride_elements, valid_rows, valid_cols);
        }
    }
}

#pragma endregion F16 Native

#pragma region E5M2 Source (widened to FP16 tiles)

/*  E5M2 Granite AMX kernels: same F16 tile shapes, same TDPFP16PS compute body as the F16 path.
 *  The only difference is byte-to-word widening during A-load and B-pack: E5M2 shares F16's
 *  exponent bias (15), so `(byte << 8)` is the exact F16 bit pattern for every E5M2 value,
 *  including zero/subnormals/Inf/NaN. Tile buffers hold F16 after widen — "e5m2" in the
 *  typedefs refers to the source dtype, not the on-tile representation.
 *
 *  The tile types below alias the F16 types (identical memory layout) so we can reuse the
 *  F16 init/store/output2x2/update internal helpers by pointer cast at the boundary. Only the
 *  public entry points and the load-A widen helper are new.
 */

typedef nk_dots_f16_a16x32_graniteamx_t nk_dots_e5m2_a16x32_graniteamx_t;
typedef nk_dots_f16_b32x16_graniteamx_t nk_dots_e5m2_b32x16_graniteamx_t;
typedef nk_dots_f16_state_graniteamx_t nk_dots_e5m2_state_graniteamx_t;
typedef nk_dots_f16_state2x2_graniteamx_t nk_dots_e5m2_state2x2_graniteamx_t;

/* Load A tile from E5M2 row-major source, widen to F16 via `(byte << 8)` into the F16 tile buffer. */
NK_INTERNAL void nk_dots_e5m2_load_a_graniteamx_(        //
    nk_dots_e5m2_a16x32_graniteamx_t *a_tile,            //
    nk_e5m2_t const *src, nk_size_t src_stride_elements, //
    nk_size_t valid_rows, nk_size_t valid_cols) {

    __mmask32 column_mask = (valid_cols >= 32) ? 0xFFFFFFFF : ((__mmask32)1 << valid_cols) - 1;
    __m512i zero_i16x32 = _mm512_setzero_si512();

    for (nk_size_t row_idx = 0; row_idx < 16; row_idx++) {
        if (row_idx < valid_rows) {
            __m256i e5m2_u8x32 = _mm256_maskz_loadu_epi8(column_mask, src + row_idx * src_stride_elements);
            __m512i word_u16x32 = _mm512_cvtepu8_epi16(e5m2_u8x32);
            __m512i f16_bits_i16x32 = _mm512_slli_epi16(word_u16x32, 8);
            _mm512_store_si512((__m512i *)a_tile->data[row_idx], f16_bits_i16x32);
        }
        else { _mm512_store_si512((__m512i *)a_tile->data[row_idx], zero_i16x32); }
    }
    nk_compiler_barrier_sapphireamx_();
}

NK_PUBLIC nk_size_t nk_dots_packed_size_e5m2_graniteamx(nk_size_t column_count, nk_size_t depth) {
    nk_size_t const tmm_rows = 16;
    nk_size_t const tmm_cols = 32;
    nk_size_t const tile_bytes = 512 * sizeof(nk_f16_t); // Tiles hold F16 after widen: same 1KB as F16.

    nk_size_t const full_column_tiles = column_count / tmm_rows;
    nk_size_t const tiles_along_depth = nk_size_divide_round_up_(depth, tmm_cols);
    nk_size_t const column_remainder_count = column_count - full_column_tiles * tmm_rows;

    nk_size_t size = sizeof(nk_dots_amx_packed_header_t);
    size += full_column_tiles * tiles_along_depth * tile_bytes;
    if (column_remainder_count > 0) size += column_remainder_count * depth * sizeof(nk_f16_t);
    size += column_count * sizeof(nk_f32_t);

    return size;
}

NK_PUBLIC void nk_dots_pack_e5m2_graniteamx(                     //
    nk_e5m2_t const *b, nk_size_t column_count, nk_size_t depth, //
    nk_size_t b_stride_in_bytes, void *b_packed) {

    nk_size_t const tmm_rows = 16;
    nk_size_t const tmm_cols = 32;
    nk_size_t const tile_elements = 512;
    nk_size_t const tile_bytes = tile_elements * sizeof(nk_f16_t);
    nk_size_t const b_stride_elements = b_stride_in_bytes; // E5M2: 1 byte per element

    nk_size_t const column_tiles_count = column_count / tmm_rows;
    nk_size_t const depth_tiles_count = nk_size_divide_round_up_(depth, tmm_cols);
    nk_size_t const column_remainder_count = column_count - column_tiles_count * tmm_rows;
    nk_size_t const total_tiles = column_tiles_count * depth_tiles_count;

    nk_dots_amx_packed_header_t *header = (nk_dots_amx_packed_header_t *)b_packed;
    header->full_column_tiles = (nk_u32_t)column_tiles_count;
    header->full_depth_tiles = (nk_u32_t)depth_tiles_count;
    header->column_remainder_count = (nk_u32_t)column_remainder_count;

    nk_size_t const tiles_offset = sizeof(nk_dots_amx_packed_header_t);
    nk_size_t const column_edge_offset = tiles_offset + total_tiles * tile_bytes;
    header->column_edge_offset = (nk_u32_t)column_edge_offset;

    nk_f16_t *tiles_ptr = (nk_f16_t *)((char *)b_packed + tiles_offset);
    nk_f16_t *column_edge_ptr = (nk_f16_t *)((char *)b_packed + column_edge_offset);

    // Pack tiles: widen-and-gather 16 strided E5M2 rows into an F16-aligned source tile,
    // then transpose via the shared BF16/F16 transposer (same 16-bit lane layout).
    for (nk_size_t column_tile_idx = 0; column_tile_idx < column_tiles_count; column_tile_idx++) {
        for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {

            nk_size_t const tile_index = column_tile_idx * depth_tiles_count + depth_tile_idx;
            nk_f16_t *tile_output = tiles_ptr + tile_index * tile_elements;

            nk_size_t const src_row_start = column_tile_idx * tmm_rows;
            nk_size_t const src_column_start = depth_tile_idx * tmm_cols;
            nk_size_t const columns_to_pack = (src_column_start + tmm_cols <= depth) ? tmm_cols
                                                                                     : (depth - src_column_start);

            nk_dots_bf16_a16x32_sapphireamx_t source_tile;
            __mmask32 depth_mask = (__mmask32)((columns_to_pack < 32) ? ((1U << columns_to_pack) - 1) : ~0U);
            for (nk_size_t row_idx = 0; row_idx < tmm_rows; row_idx++) {
                nk_e5m2_t const *source_row = b + (src_row_start + row_idx) * b_stride_elements + src_column_start;
                __m256i e5m2_u8x32 = _mm256_maskz_loadu_epi8(depth_mask, source_row);
                __m512i word_u16x32 = _mm512_cvtepu8_epi16(e5m2_u8x32);
                __m512i f16_bits_i16x32 = _mm512_slli_epi16(word_u16x32, 8);
                _mm512_store_si512(&source_tile.data[row_idx][0], f16_bits_i16x32);
            }

            nk_dots_bf16_b32x16_sapphireamx_t transposed_tile;
            nk_dots_pack_bf16_transposed_sapphireamx_(&source_tile, &transposed_tile);
            for (nk_size_t i = 0; i < tile_bytes; i += 64)
                _mm512_storeu_si512((char *)tile_output + i, _mm512_load_si512((char const *)&transposed_tile + i));
        }
    }

    // Column-remainder: widen and store row-major (depth contiguous per column)
    if (column_remainder_count > 0) {
        nk_size_t const remainder_start_row = column_tiles_count * tmm_rows;
        for (nk_size_t row_idx = 0; row_idx < column_remainder_count; row_idx++) {
            nk_e5m2_t const *src_row = b + (remainder_start_row + row_idx) * b_stride_elements;
            nk_f16_t *dst_row = column_edge_ptr + row_idx * depth;
            nk_size_t column_idx = 0;
            for (; column_idx + 32 <= depth; column_idx += 32) {
                __m256i e5m2_u8x32 = _mm256_loadu_si256((__m256i const *)(src_row + column_idx));
                __m512i word_u16x32 = _mm512_cvtepu8_epi16(e5m2_u8x32);
                __m512i f16_bits_i16x32 = _mm512_slli_epi16(word_u16x32, 8);
                _mm512_storeu_si512(dst_row + column_idx, f16_bits_i16x32);
            }
            if (column_idx < depth) {
                __mmask32 tail_mask = (__mmask32)((1U << (depth - column_idx)) - 1);
                __m256i e5m2_u8x32 = _mm256_maskz_loadu_epi8(tail_mask, src_row + column_idx);
                __m512i word_u16x32 = _mm512_cvtepu8_epi16(e5m2_u8x32);
                __m512i f16_bits_i16x32 = _mm512_slli_epi16(word_u16x32, 8);
                _mm512_mask_storeu_epi16(dst_row + column_idx, tail_mask, f16_bits_i16x32);
            }
        }
    }

    nk_size_t norms_offset = column_edge_offset +
                             (column_remainder_count > 0 ? column_remainder_count * depth * sizeof(nk_f16_t) : 0);
    header->norms_byte_offset = (nk_u32_t)norms_offset;
    nk_f32_t *norms = (nk_f32_t *)((char *)b_packed + norms_offset);
    for (nk_size_t col = 0; col < column_count; col++)
        norms[col] = nk_dots_reduce_sumsq_e5m2_(b + col * b_stride_elements, depth);
}

NK_PUBLIC void nk_dots_packed_e5m2_graniteamx(             //
    nk_e5m2_t const *a, void const *b_packed, nk_f32_t *c, //
    nk_size_t rows_count, nk_size_t cols_count, nk_size_t depth, nk_size_t a_stride_bytes, nk_size_t c_stride_bytes) {
    nk_unused_(cols_count);

    nk_dots_amx_packed_header_t const *header = (nk_dots_amx_packed_header_t const *)b_packed;
    nk_size_t const column_tiles_count = header->full_column_tiles;
    nk_size_t const depth_tiles_count = header->full_depth_tiles;
    nk_size_t const column_remainder_count = header->column_remainder_count;

    nk_f16_t const *b_tiles_base = (nk_f16_t const *)((char const *)b_packed + sizeof(nk_dots_amx_packed_header_t));
    nk_f16_t const *col_edge_ptr = (nk_f16_t const *)((char const *)b_packed + header->column_edge_offset);

    nk_size_t const a_stride_elements = a_stride_bytes; // E5M2: 1 byte per element
    nk_size_t const c_stride_elements = c_stride_bytes / sizeof(nk_f32_t);

    nk_size_t const tile_depth = 32;
    nk_size_t const tile_size = 512;
    nk_size_t const full_cols = column_tiles_count * 16;

    nk_size_t const row_blocks_count = nk_size_divide_round_up_(rows_count, 32);
    nk_size_t const col_blocks_count = column_tiles_count / 2;

    if (depth_tiles_count == 0) return;

    nk_dots_e5m2_a16x32_graniteamx_t a_tile_top, a_tile_bottom;
    nk_dots_e5m2_state2x2_graniteamx_t c_accum_buffer;

    nk_size_t const full_depth_tiles_count = depth / tile_depth;
    nk_size_t const depth_remainder = depth % tile_depth;

    nk_amx_tile_configure_sapphireamx_();

    for (nk_size_t row_block_idx = 0; row_block_idx < row_blocks_count; row_block_idx++) {
        nk_size_t const row_block_start = row_block_idx * 32;
        nk_size_t const valid_rows_count = (row_block_start + 32 <= rows_count) ? 32 : (rows_count - row_block_start);
        nk_size_t const is_full_row_block = (valid_rows_count == 32);

        for (nk_size_t column_block_idx = 0; column_block_idx < col_blocks_count; column_block_idx++) {
            nk_size_t const col_block_start = column_block_idx * 32;
            nk_size_t const b_column_left_base = (column_block_idx * 2) * depth_tiles_count;
            nk_size_t const b_column_right_base = (column_block_idx * 2 + 1) * depth_tiles_count;

            _tile_zero(4);
            _tile_zero(5);
            _tile_zero(6);
            _tile_zero(7);

            // E5M2 A needs widen-then-tile-load every time (no fast in-place path), so we always
            // route through the widen buffer a_tile_{top,bottom} rather than `_tile_loadd` from A.
            for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {
                nk_size_t const depth_offset = depth_tile_idx * tile_depth;
                nk_size_t const valid_depth = (depth_tile_idx < full_depth_tiles_count) ? tile_depth : depth_remainder;
                nk_size_t const rows_in_high_tile = is_full_row_block
                                                        ? 16
                                                        : ((valid_rows_count > 16) ? 16 : valid_rows_count);
                nk_size_t const rows_in_low_tile = is_full_row_block
                                                       ? 16
                                                       : ((valid_rows_count > 16) ? valid_rows_count - 16 : 0);

                nk_dots_e5m2_load_a_graniteamx_(&a_tile_top, a + row_block_start * a_stride_elements + depth_offset,
                                                a_stride_elements, rows_in_high_tile, valid_depth);
                if (rows_in_low_tile > 0) {
                    nk_dots_e5m2_load_a_graniteamx_(&a_tile_bottom,
                                                    a + (row_block_start + 16) * a_stride_elements + depth_offset,
                                                    a_stride_elements, rows_in_low_tile, valid_depth);
                }

                nk_dots_f16_b32x16_graniteamx_t const *b_tile_left = //
                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base +
                                                              (b_column_left_base + depth_tile_idx) * tile_size);
                nk_dots_f16_b32x16_graniteamx_t const *b_tile_right = //
                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base +
                                                              (b_column_right_base + depth_tile_idx) * tile_size);

                _tile_loadd(0, a_tile_top.data, 64);
                if (rows_in_low_tile > 0) _tile_loadd(1, a_tile_bottom.data, 64);
                _tile_loadd(2, b_tile_left->data, 64);
                _tile_loadd(3, b_tile_right->data, 64);

                _tile_dpfp16ps(4, 0, 2);
                _tile_dpfp16ps(5, 0, 3);
                if (rows_in_low_tile > 0) {
                    _tile_dpfp16ps(6, 1, 2);
                    _tile_dpfp16ps(7, 1, 3);
                }
            }

            if (is_full_row_block) {
                nk_f32_t *c_block = c + row_block_start * c_stride_elements + col_block_start;
                _tile_stored(4, c_block, c_stride_bytes);
                _tile_stored(5, c_block + 16, c_stride_bytes);
                _tile_stored(6, (nk_f32_t *)((char *)c_block + 16 * c_stride_bytes), c_stride_bytes);
                _tile_stored(7, (nk_f32_t *)((char *)c_block + 16 * c_stride_bytes) + 16, c_stride_bytes);
            }
            else {
                _tile_stored(4, c_accum_buffer.c[0][0].data, 64);
                _tile_stored(5, c_accum_buffer.c[0][1].data, 64);
                _tile_stored(6, c_accum_buffer.c[1][0].data, 64);
                _tile_stored(7, c_accum_buffer.c[1][1].data, 64);
                nk_dots_f16_output2x2_graniteamx_(&c_accum_buffer,
                                                  c + row_block_start * c_stride_elements + col_block_start,
                                                  c_stride_elements, valid_rows_count, 32);
            }
        }
    }

    // Odd column tile
    if (column_tiles_count % 2 == 1) {
        nk_size_t const column_tile_idx = column_tiles_count - 1;
        nk_size_t const col_start = column_tile_idx * 16;
        nk_size_t const b_column_base = column_tile_idx * depth_tiles_count;

        for (nk_size_t row_block_idx = 0; row_block_idx < row_blocks_count; row_block_idx++) {
            nk_size_t const row_block_start = row_block_idx * 32;
            nk_size_t const valid_rows_count = (row_block_start + 32 <= rows_count) ? 32
                                                                                    : (rows_count - row_block_start);
            nk_size_t const rows_in_high_tile = (valid_rows_count > 16) ? 16 : valid_rows_count;
            nk_size_t const rows_in_low_tile = (valid_rows_count > 16) ? valid_rows_count - 16 : 0;

            nk_dots_f16_state_graniteamx_t c_high_state, c_low_state;

            _tile_zero(4);
            _tile_zero(6);

            for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {
                nk_size_t const depth_offset = depth_tile_idx * tile_depth;
                nk_size_t const valid_depth = (depth_tile_idx < full_depth_tiles_count) ? tile_depth : depth_remainder;

                nk_dots_e5m2_load_a_graniteamx_(&a_tile_top, a + row_block_start * a_stride_elements + depth_offset,
                                                a_stride_elements, rows_in_high_tile, valid_depth);
                if (rows_in_low_tile > 0) {
                    nk_dots_e5m2_load_a_graniteamx_(&a_tile_bottom,
                                                    a + (row_block_start + 16) * a_stride_elements + depth_offset,
                                                    a_stride_elements, rows_in_low_tile, valid_depth);
                }

                nk_dots_f16_b32x16_graniteamx_t const *b_tile = //
                    (nk_dots_f16_b32x16_graniteamx_t const *)(b_tiles_base +
                                                              (b_column_base + depth_tile_idx) * tile_size);

                _tile_loadd(0, a_tile_top.data, 64);
                if (rows_in_low_tile > 0) _tile_loadd(1, a_tile_bottom.data, 64);
                _tile_loadd(2, b_tile->data, 64);

                _tile_dpfp16ps(4, 0, 2);
                if (rows_in_low_tile > 0) _tile_dpfp16ps(6, 1, 2);
            }

            _tile_stored(4, c_high_state.data, 64);
            if (rows_in_low_tile > 0) _tile_stored(6, c_low_state.data, 64);

            nk_dots_f16_store_graniteamx_(&c_high_state, c + row_block_start * c_stride_elements + col_start,
                                          c_stride_elements, rows_in_high_tile, 16);
            if (rows_in_low_tile > 0) {
                nk_dots_f16_store_graniteamx_(&c_low_state, c + (row_block_start + 16) * c_stride_elements + col_start,
                                              c_stride_elements, rows_in_low_tile, 16);
            }
        }
    }

    // Column edge (fewer than 16 extra columns, stored row-major as F16 after pack)
    if (column_remainder_count > 0) {
        for (nk_size_t row_block_idx = 0; row_block_idx < row_blocks_count; row_block_idx++) {
            nk_size_t const row_block_start = row_block_idx * 32;
            nk_size_t const valid_rows_count = (row_block_start + 32 <= rows_count) ? 32
                                                                                    : (rows_count - row_block_start);
            nk_size_t const rows_in_high_tile = (valid_rows_count > 16) ? 16 : valid_rows_count;
            nk_size_t const rows_in_low_tile = (valid_rows_count > 16) ? valid_rows_count - 16 : 0;

            nk_dots_f16_state_graniteamx_t c_high_state, c_low_state;
            nk_dots_bf16_a16x32_sapphireamx_t b_as_a;
            nk_dots_bf16_b32x16_sapphireamx_t b_tile;

            _tile_zero(4);
            _tile_zero(6);

            for (nk_size_t depth_tile_idx = 0; depth_tile_idx < depth_tiles_count; depth_tile_idx++) {
                nk_size_t const depth_offset = depth_tile_idx * tile_depth;
                nk_size_t const valid_depth = (depth_tile_idx < full_depth_tiles_count) ? tile_depth : depth_remainder;

                nk_dots_e5m2_load_a_graniteamx_(&a_tile_top, a + row_block_start * a_stride_elements + depth_offset,
                                                a_stride_elements, rows_in_high_tile, valid_depth);
                if (rows_in_low_tile > 0) {
                    nk_dots_e5m2_load_a_graniteamx_(&a_tile_bottom,
                                                    a + (row_block_start + 16) * a_stride_elements + depth_offset,
                                                    a_stride_elements, rows_in_low_tile, valid_depth);
                }

                nk_dots_bf16_load_a_sapphireamx_(&b_as_a, (nk_bf16_t const *)(col_edge_ptr + depth_offset), depth,
                                                 column_remainder_count, valid_depth);
                nk_dots_pack_bf16_transposed_sapphireamx_(&b_as_a, &b_tile);

                _tile_loadd(0, a_tile_top.data, 64);
                if (rows_in_low_tile > 0) _tile_loadd(1, a_tile_bottom.data, 64);
                _tile_loadd(2, b_tile.data, 64);

                _tile_dpfp16ps(4, 0, 2);
                if (rows_in_low_tile > 0) _tile_dpfp16ps(6, 1, 2);
            }

            _tile_stored(4, c_high_state.data, 64);
            if (rows_in_low_tile > 0) _tile_stored(6, c_low_state.data, 64);

            nk_dots_f16_store_graniteamx_(&c_high_state, c + row_block_start * c_stride_elements + full_cols,
                                          c_stride_elements, rows_in_high_tile, column_remainder_count);
            if (rows_in_low_tile > 0) {
                nk_dots_f16_store_graniteamx_(&c_low_state, c + (row_block_start + 16) * c_stride_elements + full_cols,
                                              c_stride_elements, rows_in_low_tile, column_remainder_count);
            }
        }
    }

    _tile_release();
}

NK_PUBLIC void nk_dots_symmetric_e5m2_graniteamx(                                  //
    nk_e5m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth,            //
    nk_size_t stride_in_bytes, nk_f32_t *result, nk_size_t result_stride_in_bytes, //
    nk_size_t row_start, nk_size_t row_count) {

    nk_size_t const stride_elements = stride_in_bytes; // E5M2: 1 byte per element
    nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f32_t);

    nk_size_t const row_end = (row_count == 0)
                                  ? vectors_count
                                  : (row_start + row_count < vectors_count ? row_start + row_count : vectors_count);

    nk_size_t const depth_tiles = nk_size_divide_round_up_(depth, 32);
    nk_size_t const depth_tile_groups = nk_size_divide_round_up_(depth_tiles, 3);

    nk_dots_e5m2_a16x32_graniteamx_t a_tiles[3];
    nk_dots_e5m2_a16x32_graniteamx_t b_src_tiles[3];
    nk_dots_f16_b32x16_graniteamx_t b_tiles[3];
    nk_dots_f16_state_graniteamx_t state;

    nk_amx_tile_configure_sapphireamx_();

    for (nk_size_t row_tile = row_start; row_tile < row_end; row_tile += 16) {
        nk_size_t const valid_rows = (row_tile + 16 <= row_end) ? 16 : (row_end - row_tile);

        for (nk_size_t col_tile = 0; col_tile < vectors_count; col_tile += 16) {
            nk_size_t const valid_cols = (col_tile + 16 <= vectors_count) ? 16 : (vectors_count - col_tile);

            nk_dots_f16_init_graniteamx_(&state);

            for (nk_size_t depth_group_idx = 0; depth_group_idx < depth_tile_groups; depth_group_idx++) {
                nk_size_t const depth_base = depth_group_idx * 96;

                for (int tile_idx = 0; tile_idx < 3; tile_idx++) {
                    nk_size_t const depth_start = depth_base + tile_idx * 32;
                    nk_size_t const valid_depth = (depth_start + 32 <= depth)
                                                      ? 32
                                                      : (depth > depth_start ? depth - depth_start : 0);

                    nk_dots_e5m2_load_a_graniteamx_(                        //
                        &a_tiles[tile_idx],                                 //
                        vectors + row_tile * stride_elements + depth_start, //
                        stride_elements, valid_rows, valid_depth);

                    if (row_tile == col_tile) {
                        nk_dots_pack_bf16_transposed_sapphireamx_(
                            (nk_dots_bf16_a16x32_sapphireamx_t const *)&a_tiles[tile_idx],
                            (nk_dots_bf16_b32x16_sapphireamx_t *)&b_tiles[tile_idx]);
                    }
                    else {
                        nk_dots_e5m2_load_a_graniteamx_(                        //
                            &b_src_tiles[tile_idx],                             //
                            vectors + col_tile * stride_elements + depth_start, //
                            stride_elements, valid_cols, valid_depth);
                        nk_dots_pack_bf16_transposed_sapphireamx_(
                            (nk_dots_bf16_a16x32_sapphireamx_t const *)&b_src_tiles[tile_idx],
                            (nk_dots_bf16_b32x16_sapphireamx_t *)&b_tiles[tile_idx]);
                    }
                }

                nk_dots_f16_update_graniteamx_(                    //
                    &state, &a_tiles[0], &a_tiles[1], &a_tiles[2], //
                    &b_tiles[0], &b_tiles[1], &b_tiles[2]);
            }

            nk_dots_f16_store_graniteamx_(                                     //
                &state, result + row_tile * result_stride_elements + col_tile, //
                result_stride_elements, valid_rows, valid_cols);
        }
    }
}

#pragma endregion E5M2 Source

#if defined(__clang__)
#pragma clang attribute pop
#elif defined(__GNUC__)
#pragma GCC pop_options
#endif

#if defined(__cplusplus)
} // extern "C"
#endif

#endif // NK_TARGET_GRANITEAMX
#endif // NK_TARGET_X8664_
#endif // NK_DOTS_GRANITEAMX_H
