Threading support for Hybrid core architecture (#6728)

onnxruntime/core/common/cpuid_info.cc

@@ -68,6 +68,7 @@ CPUIDInfo::CPUIDInfo() noexcept {
         // Add check for AVX512 Skylake since tensorization GEMM need intrinsics from avx512bw/avx512dq.
         // avx512_skylake = avx512f | avx512vl | avx512cd | avx512bw | avx512dq
         has_avx512_skylake_ = has_avx512 && (data[1] & ((1 << 16) | (1 << 17) | (1 << 28) | (1 << 30) | (1 << 31)));
+        is_hybrid_ = (data[3] & (1 << 15));
       }
     }
   }

onnxruntime/core/common/cpuid_info.h

@@ -18,6 +18,7 @@ class CPUIDInfo {
   bool HasAVX512Skylake() const { return has_avx512_skylake_; }
   bool HasF16C() const { return has_f16c_; }
   bool HasSSE3() const { return has_sse3_; }
+  bool IsHybrid() const { return is_hybrid_; }
 
  private:
   CPUIDInfo() noexcept;
@@ -27,6 +28,7 @@ class CPUIDInfo {
   bool has_avx512_skylake_{false};
   bool has_f16c_{false};
   bool has_sse3_{false};
+  bool is_hybrid_{false};
 };
 
 }  // namespace onnxruntime

onnxruntime/core/common/threadpool.cc

@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "core/platform/threadpool.h"
 #include "core/common/common.h"
+#include "core/common/cpuid_info.h"
 #include "core/common/eigen_common_wrapper.h"
 #include "core/platform/EigenNonBlockingThreadPool.h"
 #include "core/platform/ort_mutex.h"
