Add cmake options and fix bad ifdef

Author: jslap-ubi

cmake/CMakeLists.txt

@@ -104,6 +104,15 @@ option(onnxruntime_USE_PREINSTALLED_EIGEN "Use pre-installed EIGEN. Need to prov
 option(onnxruntime_BUILD_BENCHMARKS "Build ONNXRuntime micro-benchmarks" OFF)
 option(onnxruntime_USE_VSINPU "Build with VSINPU support" OFF)
 
+#  Usually, SIMD instructions in kernel compile fine, and we detect at run-time if they are supported
+#   But on some platforms, even compiling these SIMD specific instructions could fail.
+#   So we provide options to disable compiling these SIMD instructions.
+option(onnxruntime_DISABLE_SSE4 "Disable compiling kernel with SSE4 instructions" OFF)
+option(onnxruntime_DISABLE_AVX "Disable compiling kernel with AVX instructions" OFF)
+option(onnxruntime_DISABLE_AVX2 "Disable compiling kernel with AVX2 instructions" OFF)
+option(onnxruntime_DISABLE_AVX512 "Disable compiling kernel with AVX512 instructions" OFF)
+option(onnxruntime_DISABLE_AMX "Disable compiling kernel with AMX instructions" OFF)
+
 cmake_dependent_option(onnxruntime_USE_FLASH_ATTENTION "Build flash attention kernel for scaled dot product attention" ON "onnxruntime_USE_CUDA" OFF)
 cmake_dependent_option(onnxruntime_USE_LEAN_ATTENTION "Build lean attention kernel for scaled dot product attention" ON "onnxruntime_USE_CUDA; NOT WIN32" OFF)
 option(onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION "Build memory efficient attention kernel for scaled dot product attention" ON)
@@ -976,6 +985,26 @@ if (onnxruntime_FORCE_GENERIC_ALGORITHMS)
     add_compile_definitions(FORCE_GENERIC_ALGORITHMS)
 endif()
 
+if (onnxruntime_DISABLE_SSE4)
+  add_compile_definitions(DISABLE_SSE4)
+endif()
+
+if (onnxruntime_DISABLE_AVX)
+  add_compile_definitions(DISABLE_AVX)
+endif()
+
+if (onnxruntime_DISABLE_AVX2)
+  add_compile_definitions(DISABLE_AVX2)
+endif()
+
+if (onnxruntime_DISABLE_AVX512)
+  add_compile_definitions(DISABLE_AVX512)
+endif()
+
+if (onnxruntime_DISABLE_AMX)
+  add_compile_definitions(DISABLE_AMX)
+endif()
+
 if (onnxruntime_ENABLE_LAZY_TENSOR)
   # To support LazyTensor, ORT needs to call Python function from C/C++.
   # so onnxruntime_ENABLE_PYTHON is required.

cmake/onnxruntime_mlas.cmake

@@ -542,6 +542,11 @@ else()
         endif()
         set_source_files_properties(${mlas_platform_srcs_sse2} PROPERTIES COMPILE_FLAGS "-msse2")
 
+        set(mlas_platform_srcs_sse41
+          ${MLAS_SRC_DIR}/qgemm_kernel_sse41.cpp
+        )
+        set_source_files_properties(${mlas_platform_srcs_sse41} PROPERTIES COMPILE_FLAGS "-msse4.1")
+
         set(mlas_platform_srcs_avx
           ${MLAS_SRC_DIR}/x86_64/DgemmKernelAvx.S
           ${MLAS_SRC_DIR}/x86_64/SgemmKernelAvx.S
@@ -615,36 +620,65 @@ endif()
         )
         set_source_files_properties(${mlas_platform_srcs_avx512vnni} PROPERTIES COMPILE_FLAGS "-mfma -mavx512vnni -mavx512bw -mavx512dq -mavx512vl -mavx512f")
 
+        set(mlas_platform_srcs_amx
+          ${MLAS_SRC_DIR}/x86_64/QgemmU8S8KernelAmxCommon.S
+          ${MLAS_SRC_DIR}/qgemm_kernel_amx.cpp
+          ${MLAS_SRC_DIR}/x86_64/QgemmU8S8KernelAmx.S
+        )
+        set_source_files_properties(${MLAS_SRC_DIR}/qgemm_kernel_amx.cpp PROPERTIES COMPILE_FLAGS "-mavx2 -mavx512bw -mavx512dq -mavx512vl -mavx512f")
+        set_source_files_properties(${MLAS_SRC_DIR}/x86_64/QgemmU8S8KernelAmx.S PROPERTIES COMPILE_FLAGS "-mavx2 -mavx512bw -mavx512dq -mavx512vl -mavx512f")
+
         set(mlas_platform_srcs
           ${MLAS_SRC_DIR}/activate_fp16.cpp
           ${MLAS_SRC_DIR}/dwconv.cpp
           ${MLAS_SRC_DIR}/dgemm.cpp
           ${MLAS_SRC_DIR}/pooling_fp16.cpp
-          ${MLAS_SRC_DIR}/qgemm_kernel_avx2.cpp
           ${mlas_platform_srcs_sse2}
-          ${mlas_platform_srcs_avx}
-          ${mlas_platform_srcs_avx2}
-          ${mlas_platform_srcs_avx512f}
-          ${mlas_platform_srcs_avx512core}
-          ${mlas_platform_srcs_avx512vnni}
         )
 
-        if (NOT onnxruntime_ORT_MINIMAL_BUILD)
+        if (NOT onnxruntime_DISABLE_SSE4)
+          set(mlas_platform_srcs
+            ${mlas_platform_srcs}
+            ${mlas_platform_srcs_sse41}
+          )
+        endif()
+
+        if (NOT onnxruntime_DISABLE_AVX)
+          set(mlas_platform_srcs
+            ${mlas_platform_srcs}
+            ${mlas_platform_srcs_avx}
+          )
+        endif()
+
+        if (NOT onnxruntime_DISABLE_AVX2)
+          set(mlas_platform_srcs
+            ${mlas_platform_srcs}
+            ${mlas_platform_srcs_avx2}
+            ${MLAS_SRC_DIR}/qgemm_kernel_avx2.cpp
+          )
+        endif()
+
+        if (NOT onnxruntime_DISABLE_AVX512)
+          set(mlas_platform_srcs
+            ${mlas_platform_srcs}
+            ${mlas_platform_srcs_avx512f}
+            ${mlas_platform_srcs_avx512core}
+            ${mlas_platform_srcs_avx512vnni}
+          )
+        endif()
+
+        if (NOT onnxruntime_ORT_MINIMAL_BUILD AND NOT onnxruntime_DISABLE_AVX512)
           set(mlas_platform_srcs
             ${mlas_platform_srcs}
             ${MLAS_SRC_DIR}/q4gemm_avx512.cpp
           )
           set_source_files_properties(${MLAS_SRC_DIR}/q4gemm_avx512.cpp PROPERTIES COMPILE_FLAGS "-mfma -mavx512vnni -mavx512bw -mavx512dq -mavx512vl -mavx512f")
         endif()
-        if(NOT APPLE)
+        if(NOT APPLE AND NOT onnxruntime_DISABLE_AMX)
           set(mlas_platform_srcs
             ${mlas_platform_srcs}
-	        ${MLAS_SRC_DIR}/x86_64/QgemmU8S8KernelAmxCommon.S
-            ${MLAS_SRC_DIR}/qgemm_kernel_amx.cpp
-            ${MLAS_SRC_DIR}/x86_64/QgemmU8S8KernelAmx.S
-            )
-          set_source_files_properties(${MLAS_SRC_DIR}/qgemm_kernel_amx.cpp PROPERTIES COMPILE_FLAGS "-mavx2 -mavx512bw -mavx512dq -mavx512vl -mavx512f")
-          set_source_files_properties(${MLAS_SRC_DIR}/x86_64/QgemmU8S8KernelAmx.S PROPERTIES COMPILE_FLAGS "-mavx2 -mavx512bw -mavx512dq -mavx512vl -mavx512f")
+            ${mlas_platform_srcs_amx}
+          )
         endif()
 
         if(ONNXRUNTIME_MLAS_MULTI_ARCH)

