[experimental][kleidi] rebase fixes with int to size_t

pytorch · Oct 11, 2024 · d4bb3ed · d4bb3ed
1 parent 3049ded
commit d4bb3ed
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 10 deletions.
diff --git a/...nels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p4x8_1x4x32_neon_dotprod.h b/...nels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p4x8_1x4x32_neon_dotprod.h
@@ -40,7 +40,7 @@ const Ukernel get_ukernel() {
           kai_run_matmul_clamp_f32_qai8dxp1x8_qsi4c32p4x8_1x4x32_neon_dotprod};
 }
 
-int activation_data_size(int m, int k, int group_size) {
+size_t activation_data_size(int m, int k, int group_size) {
   (void)group_size; // unused
   return kai_matmul_clamp_f32_qai8dxp_qsi4c32p::activation_data_size(
       get_ukernel(), m, k);
@@ -57,7 +57,7 @@ void prepare_activation_data(
       get_ukernel(), activation_data, m, k, activations);
 }
 
-int weight_data_size(int n, int k, int group_size) {
+size_t weight_data_size(int n, int k, int group_size) {
   return kai_matmul_clamp_f32_qai8dxp_qsi4c32p::weight_data_size(
       get_ukernel(), n, k, group_size);
 }
@@ -115,7 +115,7 @@ void kernel(
       clamp_max);
 }
 
-size_t get_alignement() {
+size_t get_preferred_alignement() {
   return 16;
 }
 } // namespace neon_dotprod_1x4x32

diff --git a/...nels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p8x8_1x8x32_neon_dotprod.h b/...nels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p8x8_1x8x32_neon_dotprod.h
@@ -39,7 +39,7 @@ const Ukernel get_ukernel() {
           kai_run_matmul_clamp_f32_qai8dxp1x8_qsi4c32p8x8_1x8x32_neon_dotprod};
 }
 
-int activation_data_size(int m, int k, int group_size) {
+size_t activation_data_size(int m, int k, int group_size) {
   (void) group_size; // unused
   return kai_matmul_clamp_f32_qai8dxp_qsi4c32p::activation_data_size(get_ukernel(), m, k);
 }
@@ -59,7 +59,7 @@ void prepare_activation_data(
       activations);
 }
 
-int weight_data_size(int n, int k, int group_size) {
+size_t weight_data_size(int n, int k, int group_size) {
   return kai_matmul_clamp_f32_qai8dxp_qsi4c32p::weight_data_size(get_ukernel(), n, k, group_size);
 }
 
@@ -116,7 +116,7 @@ void kernel(
       clamp_max);
 }
 
-size_t get_alignement() {
+size_t get_preferred_alignement() {
   return 16;
 }
 } // namespace neon_dotprod_1x4x32

diff --git a/torchao/experimental/kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp_qsi4c32p.h b/torchao/experimental/kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp_qsi4c32p.h
@@ -43,7 +43,7 @@ namespace kai_matmul_clamp_f32_qai8dxp_qsi4c32p {
 
 using Ukernel = struct kai_matmul_clamp_f32_qai8dxp_qsi4c32p_ukernel;
 
-int activation_data_size(const Ukernel ukernel, int m, int k) {
+size_t activation_data_size(const Ukernel ukernel, int m, int k) {
   auto lhs_packing = get_lhs_packing();
   return lhs_packing.get_lhs_packed_size(
       m, k, ukernel.get_mr(), ukernel.get_kr(), ukernel.get_sr());
@@ -69,7 +69,7 @@ void prepare_activation_data(
       activation_data);
 }
 
-int weight_data_size(const Ukernel ukernel, int n, int k, int group_size) {
+size_t weight_data_size(const Ukernel ukernel, int n, int k, int group_size) {
   auto rhs_pack = get_rhs_packing();
   return rhs_pack.get_rhs_packed_size(
       n,

diff --git a/torchao/experimental/kernels/cpu/aarch64/tests/test_utils.h b/torchao/experimental/kernels/cpu/aarch64/tests/test_utils.h
@@ -44,7 +44,8 @@ inline std::vector<uint8_t> get_random_lowbit_vector(int size, int nbit) {
 }
 
 // TODO move these to a common utils 
-uint16_t get_bf16_from_float(float f) {
+inline uint16_t
+get_bf16_from_float(float f) {
   uint16_t bf16;
 #if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
   memcpy(&bf16, &f, sizeof(uint16_t));
@@ -56,7 +57,8 @@ uint16_t get_bf16_from_float(float f) {
   return bf16;
 }
 
-float get_float_from_bf16(uint16_t bf16) {
+inline float
+get_float_from_bf16(uint16_t bf16) {
   float f;
   const uint32_t i32 = (bf16 << 16);
   memcpy(&f, &i32, sizeof(uint32_t));