[MLAS] Add AVX2 (+VNNI) 2-bit weight CPU kernels

cmake/onnxruntime_mlas.cmake

@@ -871,13 +871,15 @@ endif()
         )
         set_source_files_properties(${mlas_platform_srcs_avx512core} PROPERTIES COMPILE_FLAGS "-mfma -mavx512vnni -mavx512bw -mavx512dq -mavx512vl")
 
-        # Strip -mavx512vnni from the W2 scalar oracle / pack-helper TU so the
-        # compiler cannot autovectorize its int8 dot-product loops to vpdpbusd.
-        # Needed because this helper runs at model load on AVX-512-only
-        # (non-VNNI) hosts via the AVX-512 W2 dispatch. TU is pure C++ -- no
-        # AVX-512 intrinsics inside.
+        # NOTE: this TU is intrinsic-free scalar helpers reached from both the
+        # AVX2 and AVX-512 W2 dispatch tables at model load. Flags must not
+        # enable any ISA the AVX2-only host lacks, or the compiler may
+        # autovectorize the scalar loops into EVEX instructions and SIGILL.
+        # If the AVX-512 W2 pack ever becomes a perf hot spot, split this
+        # file: keep the scalar pack in a flag-less TU and put an optimized
+        # variant in a separate one only used by the AVX-512 tables.
         set_source_files_properties(${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512_2bit.cpp PROPERTIES
-          COMPILE_FLAGS "-mfma -mavx512bw -mavx512dq -mavx512vl")
+          COMPILE_FLAGS "")
 
         set(mlas_platform_srcs_avx512vnni
           ${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512vnni.cpp

onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx2.cpp

@@ -26,6 +26,12 @@ Module Name:
 #include "sqnbitgemm_kernel_avx2_int8_blklen32.h"
 #include "sqnbitgemm_kernel_avx2_int8_blklen64.h"
 
+#include "sqnbitgemm_kernel_avx512_2bit.h"
+#include "sqnbitgemm_kernel_avx2_2bit.h"
+#include "sqnbitgemm_kernel_avx2_2bit_blklen32.h"
+#include "sqnbitgemm_kernel_avx2_2bit_blklen64.h"
+#include "sqnbitgemm_kernel_avx2_2bit_blklen128.h"
+
 #include "sqnbitgemm_m1_sym_kernel_avx2_int8_blklen32.h"
 #include "sqnbitgemm_m1_sym_kernel_avx2_int8_blklen64.h"
 
@@ -1439,6 +1445,80 @@ SQ8BitGemmPackQuantBDataAndBlkSum(
         HasZeroPoint, QuantBZPBegin, PackedQuantB, ThreadPool);
 }
 
+//
+// BlkLen-routing wrappers for the W2 CompInt8 AVX2 / AVX2-VNNI dispatch
+// entries. Production code reaches these via the MLAS dispatch table; tests
+// call them directly via the namespace. Siblings of the AVX-512 variants in
+// sqnbitgemm_kernel_avx512.cpp / sqnbitgemm_kernel_avx512vnni.cpp.
+//
+namespace onnxruntime::mlas::sq2bit_avx2 {
+size_t MLASCALL
+SQ2BitGemmKernel_BlkSum_CompInt8_Avx2_Dispatch(
+    size_t BlkLen,
+    const std::byte* QuantA,
+    const float* QuantAScale,
+    const std::byte* QuantBData,
+    const float* QuantBScale,
+    const std::byte* QuantBZeroPoint,
+    float* C,
+    size_t CountM,
+    size_t CountN,
+    size_t CountK,
+    size_t BlockCountK,
+    const float* Bias,
+    size_t ldc,
+    const float* ABlockSum,
+    const float* QuantBBlkSum)
+{
+    if (BlkLen == 128) {
+        return SQ2BitGemmKernel_BlkSum_CompInt8_BlkLen128_Avx2(
+            QuantA, QuantAScale, QuantBData, QuantBScale,
+            C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
+    }
+    if (BlkLen == 32) {
+        return SQ2BitGemmKernel_BlkSum_CompInt8_BlkLen32_Avx2(
+            QuantA, QuantAScale, QuantBData, QuantBScale,
+            C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
+    }
+    return SQ2BitGemmKernel_BlkSum_CompInt8_Avx2(
+        BlkLen, QuantA, QuantAScale, QuantBData, QuantBScale, QuantBZeroPoint,
+        C, CountM, CountN, CountK, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
+}
+
+size_t MLASCALL
+SQ2BitGemmKernel_BlkSum_CompInt8_Avx2Vnni_Dispatch(
+    size_t BlkLen,
+    const std::byte* QuantA,
+    const float* QuantAScale,
+    const std::byte* QuantBData,
+    const float* QuantBScale,
+    const std::byte* QuantBZeroPoint,
+    float* C,
+    size_t CountM,
+    size_t CountN,
+    size_t CountK,
+    size_t BlockCountK,
+    const float* Bias,
+    size_t ldc,
+    const float* ABlockSum,
+    const float* QuantBBlkSum)
+{
+    if (BlkLen == 128) {
+        return SQ2BitGemmKernel_BlkSum_CompInt8_BlkLen128_Avx2Vnni(
+            QuantA, QuantAScale, QuantBData, QuantBScale,
+            C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
+    }
+    if (BlkLen == 32) {
+        return SQ2BitGemmKernel_BlkSum_CompInt8_BlkLen32_Avx2Vnni(
+            QuantA, QuantAScale, QuantBData, QuantBScale,
+            C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
+    }
+    return SQ2BitGemmKernel_BlkSum_CompInt8_Avx2Vnni(
+        BlkLen, QuantA, QuantAScale, QuantBData, QuantBScale, QuantBZeroPoint,
+        C, CountM, CountN, CountK, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
+}
+}  // namespace onnxruntime::mlas::sq2bit_avx2
+
 //
 // Kernel dispatch structure definition.
 //
@@ -1461,6 +1541,18 @@ const MLAS_QNBIT_GEMM_DISPATCH MlasSQNBitGemmDispatchAvx2 = []() {
     d.SQ8BitGemmKernel_BlkSum_CompInt8 = SQ8BitGemmKernel_BlkSum_CompInt8_avx2<false>;
     d.QuantizeARowComputeBlkSum_CompInt8 = QuantizeARow_CompInt8_avx2;
 
+    // 2-bit native CompInt8 path. Reuses the portable block-group packer and
+    // effective-K helper (shared with AVX-512); the AVX2 kernel supplies the
+    // 256-bit compute. QuantizeARow above already provides the signed-int8 A
+    // layout the W2 path consumes.
+    static_assert(
+        onnxruntime::mlas::sq2bit_avx512::kBlockGroupBlks == kSq2BitAvx512WeightKBlockGroup,
+        "kBlockGroupBlks (kernel-internal) must match kSq2BitAvx512WeightKBlockGroup (qnbitgemm.h).");
+    d.Q2BitGemmPackQuantBDataSize       = onnxruntime::mlas::sq2bit_avx512::Q2BitGemmPackQuantBDataSize_Avx512;
+    d.SQ2BitGemmPackQuantBDataAndBlkSum = onnxruntime::mlas::sq2bit_avx512::SQ2BitGemmPackQuantBDataAndBlkSum_Scalar;
+    d.SQ2BitGemmKernel_BlkSum_CompInt8  = onnxruntime::mlas::sq2bit_avx2::SQ2BitGemmKernel_BlkSum_CompInt8_Avx2_Dispatch;
+    d.Q2BitGemmEffectiveBlockCountK     = onnxruntime::mlas::sq2bit_avx512::Q2BitGemmEffectiveBlockCountK;
+
     return d;
 }();
 
@@ -1483,5 +1575,14 @@ const MLAS_QNBIT_GEMM_DISPATCH MlasSQNBitGemmDispatchAvx2vnni = []() {
     d.SQ8BitGemmKernel_BlkSum_CompInt8 = SQ8BitGemmKernel_BlkSum_CompInt8_avx2<true>;
     d.QuantizeARowComputeBlkSum_CompInt8 = QuantizeARow_CompInt8_avx2;
 
+    // 2-bit native CompInt8 path (AVX-VNNI compute). See the AVX2 table above.
+    static_assert(
+        onnxruntime::mlas::sq2bit_avx512::kBlockGroupBlks == kSq2BitAvx512WeightKBlockGroup,
+        "kBlockGroupBlks (kernel-internal) must match kSq2BitAvx512WeightKBlockGroup (qnbitgemm.h).");
+    d.Q2BitGemmPackQuantBDataSize       = onnxruntime::mlas::sq2bit_avx512::Q2BitGemmPackQuantBDataSize_Avx512;
+    d.SQ2BitGemmPackQuantBDataAndBlkSum = onnxruntime::mlas::sq2bit_avx512::SQ2BitGemmPackQuantBDataAndBlkSum_Scalar;
+    d.SQ2BitGemmKernel_BlkSum_CompInt8  = onnxruntime::mlas::sq2bit_avx2::SQ2BitGemmKernel_BlkSum_CompInt8_Avx2Vnni_Dispatch;
+    d.Q2BitGemmEffectiveBlockCountK     = onnxruntime::mlas::sq2bit_avx512::Q2BitGemmEffectiveBlockCountK;
+
     return d;
 }();

onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx2_2bit.h

@@ -0,0 +1,72 @@
+/*++
+
+Copyright (c) Microsoft Corporation. All rights reserved.
+
+Licensed under the MIT License.
+
+Module Name:
+
+    sqnbitgemm_kernel_avx2_2bit.h
+
+Abstract:
+
+    Declarations for the AVX2 / AVX2-VNNI W2 CompInt8 BlkLen-routing dispatch
+    forwarders. The forwarders are defined in sqnbitgemm_kernel_avx2.cpp and the
+    per-BlkLen kernels live in sqnbitgemm_kernel_avx2_2bit_blklen{32,64,128}.h.
+    Declared here so the dispatch tables and the unit tests can reference them,
+    mirroring the AVX-512 forwarder declarations in
+    sqnbitgemm_kernel_avx512_2bit.h.
+
+--*/
+
+#pragma once
+
+#include <cstddef>
+
+#include "mlas.h"
+
+namespace onnxruntime {
+namespace mlas {
+namespace sq2bit_avx2 {
+
+size_t MLASCALL
+SQ2BitGemmKernel_BlkSum_CompInt8_Avx2_Dispatch(
+    size_t BlkLen,
+    const std::byte* QuantA,
+    const float* QuantAScale,
+    const std::byte* QuantBData,
+    const float* QuantBScale,
+    const std::byte* QuantBZeroPoint,
+    float* C,
+    size_t CountM,
+    size_t CountN,
+    size_t CountK,
+    size_t BlockCountK,
+    const float* Bias,
+    size_t ldc,
+    const float* ABlockSum,
+    const float* QuantBBlkSum
+);
+
+size_t MLASCALL
+SQ2BitGemmKernel_BlkSum_CompInt8_Avx2Vnni_Dispatch(
+    size_t BlkLen,
+    const std::byte* QuantA,
+    const float* QuantAScale,
+    const std::byte* QuantBData,
+    const float* QuantBScale,
+    const std::byte* QuantBZeroPoint,
+    float* C,
+    size_t CountM,
+    size_t CountN,
+    size_t CountK,
+    size_t BlockCountK,
+    const float* Bias,
+    size_t ldc,
+    const float* ABlockSum,
+    const float* QuantBBlkSum
+);
+
+}  // namespace sq2bit_avx2
+}  // namespace mlas
+}  // namespace onnxruntime