usearch 2.21.4 → 2.25.1

package/numkong/bench/bench_attention.cpp

@@ -0,0 +1,113 @@
+/**
+ *  @brief Attention benchmarks - SME ISA.
+ *  @file bench/bench_attention.cpp
+ *  @author Ash Vardanian
+ *  @date February 21, 2026
+ */
+
+#if NK_TARGET_SME
+#include "numkong/attention/sme.h"
+#endif
+
+#include "bench.hpp"
+
+#if NK_TARGET_SME
+
+/**
+ *  @brief Attention benchmark template: pack KV once, then benchmark the attention kernel.
+ *
+ *  Uses matrix_height as num_heads * query_len, matrix_width as kv_len, matrix_depth as head_dim.
+ *  The benchmark reports scalar-ops as: num_heads * query_len * kv_len * head_dim * 2 (Q×K^T + P×V).
+ *
+ *  @tparam scalar_type_ The input/output element type (nk_bf16_t or nk_f16_t)
+ */
+template <typename scalar_type_>
+void measure_attention_sme(                                                                                    //
+    bm::State &state,                                                                                          //
+    nk_size_t (*packed_kv_size_fn)(nk_size_t, nk_size_t, nk_size_t),                                           //
+    void (*pack_kv_fn)(scalar_type_ const *, scalar_type_ const *, nk_size_t, nk_size_t, nk_size_t, nk_size_t, //
+                       nk_size_t, void *),                                                                     //
+    void (*attention_fn)(scalar_type_ const *, void const *, scalar_type_ *, nk_size_t, nk_size_t, nk_size_t,  //
+                         nk_size_t, nk_size_t, nk_f32_t),                                                      //
+    nk_size_t num_heads, nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim) {
+
+    nk_size_t num_kv_heads = num_heads;
+    nk_f32_t scale = 1.0f / std::sqrt((float)head_dim);
+
+    // Allocate Q, K, V, output
+    std::vector<scalar_type_> q(num_heads * query_len * head_dim);
+    std::vector<scalar_type_> k(num_kv_heads * kv_len * head_dim);
+    std::vector<scalar_type_> v(num_kv_heads * kv_len * head_dim);
+    std::vector<scalar_type_> output(num_heads * query_len * head_dim);
+
+    // Fill with random bit patterns (half-precision range)
+    std::mt19937 gen(42);
+    std::uniform_int_distribution<uint16_t> dist(0x0000, 0x3C00); // [0, 1.0] in f16 range
+    for (auto &x : q) {
+        uint16_t bits = dist(gen);
+        std::memcpy(&x, &bits, sizeof(x));
+    }
+    for (auto &x : k) {
+        uint16_t bits = dist(gen);
+        std::memcpy(&x, &bits, sizeof(x));
+    }
+    for (auto &x : v) {
+        uint16_t bits = dist(gen);
+        std::memcpy(&x, &bits, sizeof(x));
+    }
+
+    // Pack KV
+    nk_size_t kv_packed_size = packed_kv_size_fn(num_kv_heads, head_dim, kv_len);
+    std::vector<char> kv_packed(kv_packed_size, 0);
+    nk_size_t k_stride = kv_len * head_dim;
+    nk_size_t v_stride = kv_len * head_dim;
+    pack_kv_fn(k.data(), v.data(), num_kv_heads, head_dim, kv_len, k_stride, v_stride, kv_packed.data());
+
+    // Benchmark
+    std::size_t iterations = 0;
+    for (auto _ : state) {
+        attention_fn(q.data(), kv_packed.data(), output.data(), num_heads, num_kv_heads, query_len, kv_len, head_dim,
+                     scale);
+        ++iterations;
+        bm::DoNotOptimize(output.data());
+    }
+
+    // Q×K^T: num_heads * query_len * kv_len * head_dim * 2 (mul + add)
+    // P×V:   num_heads * query_len * head_dim * kv_len * 2 (mul + add)
+    // Total: 2 * 2 * num_heads * query_len * kv_len * head_dim
+    state.counters["scalar-ops"] = bm::Counter(4.0 * iterations * num_heads * query_len * kv_len * head_dim,
+                                               bm::Counter::kIsRate);
+}
+
+template <typename scalar_type_>
+void attention_sme_(                                                                                           //
+    std::string name,                                                                                          //
+    nk_size_t (*packed_kv_size_fn)(nk_size_t, nk_size_t, nk_size_t),                                           //
+    void (*pack_kv_fn)(scalar_type_ const *, scalar_type_ const *, nk_size_t, nk_size_t, nk_size_t, nk_size_t, //
+                       nk_size_t, void *),                                                                     //
+    void (*attention_fn)(scalar_type_ const *, void const *, scalar_type_ *, nk_size_t, nk_size_t, nk_size_t,  //
+                         nk_size_t, nk_size_t, nk_f32_t),                                                      //
+    nk_size_t num_heads, nk_size_t query_len, nk_size_t kv_len, nk_size_t head_dim) {
+
+    std::string bench_name = name + "<" + std::to_string(num_heads) + "h_" + std::to_string(query_len) + "q_" +
+                             std::to_string(kv_len) + "kv_" + std::to_string(head_dim) + "d>";
+    bm::RegisterBenchmark(bench_name.c_str(), measure_attention_sme<scalar_type_>, packed_kv_size_fn, pack_kv_fn,
+                          attention_fn, num_heads, query_len, kv_len, head_dim);
+}
+
+#endif // NK_TARGET_SME
+
+void bench_attention() {
+#if NK_TARGET_SME
+    // Attention benchmarks: 8 heads, head_dim=128, varying query_len and kv_len
+    // head_dim is always 128 (typical for LLaMA/Qwen); kv_len from matrix_height
+    nk_size_t attn_num_heads = 8;
+    nk_size_t attn_head_dim = 128;
+    nk_size_t attn_query_len = 64;
+    nk_size_t attn_kv_len = bench_config.matrix_height > 0 ? bench_config.matrix_height : 1024;
+    attention_sme_<nk_bf16_t>("attention_bf16_sme", nk_attention_packed_kv_size_bf16_sme, nk_attention_pack_kv_bf16_sme,
+                              nk_attention_bf16_sme, attn_num_heads, attn_query_len, attn_kv_len, attn_head_dim);
+    attention_sme_<nk_f16_t>("attention_f16_sme", nk_attention_packed_kv_size_f16_sme, nk_attention_pack_kv_f16_sme,
+                             nk_attention_f16_sme, attn_num_heads, attn_query_len, attn_kv_len, attn_head_dim);
+#endif
+}

