fix: Add vscale helpers

As `vscale` by design doesn't consider factors like accumulation it's
straightforward to use shared helpers and avoid using magic numbers.

This change introduces the constant `KAI_VSCALE_MAX` which is the
maximum allowed `vscale` factor, and `kai_get_sme_vscale()` which is the
SME `vscale` factor for current CPU.

Signed-off-by: Emil Ohlsson <emil.ohlsson@arm.com>

Reviewed-by: James Gross <james.gross@arm.com>
Reviewed-by: Felix Johnny Thomasmathibalan <felixjohnny.thomasmathibalan@arm.com>
Approved-by: Felix Johnny Thomasmathibalan <felixjohnny.thomasmathibalan@arm.com>

Author: EmilOhlssonARM

kai/kai_common.h

@@ -118,8 +118,17 @@ extern "C" {
 #define KAI_MIN(a, b) (((a) < (b)) ? (a) : (b))
 #define KAI_MAX(a, b) (((a) > (b)) ? (a) : (b))
 
-/// Largest supported SME vector length in bytes
-#define KAI_SME_VEC_LENGTH_MAX_BYTES 256  // NOLINT(cppcoreguidelines-macro-to-enum,modernize-macro-to-enum)
+/// KleidiAI shared constants
+enum {
+    /// Largest supported SME vector length in bytes
+    KAI_SME_VEC_LENGTH_MAX_BYTES = 256,
+
+    /// Size of one vscale unit, in bytes
+    KAI_VSCALE_UNIT_BYTES = 16,
+
+    /// Maximum possible VSCALE
+    KAI_VSCALE_MAX = KAI_SME_VEC_LENGTH_MAX_BYTES / KAI_VSCALE_UNIT_BYTES,
+};
 
 /// Gets the version of the project in the Major.Minor.Patch semantic versioning format.
 ///
@@ -221,6 +230,11 @@ inline static uint64_t kai_get_sme_vector_length_u32(void) {
     return kai_get_sme_vector_length_u8() / 4;
 }
 
+/// Gets the vscale scale factor for SME
+inline static uint64_t kai_get_sme_vscale(void) {
+    return kai_get_sme_vector_length_u8() / KAI_VSCALE_UNIT_BYTES;
+}
+
 /// Commit ZA to lazy save buffer
 void kai_commit_za(void);
 #endif  // defined(__ARM_FEATURE_SVE2) || defined(_M_ARM64)
@@ -244,7 +258,7 @@ inline static uint64_t kai_get_sve_vector_length_u32(void) {
 /// @return the int8_t value with sign extended
 inline static int8_t kai_ext_sign_i8_i4(int8_t value) {
     // Make sure value holds correct int4 value
-    KAI_ASSERT(value <= 0xF);
+    KAI_ASSUME(value <= 0xF);
 
     return (value ^ 0x8) - 8;  // NOLINT(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
 }

kai/ukernels/matmul/pack/kai_matmul_pack_lhs_mxk_x32p4vsx1_x32_sme.c

@@ -16,14 +16,14 @@ enum {
     MR_VSCALE = 4,
     KR = 1,
 
-    MAX_MR = MR_VSCALE * KAI_SME_VEC_LENGTH_MAX_BYTES / 16,
+    MAX_MR = MR_VSCALE * KAI_VSCALE_MAX,
 };
 
 void kai_kernel_matmul_pack_lhs_mxk_x32p4vsx1_x32_sme(
     size_t height, size_t width, const void* in, size_t row_offset, void* out);
 
 static size_t get_mr(void) {
-    return MR_VSCALE * kai_get_sme_vector_length_u8() / 16;
+    return MR_VSCALE * kai_get_sme_vscale();
 }
 
 static size_t div_ceil(size_t a, size_t b) {

kai/ukernels/matmul/pack/kai_matmul_pack_rhs_kxn_x32p4vsx1bx32_x32_x32_sme.c

@@ -15,7 +15,7 @@ enum {
     NR_VSCALE = 4,
     KR = 1,
 
-    MAX_NR = NR_VSCALE * KAI_SME_VEC_LENGTH_MAX_BYTES / 16,
+    MAX_NR = NR_VSCALE * KAI_VSCALE_MAX,
 };
 
 struct uker_args_t {
@@ -31,7 +31,7 @@ struct uker_args_t {
 void kai_kernel_matmul_pack_rhs_kxn_x32p4vsx1bx32_x32_x32_sme(const struct uker_args_t* args);
 
 static size_t get_nr(void) {
-    return NR_VSCALE * kai_get_sme_vector_length_u8() / 16;
+    return NR_VSCALE * kai_get_sme_vscale();
 }
 
 static size_t div_ceil(size_t a, size_t b) {

kai/ukernels/matmul/pack/kai_matmul_pack_rhs_nxk_x32p4vsx1bx32_x32_x32_sme.c

@@ -17,14 +17,14 @@ enum {
     NR_VSCALE = 4,
     KR = 1,
 
-    MAX_NR = NR_VSCALE * KAI_SME_VEC_LENGTH_MAX_BYTES / 16,
+    MAX_NR = NR_VSCALE * KAI_VSCALE_MAX,
 };
 
 void kai_kernel_matmul_pack_rhs_nxk_x32p4vsx1bx32_x32_x32_sme(
     size_t height, size_t width, const void* in, size_t row_offset, void* out, const void* bias);
 
 static size_t get_nr(void) {
-    return NR_VSCALE * kai_get_sme_vector_length_u8() / 16;
+    return NR_VSCALE * kai_get_sme_vscale();
 }
 
 static size_t div_ceil(size_t a, size_t b) {