-
Notifications
You must be signed in to change notification settings - Fork 183
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Experimental] Add Kleidi i8mm gemm kernels (#1295)
* Update git ignore * [experimental] Add Kleidi compile def at the top level * [Experimental] Add Kleidi i8mm gemm kernels Add kernel level tests, with basic cross compilation support. Tested with S24 + r26c ``` [----------] 6 tests from test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.k_eq_gs_32 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.k_eq_gs_32 (0 ms) [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.large_k_n_gs32 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.large_k_n_gs32 (79 ms) [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.even_n_gs32 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.even_n_gs32 (28 ms) [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.k_eq_gs128 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.k_eq_gs128 (3 ms) [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.clamp_k_eq_gs128 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.clamp_k_eq_gs128 (3 ms) [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.m_clamp_k_eq_gs128 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.m_clamp_k_eq_gs128 (5 ms) [----------] 6 tests from test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm (121 ms total) [----------] 6 tests from test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.k_eq_gs_32 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.k_eq_gs_32 (0 ms) [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.large_k_n_gs32 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.large_k_n_gs32 (79 ms) [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.even_n_gs32 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.even_n_gs32 (28 ms) [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.k_eq_gs128 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.k_eq_gs128 (3 ms) [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.clamp_k_eq_gs128 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.clamp_k_eq_gs128 (3 ms) [ RUN ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.m_clamp_k_eq_gs128 [ OK ] test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.m_clamp_k_eq_gs128 (5 ms) [----------] 6 tests from test_kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm (121 ms total) ``` * [Exeprimental] Kleidi: rename arg name for packing functions * [Experimental] Change kernel cmake_out dir to avoid conflict
- Loading branch information
1 parent
72fb597
commit ca52cdc
Showing
10 changed files
with
557 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -371,4 +371,7 @@ venv/ | |
| sweep/ | ||
|
|
||
| # Model checkpoints | ||
| checkpoints/ | ||
| checkpoints/ | ||
|
|
||
| # Experimental | ||
| torchao/experimental/cmake-out | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
120 changes: 120 additions & 0 deletions
120
...kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,120 @@ | ||
| // Copyright (c) Meta Platforms, Inc. and affiliates. | ||
| // All rights reserved. | ||
| // | ||
| // This source code is licensed under the license found in the | ||
| // LICENSE file in the root directory of this source tree. | ||
|
|
||
| #pragma once | ||
| #include <kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi4c32p/kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm.h> | ||
| #include <torchao/experimental/kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp_qsi4c32p.h> | ||
|
|
||
| namespace torchao::kernels::cpu::aarch64::kleidi { | ||
| namespace kai_matmul_clamp_f32_qai8dxp_qsi4c32p { | ||
| namespace neon_i8mm_8x4x32 { | ||
|
|
||
| const Ukernel get_ukernel() { | ||
| return Ukernel{ | ||
| .get_m_step = | ||
| kai_get_m_step_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm, | ||
| .get_n_step = | ||
| kai_get_n_step_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm, | ||
| .get_mr = | ||
| kai_get_mr_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm, | ||
| .get_nr = | ||
| kai_get_nr_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm, | ||
| .get_kr = | ||
| kai_get_kr_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm, | ||
| .get_sr = | ||
| kai_get_sr_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm, | ||
| .get_lhs_packed_offset = | ||
| kai_get_lhs_packed_offset_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm, | ||
| .get_rhs_packed_offset = | ||
| kai_get_rhs_packed_offset_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm, | ||
| .get_dst_offset = | ||
| kai_get_dst_offset_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm, | ||
| .get_dst_size = | ||
| kai_get_dst_size_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm, | ||
| .run_matmul = | ||
| kai_run_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_8x4x32_neon_i8mm}; | ||
| } | ||
|
|
||
| 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); | ||
| } | ||
|
|
||
| void prepare_activation_data( | ||
| void* prepared_activation_data, | ||
| int m, | ||
| int k, | ||
| int group_size, | ||
| const float* activations) { | ||
| (void)group_size; // unused | ||
| kai_matmul_clamp_f32_qai8dxp_qsi4c32p::prepare_activation_data( | ||
| get_ukernel(), prepared_activation_data, m, k, activations); | ||
| } | ||
|
|
||
| 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); | ||
| } | ||
|
|
||
| void prepare_weight_data( | ||
| void* prepared_weight_data, | ||
| int n, | ||
| int k, | ||
| int group_size, | ||
| const int8_t* weight_qvals, | ||
| const float* weight_scales, | ||
| const int8_t* weight_zeros, | ||
| const float* bias) { | ||
| kai_matmul_clamp_f32_qai8dxp_qsi4c32p::prepare_weight_data( | ||
| get_ukernel(), | ||
| prepared_weight_data, | ||
| n, | ||
| k, | ||
| group_size, | ||
| weight_qvals, | ||
| weight_scales, | ||
| weight_zeros, | ||
| bias); | ||
| } | ||
|
|
||
| void kernel( | ||
| float32_t* output, | ||
| int output_m_stride, | ||
| int m, | ||
| int n, | ||
| int k, | ||
| int group_size, | ||
| const void* weight_data, | ||
| const void* activation_data, | ||
| float clamp_min, | ||
| float clamp_max) { | ||
| if (clamp_min == 0 && clamp_max == 0) { | ||
| clamp_min = std::numeric_limits<float>::lowest(); | ||
| clamp_max = std::numeric_limits<float>::max(); | ||
| } | ||
|
|
||
| auto ukernel = get_ukernel(); | ||
| ukernel.run_matmul( | ||
| m, | ||
| n, | ||
| k, | ||
| group_size, | ||
| activation_data, | ||
| weight_data, | ||
| output, | ||
| /*dst_stride_row=*/n * sizeof(float), | ||
| /*dst_stride_col=*/sizeof(float), | ||
| clamp_min, | ||
| clamp_max); | ||
| } | ||
|
|
||
| size_t get_preferred_alignement() { | ||
| return 16; | ||
| } | ||
| } // namespace neon_i8mm_8x4x32 | ||
| } // namespace kai_matmul_clamp_f32_qai8dxp_qsi4c32p | ||
| } // namespace torchao::kernels::cpu::aarch64::kleidi |
122 changes: 122 additions & 0 deletions
122
...kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,122 @@ | ||
| // Copyright (c) Meta Platforms, Inc. and affiliates. | ||
| // All rights reserved. | ||
| // | ||
| // This source code is licensed under the license found in the | ||
| // LICENSE file in the root directory of this source tree. | ||
|
|
||
| #pragma once | ||
| #include <kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi4c32p/kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm.h> | ||
|
|
||
| #include <torchao/experimental/kernels/cpu/aarch64/kleidi/kai_matmul_clamp_f32_qai8dxp_qsi4c32p.h> | ||
|
|
||
| namespace torchao::kernels::cpu::aarch64::kleidi { | ||
| namespace kai_matmul_clamp_f32_qai8dxp_qsi4c32p { | ||
| namespace neon_i8mm_4x8x32 { | ||
|
|
||
| const Ukernel get_ukernel() { | ||
| return Ukernel{ | ||
| .get_m_step = | ||
| kai_get_m_step_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm, | ||
| .get_n_step = | ||
| kai_get_n_step_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm, | ||
| .get_mr = | ||
| kai_get_mr_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm, | ||
| .get_nr = | ||
| kai_get_nr_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm, | ||
| .get_kr = | ||
| kai_get_kr_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm, | ||
| .get_sr = | ||
| kai_get_sr_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm, | ||
| .get_lhs_packed_offset = | ||
| kai_get_lhs_packed_offset_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm, | ||
| .get_rhs_packed_offset = | ||
| kai_get_rhs_packed_offset_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm, | ||
| .get_dst_offset = | ||
| kai_get_dst_offset_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm, | ||
| .get_dst_size = | ||
| kai_get_dst_size_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm, | ||
| .run_matmul = | ||
| kai_run_matmul_clamp_f32_qai8dxp4x8_qsi4c32p8x8_4x8x32_neon_i8mm}; | ||
| } | ||
|
|
||
| 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); | ||
| } | ||
|
|
||
| void prepare_activation_data( | ||
| void* prepared_activation_data, | ||
| int m, | ||
| int k, | ||
| int group_size, | ||
| const float* activations) { | ||
| (void)group_size; // unused | ||
| kai_matmul_clamp_f32_qai8dxp_qsi4c32p::prepare_activation_data( | ||
| get_ukernel(), prepared_activation_data, m, k, activations); | ||
| } | ||
|
|
||
| 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); | ||
| } | ||
|
|
||
| void prepare_weight_data( | ||
| void* prepared_weight_data, | ||
| int n, | ||
| int k, | ||
| int group_size, | ||
| const int8_t* weight_qvals, | ||
| const float* weight_scales, | ||
| const int8_t* weight_zeros, | ||
| const float* bias) { | ||
| kai_matmul_clamp_f32_qai8dxp_qsi4c32p::prepare_weight_data( | ||
| get_ukernel(), | ||
| prepared_weight_data, | ||
| n, | ||
| k, | ||
| group_size, | ||
| weight_qvals, | ||
| weight_scales, | ||
| weight_zeros, | ||
| bias); | ||
| } | ||
|
|
||
| void kernel( | ||
| float32_t* output, | ||
| int output_m_stride, | ||
| int m, | ||
| int n, | ||
| int k, | ||
| int group_size, | ||
| const void* weight_data, | ||
| const void* activation_data, | ||
| float clamp_min, | ||
| float clamp_max) { | ||
| if (clamp_min == 0 && clamp_max == 0) { | ||
| clamp_min = std::numeric_limits<float>::lowest(); | ||
| clamp_max = std::numeric_limits<float>::max(); | ||
| } | ||
|
|
||
| auto ukernel = get_ukernel(); | ||
| ukernel.run_matmul( | ||
| m, | ||
| n, | ||
| k, | ||
| group_size, | ||
| activation_data, | ||
| weight_data, | ||
| output, | ||
| /*dst_stride_row=*/n * sizeof(float), | ||
| /*dst_stride_col=*/sizeof(float), | ||
| clamp_min, | ||
| clamp_max); | ||
| } | ||
|
|
||
| size_t get_preferred_alignement() { | ||
| return 16; | ||
| } | ||
|
|
||
| } // namespace neon_i8mm_4x8x32 | ||
| } // namespace kai_matmul_clamp_f32_qai8dxp_qsi4c32p | ||
| } // namespace torchao::kernels::cpu::aarch64::kleidi |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.