function deviceIndex = selectGPU(nodeInfo)
%SELECTGPU Select a GPU device for the present worker to use
%   deviceIndex = selectGPU(nodeInfo) returns in deviceIndex the one-based 
%   index of the GPU to be selected based on the information in nodeInfo.
%   The deviceIndex will be empty if no device is available. 
%
%   The nodeInfo argument is an instance of the parallel.gpu.NodeInfo class
%   that stores information about the GPUs and workers on this node. The
%   information in nodeInfo is collected at the beginning of a MATLAB
%   session and does not change. If the GPU device configuration changes
%   during a MATLAB session, selection might not work as expected.
%
%   Normally, you should not need to call this function: it is called by
%   the GPU infrastructure to determine the proper GPU to use. To supply an
%   alternative GPU selection method, create a MATLAB file called
%   'selectGPU.m' and place it higher on the MATLAB path.
% 
%   If you write your own version of this function, keep in mind the
%   following guidelines.
%     * Your implementation should return a one-based index. It is 
%       acceptable to return an empty array ([]) if no device is available 
%       or no device should be used.
%     * Your implementation should not actually create any GPU data or
%       reserve a GPU device.
%     * Consider using parallel.gpu.GPUDevice.isAvailable to determine
%       whether a device is available without selecting it.
%     * The implementation we provide does not perform collective 
%       operations that communicate between workers. Collective 
%       operations would require that each worker call this function at
%       the same time, which would cause a synchronization point.
%     * Before we call this function, we verify that at least one device
%       exists and is in either default compute mode or exclusive compute
%       mode.
%     * In your implementation, if multiple workers try to use the same
%       exclusive-mode GPU device, only one of them will be able to use it.
%       The implementation we provide allows only one worker to try to use
%       each exclusive-mode device.
%
%   Example:
%     function deviceIndex = selectGPU(nodeInfo)
%         % This function assigns each worker to use one of the GPU devices
%         % that are in 'default' compute mode in a round-robin fashion,
%         % that is, worker 1 uses device 1, worker 2 device 2, and so on.
%         % Devices in 'exclusive' compute mode are ignored.
%         deviceIndex = [];
%         defaultDevices = nodeInfo.DefaultModeDevices;
%         myWorkerIndex = nodeInfo.MyWorkerIndex;
%
%         nD = numel(defaultDevices);
%         if nD > 0
%             % Index into the list of default mode devices using
%             % my worker index. If there are more workers than devices,
%             % multiple workers share devices. 
%             deviceIndex = defaultDevices(mod(myWorkerIndex-1, nD) + 1);
%         end
%     end
%
%   See also parallel.gpu.NodeInfo, parallel.gpu.GPUDevice.isAvailable.

% Copyright 2012-2013 The MathWorks, Inc.

exclusiveDevices = nodeInfo.ExclusiveModeDevices;
defaultDevices = nodeInfo.DefaultModeDevices;
myWorkerIndex = nodeInfo.MyWorkerIndex;
numWorkers = nodeInfo.NumNodeWorkers;

nE = numel(exclusiveDevices);
nD = numel(defaultDevices);

if myWorkerIndex <= nE
    % In this implementation, one and only one worker will try to access a
    % given exclusive-mode device: the first nE workers will each try one
    % exclusive-mode device. This minimizes the possibility of MATLAB
    % causing race conditions. If any of the exclusive mode devices is not
    % available, use a default mode device instead, or error if there are
    % none.
    
    if numWorkers == 1
        % If I'm the only worker on the node, no race conditions are 
        % possible: try all the devices.
        idx = 0;
        deviceFound = false;
        while ~deviceFound && (idx < nE)
            idx = idx+1;
            deviceIndex = exclusiveDevices(idx);
            deviceFound = parallel.gpu.GPUDevice.isAvailable(deviceIndex);
        end
        if ~deviceFound
            % As long as at least one default device is available this call
            % will succeed.
            errId = 'parallel:gpu:device:UnavailableDevice';
            deviceIndex = iSelectOrError(nD, myWorkerIndex, ...
                defaultDevices,...
                MException(errId, '%s', getString(message(errId))));
        end
    else
        % In a multi-worker situation, try only one exclusive-mode device,
        % based on my worker index.  This is not completely free from race
        % conditions. If device availability changes between now and when
        % the device is actually used, then an idle device or an error
        % could result. But it does not cause additional race conditions.
        deviceIndex = exclusiveDevices(myWorkerIndex);

        if ~parallel.gpu.GPUDevice.isAvailable(deviceIndex)
            errId = 'parallel:gpu:device:WorkerUnavailableDevice';
            deviceIndex = iSelectOrError(nD, ...
                myWorkerIndex, ...
                defaultDevices, ...
                MException(errId, '%s', ...
                    getString(message(...
                    errId, myWorkerIndex, deviceIndex))));
        end
    end
else
    % The more typical case. If there are no exclusive mode devices or the
    % number of workers is greater than the number of exclusive mode
    % devices, share one of the default mode devices. If there aren't any
    % default mode devices, a worker with device index greater than the
    % number of exclusive mode devices will throw an error.
    errId = 'parallel:gpu:device:UnavailableDevice';
    deviceIndex = iSelectOrError(nD, myWorkerIndex, ...
        defaultDevices, ...
        MException(errId, '%s', getString(message(errId))));
end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function deviceIndex = iSelectOrError(nD, myWorkerIndex, ...
    defaultDevices, exceptionToThrow)
if nD == 0
    % Exclusive device(s) not available and no default devices exist: issue
    % an error.
    throwAsCaller(exceptionToThrow);
else
    % Couldn't get the exclusive device, revert to a default mode device.
    % We decided against issuing a warning in this case.
    deviceIndex = defaultDevices(mod(myWorkerIndex-1, nD) + 1);
end
end

