numkong 7.4.5 → 7.6.0

package/include/numkong/spatials/smef64.h

@@ -13,6 +13,7 @@
 #if NK_TARGET_SME
 
 #include "numkong/dots/serial.h"
+#include "numkong/reduce/sve.h" // `nk_svaddv_f64_`
 #include "numkong/dots/smef64.h"
 
 #if defined(__cplusplus)
@@ -44,7 +45,7 @@ NK_PUBLIC nk_f64_t nk_dots_reduce_sumsq_f32_ssve_(nk_f32_t const *data, nk_size_
         svfloat64_t values_odd_f64x = svcvtlt_f64_f32_x(predicate_odd_b64x, values_f32x);
         accumulator_odd_f64x = svmla_f64_m(predicate_odd_b64x, accumulator_odd_f64x, values_odd_f64x, values_odd_f64x);
     }
-    return svaddv_f64(svptrue_b64(), accumulator_even_f64x) + svaddv_f64(svptrue_b64(), accumulator_odd_f64x);
+    return nk_svaddv_f64_(svptrue_b64(), accumulator_even_f64x) + nk_svaddv_f64_(svptrue_b64(), accumulator_odd_f64x);
 }
 
 NK_PUBLIC nk_f64_t nk_dots_reduce_sumsq_f64_ssve_(nk_f64_t const *data, nk_size_t count) NK_STREAMING_ {
@@ -55,7 +56,7 @@ NK_PUBLIC nk_f64_t nk_dots_reduce_sumsq_f64_ssve_(nk_f64_t const *data, nk_size_
         svfloat64_t values_f64x = svld1_f64(predicate_b64x, data + i);
         accumulator_f64x = svmla_f64_m(predicate_b64x, accumulator_f64x, values_f64x, values_f64x);
     }
-    return svaddv_f64(svptrue_b64(), accumulator_f64x);
+    return nk_svaddv_f64_(svptrue_b64(), accumulator_f64x);
 }
 
 NK_PUBLIC svfloat64_t nk_angulars_from_dot_f64x_ssvef64_(svbool_t predicate_b64x, svfloat64_t dots_f64x,
@@ -85,10 +86,9 @@ NK_PUBLIC svfloat64_t nk_euclideans_from_dot_f64x_ssvef64_(svbool_t predicate_b6
 
 #pragma region F32 Packed Angular
 
-__arm_locally_streaming static void nk_angulars_packed_f32_smef64_finalize_streaming_( //
-    nk_f32_t const *a, void const *b_packed, nk_f64_t *c,                              //
-    nk_size_t rows, nk_size_t columns, nk_size_t depth,                                //
-    nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
+static void nk_angulars_packed_f32_smef64_finalize_ssve_( //
+    nk_f32_t const *a, void const *b_packed, nk_f64_t *c, nk_size_t rows, nk_size_t columns, nk_size_t depth,
+    nk_size_t a_stride_elements, nk_size_t c_stride_elements) NK_STREAMING_ {
 
     nk_dots_sme_packed_header_t const *header = (nk_dots_sme_packed_header_t const *)b_packed;
     nk_f64_t const *b_norms = (nk_f64_t const *)((char const *)b_packed + header->norms_offset);
@@ -110,26 +110,26 @@ __arm_locally_streaming static void nk_angulars_packed_f32_smef64_finalize_strea
     }
 }
 
-NK_PUBLIC void nk_angulars_packed_f32_smef64(             //
-    nk_f32_t const *a, void const *b_packed, nk_f64_t *c, //
-    nk_size_t rows, nk_size_t columns, nk_size_t depth,   //
+NK_PUBLIC void nk_angulars_packed_f32_smef64( //
+    nk_f32_t const *a, void const *b_packed, nk_f64_t *c, nk_size_t rows, nk_size_t columns, nk_size_t depth,
     nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
 
     nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_f32_t);
     nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f64_t);
 
+    nk_sme_start_streaming_();
     nk_dots_packed_f32_smef64_streaming_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
-    nk_angulars_packed_f32_smef64_finalize_streaming_(a, b_packed, c, rows, columns, depth, a_stride_elements,
-                                                      c_stride_elements);
+    nk_angulars_packed_f32_smef64_finalize_ssve_(a, b_packed, c, rows, columns, depth, a_stride_elements,
+                                                 c_stride_elements);
+    nk_sme_stop_streaming_();
 }
 
 #pragma endregion F32 Packed Angular
 #pragma region F32 Packed Euclidean
 
-__arm_locally_streaming static void nk_euclideans_packed_f32_smef64_finalize_streaming_( //
-    nk_f32_t const *a, void const *b_packed, nk_f64_t *c,                                //
-    nk_size_t rows, nk_size_t columns, nk_size_t depth,                                  //
-    nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
+static void nk_euclideans_packed_f32_smef64_finalize_ssve_( //
+    nk_f32_t const *a, void const *b_packed, nk_f64_t *c, nk_size_t rows, nk_size_t columns, nk_size_t depth,
+    nk_size_t a_stride_elements, nk_size_t c_stride_elements) NK_STREAMING_ {
 
     nk_dots_sme_packed_header_t const *header = (nk_dots_sme_packed_header_t const *)b_packed;
     nk_f64_t const *b_norms = (nk_f64_t const *)((char const *)b_packed + header->norms_offset);
@@ -151,25 +151,26 @@ __arm_locally_streaming static void nk_euclideans_packed_f32_smef64_finalize_str
     }
 }
 
-NK_PUBLIC void nk_euclideans_packed_f32_smef64(           //
-    nk_f32_t const *a, void const *b_packed, nk_f64_t *c, //
-    nk_size_t rows, nk_size_t columns, nk_size_t depth,   //
+NK_PUBLIC void nk_euclideans_packed_f32_smef64( //
+    nk_f32_t const *a, void const *b_packed, nk_f64_t *c, nk_size_t rows, nk_size_t columns, nk_size_t depth,
     nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
 
     nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_f32_t);
     nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f64_t);
 
+    nk_sme_start_streaming_();
     nk_dots_packed_f32_smef64_streaming_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
-    nk_euclideans_packed_f32_smef64_finalize_streaming_(a, b_packed, c, rows, columns, depth, a_stride_elements,
-                                                        c_stride_elements);
+    nk_euclideans_packed_f32_smef64_finalize_ssve_(a, b_packed, c, rows, columns, depth, a_stride_elements,
+                                                   c_stride_elements);
+    nk_sme_stop_streaming_();
 }
 
 #pragma endregion F32 Packed Euclidean
 #pragma region F32 Symmetric Angular
 
