gpuArray support for Optimization Toolbox solvers

Question

1 个投票

lsqnonlin_Simplified.m

I have wondered for a while why the OptimizationToolbox doesn't currently support gpuArray data types in its solvers (lsqnonlin, fmincon, etc...), so that the solver algorithms could run completely on the GPU. I happen to know that it is something the Mathworks has expressly ruled out as a worthwhile enhancement. In previous correspondence, my TMW contact said:

One of our developers recently worked on a Proof of Concept for a customer solving a large set of nonlinear equations using GPU arrays and the speed up was marginal, at best. Elsewhere, there is no evidence of GPU use by the usual competitors (Gurobi, CPLEX etc) and there seems to be similar conclusions in the open source world.

As a workaround, it is of course possible for the user-provided objective and constraint functions to send its work to the GPU. However, this workaround requries that the user-provided functions gather() the results back to the CPU, to ensure that non-gpuArray data gets handed back to the parent solver.

It seems abundantly obvious that forcing users to do this incurs two sources of overhead:

(1) CPU-GPU transfers must be done every iteration.

(2) Computations done in the Mathworks-provided portion of the algorithm are always done on the CPU where they cannot benefit from GPU-accelerated fast linear algebra routines.

I have attached a simplified version of lsqnonlin with full gpuArray compatibility, and demonstrated in the timing test below that overhead from (1) and (2) lead to about a 50% slow-down, even for modest-sized problems. This test was done on the NVIDIA RTX 5000 Ada Generation. Is there any reason why this is not a convincing example?

gd=gpuDevice(1); %
m = 5e4;
n = 50;
Atrue = randn(m,n,'single');
xtrue = randn(n,1,'single');
A = gpuArray(Atrue);
y = exp(-A*xtrue);  %Also always gpuArray
y = y + 0.01*randn(size(y),'like',y);
opts.MaxIter = 10;
opts.CGMaxIter = 50;
opts.CGTol = 1e-3;
opts=namedargs2cell(opts);
%%% GPU in objective function only
x0 = 0.1*xtrue;
fun=@(x)resFcn(x,A,y,1);
jacobmult = @(x,W,flag) expJacobianMultiply(A,x,W,flag, 1);
tic;
    [x,resnorm,residual,exitflag,output] =...
        lsqnonlin_Simplified(fun,x0,jacobmult,opts{:});
toc

Elapsed time is 0.304361 seconds.

%%% Fully GPU
x0=gpuArray(x0);
fun=@(x)resFcn(x,A,y,0);
jacobmult = @(x,W,flag) expJacobianMultiply(A,x,W,flag, 0);
tic
 [xhat,info] = lsqnonlin_Simplified(fun,x0,jacobmult,opts{:});
wait(gd)
toc

Elapsed time is 0.136295 seconds.

function [r,jminfo] = resFcn(x,A,y,doTransfer)
    if doTransfer
         x=gpuArray(x);
    end
    r = exp(-A*x) - y;
   jminfo=r.*x(:).';   %Fake placeholder
  
    if doTransfer
         r=gather(r);
         jminfo=gather(jminfo);
    end
end
function Z = expJacobianMultiply(A,x,W,flag,doTransfer)
    if doTransfer
        x = gpuArray(x);
        W = gpuArray(W);
    end
    f = exp(-A*x);
    switch flag
        case 'notransp'
            Z = -f .* (A*W);
        case 'transp'
            Z = -A' * (f .* W);
        otherwise
            error('Unknown flag.');
    end
    if doTransfer
        Z = gather(Z);
    end
end