I have a 3D GPU array A
with dimensions K x M x N
and an int
vector v
of length M
and want to construct 2D GPU arrays of the form
X = [A(:,1,v(1)), A(:,2,v(2)),..., A(:,M,v(M))]
(depending on v
)
in the most time-efficient way. Since all these are GPU arrays, I was wondering if there is a faster way to accomplish this than pre-allocating X
and using the obvious for
loop. My code needs to invoke several millions of these instances, so this becomes quite the bottleneck. Typical oders of magnitude would be K = 350 000, 2<=M<=15, N<=2000
, if that matters.
EDIT: Here is a minimal working version of the original bottleneck code I am trying to improve. Conversion to the 3D array A
has been commented out. Adjust the array size parameters as you see fit.
% generate test data:
K = 4000; M = 2; % N = 100
A_cell = cell(1,M);
s = zeros(1,M,'uint16');
for m=1:M
s(m) = m*50; % defines some widths for the matrices in the cells
A_cell{m} = cast(randi([0 1],K,s(m),'gpuArray'),'logical');
end
N = max(s,[],2);
% % A_cell can be replaced by a 3D array A of dimensions K x M x N:
% A = true(K,M,N,'gpuArray');
% for m=1:M
% A(:,m,1:s(m)) = permute(A_cell{m},[1 3 2]);
% end
% bottleneck code starts here and has M = 2 nested loops:
I_true = true(K,1,'gpuArray');
I_01 = false(K,1,'gpuArray');
I_02 = false(K,1,'gpuArray');
for j_1=1:s(1)
for j_2=1:s(2)
v = [j_1,j_2];
I_tmp = I_true;
for m=1:M
I_tmp = I_tmp & A_cell{m}(:,v(m));
end
I_02 = I_02 | I_tmp;
end
I_01 = I_01 | I_02;
end
Out = gather(I_01);
% A_cell can be replaced by 3D array A above
MATLAB allows you to index multiple dimensions at once. This allows you to give a linear indexing vector h
which indexes both the second and third dimension at the same time:
% Some example data
k=2;
m=3;
n=4;
v=[2,3,1];
A=rand(k,m,n);
X=[A(:,1,v(1)),A(:,2,v(2)),A(:,3,v(3))]
%solution
h=sub2ind([m,n],[1:m],v);
Y=A(:,h)
Further reading: Linear indexing, logical indexing, and all that