Ruy: Permit GEMV code to thread if thread_count above 1

PiperOrigin-RevId: 280708178
Change-Id: I519ca15e4b52d2e4483dae63418ebfb7b6b8e31f
This commit is contained in:
T.J. Alumbaugh 2019-11-15 12:06:46 -08:00 committed by TensorFlower Gardener
parent 4608b40a30
commit 0c04770ad1

View File

@ -261,11 +261,13 @@ int GetThreadCount(Context* context, int rows, int cols, int depth) {
LoopStructure GetLoopStructure(int tentative_thread_count, int rows, int cols, LoopStructure GetLoopStructure(int tentative_thread_count, int rows, int cols,
int depth, int depth,
int cache_friendly_traversal_threshold) { int cache_friendly_traversal_threshold) {
if (cols == 1) { // Use a simple loop for the GEMV case. if (tentative_thread_count == 1) {
return LoopStructure::kSimple; // If we are in the GEMV case or the size is below the
} else if (tentative_thread_count == 1 && // threshold, stay with the simple loop structure.
(rows + cols) * depth < cache_friendly_traversal_threshold) { if ((cols == 1) ||
return LoopStructure::kSimple; (rows + cols) * depth < cache_friendly_traversal_threshold) {
return LoopStructure::kSimple;
}
} }
return LoopStructure::kGeneral; return LoopStructure::kGeneral;
} }