function varargout = gpuDeviceCount(mode)
%GPUDEVICECOUNT Count GPU devices present in system
%   N = GPUDEVICECOUNT returns the number of GPU devices in your system as
%   reported by the device driver. gpuDeviceCount counts all detected
%   devices, including any unsupported devices or devices that are not
%   available for use in this MATLAB session.
%
%   N = GPUDEVICECOUNT(COUNTMODE) returns the number of GPU devices counted
%   according to the COUNTMODE input. COUNTMODE must be one of:
%      'all'       - (default) counts all GPU devices reported by the
%                    driver. 
%      'supported' - counts only supported GPU devices, i.e., those
%                    which have DeviceSupported property equal to true. 
%      'available' - counts only GPU devices that are supported and are
%                    currently available for use in this MATLAB session.
%
%   [N,IDX] = GPUDEVICECOUNT(...) also returns a vector of indices for
%   the counted devices.
%
%   % Example: Count and return supported devices
%   [n,indx] = gpuDeviceCount("supported");
%   for ii=1:n
%       device = parallel.gpu.GPUDevice.getDevice(indx(ii)); 
%       disp(device.Name);
%   end
%   % Available devices will be displayed
%   % Quadro P620
%   % Tesla K20c
%
%   See also gpuDevice, parallel.gpu.GPUDevice, parallel.gpu.CUDADevice.

%   Copyright 2010-2022 The MathWorks, Inc.

arguments
    mode (1,1) string = "all"
end

nargoutchk(0,2);

% Deal with partial case-insensitive matches for valid modes (error if not)
mode = mustBeValidCountMode(mode);

varargout = cell(1,max(1,nargout));
switch mode
    case "all"
        [varargout{:}] = iCountAll();
    case "available"
        [varargout{:}] = iCountAvailable();
    case "supported"
        [varargout{:}] = iCountSupported();
end
end

function mode = mustBeValidCountMode(mode)
allowedModes = ["all", "available", "supported"];
ii = startsWith(allowedModes, mode, "IgnoreCase", true);
if nnz(ii)~=1
    error(message("parallel:gpu:device:DeviceCountMode"));
end
mode = allowedModes(ii);
end

function [n, idx] = iCountAll()
% Count all devices that we can find, supported or not.

% This can't be allowed to fail. If CUDA isn't available or there's some
% other problem we catch it and return 0
try
    n = parallel.internal.gpudeviceinfo.gpuDeviceCount;
catch err %#ok<NASGU>
    n = 0;
end
if nargout>1
    idx = 1:n;
end
end

function [n,idx] = iCountAvailable()
% Count only devices where isAvailable==true
nGPUs = iCountAll();
availableGPUs = parallel.gpu.GPUDevice.isAvailable(1:nGPUs);
n = nnz(availableGPUs);
if nargout>1
    idx = find(availableGPUs);
end
end

function [n,idx] = iCountSupported()
% Count only devices where isSupported==true
nGPUs = iCountAll();
supportedGPUs = parallel.gpu.GPUDevice.isSupported(1:nGPUs);
n = nnz(supportedGPUs);
if nargout>1
    idx = find(supportedGPUs);
end
end
