classdef NodeInfo
%NodeInfo Stores information about GPU devices on this node
%   Use the information stored in the NodeInfo object to help decide how to
%   assign workers to GPUs in a parallel pool or communicating job. An object
%   of this class is passed as an input to the function selectGPU. You
%   can write your own version of selectGPU if the default assignment does
%   not meet your needs.
% 
%   NodeInfo properties:
%
%   DefaultModeDevices   - indices of GPU devices in default compute mode
%   ExclusiveModeDevices - indices of GPU devices in exclusive compute mode
%   MyWorkerIndex        - relative index of this worker
%   NumNodeWorkers       - number of workers sharing this node
%   SupportedDevices     - indices of supported GPU devices on this node
%
%   See also selectGPU, parallel.gpu.GPUDevice.isAvailable.

%   Copyright 2012-2022 The MathWorks, Inc.

properties (GetAccess = public, SetAccess = immutable)
    %DefaultModeDevices indices of GPU devices in default compute mode
    %   A vector of indices of supported GPU devices for which the
    %   ComputeMode property is 'Default'. Such devices can be used by all
    %   workers.
    DefaultModeDevices
    %ExclusiveModeDevices indices of GPU devices in exclusive compute mode
    %   A vector of indices of supported GPU devices for which the
    %   ComputeMode property is 'Exclusive thread' or 'Exclusive process'.
    %   Such devices can be used by only one worker. The devices are not
    %   necessarily available. Availability can be checked using the
    %   function parallel.gpu.GPUDevice.isAvailable.
    ExclusiveModeDevices
    %MyWorkerIndex relative index of this worker
    %   The index of this worker relative to other workers on this node
    %   participating in the current communicating job. The index might not
    %   correspond to the order of spmdIndex for the workers. If there is
    %   no communicating job, this returns 1.
    MyWorkerIndex
    %NumNodeWorkers number of workers sharing this node
    %   The number of workers on this node participating in the current
    %   communicating job. If there is no communicating job, this property
    %   has the value 1.
    NumNodeWorkers
    %SupportedDevices indices of supported GPU devices on this node
    %   A vector of indices of supported GPU devices. They are not
    %   necessarily available. Availability can be checked using the
    %   function parallel.gpu.GPUDevice.isAvailable.
    SupportedDevices
end

methods (Access = public)
    function obj = NodeInfo(mediator, ...
            defaultModeDevices, ...
            exclusiveModeDevices)
        % NodeInfo constructor stores information about the available
        % devices. It takes the machine-to-worker mapping (provided by the
        % mediator) and determines information about how many workers share
        % this node. It is an input argument to selectGPU.
        
        obj.DefaultModeDevices = defaultModeDevices;
        obj.ExclusiveModeDevices = exclusiveModeDevices;
        obj.SupportedDevices = sort([defaultModeDevices, exclusiveModeDevices]);
        % Extract information from the mediator
        machineToWorkerMapping = mediator.MachineToWorkerMapping;
        hostName = mediator.HostName;
        storedWorkerIndex = mediator.StoredWorkerIndex;
        if isempty(machineToWorkerMapping) || isempty(hostName) ...
                || isempty(storedWorkerIndex)
            % Missing information: assume I am alone on the node.
            obj.MyWorkerIndex = 1;
            obj.NumNodeWorkers = 1;
        else
            % Sort the mapping cell array by hostName.
            machineToWorkerMapping = sortrows(machineToWorkerMapping.',1);
            
            % Detect all instances of my hostName.
            indices = strcmp(hostName, ...
                machineToWorkerMapping(:,1))==true;
            
            % Retrieve indices of the workers that match my hostName.
            myRange = cell2mat(machineToWorkerMapping(indices, 2));
            obj.NumNodeWorkers = numel(myRange);
            
            % Find my location in the resulting continuous range.
            % Use the worker index stored in the mediator, because if we
            % are in a parfor loop, we no longer know the spmdIndex.
            obj.MyWorkerIndex = find(myRange==storedWorkerIndex);
        end
    end
end
end
