Simultaneous sortrows() for all slices of a 3D matrix?

I need to apply sortrows to every page of a 3D matrix and I was wondering if there is a faster way of doing this than simply looping through all pages. I have the option to execute on a GPU. While I am aware that sortrows also works with GPU arrays, the code appears to execute much faster on the CPU than on the GPU, probably because of the very long loop.

Here is a minimal working example:

% Parameters:

N = 20;

n = 40;

m = 2;

d = 200000;

% Generate an example matrix 'A':

A = randi(N,n,m,d,'uint16'); % generate a 3D matrix with above parameters

% To be vectorized?

tic;

B = zeros(n,d,'double');

for i=1:d

[~,b] = sortrows(A(:,:,i),[1 2]);

B(:,i) = b;

end

toc;

[Bruno Luong]

If N is small (e.g, N < 256) and m is small (m <= 8), one can combine rows of A in single number and use sort instead of sortrows.

[Marin Genov]

Dear @Bruno Luong, yes, N is definitely much smaller than 256 and m is fixed equal 2. What would the sort solution look like?

[Bruno Luong]

I provide the answer bellow.

Answer 1

[Bruno Luong]

Code applicable for n == 2 and N < 256 only:

% Parameters:
N = 20;
n = 40;
m = 2;
d = 200000;
% Generate an example matrix 'A':
A = randi(N,n,m,d,'uint16'); % generate a 3D matrix with above parameters

% Original method: To be vectorized?
tic;
B = zeros(n,d,'double');
for i=1:d
    [~,b] = sortrows(A(:,:,i),[1 2]);
    B(:,i) = b;
end
toc

Elapsed time is 1.019606 seconds.

% Code applicable for n == 2 and N < 256 only
tic
[m,n,p] = size(A);
% Uncomment this sanity test if you are not sure A meets the requirement 
% if n ~= 2 && ~all(A<256,'all')
%     error('Method 3 not applicable');
% end
C = A(:,1,:)*feval(class(A),2^8) + A(:,2,:);
[~,B3] = sort(C,1);
B3 = reshape(B3,m,d);
toc

Elapsed time is 0.129230 seconds.

isequal(B,B3)

ans = logical
   1

[Marin Genov]

Thanks for posting this method as well! This helps a lot!

Answer 2

[Bruno Luong]

Somewhat faster

% Parameters:
N = 20;
n = 40;
m = 2;
d = 200000;
% Generate an example matrix 'A':
A = randi(N,n,m,d,'uint16'); % generate a 3D matrix with above parameters

% Original method: To be vectorized?
tic;
B = zeros(n,d,'double');
for i=1:d
    [~,b] = sortrows(A(:,:,i),[1 2]);
    B(:,i) = b;
end
toc

Elapsed time is 1.090424 seconds.

% No loop but a big permute
tic
[m,n,p] = size(A);
q = repelem((1:p)',m,1);
A3 = reshape(permute(A,[1 3 2]),m*p,n);
A3 = [q,double(A3)];
[~,B2] = sortrows(A3);
B2 = reshape(B2,m,d)-(0:d-1)*m;
toc

Elapsed time is 0.667823 seconds.

isequal(B,B2)

ans = logical
   1

[Marin Genov]

Thanks, this looks great!