这个包含不同索引的循环可以向量化或加速吗？

Question

我有一个代码对 3D 矩阵中的每个点进行一些处理。数组input_vec_1D 由不寻常的索引ind_prime 访问，该索引取决于循环变量（对于上下文，索引由我在this paper 的方程式42e 中使用的算法确定，我的完整代码是@ 987654322@)。我首先将矩阵转换为 1D 数组，计算正确的索引，进行处理，然后重新整形为 3D，从而使其正常工作：

Nx = 8; Ny = 6; Nz = 4; Ntot = Nx*Ny*Nz;                    % Number of points
xvals = rand(1,Nx); yvals = rand(1,Ny); zvals = rand(1,Nz); % Grid vectors

input_vec_3D = rand(Ny,Nx,Nz); % Dummy 3D array

factor1 = 3.6*xvals; % some constant times xvals
factor2 = 1.2*yvals;
factor3 = 8.5*zvals;

input_vec_1D = reshape( permute(input_vec_3D,[3,1,2]) , [Ntot 1]); % Reshape to 1D for loop
output_vec = zeros(Ntot,1);
for ind = 1:Ntot
    j1 = floor( floor( (ind-1)/Nz ) /Ny ) + 1;   
    j2 = mod( floor( (ind-1)/Nz ) , Ny ) + 1; 
    j3 = mod( (ind-1) , Nz ) + 1;
    n1 = mod( 5*(j1-1) ,Nx);
    n2 = mod( 3*(j2-1) ,Ny);
    n3 = mod( 2*(j3-1) ,Nz);
    ind_prime = mod( ( n3 + Nz*(n2 + Ny*n1) ) , Ntot ) + 1; % a different index for input_vec
    output_vec(ind) = output_vec(ind) + input_vec_1D(ind_prime) * factor1(j1)*factor2(j2)*factor3(j3);
end
output_vec = permute( reshape( output_vec, [Nz,Ny,Nx] ) , [2,3,1] );  % Reshape back to 3D

所有元素的循环是我代码中最慢的部分，所以我想加快它的速度——通过矢量化或其他方式。

我的数组通常是 512x512x1024 复数双精度数，因此对于我的应用程序来说至关重要的是，由于 RAM 有限（大约 6 GB），我没有存储任何临时超大矩阵，这排除了使用 meshgrid() 生成因子（请注意，factor1、factor2、factor3 只是一维向量，因此它们的内存使用量很小）。

我得到了一个非常相似的循环here 的帮助，在这种情况下使用 Matlab 的隐式扩展解决了这个问题。但是，这更复杂，因为在处理行中使用了不同的索引 ind_prime、ind 和 j。

Answer 1

您可以将其完全矢量化，这会更快（根据我对 512x512x10 输入矩阵的测试，大约 50%）。但这涉及到创建几个数组，我们可以通过两种方式减小它们的大小

1. 对索引使用整数数据类型（例如uint32）。 uint32 是每个索引 4 个字节，而 double 是 8 个字节，所以这是一个不错的节省，尤其是当索引始终是整数时。请注意，要充分利用这一点，您必须使用 idivide 而不是 ./ 以避免 MATLAB 在内部转换为 double，并且您必须将 Nx/Ny/Nz 转换为 uint32 相同原因。

   我们不能使用 uint8 或 uint16，因为它们的最大值太小，无法满足您的大型数组。

   还请注意（默认情况下）idivide 使用 fix 舍入，即向 0 舍入，因此您可以跳过使用 floor 并可能在那里弥补少量性能。

2. 回收您的阵列。除了使用j1..3和n1..3作为6个索引数组，我们可以稍微重新排序操作并回收j1..3。

这样组合在一起：

IND = uint32(1:Ntot); % Shorthand, could skip defining this and write each tiem if memory was tight
Nx = uint32(Nx); Ny = uint32(Ny); Nz = uint32(Nz); Ntot = Nx*Ny*Nz; % uint8 conversion
J1 = idivide( idivide(IND-1,Nz), Ny ) + 1; % idivide to avoid "double" casting, does "fix" rounding
J2 = mod( idivide( IND-1, Nz ), Ny ) + 1;  % idivide to avoid "double" casting, does "fix" rounding
J3 = mod( IND-1, Nz ) + 1;
FAC = (factor1(J1).*factor2(J2).*factor3(J3)).'; % done here so we can recycle J1..3
J1 = mod( 5*(J1-1), Nx );
J2 = mod( 3*(J2-1), Ny );
J3 = mod( 2*(J3-1), Nz );
ind_prime2 = (mod( ( J3 + Nz*(J2 + Ny*J1) ) , Ntot ) + 1).';
output_vec3 = input_vec_1D(ind_prime2) .* FAC;