How to construct a matrix from selected columns of a 3D array?

Question

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

Answer 1

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