-__arm_locally_streaming static void nk_angulars_symmetric_f32_smef64_finalize_streaming_(         //
-    nk_f32_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_elements, //
-    nk_f64_t *result, nk_size_t result_stride_elements, nk_size_t row_start, nk_size_t row_count) {
+static void nk_angulars_symmetric_f32_smef64_finalize_ssve_( //
+    nk_f32_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_elements, nk_f64_t *result,
+    nk_size_t result_stride_elements, nk_size_t row_start, nk_size_t row_count) NK_STREAMING_ {
     // Phase 1: cache row norms on diagonal
     for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
         nk_f32_t const *row_vector = vectors + row_index * stride_elements;
@@ -204,25 +205,27 @@ __arm_locally_streaming static void nk_angulars_symmetric_f32_smef64_finalize_st
         result[row_index * result_stride_elements + row_index] = 0;
 }
 
-NK_PUBLIC void nk_angulars_symmetric_f32_smef64(                                                  //
-    nk_f32_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
-    nk_f64_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
+NK_PUBLIC void nk_angulars_symmetric_f32_smef64( //
+    nk_f32_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, nk_f64_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_f32_t);
     nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f64_t);
 
+    nk_sme_start_streaming_();
     nk_dots_symmetric_f32_smef64_streaming_(vectors, vectors_count, depth, stride_elements, result,
                                             result_stride_elements, row_start, row_count);
-    nk_angulars_symmetric_f32_smef64_finalize_streaming_(vectors, vectors_count, depth, stride_elements, result,
-                                                         result_stride_elements, row_start, row_count);
+    nk_angulars_symmetric_f32_smef64_finalize_ssve_(vectors, vectors_count, depth, stride_elements, result,
+                                                    result_stride_elements, row_start, row_count);
+    nk_sme_stop_streaming_();
 }
 
 #pragma endregion F32 Symmetric Angular
 #pragma region F32 Symmetric Euclidean
 