@@ -41,6 +42,10 @@ namespace concurrency {
 static constexpr int CACHE_LINE_BYTES = 64;
 static constexpr unsigned MAX_SHARDS = 8;
 
+#ifndef _OPENMP
+static constexpr int TaskGranularityFactor = 4;
+#endif
+
 struct alignas(CACHE_LINE_BYTES) LoopCounterShard {
   ::std::atomic<uint64_t> _next{0};
   uint64_t _end{0};
@@ -390,7 +395,15 @@ int ThreadPool::DegreeOfParallelism(const concurrency::ThreadPool* tp) {
 #else
   // When not using OpenMP, we parallelise over the N threads created by the pool
   // tp, plus 1 for the thread entering a loop.
-  return tp ? (tp->NumThreads()+1) : 1;
+  if (tp) {
+    if (CPUIDInfo::GetCPUIDInfo().IsHybrid()) {
+      return ((tp->NumThreads() + 1)) * TaskGranularityFactor;
+    } else {
+      return ((tp->NumThreads() + 1));
+    }
+  } else {
+    return 1;
+  }
 #endif
 }
 

onnxruntime/core/mlas/lib/dgemm.cpp

@@ -866,10 +866,10 @@ Return Value:
 
     int32_t TargetThreadCount;
 
-    if (Complexity < double(MLAS_DGEMM_THREAD_COMPLEXITY * MLAS_MAXIMUM_THREAD_COUNT)) {
+    if (Complexity < double(MLAS_DGEMM_THREAD_COMPLEXITY * MlasPlatform.MaximumThreadCount)) {
         TargetThreadCount = int32_t(Complexity / double(MLAS_DGEMM_THREAD_COMPLEXITY)) + 1;
     } else {
-        TargetThreadCount = MLAS_MAXIMUM_THREAD_COUNT;
+        TargetThreadCount = MlasPlatform.MaximumThreadCount;
     }
 
     int32_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool);

onnxruntime/core/mlas/lib/mlasi.h

@@ -796,6 +796,9 @@ struct MLAS_PLATFORM {
     PMLAS_QUANTIZE_LINEAR_U8_KERNEL QuantizeLinearU8Kernel;
     uint32_t NchwcBlockSize;
     uint32_t PreferredBufferAlignment;
+    uint32_t MaximumThreadCount;
+#else
+    static constexpr uint32_t MaximumThreadCount = MLAS_MAXIMUM_THREAD_COUNT;
 #endif
 };
 

onnxruntime/core/mlas/lib/platform.cpp

@@ -148,6 +148,8 @@ Return Value:
     this->NchwcBlockSize = 8;
     this->PreferredBufferAlignment = MLAS_DEFAULT_PREFERRED_BUFFER_ALIGNMENT;
 
+    this->MaximumThreadCount = MLAS_MAXIMUM_THREAD_COUNT;
+
 #endif
 
     //
@@ -227,6 +229,14 @@ Return Value:
                 this->ConvDepthwiseU8S8Kernel = MlasConvDepthwiseKernelAvx2<int8_t>;
                 this->ConvDepthwiseU8U8Kernel = MlasConvDepthwiseKernelAvx2<uint8_t>;
                 this->ComputeSumExpF32Kernel = MlasComputeSumExpF32KernelFma3;
+                
+                //
+                // Check if the processor supports Hybrid core architecture.
+                //
+
+                if ((Cpuid7[3] & 0x8000) != 0) {
+                    this->MaximumThreadCount = MLAS_MAXIMUM_THREAD_COUNT * 4;
+                }
 
                 //
                 // Check if the processor supports AVXVNNI features.

onnxruntime/core/mlas/lib/qgemm.cpp

@@ -2244,10 +2244,10 @@ Return Value:
 
     int32_t TargetThreadCount;
 
-    if (Complexity < double(MLAS_QGEMM_THREAD_COMPLEXITY * MLAS_MAXIMUM_THREAD_COUNT)) {
+    if (Complexity < double(MLAS_QGEMM_THREAD_COMPLEXITY * MlasPlatform.MaximumThreadCount)) {
         TargetThreadCount = int32_t(Complexity / double(MLAS_QGEMM_THREAD_COMPLEXITY)) + 1;
     } else {
-        TargetThreadCount = MLAS_MAXIMUM_THREAD_COUNT;
+        TargetThreadCount = MlasPlatform.MaximumThreadCount;
     }
 
     int32_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool);

onnxruntime/core/mlas/lib/sgemm.cpp

@@ -1333,10 +1333,10 @@ Return Value:
 
     int32_t TargetThreadCount;
 
-    if (Complexity < double(MLAS_SGEMM_THREAD_COMPLEXITY * MLAS_MAXIMUM_THREAD_COUNT)) {
+    if (Complexity < double(MLAS_SGEMM_THREAD_COMPLEXITY * MlasPlatform.MaximumThreadCount)) {
         TargetThreadCount = int32_t(Complexity / double(MLAS_SGEMM_THREAD_COMPLEXITY)) + 1;
     } else {
-        TargetThreadCount = MLAS_MAXIMUM_THREAD_COUNT;
+        TargetThreadCount = MlasPlatform.MaximumThreadCount;
     }
 
     int32_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool);

onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvxVnni.S

@@ -103,9 +103,9 @@ Implicit Arguments:
         EmitIfCountGE \RowCount\(), 3, "MultiplyAccumulateRow \ColumnCount\(), ymm8, ymm9"
         EmitIfCountGE \RowCount\(), 4, "vpbroadcastd ymm2,DWORD PTR [rbx+\BroadcastOffset\()]"
         EmitIfCountGE \RowCount\(), 4, "MultiplyAccumulateRow \ColumnCount\(), ymm10, ymm11"
-        EmitIfCountGE \RowCount\(), 5, "vpbroadcastd ymm2,DWORD PTR [rbx+r9+\BroadcastOffset\()]"
+        EmitIfCountGE \RowCount\(), 5, "vpbroadcastd ymm2,DWORD PTR [rbx+rcx+\BroadcastOffset\()]"
         EmitIfCountGE \RowCount\(), 5, "MultiplyAccumulateRow \ColumnCount\(), ymm12, ymm13"
-        EmitIfCountGE \RowCount\(), 6, "vpbroadcastd ymm2,DWORD PTR [rbx+r9*2+\BroadcastOffset\()]"
+        EmitIfCountGE \RowCount\(), 6, "vpbroadcastd ymm2,DWORD PTR [rbx+rcx*2+\BroadcastOffset\()]"
         EmitIfCountGE \RowCount\(), 6, "MultiplyAccumulateRow \ColumnCount\(), ymm14, ymm15"
 
         .endm

onnxruntime/core/providers/cpu/nn/qlinearconv.cc

@@ -3,6 +3,7 @@
 
 #include "core/framework/op_kernel.h"
 #include "core/providers/cpu/nn/conv_attributes.h"
+#include "core/common/cpuid_info.h"
 #include "core/common/safeint.h"
 #include "core/providers/common.h"
 #include "core/util/math.h"
@@ -360,7 +361,12 @@ Status QLinearConv::Compute(OpKernelContext* context) const {
 
   // Replicate the logic from MlasGemmU8X8Schedule to control the number of
   // worker threads used for the convolution.
-  constexpr int32_t maximum_thread_count = 16;
+  int32_t maximum_thread_count;
+  if (CPUIDInfo::GetCPUIDInfo().IsHybrid()) {
+    maximum_thread_count = 64;
+  } else {
+    maximum_thread_count = 16;
+  }
   constexpr double thread_complexity = static_cast<double>(64 * 1024);
 
   const double complexity = static_cast<double>(output_image_size) *