gpuArray support for Optimization Toolbox solvers

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

Answer 1

[Catalytic]

Convincing yes if the optimization loop in MATLAB's lsqnonlin were MCoded, like in your example, but it is not. It is coded in optimized binaries, maybe even from 3rd party soruces. To achieve the enhancement you're looking for, and for it to be competitive with the current solvers at "modest" problem sizes, the developers would have to develop and integrate new binaries for gpuArray inputs into the Optimization Toolbox.

Their calculation might be that it isn't worth the effort as long as the current lsqnonlin can outperform MCoded alternatives like your lsqnonllin_Simplified.The following test shows that for problem sizes n<10000, that is the case --

Results =compare_times(1e4,[5e1, 5e2, 1000, 5e3, 7e3, 1e4])

gd =

CUDADevice with properties:

Name: 'NVIDIA GeForce RTX 4080 SUPER'

Index: 1 (of 2)

ComputeCapability: '8.9'

DriverModel: 'WDDM'

TotalMemory: 17170956288 (17.17 GB)

AvailableMemory: 15760584704 (15.76 GB)

Results =

6×3 table

Problem Size Fully GPU (MCode) Native Lsqnonlin

____________ _________________ ________________

50 0.27738 0.12965

500 0.046878 0.028637

1000 0.07907 0.041182

5000 0.74005 0.4737

7000 1.4157 1.1938

10000 2.2006 3.3461

Of course, this also makes the case for abandoning the native lsqnonlin when for n>10000. But maybe MathWorks considers that problem size unrealistically large...

function res=compare_times(m,N)

    gd=gpuDevice(1)

    for i=1:length(N)
    
        n=N(i);

    Atrue = 0.001*rand(m,n);
    xtrue = randn(n,1);
    
    A = gpuArray(Atrue);
    
    y = exp(-A*xtrue);  %Also always gpuArray
    y = y + 0.01*randn(size(y),'like',y);
    
    
    
    opts.MaxIter = 5;
    opts.CGMaxIter = 50;
    opts.CGTol = 1e-3;
    optsCell=namedargs2cell(opts);
    
    x0=0.1*xtrue;
    
    
    %%% 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,optsCell{:});
    wait(gd)
    t1=toc;
    
    
    %%% Built-in lsqnonlin, GPU objective only via JacobianMultiplyFcn
    
    x0 = gather(x0);   % CPU initial point
    
    fun_builtin = @(x) resFcn(x,A,y,1);
    
    jacobmult_builtin = @(Jinfo,W,flag) expJacobianMultiplyBuiltin(Jinfo,W,flag);
    
    opts_builtin = optimoptions('lsqnonlin', ...
        'Algorithm','trust-region-reflective', ...
        'SpecifyObjectiveGradient',true, ...
        'JacobianMultiplyFcn',jacobmult_builtin, ...
        'MaxIterations',opts.MaxIter, ...
        'MaxPCGIter',opts.CGMaxIter, ...
        'TolPCG',opts.CGTol, ...
        'Display','off');
    
    tic
    [x_builtin,resnorm_builtin,residual_builtin,exitflag_builtin,output_builtin] = ...
        lsqnonlin(fun_builtin,x0,[],[],opts_builtin);
    wait(gd)
    t2=toc;

      res{i}=table(n, t1,t2,Var={'Problem Size', 'Fully GPU (MCode)','Native Lsqnonlin'});

    end

    res=vertcat(res{:});

end


function [r,jminfo] = resFcn(x,A,y,doTransfer)

    if doTransfer
         x=gpuArray(x);
    end

    r = exp(-A*x) - y;


   jminfo.x=x;
   jminfo.A=A;
   jminfo.doTransfer=doTransfer;

    if doTransfer
         r=gather(r);
    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


function W = expJacobianMultiplyBuiltin(Jinfo,Y,flag)

    A = Jinfo.A;
    x = Jinfo.x;
    doTransfer = Jinfo.doTransfer;

    if doTransfer
        x = gpuArray(x);
        Y = gpuArray(Y);
    end

    f = exp(-A*x);

    if flag == 0
        % W = J'*(J*Y), Y is n-by-2
        JY = -f .* (A*Y);
        W  = -A' * (f .* JY);

    elseif flag > 0
        % W = J*Y, Y is n-by-1
        W = -f .* (A*Y);

    else
        % W = J'*Y, Y is m-by-1
        W = -A' * (f .* Y);
    end

    if doTransfer
        W = gather(W);
    end

end