onnxruntime/core/mlas/lib/platform.cpp

@@ -320,7 +320,7 @@ Return Value:
         this->GemmU8S8Dispatch = &MlasGemmU8S8DispatchSse41;
     }
 
-#endif
+#endif // defined(_MSC_VER) && !defined(ORT_DISABLE_SSE4)
 
     //
     // Check if the processor supports the AVX and OSXSAVE features.
@@ -502,7 +502,7 @@ Return Value:
                     this->GemmS8U8Kernel = MlasGemmS8U8KernelAvx2Vnni;
                 }
 
-#if !defined(__APPLE__) && !defined(ORT_DISABLE_AMX)
+#if !defined(__APPLE__)
 #if (defined(_MSC_VER) && (_MSC_VER >= 1933)) || (defined(__GNUC__) && (__GNUC__ >= 13))
                 //
                 // Check if the processor supports AVX NE CONVERT.
@@ -513,6 +513,7 @@ Return Value:
 #endif  // (defined(_MSC_VER) && (_MSC_VER >= 1933)) || (defined(__GNUC__) && (__GNUC__ >= 13))
 
 
+#if !defined(ORT_DISABLE_AMX)
                 //
                 // Check if the processor supports AMX-TILE and AMX-INT8
                 // features.
@@ -525,7 +526,8 @@ Return Value:
                         this->GemmU8S8Dispatch = &MlasGemmU8S8DispatchAmx;
                     }
                 }
-#endif // !defined(__APPLE__) && !defined(ORT_DISABLE_AMX)
+#endif // !defined(ORT_DISABLE_AMX)
+#endif // !defined(__APPLE__)
 
 #endif // ORT_MINIMAL_BUILD