|
| 1 | +--- |
| 2 | +title: Integration of KleidiAI to ORT MLAS |
| 3 | +weight: 3 |
| 4 | + |
| 5 | +### FIXED, DO NOT MODIFY |
| 6 | +layout: learningpathall |
| 7 | +--- |
| 8 | + |
| 9 | +## Integration of KleidiAI to ONNX runtime MLAS |
| 10 | +ONNX runtime is built with KleidiAI support: |
| 11 | +1. Detection: At runtime, MLAS checks the CPU capabilities for SME2 support. |
| 12 | +2. Dispatch: If SME2 is detected, MLAS overrides its default kernels. For example, a Gemm (General Matrix Multiplication) operation that would normally use standard vector instructions (such as NEON) is dispatched to a KleidiAI SME2 micro-kernel. |
| 13 | + |
| 14 | +Currently, KleidiAI in MLAS provides ArmKleidiAI::MlasConv, ArmKleidiAI::MlasGemmBatch and ArmKleidiAI::MlasDynamicQGemmBatch kernels. |
| 15 | + |
| 16 | +### The ArmKleidiAI::MlasConv kernel |
| 17 | +Usually, 2D fp32 convolution operators with batch_size=1 and multiple filters (filter kernel is equal or greater than (3,3)) are dispatched to the ArmKleidiAI::MlasConv kernel. |
| 18 | + |
| 19 | +For example, the figure below shows a (7,7) Conv node. |
| 20 | + |
| 21 | + Conv node") |
| 22 | + |
| 23 | +ArmKleidiAI::MlasConv kernel makes use of KleidiAI’s indirect matrix multiplication (imatmul) micro kernel to accelerate the convolution. |
| 24 | + |
| 25 | +The function calls are shown as below. |
| 26 | +```text |
| 27 | +onnxruntime::InferenceSession::Run |
| 28 | +|--onnxruntime::utils::ExecuteGraph |
| 29 | +| |--onnxruntime::utils::ExecuteGraphImp |
| 30 | +| | |--onnxruntime::ExecuteThePlan |
| 31 | +| | | |--onnxruntime::concurrency::ThreadPool::Schedule |
| 32 | +| | | | |--onnxruntime::RunSince |
| 33 | +| | | | | |--onnxruntime::LaunchKernelStep::Execute |
| 34 | +| | | | | | |--onnxruntime::ExecuteKernel |
| 35 | +| | | | | | | |--onnxruntime::Conv<float>::Compute |
| 36 | +| | | | | | | | |--MlasConv |
| 37 | +| | | | | | | | | |--ArmKleidiAI::MlasConv |
| 38 | +| | | | | | | | | | |--ConvolveSme |
| 39 | +| | | | | | | | | | | |--MlasTrySimpleParallel |
| 40 | +| | | | | | | | | | | | |--kai_run_lhs_imatmul_pack_x32p2vlx1_x32p_sme |
| 41 | +| | | | | | | | | | | | | |--kai_kernel_lhs_imatmul_pack_x32p2vlx1_x32p_sme |
| 42 | +| | | | | | | | | | | | |--kai_run_rhs_imatmul_pack_kxn_x32p2vlx1b_x32_x32_sme |
| 43 | +| | | | | | | | | | | | | |--kai_kernel_rhs_imatmul_pack_kxn_x32p2vlx1b_x32_x32_sme |
| 44 | +| | | | | | | | | | | | |--kai_run_imatmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme2_mopa |
| 45 | +| | | | | | | | | | | | | |--kai_kernel_imatmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme2_mopa |
| 46 | +``` |
| 47 | + |
| 48 | +### The ArmKleidiAI::MlasGemmBatch kernel |
| 49 | +It performs a batched fp32 matrix multiplication (GEMM or GemV) operation using KleidiAI matmul micro kernels. fp32 Conv operators with (1,1) filter kernels also use this kernel. |
| 50 | + |
| 51 | +For example, the figure below shows a (1,1) Conv node. |
| 52 | + |
| 53 | + FusedConv node") |
| 54 | + |
| 55 | +The function calls of fp32 Conv operators with (1,1) filter kernels are shown below. |
| 56 | + |
| 57 | +```text |
| 58 | +onnxruntime::InferenceSession::Run |
| 59 | +|--onnxruntime::utils::ExecuteGraph |
| 60 | +| |--onnxruntime::utils::ExecuteGraphImp |
| 61 | +| | |--onnxruntime::ExecuteThePlan |
| 62 | +| | | |--onnxruntime::concurrency::ThreadPool::Schedule |
| 63 | +| | | | |--onnxruntime::RunSince |
| 64 | +| | | | | |--onnxruntime::LaunchKernelStep::Execute |
| 65 | +| | | | | | |--onnxruntime::ExecuteKernel |
| 66 | +| | | | | | | |--onnxruntime::Conv<float>::Compute |
| 67 | +| | | | | | | | |--MlasConv |
| 68 | +| | | | | | | | | |--MlasGemmBatch |
| 69 | +| | | | | | | | | | |--ArmKleidiAI::MlasGemmBatch |
| 70 | +| | | | | | | | | | | |--MlasTrySimpleParallel |
| 71 | +| | | | | | | | | | | | |--kai_run_lhs_pack_f32p2vlx1_f32_sme |
| 72 | +| | | | | | | | | | | | | |--kai_kernel_lhs_pack_f32p2vlx1_f32_sme |
| 73 | +| | | | | | | | | | | | |--ArmKleidiAI::MlasGemmPackB |
| 74 | +| | | | | | | | | | | | | |--kai_run_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme |
| 75 | +| | | | | | | | | | | | | | |--kai_kernel_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme |
| 76 | +| | | | | | | | | | | | |--kai_run_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa |
| 77 | +| | | | | | | | | | | | | |--kai_kernel_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa |
| 78 | +``` |
| 79 | + |
| 80 | +For example, the figure below shows a Gemm node. |
| 81 | + |
| 82 | + |
| 83 | + |
| 84 | +The function calls of fp32 Gemm operators are shown below. |
| 85 | +```text |
| 86 | +onnxruntime::InferenceSession::Run |
| 87 | +|--onnxruntime::utils::ExecuteGraph |
| 88 | +| |--onnxruntime::utils::ExecuteGraphImp |
| 89 | +| | |--onnxruntime::ExecuteThePlan |
| 90 | +| | | |--onnxruntime::concurrency::ThreadPool::Schedule |
| 91 | +| | | | |--onnxruntime::RunSince |
| 92 | +| | | | | |--onnxruntime::LaunchKernelStep::Execute |
| 93 | +| | | | | | |--onnxruntime::ExecuteKernel |
| 94 | +| | | | | | | |--onnxruntime::Gemm<float>::Compute |
| 95 | +| | | | | | | | |--MlasGemm |
| 96 | +| | | | | | | | | |--MlasGemmBatch |
| 97 | +| | | | | | | | | | |--ArmKleidiAI::MlasGemmBatch |
| 98 | +| | | | | | | | | | | |--MlasTrySimpleParallel |
| 99 | +| | | | | | | | | | | | |--kai_run_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa |
| 100 | +| | | | | | | | | | | | | |--kai_kernel_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa |
| 101 | +``` |
| 102 | + |
| 103 | +### The ArmKleidiAI::MlasDynamicQGemmBatch kernel |
| 104 | +This kernel is for Matmul with float output of dynamic quantized A and symmetric quantized B. |
| 105 | +It uses KleidiAI *kai_kernel_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa* micro kernel. |
0 commit comments