parfor nested loop ind2sub

I have three nested loops and I would like to parallelize all of them using parfor. I know that it is not possible to nest a parfor inside another parfor so first I flatten the nested loop switching from subscripts to linear index. See the example code below:

clear
clc
close all
%% Set up fake data
alpha = 0.36;
delta = 0.06;
sigma = 1.5;
nk = 100;
nkp = 500;
nz = 7;
k_grid = linspace(1e-6,100,nk)';
kp_grid = linspace(1e-6,100,nkp)';
z_grid = linspace(0.5,1.8,nz)';
%% Original problem
disp('Original problem')

Original problem

tic
Umat = -inf(nkp,nk,nz);
for z_c=1:nz
    for k_c=1:nk
        for kp_c=1:nkp
            z  = z_grid(z_c);
            k  = k_grid(k_c);
            kp = kp_grid(kp_c);
            cons = z*k^alpha+(1-delta)*k-kp;
            if cons>0
                Umat(kp_c,k_c,z_c) = cons^(1-sigma)/(1-sigma);
            end
        end
    end
end
toc

Elapsed time is 0.068098 seconds.

%% Flatten the nested loops
disp('Flatten the nested loops')

Flatten the nested loops

siz = [nkp,nk,nz];
nn = prod(siz);
tic
Umat1 = -inf(nkp,nk,nz);
for n_c=1:nn % Cannot use parfor here!!
    [kp_c,k_c,z_c] = ind2sub(siz,n_c);
    z  = z_grid(z_c);
    k  = k_grid(k_c);
    kp = kp_grid(kp_c);
    cons = z*k^alpha+(1-delta)*k-kp;
    if cons>0
        Umat1(kp_c,k_c,z_c) = cons^(1-sigma)/(1-sigma);
    end
end
toc

Elapsed time is 0.702280 seconds.

err = max(abs(Umat1(:)-Umat(:)))

err = 0

The flattened loops are slower, as expected since I'm calling another function (ind2sub). However, if I replace the for with parfor I get a mistake and it says that "parfor cannot be used due to the way Umat1 is used". But the iterations are obviously independent, Indeed I have parallelize the above code in Fortran with OpenMP and it works well. Any help/suggestion is greatly appreciated!

Answer 1

[Walter Roberson]

You cannot do that. parfor requires that the indexing of output objects must involve a simple linear transformation of the loop control variable for one index, and constants or : for the other indices.

Use your ind2sub version but assign to Umat1(n_c)

[Alessandro Maria Marco]

Thanks!

[Alessandro Maria Marco]

Following your suggestion I rewrote the code as follows

siz = [nkp,nk,nz];
nn = prod(siz);
tic
Umat1 = -inf(nkp,nk,nz);
for n_c=1:nn % Cannot use parfor here!!
    [kp_c,k_c,z_c] = ind2sub(siz,n_c);
    z  = z_grid(z_c);
    k  = k_grid(k_c);
    kp = kp_grid(kp_c);
    cons = z*k^alpha+(1-delta)*k-kp;
    if cons>0
        Umat1(n_c) = cons^(1-sigma)/(1-sigma);
    end
end
toc

[Alessandro Maria Marco]

The code above works but I get a warning for

z_grid(z_c)

and says that "z_grid is a broadcast variable this might result in unnecessary overhead". I can eliminate the warning in this way, which allows me to get rid of the call to ind2sub:

siz = [nkp,nk,nz];
[kp_mat,k_mat,z_mat] = ndgrid(kp_grid,k_grid,z_grid);
tic
Umat3 = -inf(nkp,nk,nz);
parfor n_c=1:prod(siz)
    z  = z_mat(n_c);
    k  = k_mat(n_c);
    kp = kp_mat(n_c);
    cons = z*k^alpha+(1-delta)*k-kp;
    if cons>0
        Umat3(n_c) = cons^(1-sigma)/(1-sigma);
        %Umat2(n_c) = cons^(1-sigma)/(1-sigma);
    end
end
time_parallel=toc
err = max(abs(Umat3(:)-Umat(:)))

I'm trying to write the code following the best practices and I was wondering if it is worth to define the extra arrays kp_mat,k_mat,z_mat and access them with the linear index n_c. Any opinion on this is appreciated!