-__arm_locally_streaming static void nk_euclideans_symmetric_f32_smef64_finalize_streaming_(       //
-    nk_f32_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_elements, //
-    nk_f64_t *result, nk_size_t result_stride_elements, nk_size_t row_start, nk_size_t row_count) {
+static void nk_euclideans_symmetric_f32_smef64_finalize_ssve_( //
+    nk_f32_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_elements, nk_f64_t *result,
+    nk_size_t result_stride_elements, nk_size_t row_start, nk_size_t row_count) NK_STREAMING_ {
     // Phase 1: cache row norms on diagonal
     for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
         nk_f32_t const *row_vector = vectors + row_index * stride_elements;
@@ -257,26 +260,27 @@ __arm_locally_streaming static void nk_euclideans_symmetric_f32_smef64_finalize_
         result[row_index * result_stride_elements + row_index] = 0;
 }
 
-NK_PUBLIC void nk_euclideans_symmetric_f32_smef64(                                                //
-    nk_f32_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
-    nk_f64_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
+NK_PUBLIC void nk_euclideans_symmetric_f32_smef64( //
+    nk_f32_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, nk_f64_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_f32_t);
     nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f64_t);
 
+    nk_sme_start_streaming_();
     nk_dots_symmetric_f32_smef64_streaming_(vectors, vectors_count, depth, stride_elements, result,
                                             result_stride_elements, row_start, row_count);
-    nk_euclideans_symmetric_f32_smef64_finalize_streaming_(vectors, vectors_count, depth, stride_elements, result,
-                                                           result_stride_elements, row_start, row_count);
+    nk_euclideans_symmetric_f32_smef64_finalize_ssve_(vectors, vectors_count, depth, stride_elements, result,
+                                                      result_stride_elements, row_start, row_count);
+    nk_sme_stop_streaming_();
 }
 
 #pragma endregion F32 Symmetric Euclidean
 #pragma region F64 Packed Angular
 
-__arm_locally_streaming static void nk_angulars_packed_f64_smef64_finalize_streaming_( //
-    nk_f64_t const *a, void const *b_packed, nk_f64_t *c,                              //
-    nk_size_t rows, nk_size_t columns, nk_size_t depth,                                //
-    nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
+static void nk_angulars_packed_f64_smef64_finalize_ssve_( //
+    nk_f64_t const *a, void const *b_packed, nk_f64_t *c, nk_size_t rows, nk_size_t columns, nk_size_t depth,
+    nk_size_t a_stride_elements, nk_size_t c_stride_elements) NK_STREAMING_ {
 
     nk_dots_sme_packed_header_t const *header = (nk_dots_sme_packed_header_t const *)b_packed;
     nk_f64_t const *b_norms = (nk_f64_t const *)((char const *)b_packed + header->norms_offset);
@@ -298,26 +302,26 @@ __arm_locally_streaming static void nk_angulars_packed_f64_smef64_finalize_strea
     }
 }
 
-NK_PUBLIC void nk_angulars_packed_f64_smef64(             //
-    nk_f64_t const *a, void const *b_packed, nk_f64_t *c, //
-    nk_size_t rows, nk_size_t columns, nk_size_t depth,   //
+NK_PUBLIC void nk_angulars_packed_f64_smef64( //
+    nk_f64_t const *a, void const *b_packed, nk_f64_t *c, nk_size_t rows, nk_size_t columns, nk_size_t depth,
     nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
 
     nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_f64_t);
     nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f64_t);
 
+    nk_sme_start_streaming_();
     nk_dots_packed_f64_smef64_streaming_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
