Как построить матрицу из выбранных столбцов трехмерного массива?

Question

У меня есть массив 3D GPU A с размерами K x M x N и int вектор v длины M и хотите построить 2D-массивы графического процессора вида

X = [A(:,1,v(1)), A(:,2,v(2)),..., A(:,M,v(M))] (в зависимости от v)

самым быстрым способом. Поскольку все это массивы графического процессора, мне было интересно, есть ли более быстрый способ сделать это, чем предварительное выделениеX и используя очевидное forпетля. Мой код должен вызывать несколько миллионов этих экземпляров, поэтому это становится узким местом. Типичные приказы величины были быK = 350 000, 2<=M<=15, N<=2000, если это важно.

EDIT: вот минимальная рабочая версия исходного кода узкого места, который я пытаюсь улучшить. Преобразование в 3D массивAбыл закомментирован. Настройте параметры размера массива по своему усмотрению.

% 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 позволяет индексировать сразу несколько измерений. Это позволяет задать вектор линейной индексацииh который индексирует и второе, и третье измерение одновременно:

% 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)

Дополнительная литература: линейное индексирование, логическое индексирование и все такое

Answer 2

Что касается кода, который я опубликовал выше, оказалось, что быстрее использовать 2D gpuAarray вместо 3D gpuArray вместо ячейки. Это обеспечивает очень простой выбор столбцов и векторизацию самого дальнего внутреннего цикла. Точнее:

% generate test data:
K = 4000; M = 2;

A_cell = cell(1,M); % this is given externally
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)),'logical'); % cell2mat doesn't work with cells of gpuArrays
end

% conversion of A_cell into an appropriate 2D array is straightforward:
A_hor = gpuArray(cell2mat(A_cell)); % horizontal concatenation of the cells

% bottleneck code starts here and has M = 2 nested loops:
I_01 = false(K,1,'gpuArray');
I_02 = false(K,1,'gpuArray');

t = [1,s]; t = t(1:M); % vector of the starting indices of the old matrices inside A_hor

for j_1=1:s(1)
    for j_2=1:s(2)

        j = [j_1,j_2];

        k = t-1+j; % vector of the positions of the needed columns

        I_02 = I_02 | all(A_hor(:,k),2);
    end

    I_01 = I_01 | I_02;
end

Out = gather(I_01);