Matlab依赖索引网格的循环向量化优化：加速三维数组最大化计算

需要对三维数组V(e,b,d)沿第一维度执行最大化操作，得到b×d维度的矩阵e_opt，定义如下：
e_opt(b,d) = argmaxₑ {V(e,b,d)}

特殊点在于e的网格范围依赖于d，因此无法消除d维度的循环。现有代码单次运行约1秒，但因需多次执行，速度无法满足需求，寻求加速方案。以下为最小工作示例（MWE）：

clear
close all
clc
rng("default")

% Size of grids
n_b = 1000;
n_d = 1000;
n_e = 5000;

% Parameters
delta   = 0.2696037674949296; 
%omega_0 = 0.0015534835229589;
%omega_1 = 0.0013023827114628;

% Generate fake data
cost_mat = rand(n_e,n_d);
prob_mat = rand(n_e,n_d);
option_nob = rand(n_b,1);
d_grid = linspace(0,2.3,n_d)';

%% Create a grid for "e". The non-standard feature is that the upper bound
% of the e_grid depends on "d"
e_grid_mat = zeros(n_e,n_d);
e_min = 0;
space = 1.5;
for d_c=1:n_d
    d_val = d_grid(d_c);
    e_max = 0.1+0.2*d_val; %silly example, don't take this literally
    % e_grid is NOT equally spaced
    e_grid_mat(:,d_c) = e_min+(e_max-e_min)*(linspace(0,1,n_e).^space)';
end

%% The code below is the part that I'd like to speed up
tic
e_opt = zeros(n_b,n_d);
for d_c = 1:n_d
    % Effort grid depends on d (upper bound changes, but same no. of elements)
    e_grid = e_grid_mat(:,d_c); %each column of e_grid_mat is different!
    % V(e,b) has dim: (n_e,n_b) and I maximize with respect to the first dimension, e
    V = -cost_mat(:,d_c)-delta*prob_mat(:,d_c).*option_nob';
    [~,max_ind] = max(V,[],1); %maxind is (1,n_b) vector
    e_opt(:,d_c)  = e_grid(max_ind);
end %end d
toc

%To check results
disp(mean(mean(e_opt)))

1. 优化核心运算的向量化与预计算 #

原循环内重复计算V矩阵时，可利用矩阵乘法替代逐元素广播（MATLAB对BLAS优化的矩阵乘法效率更高），同时预计算固定部分减少重复运算：

% 预计算-cost_mat，避免循环内重复取负
pre_cost = -cost_mat;

tic
e_opt = zeros(n_b, n_d);
for d_c = 1:n_d
    e_grid = e_grid_mat(:,d_c);
    % 用矩阵乘法替代广播，提升运算速度
    V = pre_cost(:,d_c) - delta * prob_mat(:,d_c) * option_nob';
    [~,max_ind] = max(V,[],1);
    e_opt(:,d_c) = e_grid(max_ind);
end
toc

2. 并行化循环（Parfor） #

由于各d_c对应的运算完全独立，可使用parfor开启多核心并行计算：

% 启动并行池（需提前配置MATLAB并行工具箱）
parpool;

tic
e_opt = zeros(n_b, n_d);
parfor d_c = 1:n_d
    e_grid = e_grid_mat(:,d_c);
    V = -cost_mat(:,d_c) - delta * prob_mat(:,d_c) * option_nob';
    [~,max_ind] = max(V,[],1);
    e_opt(:,d_c) = e_grid(max_ind);
end
toc

% 关闭并行池
delete(gcp);

3. GPU加速（若有可用GPU） #

将数组转移至GPU执行运算，MATLAB会自动利用GPU的并行能力加速矩阵操作：

% 将所有数组转移到GPU
cost_mat_gpu = gpuArray(cost_mat);
prob_mat_gpu = gpuArray(prob_mat);
option_nob_gpu = gpuArray(option_nob);
e_grid_mat_gpu = gpuArray(e_grid_mat);
e_opt_gpu = gpuArray.zeros(n_b, n_d);
delta_gpu = gpuArray(delta);

tic
for d_c = 1:n_d
    e_grid = e_grid_mat_gpu(:,d_c);
    V = -cost_mat_gpu(:,d_c) - delta_gpu * prob_mat_gpu(:,d_c) * option_nob_gpu';
    [~,max_ind] = max(V,[],1);
    e_opt_gpu(:,d_c) = e_grid(max_ind);
end
toc

% 将结果转回CPU
e_opt = gather(e_opt_gpu);

4. 优化e_grid_mat的生成 #

原代码中每次循环都计算linspace(0,1,n_e).^space，可提前预计算该基础序列，减少重复运算：

e_base = linspace(0,1,n_e).^space;
e_grid_mat = zeros(n_e,n_d);
e_min = 0;
space = 1.5;
for d_c=1:n_d
    d_val = d_grid(d_c);
    e_max = 0.1+0.2*d_val;
    e_grid_mat(:,d_c) = e_min + (e_max - e_min) * e_base;
end

内容的提问来源于stack exchange，提问作者Alessandro