使用 SSE 和 AVX 查找矩阵中的最大元素及其列和行索引
Posted
技术标签:
【中文标题】使用 SSE 和 AVX 查找矩阵中的最大元素及其列和行索引【英文标题】：Find largest element in matrix and its column and row indexes using SSE and AVX 【发布时间】：2015-11-21 14:29:39 【问题描述】：我需要找到一维矩阵中最大的元素及其列和行索引。
我使用一维矩阵，所以首先需要找到最大元素的索引，然后很容易得到行和列。
我的问题是我无法获得该索引。
我有一个可以找到最大元素并使用 SSE 的工作函数，这里是：
float find_largest_element_in_matrix_SSE(float* m, unsigned const int dims)
size_t i;
int index = -1;
__m128 max_el = _mm_loadu_ps(m);
__m128 curr;
for (i = 4; i < dims * dims; i += 4)
curr = _mm_loadu_ps(m + i);
max_el = _mm_max_ps(max_el, curr);
__declspec(align(16))float max_v[4] = 0 ;
_mm_store_ps(max_v, max_el);
return max(max(max(max_v[0], max_v[1]), max_v[2]), max_v[3]);
而且我还有一个使用 AVX 的非工作功能：
float find_largest_element_in_matrix_AVX(float* m, unsigned const int dims)
size_t i;
int index = -1;
__m256 max_el = _mm256_loadu_ps(m);
__m256 curr;
for (i = 8; i < dims * dims; i += 8)
curr = _mm256_loadu_ps(m + i);
max_el = _mm256_max_ps(max_el, curr);
__declspec(align(32))float max_v[8] = 0 ;
_mm256_store_ps(max_v, max_el);
__m256 y = _mm256_permute2f128_ps(max_el, max_el, 1);
__m256 m1 = _mm256_max_ps(max_el, y);m1[1] = max(max_el[1], max_el[3])
__m256 m2 = _mm256_permute_ps(m1, 5);
__m256 m_res = _mm256_max_ps(m1, m2);
return m[0];
谁能帮我找到最大元素的索引并让我的 AVX 版本工作？
【问题讨论】：
我没有在您的 AVX 中查看 et，但您的 SSE 功能存在问题。这可能是您的 AVX 的相同原因： for (i = 4; i @user3545806 保证dims 始终是 8 的倍数。
除了 return max(max(max(max_v[0], max_v[1]), max(max_v[2],max_v[3] )), .. .);在 _m256 y 行之前？
@user3545806 不，我在此处发布的代码与我在计算机上运行的代码完全相同。
要使查找最大元素 AVX 函数工作（返回最大元素的值），您需要将最后 5 行替换为 return max(max(max(max_v[0], max_v[1] ), max(max_v[2], max_v[3])), max(max(max_v[4], max_v[5]), max(max_v[6], max_v[7])));查找索引比较棘手，因为您需要进入数组并再次查找它。根据我过去的经验，如果您编写的代码没有显式 AVX/SSE，您可能会获得更快的代码，因为编译器可以为您优化 ti。
【参考方案1】：
这是一个有效的 SSE (SSE 4) 实现，它返回最大 val 和相应的索引，以及一个标量参考实现和测试工具：
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <time.h>
#include <smmintrin.h> // SSE 4.1
float find_largest_element_in_matrix_ref(const float* m, int dims, int *maxIndex)
float maxVal = m[0];
int i;
*maxIndex = 0;
for (i = 1; i < dims * dims; ++i)
if (m[i] > maxVal)
maxVal = m[i];
*maxIndex = i;
return maxVal;
float find_largest_element_in_matrix_SSE(const float* m, int dims, int *maxIndex)
float maxVal = m[0];
float aMaxVal[4];
int32_t aMaxIndex[4];
int i;
*maxIndex = 0;
const __m128i vIndexInc = _mm_set1_epi32(4);
__m128i vMaxIndex = _mm_setr_epi32(0, 1, 2, 3);
__m128i vIndex = vMaxIndex;
__m128 vMaxVal = _mm_loadu_ps(m);
for (i = 4; i < dims * dims; i += 4)
__m128 v = _mm_loadu_ps(&m[i]);
__m128 vcmp = _mm_cmpgt_ps(v, vMaxVal);
vIndex = _mm_add_epi32(vIndex, vIndexInc);
vMaxVal = _mm_max_ps(vMaxVal, v);
vMaxIndex = _mm_blendv_epi8(vMaxIndex, vIndex, _mm_castps_si128(vcmp));
_mm_storeu_ps(aMaxVal, vMaxVal);
_mm_storeu_si128((__m128i *)aMaxIndex, vMaxIndex);
maxVal = aMaxVal[0];
*maxIndex = aMaxIndex[0];
for (i = 1; i < 4; ++i)
if (aMaxVal[i] > maxVal)
maxVal = aMaxVal[i];
*maxIndex = aMaxIndex[i];
return maxVal;
int main()
const int dims = 1024;
float m[dims * dims];
float maxVal_ref, maxVal_SSE;
int maxIndex_ref = -1, maxIndex_SSE = -1;
int i;
srand(time(NULL));
for (i = 0; i < dims * dims; ++i)
m[i] = (float)rand() / RAND_MAX;
maxVal_ref = find_largest_element_in_matrix_ref(m, dims, &maxIndex_ref);
maxVal_SSE = find_largest_element_in_matrix_SSE(m, dims, &maxIndex_SSE);
if (maxVal_ref == maxVal_SSE && maxIndex_ref == maxIndex_SSE)
printf("PASS: maxVal = %f, maxIndex = %d\n",
maxVal_ref, maxIndex_ref);
else
printf("FAIL: maxVal_ref = %f, maxVal_SSE = %f, maxIndex_ref = %d, maxIndex_SSE = %d\n",
maxVal_ref, maxVal_SSE, maxIndex_ref, maxIndex_SSE);
return 0;
编译运行：
$ gcc -Wall -msse4 Yakovenko.c && ./a.out
PASS: maxVal = 0.999999, maxIndex = 120409
显然，如果需要，您可以获取行和列索引：
int rowIndex = maxIndex / dims;
int colIndex = maxIndex % dims;
从这里开始编写 AVX2 实现应该相当简单。
【讨论】：
@stgatilov：谢谢 - 你可能是对的 - 我没有收到关于 clang 的警告，但其他编译器可能会抱怨 - 你是否尝试使用特定的编译器并看到警告或错误？ 是的，MSVC2013 产生错误：_mm_blendv_epi8 不能接受 __m128 参数。
是的，它有固定的编译。
_mm_blendv_ps 可以替换为_mm_max_ps。
最近关于同一问题的问答（最后没有矩阵 2D 索引分解）：Efficient C vectors for generic SIMD (SSE, AVX, NEON) test for zero matches. (find FP max absolute value and index) 一些答案是使用英特尔内在函数，而不是通用的，与您拥有的基本相同的内部循环现在。【参考方案2】：
一种方法是在第一遍中计算最大值，并在第二遍中通过线性搜索找到索引。以下是 SSE2 中的示例实现：
#define anybit __builtin_ctz //or lookup table with 16 entries...
float find_largest_element_in_matrix_SSE(const float* m, int dims, int *maxIndex)
//first pass: calculate maximum as usual
__m128 vMaxVal = _mm_loadu_ps(m);
for (int i = 4; i < dims * dims; i += 4)
vMaxVal = _mm_max_ps(vMaxVal, _mm_loadu_ps(&m[i]));
//perform in-register reduction
vMaxVal = _mm_max_ps(vMaxVal, _mm_shuffle_ps(vMaxVal, vMaxVal, _MM_SHUFFLE(2, 3, 0, 1)));
vMaxVal = _mm_max_ps(vMaxVal, _mm_shuffle_ps(vMaxVal, vMaxVal, _MM_SHUFFLE(1, 0, 3, 2)));
//second pass: search for maximal value
for (int i = 0; i < dims * dims; i += 4)
__m128 vIsMax = _mm_cmpeq_ps(vMaxVal, _mm_loadu_ps(&m[i]));
if (int mask = _mm_movemask_ps(vIsMax))
*maxIndex = i + anybit(mask);
return _mm_cvtss_f32(vMaxVal);
请注意，除非您的输入数据非常小，否则第二个循环中的分支应该几乎可以完美预测。
该解决方案存在几个问题，特别是：
它可能会在出现奇怪的浮点值时无法正常工作，例如使用 NaN。
如果您的矩阵不适合 CPU 缓存，那么代码将从主内存中读取矩阵两次，因此它会比单次通过方法慢两倍。这可以通过分块处理来解决大型矩阵。
在第一个循环中，每个迭代都依赖于前一个（vMaxVal 被修改和读取），所以它会因_mm_max_ps 的延迟而减慢。也许将第一个循环展开一点（2x 或 4x）会很好，同时为vMaxVal 提供 4 个独立的寄存器（实际上，第二个循环也会受益于展开）。
移植到 AVX 应该非常简单，除了寄存器内减少：
vMaxVal = _mm256_max_ps(vMaxVal, _mm256_shuffle_ps(vMaxVal, vMaxVal, _MM_SHUFFLE(2, 3, 0, 1)));
vMaxVal = _mm256_max_ps(vMaxVal, _mm256_shuffle_ps(vMaxVal, vMaxVal, _MM_SHUFFLE(1, 0, 3, 2)));
vMaxVal = _mm256_max_ps(vMaxVal, _mm256_permute2f128_ps(vMaxVal, vMaxVal, 1));
【讨论】：
【参考方案3】：另一种方法：
void find_largest_element_in_matrix_SSE(float * matrix, size_t n, int * row, int * column, float * v)
__m128 indecies = _mm_setr_ps(0, 1, 2, 3);
__m128 update = _mm_setr_ps(4, 4, 4, 4);
__m128 max_indecies = _mm_setr_ps(0, 1, 2, 3);
__m128 max = _mm_load_ps(matrix);
for (int i = 4; i < n * n; i+=4)
indecies = _mm_add_ps(indecies, update);
__m128 pm2 = _mm_load_ps(&matrix[i]);
__m128 mask = _mm_cmpge_ps(max, pm2);
max = _mm_max_ps(max, pm2);
max_indecies = _mm_or_ps(_mm_and_ps(max_indecies, mask), _mm_andnot_ps(mask, indecies));
__declspec (align(16)) int max_ind[4];
__m128i maxi = _mm_cvtps_epi32(max_indecies);
_mm_store_si128((__m128i *) max_ind, maxi);
int c = max_ind[0];
for (int i = 1; i < 4; i++)
if (matrix[max_ind[i]] >= matrix[c] && max_ind[i] < c)
c = max_ind[i];
*v = matrix[c];
*row = c / n;
*column = c % n;
void find_largest_element_in_matrix_AVX(float * matrix, size_t n, int * row, int * column, float * v)
__m256 indecies = _mm256_setr_ps(0, 1, 2, 3, 4, 5, 6, 7);
__m256 update = _mm256_setr_ps(8, 8, 8, 8, 8, 8, 8, 8);
__m256 max_indecies = _mm256_setr_ps(0, 1, 2, 3, 4, 5, 6, 7);
__m256 max = _mm256_load_ps(matrix);
for (int i = 8; i < n * n; i += 8)
indecies = _mm256_add_ps(indecies, update);
__m256 pm2 = _mm256_load_ps(&matrix[i]);
__m256 mask = _mm256_cmp_ps(max, pm2, _CMP_GE_OQ);
max = _mm256_max_ps(max, pm2);
max_indecies = _mm256_or_ps(_mm256_and_ps(max_indecies, mask), _mm256_andnot_ps(mask, indecies));
__declspec (align(32)) int max_ind[8];
__m256i maxi = _mm256_cvtps_epi32(max_indecies);
_mm256_store_si256((__m256i *) max_ind, maxi);
int c = max_ind[0];
for (int i = 1; i < 8; i++)
if (matrix[max_ind[i]] >= matrix[c] && max_ind[i] < c)
c = max_ind[i];
*v = matrix[c];
*row = c / n;
*column = c % n;
【讨论】：
这主要是@PaulR 的解决方案，除了两个更改：1) 一个blendv 被max 替换，2) 另一个blendv 被分解为三个更简单的指令（事实上，这就是 blendv 当前在处理器中的实现方式）。
_mm_add_ps 的索引非常糟糕。对整数数据使用_mm_add_epi32；具有较低的延迟，并且不需要转换。并且工作到 2^32-1，而不是在 2^24 之后四舍五入。使用_mm_castsi128_ps 将其与SSE4 _mm_blendv_ps 一起使用，或使用_mm_castps_si128 将FP 比较结果与整数内容一起使用。以上是关于使用 SSE 和 AVX 查找矩阵中的最大元素及其列和行索引的主要内容，如果未能解决你的问题，请参考以下文章