function x = gop(F, x, varargin)
%GOP  Global collective operation on gpuArray
%   RES = GOP(FUN,X)
%   RES = GOP(FUN,X,LABTARGET)
%   RES = GOP(FUN,X,CLASSNAME)
%   RES = GOP(FUN,X,LABTARGET,CLASSNAME)
%   
%   GOP has been optimized for gpuArrays under certain conditions:
%   - You specify the CLASSNAME 'gpuArray'.
%   - Your parallel pool is running under Linux.
%   - All GPUs in the pool have compute capability 3.5 or above.
%   - All data are non-empty full real gpuArrays of type single, double, uint8,
%     int32, int64 or uint64.
%   - All data are the same size in every dimension (no dimension expansion).
%   - FUN is one of @plus, @times, @min or @max.
%   Under these circumstances, peer-to-peer communication is used to
%   improve performance.
%   
%   If these conditions are not met, the behaviour of GOP for gpuArrays is
%   unchanged.
%   
%   Example:
%   spmd
%       N = 1000;
%       x = gpuArray.randn(N);
%       SumX = gop(@plus,x,'gpuArray');        % Matrix addition of each x
%       MaxX = gop(@max,max(x(:)),'gpuArray'); % Maximum value over all workers
%   end
%   
%   See also GOP, GPUARRAY, GPLUS, GCAT.
%   


%   Copyright 2016-2020 The MathWorks, Inc.

narginchk(2,4);

[labTarget, classname] = parallel.internal.mpi.gopValidate(F, x, varargin{:});

% Early out for one-worker pools
if numlabs == 1
    return;
end

% Only attempt to continue with gpuArray specialization if classname is set
% correctly, otherwise divert to the generic reduction code
if ~strcmp(classname, 'gpuArray')
    x = parallel.internal.mpi.gopReduce(F, x, labTarget);
    return;
end

% If we reached here, then we are implementing a gpuArray-specific
% reduction

% Get some pool attributes with a single global communication
attributes = iGetWorkerAttributes(F, x);
isOptimized = iIsOptimized(attributes);

if ~isOptimized
    % NCCL is not supported so divert to the normal gopReduce function 
    % which uses the GPU optimised peer to peer communications
    x = parallel.internal.mpi.gopReduce(F, x, labTarget);
    return;
    
end

% If we reached here, then we can do at least some of the reduction using
% peer-to-peer DMA

% Make sure every worker has data of the same type
x = cast(x, iOutputDatatype(attributes));

% On a single node or multinode peer-to-peer-enabled cluster with a
% GPU-per-worker, we can get the output in a single reduction operation.
world_labindex = labindex;
[~, order, ids] = unique(attributes.HostName, 'sorted');
isSingleNode = numel(order) == 1;
useMultinodeP2P = all(attributes.UseMultinodeP2P);
if (isSingleNode || useMultinodeP2P) && ~iNodesHaveSharedGPUs(attributes, ids)
    x = iNcclReduce(F, x, labTarget);
    return;
end

% The algorithm that follows can be read sequentially. The nested function
% syntax is used to ensure exception safety for the CommSplitter
% communicators, and provide some meaningful indentation.
hasResult = true;
    function reduceNode
        % Split the communicator so we are working separately on each node
        hostnameID = ids(world_labindex);
        nodeComm = distributedutil.CommSplitter(hostnameID, world_labindex); %#ok<NASGU>

        function reduceOnEachDevice
            % Split the communicator again to work only on all the labs using the
            % same device
            gpuIndex = parallel.internal.gpu.currentDeviceIndex();
            gpuComm = distributedutil.CommSplitter(gpuIndex, world_labindex); %#ok<NASGU>
            
            % Reduce to a single gpuArray per device using generic method
            x = parallel.internal.mpi.gopReduce(F, x, 1);
            hasResult = double(labindex == 1);
        end
        reduceOnEachDevice();
        
        function reduceBetweenDevices
            % Create a new communicator between the workers which own the result
            % for each device
            % hasResult+1 is needed because color=0 is ignored by the
            % CommSplitter.
            ncclComm = distributedutil.CommSplitter(hasResult+1, world_labindex); %#ok<NASGU>
            
            % Do an optimized peer-to-peer reduction between those GPUs
            if hasResult
                x = iNcclReduce(F, x, 1);
            end
            % Result is now on the worker on each node with the lowest labindex
        end
        if ~useMultinodeP2P
            reduceBetweenDevices();
            hasResult = double(labindex == 1);
        end
    end
reduceNode();

% Finally, we need to reduce between the single worker on each node
% that has the result. If we're using multi-node peer-to-peer, this will
% reduce between all workers in the cluster with unique GPUs.
    function reduceCluster
        % Create communicator
        finalComm = distributedutil.CommSplitter(hasResult+1, world_labindex); %#ok<NASGU>
        % Reduce
        if hasResult
            if useMultinodeP2P
                x = iNcclReduce(F, x, 1);
            else
                x = parallel.internal.mpi.gopReduce(F, x, 1);
            end
        end
    end