package/numkong/bench/bench_cast.cpp

@@ -0,0 +1,264 @@
+/**
+ *  @brief Type casting benchmarks.
+ *  @file bench/bench_cast.cpp
+ *  @author Ash Vardanian
+ *  @date March 14, 2023
+ */
+
+#include "numkong/cast.h"
+
+#include "bench.hpp"
+
+using cast_kernel_t = void (*)(void const *, nk_dtype_t, nk_size_t, void *, nk_dtype_t);
+
+/**
+ *  @brief Measures the performance of type casting operations using Google Benchmark.
+ *  @param state The benchmark state object provided by Google Benchmark.
+ *  @param kernel The cast kernel function to benchmark.
+ *  @param count The number of elements to cast.
+ */
+template <nk_dtype_t input_dtype_, nk_dtype_t output_dtype_>
+void measure_cast(bm::State &state, cast_kernel_t kernel, std::size_t count) {
+
+    using input_t = typename nk::type_for<input_dtype_>::type;
+    using output_t = typename nk::type_for<output_dtype_>::type;
+
+    auto input = make_vector<input_t>(count);
+    auto output = make_vector<output_t>(count);
+
+    // Initialize input with random values
+    auto generator = make_random_engine();
+    nk::fill_uniform(generator, input.values_data(), count);
+
+    // Benchmark loop
+    std::size_t iterations = 0;
+    for (auto _ : state) {
+        kernel(input.values_data(), input_dtype_, count, output.values_data(), output_dtype_);
+        bm::ClobberMemory();
+        iterations++;
+    }
+
+    std::size_t const bytes_per_call = count * (sizeof(input_t) + sizeof(output_t));
+    state.counters["bytes"] = bm::Counter(1.0 * iterations * bytes_per_call, bm::Counter::kIsRate);
+    state.counters["calls"] = bm::Counter(iterations, bm::Counter::kIsRate);
+}
+
+template <nk_dtype_t input_dtype_, nk_dtype_t output_dtype_>
+void run_cast(std::string name, cast_kernel_t kernel) {
+    std::string bench_name = name + "<" + std::to_string(bench_config.dense_dimensions) + "d>";
+    bm::RegisterBenchmark(bench_name.c_str(), measure_cast<input_dtype_, output_dtype_>, kernel,
+                          bench_config.dense_dimensions);
+}
+
+void bench_cast() {
+
+#if NK_TARGET_HASWELL
+    run_cast<nk_f32_k, nk_f16_k>("cast_f32_to_f16_haswell", nk_cast_haswell);
+    run_cast<nk_f16_k, nk_f32_k>("cast_f16_to_f32_haswell", nk_cast_haswell);
+    run_cast<nk_f32_k, nk_bf16_k>("cast_f32_to_bf16_haswell", nk_cast_haswell);
+    run_cast<nk_bf16_k, nk_f32_k>("cast_bf16_to_f32_haswell", nk_cast_haswell);
+    run_cast<nk_f32_k, nk_e4m3_k>("cast_f32_to_e4m3_haswell", nk_cast_haswell);
+    run_cast<nk_e4m3_k, nk_f32_k>("cast_e4m3_to_f32_haswell", nk_cast_haswell);
+    run_cast<nk_f32_k, nk_e5m2_k>("cast_f32_to_e5m2_haswell", nk_cast_haswell);
+    run_cast<nk_e5m2_k, nk_f32_k>("cast_e5m2_to_f32_haswell", nk_cast_haswell);
+    run_cast<nk_f32_k, nk_e2m3_k>("cast_f32_to_e2m3_haswell", nk_cast_haswell);
+    run_cast<nk_e2m3_k, nk_f32_k>("cast_e2m3_to_f32_haswell", nk_cast_haswell);
+    run_cast<nk_f32_k, nk_e3m2_k>("cast_f32_to_e3m2_haswell", nk_cast_haswell);
+    run_cast<nk_e3m2_k, nk_f32_k>("cast_e3m2_to_f32_haswell", nk_cast_haswell);
+    run_cast<nk_i8_k, nk_f32_k>("cast_i8_to_f32_haswell", nk_cast_haswell);
+    run_cast<nk_f32_k, nk_i8_k>("cast_f32_to_i8_haswell", nk_cast_haswell);
+    run_cast<nk_i16_k, nk_f32_k>("cast_i16_to_f32_haswell", nk_cast_haswell);
+    run_cast<nk_f32_k, nk_i16_k>("cast_f32_to_i16_haswell", nk_cast_haswell);
+    run_cast<nk_u16_k, nk_f32_k>("cast_u16_to_f32_haswell", nk_cast_haswell);
+    run_cast<nk_f32_k, nk_u16_k>("cast_f32_to_u16_haswell", nk_cast_haswell);
+    run_cast<nk_u8_k, nk_f32_k>("cast_u8_to_f32_haswell", nk_cast_haswell);
+    run_cast<nk_f32_k, nk_u8_k>("cast_f32_to_u8_haswell", nk_cast_haswell);
+#endif
+
+#if NK_TARGET_SKYLAKE
+    // float ↔ half/brain/MX
+    run_cast<nk_f32_k, nk_f16_k>("cast_f32_to_f16_skylake", nk_cast_skylake);
+    run_cast<nk_f16_k, nk_f32_k>("cast_f16_to_f32_skylake", nk_cast_skylake);
+    run_cast<nk_f32_k, nk_bf16_k>("cast_f32_to_bf16_skylake", nk_cast_skylake);
+    run_cast<nk_bf16_k, nk_f32_k>("cast_bf16_to_f32_skylake", nk_cast_skylake);
+    run_cast<nk_f32_k, nk_e4m3_k>("cast_f32_to_e4m3_skylake", nk_cast_skylake);
+    run_cast<nk_e4m3_k, nk_f32_k>("cast_e4m3_to_f32_skylake", nk_cast_skylake);
+    run_cast<nk_f32_k, nk_e5m2_k>("cast_f32_to_e5m2_skylake", nk_cast_skylake);
+    run_cast<nk_e5m2_k, nk_f32_k>("cast_e5m2_to_f32_skylake", nk_cast_skylake);
+    run_cast<nk_f32_k, nk_e2m3_k>("cast_f32_to_e2m3_skylake", nk_cast_skylake);
+    run_cast<nk_e2m3_k, nk_f32_k>("cast_e2m3_to_f32_skylake", nk_cast_skylake);
+    run_cast<nk_f32_k, nk_e3m2_k>("cast_f32_to_e3m2_skylake", nk_cast_skylake);
+    run_cast<nk_e3m2_k, nk_f32_k>("cast_e3m2_to_f32_skylake", nk_cast_skylake);
+    // float ↔ double, integer ↔ float
+    run_cast<nk_f64_k, nk_f32_k>("cast_f64_to_f32_skylake", nk_cast_skylake);
+    run_cast<nk_f32_k, nk_f64_k>("cast_f32_to_f64_skylake", nk_cast_skylake);
+    run_cast<nk_i32_k, nk_f64_k>("cast_i32_to_f64_skylake", nk_cast_skylake);
+    run_cast<nk_f64_k, nk_i32_k>("cast_f64_to_i32_skylake", nk_cast_skylake);
+    run_cast<nk_i8_k, nk_i32_k>("cast_i8_to_i32_skylake", nk_cast_skylake);
+    run_cast<nk_i32_k, nk_i8_k>("cast_i32_to_i8_skylake", nk_cast_skylake);
+    run_cast<nk_i16_k, nk_f32_k>("cast_i16_to_f32_skylake", nk_cast_skylake);
+    run_cast<nk_f32_k, nk_i16_k>("cast_f32_to_i16_skylake", nk_cast_skylake);
+    run_cast<nk_u16_k, nk_f32_k>("cast_u16_to_f32_skylake", nk_cast_skylake);
+    run_cast<nk_f32_k, nk_u16_k>("cast_f32_to_u16_skylake", nk_cast_skylake);
+    run_cast<nk_u8_k, nk_f32_k>("cast_u8_to_f32_skylake", nk_cast_skylake);
+    run_cast<nk_f32_k, nk_u8_k>("cast_f32_to_u8_skylake", nk_cast_skylake);
+    run_cast<nk_i64_k, nk_f64_k>("cast_i64_to_f64_skylake", nk_cast_skylake);
+    run_cast<nk_f64_k, nk_i64_k>("cast_f64_to_i64_skylake", nk_cast_skylake);
+    run_cast<nk_u64_k, nk_f64_k>("cast_u64_to_f64_skylake", nk_cast_skylake);
+    run_cast<nk_f64_k, nk_u64_k>("cast_f64_to_u64_skylake", nk_cast_skylake);
+    run_cast<nk_u32_k, nk_f64_k>("cast_u32_to_f64_skylake", nk_cast_skylake);
+    run_cast<nk_f64_k, nk_u32_k>("cast_f64_to_u32_skylake", nk_cast_skylake);
+#endif
+
+#if NK_TARGET_ICELAKE
+    // float ↔ half/brain/MX
+    run_cast<nk_f32_k, nk_f16_k>("cast_f32_to_f16_icelake", nk_cast_icelake);
+    run_cast<nk_f16_k, nk_f32_k>("cast_f16_to_f32_icelake", nk_cast_icelake);
+    run_cast<nk_f32_k, nk_bf16_k>("cast_f32_to_bf16_icelake", nk_cast_icelake);
+    run_cast<nk_bf16_k, nk_f32_k>("cast_bf16_to_f32_icelake", nk_cast_icelake);
+    run_cast<nk_f32_k, nk_e4m3_k>("cast_f32_to_e4m3_icelake", nk_cast_icelake);
+    run_cast<nk_e4m3_k, nk_f32_k>("cast_e4m3_to_f32_icelake", nk_cast_icelake);
+    run_cast<nk_f32_k, nk_e5m2_k>("cast_f32_to_e5m2_icelake", nk_cast_icelake);
+    run_cast<nk_e5m2_k, nk_f32_k>("cast_e5m2_to_f32_icelake", nk_cast_icelake);
+    run_cast<nk_f32_k, nk_e2m3_k>("cast_f32_to_e2m3_icelake", nk_cast_icelake);
+    run_cast<nk_e2m3_k, nk_f32_k>("cast_e2m3_to_f32_icelake", nk_cast_icelake);
+    run_cast<nk_f32_k, nk_e3m2_k>("cast_f32_to_e3m2_icelake", nk_cast_icelake);
+    run_cast<nk_e3m2_k, nk_f32_k>("cast_e3m2_to_f32_icelake", nk_cast_icelake);
+    // integer ↔ float
+    run_cast<nk_i8_k, nk_f32_k>("cast_i8_to_f32_icelake", nk_cast_icelake);
+    run_cast<nk_f32_k, nk_i8_k>("cast_f32_to_i8_icelake", nk_cast_icelake);
+#endif
+
+#if NK_TARGET_SAPPHIRE
+    // float ↔ half/brain/MX
+    run_cast<nk_f32_k, nk_f16_k>("cast_f32_to_f16_sapphire", nk_cast_sapphire);
+    run_cast<nk_f16_k, nk_f32_k>("cast_f16_to_f32_sapphire", nk_cast_sapphire);
+    run_cast<nk_f32_k, nk_bf16_k>("cast_f32_to_bf16_sapphire", nk_cast_sapphire);
+    run_cast<nk_bf16_k, nk_f32_k>("cast_bf16_to_f32_sapphire", nk_cast_sapphire);
+    run_cast<nk_f32_k, nk_e4m3_k>("cast_f32_to_e4m3_sapphire", nk_cast_sapphire);
+    run_cast<nk_e4m3_k, nk_f32_k>("cast_e4m3_to_f32_sapphire", nk_cast_sapphire);
+    run_cast<nk_f32_k, nk_e5m2_k>("cast_f32_to_e5m2_sapphire", nk_cast_sapphire);
+    run_cast<nk_e5m2_k, nk_f32_k>("cast_e5m2_to_f32_sapphire", nk_cast_sapphire);
+    run_cast<nk_f32_k, nk_e2m3_k>("cast_f32_to_e2m3_sapphire", nk_cast_sapphire);
+    run_cast<nk_e2m3_k, nk_f32_k>("cast_e2m3_to_f32_sapphire", nk_cast_sapphire);
+    run_cast<nk_f32_k, nk_e3m2_k>("cast_f32_to_e3m2_sapphire", nk_cast_sapphire);
+    run_cast<nk_e3m2_k, nk_f32_k>("cast_e3m2_to_f32_sapphire", nk_cast_sapphire);
+    // integer ↔ float
+    run_cast<nk_i8_k, nk_f32_k>("cast_i8_to_f32_sapphire", nk_cast_sapphire);
+    run_cast<nk_f32_k, nk_i8_k>("cast_f32_to_i8_sapphire", nk_cast_sapphire);
+#endif
+
+#if NK_TARGET_NEON
+    // NEON — float ↔ half/brain/MX
+    run_cast<nk_f32_k, nk_f16_k>("cast_f32_to_f16_neon", nk_cast_neon);
+    run_cast<nk_f16_k, nk_f32_k>("cast_f16_to_f32_neon", nk_cast_neon);
+    run_cast<nk_f32_k, nk_bf16_k>("cast_f32_to_bf16_neon", nk_cast_neon);
+    run_cast<nk_bf16_k, nk_f32_k>("cast_bf16_to_f32_neon", nk_cast_neon);
+    run_cast<nk_f32_k, nk_e4m3_k>("cast_f32_to_e4m3_neon", nk_cast_neon);
+    run_cast<nk_e4m3_k, nk_f32_k>("cast_e4m3_to_f32_neon", nk_cast_neon);
+    run_cast<nk_f32_k, nk_e5m2_k>("cast_f32_to_e5m2_neon", nk_cast_neon);
+    run_cast<nk_e5m2_k, nk_f32_k>("cast_e5m2_to_f32_neon", nk_cast_neon);
+    run_cast<nk_f32_k, nk_e2m3_k>("cast_f32_to_e2m3_neon", nk_cast_neon);
+    run_cast<nk_e2m3_k, nk_f32_k>("cast_e2m3_to_f32_neon", nk_cast_neon);
+    run_cast<nk_f32_k, nk_e3m2_k>("cast_f32_to_e3m2_neon", nk_cast_neon);
+    run_cast<nk_e3m2_k, nk_f32_k>("cast_e3m2_to_f32_neon", nk_cast_neon);
+    // NEON — float ↔ double
+    run_cast<nk_f64_k, nk_f32_k>("cast_f64_to_f32_neon", nk_cast_neon);
+    run_cast<nk_f32_k, nk_f64_k>("cast_f32_to_f64_neon", nk_cast_neon);
+    // NEON — integer ↔ float
+    run_cast<nk_i8_k, nk_f32_k>("cast_i8_to_f32_neon", nk_cast_neon);
+    run_cast<nk_f32_k, nk_i8_k>("cast_f32_to_i8_neon", nk_cast_neon);
+    run_cast<nk_u8_k, nk_f32_k>("cast_u8_to_f32_neon", nk_cast_neon);
+    run_cast<nk_f32_k, nk_u8_k>("cast_f32_to_u8_neon", nk_cast_neon);
+    run_cast<nk_i16_k, nk_f32_k>("cast_i16_to_f32_neon", nk_cast_neon);
+    run_cast<nk_f32_k, nk_i16_k>("cast_f32_to_i16_neon", nk_cast_neon);
+    run_cast<nk_u16_k, nk_f32_k>("cast_u16_to_f32_neon", nk_cast_neon);
+    run_cast<nk_f32_k, nk_u16_k>("cast_f32_to_u16_neon", nk_cast_neon);
+    // NEON — integer ↔ double
+    run_cast<nk_i32_k, nk_f64_k>("cast_i32_to_f64_neon", nk_cast_neon);
+    run_cast<nk_f64_k, nk_i32_k>("cast_f64_to_i32_neon", nk_cast_neon);
+    run_cast<nk_u32_k, nk_f64_k>("cast_u32_to_f64_neon", nk_cast_neon);
+    run_cast<nk_f64_k, nk_u32_k>("cast_f64_to_u32_neon", nk_cast_neon);
+    run_cast<nk_i64_k, nk_f64_k>("cast_i64_to_f64_neon", nk_cast_neon);
+    run_cast<nk_f64_k, nk_i64_k>("cast_f64_to_i64_neon", nk_cast_neon);
+    run_cast<nk_u64_k, nk_f64_k>("cast_u64_to_f64_neon", nk_cast_neon);
+    run_cast<nk_f64_k, nk_u64_k>("cast_f64_to_u64_neon", nk_cast_neon);
+#endif
+
+#if NK_TARGET_POWERVSX
+    run_cast<nk_f32_k, nk_f16_k>("cast_f32_to_f16_powervsx", nk_cast_powervsx);
+    run_cast<nk_f16_k, nk_f32_k>("cast_f16_to_f32_powervsx", nk_cast_powervsx);
+    run_cast<nk_f32_k, nk_bf16_k>("cast_f32_to_bf16_powervsx", nk_cast_powervsx);
+    run_cast<nk_bf16_k, nk_f32_k>("cast_bf16_to_f32_powervsx", nk_cast_powervsx);
+    run_cast<nk_i8_k, nk_f32_k>("cast_i8_to_f32_powervsx", nk_cast_powervsx);
+    run_cast<nk_f32_k, nk_i8_k>("cast_f32_to_i8_powervsx", nk_cast_powervsx);
+    run_cast<nk_u8_k, nk_f32_k>("cast_u8_to_f32_powervsx", nk_cast_powervsx);
+    run_cast<nk_f32_k, nk_u8_k>("cast_f32_to_u8_powervsx", nk_cast_powervsx);
+    run_cast<nk_i16_k, nk_f32_k>("cast_i16_to_f32_powervsx", nk_cast_powervsx);
+    run_cast<nk_f32_k, nk_i16_k>("cast_f32_to_i16_powervsx", nk_cast_powervsx);
+    run_cast<nk_u16_k, nk_f32_k>("cast_u16_to_f32_powervsx", nk_cast_powervsx);
+    run_cast<nk_f32_k, nk_u16_k>("cast_f32_to_u16_powervsx", nk_cast_powervsx);
+#endif
+
+#if NK_TARGET_V128RELAXED
+    run_cast<nk_f32_k, nk_f16_k>("cast_f32_to_f16_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_f16_k, nk_f32_k>("cast_f16_to_f32_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_f32_k, nk_bf16_k>("cast_f32_to_bf16_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_bf16_k, nk_f32_k>("cast_bf16_to_f32_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_f32_k, nk_e4m3_k>("cast_f32_to_e4m3_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_e4m3_k, nk_f32_k>("cast_e4m3_to_f32_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_f32_k, nk_e5m2_k>("cast_f32_to_e5m2_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_e5m2_k, nk_f32_k>("cast_e5m2_to_f32_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_f32_k, nk_e2m3_k>("cast_f32_to_e2m3_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_e2m3_k, nk_f32_k>("cast_e2m3_to_f32_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_f32_k, nk_e3m2_k>("cast_f32_to_e3m2_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_e3m2_k, nk_f32_k>("cast_e3m2_to_f32_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_i8_k, nk_f32_k>("cast_i8_to_f32_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_f32_k, nk_i8_k>("cast_f32_to_i8_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_u8_k, nk_f32_k>("cast_u8_to_f32_v128relaxed", nk_cast_v128relaxed);
+    run_cast<nk_f32_k, nk_u8_k>("cast_f32_to_u8_v128relaxed", nk_cast_v128relaxed);
+#endif
+
+    // Serial — float ↔ half/brain/MX
+    run_cast<nk_f32_k, nk_f16_k>("cast_f32_to_f16_serial", nk_cast_serial);
+    run_cast<nk_f16_k, nk_f32_k>("cast_f16_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_bf16_k>("cast_f32_to_bf16_serial", nk_cast_serial);
+    run_cast<nk_bf16_k, nk_f32_k>("cast_bf16_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_e4m3_k>("cast_f32_to_e4m3_serial", nk_cast_serial);
+    run_cast<nk_e4m3_k, nk_f32_k>("cast_e4m3_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_e5m2_k>("cast_f32_to_e5m2_serial", nk_cast_serial);
+    run_cast<nk_e5m2_k, nk_f32_k>("cast_e5m2_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_e2m3_k>("cast_f32_to_e2m3_serial", nk_cast_serial);
+    run_cast<nk_e2m3_k, nk_f32_k>("cast_e2m3_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_e3m2_k>("cast_f32_to_e3m2_serial", nk_cast_serial);
+    run_cast<nk_e3m2_k, nk_f32_k>("cast_e3m2_to_f32_serial", nk_cast_serial);
+    // Serial — float ↔ double
+    run_cast<nk_f64_k, nk_f32_k>("cast_f64_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_f64_k>("cast_f32_to_f64_serial", nk_cast_serial);
+    // Serial — integer ↔ float
+    run_cast<nk_i8_k, nk_f32_k>("cast_i8_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_i8_k>("cast_f32_to_i8_serial", nk_cast_serial);
+    run_cast<nk_u8_k, nk_f32_k>("cast_u8_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_u8_k>("cast_f32_to_u8_serial", nk_cast_serial);
+    run_cast<nk_i16_k, nk_f32_k>("cast_i16_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_i16_k>("cast_f32_to_i16_serial", nk_cast_serial);
+    run_cast<nk_u16_k, nk_f32_k>("cast_u16_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_u16_k>("cast_f32_to_u16_serial", nk_cast_serial);
+    run_cast<nk_i32_k, nk_f32_k>("cast_i32_to_f32_serial", nk_cast_serial);
+    run_cast<nk_f32_k, nk_i32_k>("cast_f32_to_i32_serial", nk_cast_serial);
+    // Serial — integer ↔ double
+    run_cast<nk_i8_k, nk_f64_k>("cast_i8_to_f64_serial", nk_cast_serial);
+    run_cast<nk_u8_k, nk_f64_k>("cast_u8_to_f64_serial", nk_cast_serial);
+    run_cast<nk_i32_k, nk_f64_k>("cast_i32_to_f64_serial", nk_cast_serial);
+    run_cast<nk_f64_k, nk_i32_k>("cast_f64_to_i32_serial", nk_cast_serial);
+    run_cast<nk_u32_k, nk_f64_k>("cast_u32_to_f64_serial", nk_cast_serial);
+    run_cast<nk_f64_k, nk_u32_k>("cast_f64_to_u32_serial", nk_cast_serial);
+    run_cast<nk_i64_k, nk_f64_k>("cast_i64_to_f64_serial", nk_cast_serial);
+    run_cast<nk_f64_k, nk_i64_k>("cast_f64_to_i64_serial", nk_cast_serial);
+    run_cast<nk_u64_k, nk_f64_k>("cast_u64_to_f64_serial", nk_cast_serial);
+    run_cast<nk_f64_k, nk_u64_k>("cast_f64_to_u64_serial", nk_cast_serial);
+    // Serial — integer ↔ integer
+    run_cast<nk_i8_k, nk_i32_k>("cast_i8_to_i32_serial", nk_cast_serial);
+    run_cast<nk_i32_k, nk_i8_k>("cast_i32_to_i8_serial", nk_cast_serial);
+    run_cast<nk_i16_k, nk_i64_k>("cast_i16_to_i64_serial", nk_cast_serial);
+}

package/numkong/bench/bench_cross_amx.cpp

@@ -0,0 +1,62 @@
+/**
+ *  @brief Batch operation benchmarks - AMX ISA family (Sapphire Rapids AMX).
+ *  @file bench/bench_cross_amx.cpp
+ *  @author Ash Vardanian
+ *  @date February 6, 2026
+ */
+
+#include "numkong/dot.h"
+#include "numkong/dots.h"
+#include "numkong/spatials.h"
+
+#include "bench.hpp"
+
+void bench_cross_amx() {
+    constexpr nk_dtype_t bf16_k = nk_bf16_k;
+    constexpr nk_dtype_t f16_k = nk_f16_k;
+    constexpr nk_dtype_t e5m2_k = nk_e5m2_k;
+    constexpr nk_dtype_t e4m3_k = nk_e4m3_k;
+    constexpr nk_dtype_t e3m2_k = nk_e3m2_k;
+    constexpr nk_dtype_t e2m3_k = nk_e2m3_k;
+    constexpr nk_dtype_t i8_k = nk_i8_k;
+    constexpr nk_dtype_t u8_k = nk_u8_k;
+
+#if NK_TARGET_SAPPHIREAMX
+    run_dots_packed<bf16_k>("dots_packed_bf16_sapphireamx", nk_dots_packed_size_bf16_sapphireamx,
+                            nk_dots_pack_bf16_sapphireamx, nk_dots_packed_bf16_sapphireamx);
+    run_dots_packed<e5m2_k>("dots_packed_e5m2_sapphireamx", nk_dots_packed_size_e5m2_sapphireamx,
+                            nk_dots_pack_e5m2_sapphireamx, nk_dots_packed_e5m2_sapphireamx);
+    run_dots_packed<e4m3_k>("dots_packed_e4m3_sapphireamx", nk_dots_packed_size_e4m3_sapphireamx,
+                            nk_dots_pack_e4m3_sapphireamx, nk_dots_packed_e4m3_sapphireamx);
+    run_dots_packed<e3m2_k>("dots_packed_e3m2_sapphireamx", nk_dots_packed_size_e3m2_sapphireamx,
+                            nk_dots_pack_e3m2_sapphireamx, nk_dots_packed_e3m2_sapphireamx);
+    run_dots_packed<e2m3_k>("dots_packed_e2m3_sapphireamx", nk_dots_packed_size_e2m3_sapphireamx,
+                            nk_dots_pack_e2m3_sapphireamx, nk_dots_packed_e2m3_sapphireamx);
+    run_dots_packed<i8_k>("dots_packed_i8_sapphireamx", nk_dots_packed_size_i8_sapphireamx, nk_dots_pack_i8_sapphireamx,
+                          nk_dots_packed_i8_sapphireamx);
+    run_dots_packed<u8_k>("dots_packed_u8_sapphireamx", nk_dots_packed_size_u8_sapphireamx, nk_dots_pack_u8_sapphireamx,
+                          nk_dots_packed_u8_sapphireamx);
+
+    run_dots_symmetric<bf16_k>("dots_symmetric_bf16_sapphireamx", nk_dots_symmetric_bf16_sapphireamx);
+    run_dots_symmetric<e5m2_k>("dots_symmetric_e5m2_sapphireamx", nk_dots_symmetric_e5m2_sapphireamx);
+    run_dots_symmetric<e4m3_k>("dots_symmetric_e4m3_sapphireamx", nk_dots_symmetric_e4m3_sapphireamx);
+    run_dots_symmetric<e3m2_k>("dots_symmetric_e3m2_sapphireamx", nk_dots_symmetric_e3m2_sapphireamx);
+    run_dots_symmetric<e2m3_k>("dots_symmetric_e2m3_sapphireamx", nk_dots_symmetric_e2m3_sapphireamx);
+    run_dots_symmetric<i8_k>("dots_symmetric_i8_sapphireamx", nk_dots_symmetric_i8_sapphireamx);
+    run_dots_symmetric<u8_k>("dots_symmetric_u8_sapphireamx", nk_dots_symmetric_u8_sapphireamx);
+#endif
+
+#if NK_TARGET_GRANITEAMX
+    run_dots_packed<f16_k>("dots_packed_f16_graniteamx", nk_dots_packed_size_f16_graniteamx,
+                           nk_dots_pack_f16_graniteamx, nk_dots_packed_f16_graniteamx);
+    run_dots_symmetric<f16_k>("dots_symmetric_f16_graniteamx", nk_dots_symmetric_f16_graniteamx);
+
+    run_angulars_packed<f16_k>("angulars_packed_f16_graniteamx", nk_dots_packed_size_f16_graniteamx,
+                               nk_dots_pack_f16_graniteamx, nk_angulars_packed_f16_graniteamx);
+    run_angulars_symmetric<f16_k>("angulars_symmetric_f16_graniteamx", nk_angulars_symmetric_f16_graniteamx);
+
+    run_euclideans_packed<f16_k>("euclideans_packed_f16_graniteamx", nk_dots_packed_size_f16_graniteamx,
+                                 nk_dots_pack_f16_graniteamx, nk_euclideans_packed_f16_graniteamx);
+    run_euclideans_symmetric<f16_k>("euclideans_symmetric_f16_graniteamx", nk_euclideans_symmetric_f16_graniteamx);
+#endif
+}