-    nk_angulars_packed_f64_smef64_finalize_streaming_(a, b_packed, c, rows, columns, depth, a_stride_elements,
-                                                      c_stride_elements);
+    nk_angulars_packed_f64_smef64_finalize_ssve_(a, b_packed, c, rows, columns, depth, a_stride_elements,
+                                                 c_stride_elements);
+    nk_sme_stop_streaming_();
 }
 
 #pragma endregion F64 Packed Angular
 #pragma region F64 Packed Euclidean
 
-__arm_locally_streaming static void nk_euclideans_packed_f64_smef64_finalize_streaming_( //
-    nk_f64_t const *a, void const *b_packed, nk_f64_t *c,                                //
-    nk_size_t rows, nk_size_t columns, nk_size_t depth,                                  //
-    nk_size_t a_stride_elements, nk_size_t c_stride_elements) {
+static void nk_euclideans_packed_f64_smef64_finalize_ssve_( //
+    nk_f64_t const *a, void const *b_packed, nk_f64_t *c, nk_size_t rows, nk_size_t columns, nk_size_t depth,
+    nk_size_t a_stride_elements, nk_size_t c_stride_elements) NK_STREAMING_ {
 
     nk_dots_sme_packed_header_t const *header = (nk_dots_sme_packed_header_t const *)b_packed;
     nk_f64_t const *b_norms = (nk_f64_t const *)((char const *)b_packed + header->norms_offset);
@@ -339,25 +343,26 @@ __arm_locally_streaming static void nk_euclideans_packed_f64_smef64_finalize_str
     }
 }
 