reduceCluster();

% Send result to labTarget, or broadcast to all labs
x = iDeliverResult(x, labTarget);

end

function attributes = iGetWorkerAttributes(F, x)
% Shares some essential attributes between workers

% Host and device
hostname = parallel.internal.general.HostNameUtils.getLocalHostAddress();
gpuDeviceIndex = parallel.internal.gpu.currentDeviceIndex();

% Optimization
isLinux = ~ispc() && ~ismac();
gpu = gpuDevice();
isKepler = iIsCapabilityAtLeast(gpu.ComputeCapability, 3, 5);
isNcclCompatible = ...
    isLinux && isKepler && ...
    iIsSupportedDatatype(x) && ~issparse(x) && ...
    iIsOptimizedFunction(F);

% Data
datatype = underlyingType(x);
szX = size(x);

% Create array on each worker and share
attributes = struct( ...
    'HostName', {{hostname}}, ...
    'DeviceIndex', gpuDeviceIndex, ...
    'Type', {{datatype}}, ...
    'Size', {{szX}}, ...
    'NCCLCompatible', isNcclCompatible, ...
    'UseMultinodeP2P', iIsMultinodeP2PEnabled() );
attributes = gop(@iCatStruct, attributes);

end

function tf = iIsCapabilityAtLeast(computeCapabilityStr, majorMin, minorMin)
ccAsVec = sscanf(computeCapabilityStr, '%i.%i');
tf = ccAsVec(1) > majorMin || ...
    ccAsVec(1) == majorMin && ccAsVec(2) >= minorMin;
end

function tf = iIsMultinodeP2PEnabled
envvar = getenv('PARALLEL_SERVER_FAST_MULTINODE_GPU_COMMUNICATION');
legacyenvvar = getenv('MDCS_FAST_MULTINODE_GPU_COMMUNICATION');
tf = (~isempty(envvar) && str2double(envvar) > 0) || (~isempty(legacyenvvar) && str2double(legacyenvvar) > 0);
end

function tf = iIsSupportedDatatype(x)
tf = ismember(underlyingType(x), { ...
    'single', 'double', ...
    'uint8', 'int32', 'int64', 'uint64' });
end

function s = iCatStruct(s1, s2)
% Concatenate the fields of two structures
s = s1;
f = fieldnames(s);
for i = 1:numel(f)
    fname = f{i};
    if ischar(s1.(fname))
        s.(fname) = {s1.(fname); s2.(fname)};
    else
        s.(fname) = [s1.(fname); s2.(fname)];
    end
end
end

function tf = iIsOptimized(attributes)
tf = all(attributes.NCCLCompatible) && ...
    iAreSupportedSizes(attributes.Size) && ...
    iIsSupportedDatatypeCombination(attributes.Type);
end

function tf = iAreSupportedSizes(sizes)
tf = iIsEqualOrScalar(sizes) && prod(sizes{1}) > 0;
end

function tf = iIsSupportedDatatypeCombination(types)
tf = all(ismember(types, {'single', 'double'})) || ...
    iIsEqualOrScalar(types);
end

function tf = iIsEqualOrScalar(cellData)
tf = isscalar(cellData) || isequal(cellData{:});
end

function datatype = iOutputDatatype(attributes)
if iIsEqualOrScalar(attributes.Type)
    datatype = attributes.Type{1};
else
    % Assumes types are already validated to be floating point
    datatype = iSuperiorFloat(attributes.Type);
end
end

function datatype = iSuperiorFloat(types)
if ismember('single', types)
    datatype = 'single';
else
    datatype = 'double';
end
end

function tf = iIsOptimizedFunction(F)
optimizedFuncs = { ...
    'plus', ...
    'times', ...
    'min', ...
    'max' };
tf = ismember(func2str(F), optimizedFuncs);
end

function tf = iNodesHaveSharedGPUs(attributes, hostIDs)
[~, uniqueHostGPUWorkers] = unique([hostIDs, attributes.DeviceIndex], 'rows');
tf = numel(uniqueHostGPUWorkers) < numel(hostIDs);
end

function x = iNcclReduce(F, x, labTarget)
% Call into NCCL to reduce between GPUs on a single node.
if numlabs > 1
    x = parallel.internal.gpumpi.ncclReduce(func2str(F), x, labTarget);
end
end

function x = iDeliverResult(x, labTarget)
% Send result to labTarget, or broadcast to all labs
if labTarget == 0
    if labindex == 1
        labBroadcast(1, x);
    else
        x = labBroadcast(1);
    end
else
    if labTarget ~= 1
        if labindex == 1
            labSend(x, labTarget);
            x = [];
        elseif labindex == labTarget
            x = labReceive(1);
        else
            x = [];
        end
    else
        if labindex ~= 1
            x = [];
        end
    end
end
end
