Init neon kernel for lut 2 bit gemm

Author: vraspar

cmake/onnxruntime_mlas.cmake

@@ -101,6 +101,8 @@ function(setup_mlas_source_for_windows)
         ${MLAS_SRC_DIR}/qnbitgemm_kernel_neon.cpp
         ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_fp32.cpp
         ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8.cpp
+        ${MLAS_SRC_DIR}/sqnbitgemm_lut_kernel_neon.h
+        ${MLAS_SRC_DIR}/sqnbitgemm_lut_kernel_neon.cpp
         ${MLAS_SRC_DIR}/cast_kernel_neon.cpp
         ${MLAS_SRC_DIR}/hqnbitgemm_kernel_neon_fp16.cpp
         ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon.h
@@ -470,6 +472,8 @@ else()
           ${MLAS_SRC_DIR}/qnbitgemm_kernel_neon.cpp
           ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_fp32.cpp
           ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8.cpp
+          ${MLAS_SRC_DIR}/sqnbitgemm_lut_kernel_neon.h
+          ${MLAS_SRC_DIR}/sqnbitgemm_lut_kernel_neon.cpp
           ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon.h
           ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon.cpp
           ${MLAS_SRC_DIR}/hgemm_kernel_neon.cpp

onnxruntime/core/mlas/lib/mlasi.h

@@ -1245,6 +1245,10 @@ struct MLAS_QNBIT_LUT_GEMM_DISPATCH;
 
 extern const MLAS_QNBIT_LUT_GEMM_DISPATCH MlasLutGenKernelAvx2;
 
+#if defined(MLAS_TARGET_ARM64)
+extern const MLAS_QNBIT_LUT_GEMM_DISPATCH MlasLutGenKernelNeon;
+#endif
+
 //
 // Rotary embedding dispatch structure.
 //

onnxruntime/core/mlas/lib/platform.cpp

@@ -654,6 +654,9 @@ Return Value:
     this->ArmNeonIsQuantActivationsUnsigned = HasI8MMInstructions ? false : true;
     this->QNBitGemmDispatch = &GetMlasQNBitGemmDispatchNeon(HasDotProductInstructions, HasI8MMInstructions);
 
+    // Enable LUT-based GEMM for 2-bit quantization on ARM64
+    this->LutGenKernel = &MlasLutGenKernelNeon;
+
 #if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED)
     this->CastF16ToF32Kernel = &MlasCastF16ToF32KernelNeon;
     this->CastF32ToF16Kernel = &MlasCastF32ToF16KernelNeon;

onnxruntime/core/mlas/lib/sqnbitgemm_lut_kernel_avx2.h

@@ -10,34 +10,17 @@ Module Name:
 
 Abstract:
 
-    This module implements x64 AVX2 kernel functions for LUT-based n-bit
-    quantized integer matrix multiplication.
+    This module contains the dispatch table declaration for x64 AVX2
+    LUT-based n-bit quantized integer matrix multiplication kernels.
+
 --*/
 
 #pragma once
-#include "qnbitgemm.h"
-
-void
-GenerateLUT_avx2(
-    int32_t group_size,
-    int8_t lut,
-    const float* b,
-    float* scales,
-    float* biases,
-    int K
-);
-
-void
-TMACComputeGemm_avx2(
-    const void* A,
-    const void* a_scales,
-    const void* LUT,
-    const void* LUT_Scales,
-    const void* LUT_Biases,
-    void* C,
-    int bm,
-    int K,
-    int M,
-    int N,
-    size_t BlkLen
-);
+
+#include "qlutgemm.h"
+
+//
+// External dispatch table for AVX2 LUT GEMM kernels.
+// Kernel functions are internal to the .cpp file and accessed via this dispatch.
+//
+extern const MLAS_QNBIT_LUT_GEMM_DISPATCH MlasLutGenKernelAvx2;