-NK_PUBLIC void nk_euclideans_packed_f64_smef64(           //
-    nk_f64_t const *a, void const *b_packed, nk_f64_t *c, //
-    nk_size_t rows, nk_size_t columns, nk_size_t depth,   //
+NK_PUBLIC void nk_euclideans_packed_f64_smef64( //
+    nk_f64_t const *a, void const *b_packed, nk_f64_t *c, nk_size_t rows, nk_size_t columns, nk_size_t depth,
     nk_size_t a_stride_in_bytes, nk_size_t c_stride_in_bytes) {
 
     nk_size_t const a_stride_elements = a_stride_in_bytes / sizeof(nk_f64_t);
     nk_size_t const c_stride_elements = c_stride_in_bytes / sizeof(nk_f64_t);
 
+    nk_sme_start_streaming_();
     nk_dots_packed_f64_smef64_streaming_(a, b_packed, c, rows, columns, depth, a_stride_elements, c_stride_elements);
-    nk_euclideans_packed_f64_smef64_finalize_streaming_(a, b_packed, c, rows, columns, depth, a_stride_elements,
-                                                        c_stride_elements);
+    nk_euclideans_packed_f64_smef64_finalize_ssve_(a, b_packed, c, rows, columns, depth, a_stride_elements,
+                                                   c_stride_elements);
+    nk_sme_stop_streaming_();
 }
 
 #pragma endregion F64 Packed Euclidean
 #pragma region F64 Symmetric Angular
 
-__arm_locally_streaming static void nk_angulars_symmetric_f64_smef64_finalize_streaming_(         //
-    nk_f64_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_elements, //
-    nk_f64_t *result, nk_size_t result_stride_elements, nk_size_t row_start, nk_size_t row_count) {
+static void nk_angulars_symmetric_f64_smef64_finalize_ssve_( //
+    nk_f64_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_elements, nk_f64_t *result,
+    nk_size_t result_stride_elements, nk_size_t row_start, nk_size_t row_count) NK_STREAMING_ {
     // Phase 1: cache row norms on diagonal
     for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
         nk_f64_t const *row_vector = vectors + row_index * stride_elements;
@@ -392,25 +397,27 @@ __arm_locally_streaming static void nk_angulars_symmetric_f64_smef64_finalize_st
         result[row_index * result_stride_elements + row_index] = 0;
 }
 
-NK_PUBLIC void nk_angulars_symmetric_f64_smef64(                                                  //
-    nk_f64_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
-    nk_f64_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
+NK_PUBLIC void nk_angulars_symmetric_f64_smef64( //
+    nk_f64_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, nk_f64_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_f64_t);
     nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f64_t);
 
+    nk_sme_start_streaming_();
     nk_dots_symmetric_f64_smef64_streaming_(vectors, vectors_count, depth, stride_elements, result,
                                             result_stride_elements, row_start, row_count);
-    nk_angulars_symmetric_f64_smef64_finalize_streaming_(vectors, vectors_count, depth, stride_elements, result,
-                                                         result_stride_elements, row_start, row_count);
+    nk_angulars_symmetric_f64_smef64_finalize_ssve_(vectors, vectors_count, depth, stride_elements, result,
+                                                    result_stride_elements, row_start, row_count);
+    nk_sme_stop_streaming_();
 }
 
 #pragma endregion F64 Symmetric Angular
 #pragma region F64 Symmetric Euclidean
 
-__arm_locally_streaming static void nk_euclideans_symmetric_f64_smef64_finalize_streaming_(       //
-    nk_f64_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_elements, //
-    nk_f64_t *result, nk_size_t result_stride_elements, nk_size_t row_start, nk_size_t row_count) {
+static void nk_euclideans_symmetric_f64_smef64_finalize_ssve_( //
+    nk_f64_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_elements, nk_f64_t *result,
+    nk_size_t result_stride_elements, nk_size_t row_start, nk_size_t row_count) NK_STREAMING_ {
     // Phase 1: cache row norms on diagonal
     for (nk_size_t row_index = row_start; row_index < row_start + row_count; ++row_index) {
         nk_f64_t const *row_vector = vectors + row_index * stride_elements;
@@ -445,17 +452,19 @@ __arm_locally_streaming static void nk_euclideans_symmetric_f64_smef64_finalize_
         result[row_index * result_stride_elements + row_index] = 0;
 }
 
-NK_PUBLIC void nk_euclideans_symmetric_f64_smef64(                                                //
-    nk_f64_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, //
-    nk_f64_t *result, nk_size_t result_stride_in_bytes, nk_size_t row_start, nk_size_t row_count) {
+NK_PUBLIC void nk_euclideans_symmetric_f64_smef64( //
+    nk_f64_t const *vectors, nk_size_t vectors_count, nk_size_t depth, nk_size_t stride_in_bytes, nk_f64_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_f64_t);
     nk_size_t const result_stride_elements = result_stride_in_bytes / sizeof(nk_f64_t);
 
+    nk_sme_start_streaming_();
     nk_dots_symmetric_f64_smef64_streaming_(vectors, vectors_count, depth, stride_elements, result,
                                             result_stride_elements, row_start, row_count);
-    nk_euclideans_symmetric_f64_smef64_finalize_streaming_(vectors, vectors_count, depth, stride_elements, result,
-                                                           result_stride_elements, row_start, row_count);
+    nk_euclideans_symmetric_f64_smef64_finalize_ssve_(vectors, vectors_count, depth, stride_elements, result,
+                                                      result_stride_elements, row_start, row_count);
+    nk_sme_stop_streaming_();
 }
 
 #pragma endregion F64 Symmetric Euclidean

package/include/numkong/spatials.h

@@ -739,6 +739,45 @@ NK_PUBLIC void nk_euclideans_symmetric_u8_sapphireamx(nk_u8_t const *vectors, nk
                                                       nk_size_t row_start, nk_size_t row_count);
 #endif // NK_TARGET_SAPPHIREAMX
 
+/*  Granite Rapids backends using Intel AMX-FP16.
+ *  Native FP16 spatial kernels.
+ */
+#if NK_TARGET_GRANITEAMX
+/** @copydoc nk_angulars_packed_f16 */
+NK_PUBLIC void nk_angulars_packed_f16_graniteamx(nk_f16_t const *a, void const *b_packed, nk_f32_t *result,
+                                                 nk_size_t rows, nk_size_t cols, nk_size_t depth,
+                                                 nk_size_t a_stride_in_bytes, nk_size_t r_stride_in_bytes);
+/** @copydoc nk_angulars_symmetric_f16 */
+NK_PUBLIC void nk_angulars_symmetric_f16_graniteamx(nk_f16_t const *vectors, nk_size_t vectors_count, nk_size_t depth,
+                                                    nk_size_t stride, nk_f32_t *result, nk_size_t result_stride,
+                                                    nk_size_t row_start, nk_size_t row_count);
+/** @copydoc nk_euclideans_packed_f16 */
+NK_PUBLIC void nk_euclideans_packed_f16_graniteamx(nk_f16_t const *a, void const *b_packed, nk_f32_t *result,
+                                                   nk_size_t rows, nk_size_t cols, nk_size_t depth,
+                                                   nk_size_t a_stride_in_bytes, nk_size_t r_stride_in_bytes);
+/** @copydoc nk_euclideans_symmetric_f16 */
+NK_PUBLIC void nk_euclideans_symmetric_f16_graniteamx(nk_f16_t const *vectors, nk_size_t vectors_count, nk_size_t depth,
+                                                      nk_size_t stride, nk_f32_t *result, nk_size_t result_stride,
+                                                      nk_size_t row_start, nk_size_t row_count);
+/** @copydoc nk_angulars_packed_f16 */
+NK_PUBLIC void nk_angulars_packed_e5m2_graniteamx(nk_e5m2_t const *a, void const *b_packed, nk_f32_t *result,
+                                                  nk_size_t rows, nk_size_t cols, nk_size_t depth,
+                                                  nk_size_t a_stride_in_bytes, nk_size_t r_stride_in_bytes);
+/** @copydoc nk_angulars_symmetric_f16 */
+NK_PUBLIC void nk_angulars_symmetric_e5m2_graniteamx(nk_e5m2_t const *vectors, nk_size_t vectors_count, nk_size_t depth,
+                                                     nk_size_t stride, nk_f32_t *result, nk_size_t result_stride,
+                                                     nk_size_t row_start, nk_size_t row_count);
+/** @copydoc nk_euclideans_packed_f16 */
+NK_PUBLIC void nk_euclideans_packed_e5m2_graniteamx(nk_e5m2_t const *a, void const *b_packed, nk_f32_t *result,
+                                                    nk_size_t rows, nk_size_t cols, nk_size_t depth,
+                                                    nk_size_t a_stride_in_bytes, nk_size_t r_stride_in_bytes);
+/** @copydoc nk_euclideans_symmetric_f16 */
+NK_PUBLIC void nk_euclideans_symmetric_e5m2_graniteamx(nk_e5m2_t const *vectors, nk_size_t vectors_count,
+                                                       nk_size_t depth, nk_size_t stride, nk_f32_t *result,
+                                                       nk_size_t result_stride, nk_size_t row_start,
+                                                       nk_size_t row_count);
+#endif // NK_TARGET_GRANITEAMX
+
 /*  ARM SME backends using Scalable Matrix Extension.
  *  SME provides ZA tile registers for outer product operations.
  *  F16/BF16/I8/U8/E4M3 use ZA32 tiles, F32/F64 use ZA64 tiles (FEAT_SME_F64F64).
@@ -2078,6 +2117,7 @@ NK_PUBLIC void nk_euclideans_symmetric_u8_rvv(nk_u8_t const *vectors, nk_size_t
 #include "numkong/spatials/alder.h"
 #include "numkong/spatials/sierra.h"
 #include "numkong/spatials/sapphireamx.h"
+#include "numkong/spatials/graniteamx.h"
 #include "numkong/spatials/rvv.h"
 #include "numkong/spatials/v128relaxed.h"
 #include "numkong/spatials/sme.h"
@@ -2290,7 +2330,9 @@ NK_PUBLIC void nk_euclideans_symmetric_f32(nk_f32_t const *vectors, nk_size_t ve
 NK_PUBLIC void nk_angulars_packed_f16(nk_f16_t const *a, void const *b_packed, nk_f32_t *result, nk_size_t rows,
                                       nk_size_t cols, nk_size_t depth, nk_size_t a_stride_in_bytes,
                                       nk_size_t r_stride_in_bytes) {
-#if NK_TARGET_SME
+#if NK_TARGET_GRANITEAMX
+    nk_angulars_packed_f16_graniteamx(a, b_packed, result, rows, cols, depth, a_stride_in_bytes, r_stride_in_bytes);
+#elif NK_TARGET_SME
     nk_angulars_packed_f16_sme(a, b_packed, result, rows, cols, depth, a_stride_in_bytes, r_stride_in_bytes);
 #elif NK_TARGET_NEONFHM
     nk_angulars_packed_f16_neonfhm(a, b_packed, result, rows, cols, depth, a_stride_in_bytes, r_stride_in_bytes);
@@ -2311,7 +2353,10 @@ NK_PUBLIC void nk_angulars_packed_f16(nk_f16_t const *a, void const *b_packed, n
 NK_PUBLIC void nk_angulars_symmetric_f16(nk_f16_t const *vectors, nk_size_t vectors_count, nk_size_t depth,
                                          nk_size_t stride, nk_f32_t *result, nk_size_t result_stride,
                                          nk_size_t row_start, nk_size_t row_count) {
-#if NK_TARGET_SME
+#if NK_TARGET_GRANITEAMX
+    nk_angulars_symmetric_f16_graniteamx(vectors, vectors_count, depth, stride, result, result_stride, row_start,
+                                         row_count);
+#elif NK_TARGET_SME
     nk_angulars_symmetric_f16_sme(vectors, vectors_count, depth, stride, result, result_stride, row_start, row_count);
 #elif NK_TARGET_NEONFHM
     nk_angulars_symmetric_f16_neonfhm(vectors, vectors_count, depth, stride, result, result_stride, row_start,
@@ -2337,7 +2382,9 @@ NK_PUBLIC void nk_angulars_symmetric_f16(nk_f16_t const *vectors, nk_size_t vect
 NK_PUBLIC void nk_euclideans_packed_f16(nk_f16_t const *a, void const *b_packed, nk_f32_t *result, nk_size_t rows,
                                         nk_size_t cols, nk_size_t depth, nk_size_t a_stride_in_bytes,
                                         nk_size_t r_stride_in_bytes) {
-#if NK_TARGET_SME
+#if NK_TARGET_GRANITEAMX
+    nk_euclideans_packed_f16_graniteamx(a, b_packed, result, rows, cols, depth, a_stride_in_bytes, r_stride_in_bytes);
+#elif NK_TARGET_SME
     nk_euclideans_packed_f16_sme(a, b_packed, result, rows, cols, depth, a_stride_in_bytes, r_stride_in_bytes);
 #elif NK_TARGET_NEONFHM
     nk_euclideans_packed_f16_neonfhm(a, b_packed, result, rows, cols, depth, a_stride_in_bytes, r_stride_in_bytes);
@@ -2358,7 +2405,10 @@ NK_PUBLIC void nk_euclideans_packed_f16(nk_f16_t const *a, void const *b_packed,
 NK_PUBLIC void nk_euclideans_symmetric_f16(nk_f16_t const *vectors, nk_size_t vectors_count, nk_size_t depth,
                                            nk_size_t stride, nk_f32_t *result, nk_size_t result_stride,
                                            nk_size_t row_start, nk_size_t row_count) {
-#if NK_TARGET_SME
+#if NK_TARGET_GRANITEAMX
+    nk_euclideans_symmetric_f16_graniteamx(vectors, vectors_count, depth, stride, result, result_stride, row_start,
+                                           row_count);
+#elif NK_TARGET_SME
     nk_euclideans_symmetric_f16_sme(vectors, vectors_count, depth, stride, result, result_stride, row_start, row_count);
 #elif NK_TARGET_NEONFHM
     nk_euclideans_symmetric_f16_neonfhm(vectors, vectors_count, depth, stride, result, result_stride, row_start,