Hetu/python/hetu/communicator/mpi_nccl_comm.py

from ctypes import *
from hetu import ndarray
from hetu.stream import *
from hetu.context import DeviceGroup
import numpy as np
from enum import Enum
import os
import socket


def _load_nccl_lib():
    """Load libary in build/lib."""
    curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
    lib_path = os.path.join(curr_path, '../../../build/lib/')
    path_to_so_file = os.path.join(lib_path, "lib_mpi_nccl_runtime_api.so")
    lib = CDLL(path_to_so_file, RTLD_LOCAL)
    return lib


lib_mpi_nccl = _load_nccl_lib()
# lib_mpi_nccl = CDLL("./lib_mpi_nccl_runtime_api.so", RTLD_LOCAL)


class ncclDataType_t(Enum):
    ncclInt8 = 0
    ncclChar = 0
    ncclUint8 = 1
    ncclInt32 = 2
    ncclInt = 2
    ncclUint32 = 3
    ncclInt64 = 4
    ncclUint64 = 5
    ncclFloat16 = 6
    ncclHalf = 6
    ncclFloat32 = 7
    ncclFloat = 7
    ncclFloat64 = 8
    ncclDouble = 8
    ncclNumTypes = 9


class ncclRedOp_t(Enum):
    ncclSum = 0
    ncclProd = 1
    ncclMax = 2
    ncclMin = 3
    ncclNumOps = 4


class ncclUniqueId(Structure):
    _fields_ = [("internal", (c_int8 * 128))]


class MPI_Communicator(object):
    def __init__(self, devices=None):
        '''
            mpicomm: the MPI communicator, to use in MPI_Bcast, MPI_Reduce, MPI_Scatter, etc
            nRanks: the total number of MPI threads
            myRanks: the rank in all MPI threads
            localRank: the rank among the MPI threads in this device
        '''
        self.mpicomm = c_int64(0)
        self.nRanks = c_int32(0)
        self.myRank = c_int32(0)
        self.localRank = c_int32(-1)
        self.device_id = c_int(0)

        self.MPI_Init()
        self.MPIGetComm()
        self.MPI_Comm_rank()
        self.MPI_Comm_size()
        self.hostHashs = (c_ulonglong * self.nRanks.value)()
        self.hostDevices = (c_int * self.nRanks.value)()
        self.getLocalRank()

        self.devices = devices
        self.device_id.value = self.getDeviceFromLocalRank(
            self.localRank.value)
        self.getGlobalDevice()

    @property
    def dev_id(self):
        return self.device_id.value

    @property
    def local_rank(self):
        return self.localRank.value

    @property
    def rank(self):
        return self.myRank.value

    @property
    def nrank(self):
        return self.nRanks.value

    def MPI_Init(self):
        lib_mpi_nccl.MPIInit()

    def MPI_Finalize(self):
        lib_mpi_nccl.MPIFinalize()

    def MPIGetComm(self):
        lib_mpi_nccl.MPIGetComm(ctypes.byref(self.mpicomm))

    def MPI_Comm_rank(self):
        lib_mpi_nccl.getMPICommRank(ctypes.byref(
            self.mpicomm), ctypes.byref(self.myRank))

    def MPI_Comm_size(self):
        lib_mpi_nccl.getMPICommSize(ctypes.byref(
            self.mpicomm), ctypes.byref(self.nRanks))

    def getLocalRank(self):
        lib_mpi_nccl.getLocalRank(ctypes.byref(
            self.mpicomm), self.nRanks, self.myRank, ctypes.byref(self.localRank), self.hostHashs)

    def getGlobalDevice(self):
        lib_mpi_nccl.getGlobalDevice(ctypes.byref(
            self.mpicomm), self.nRanks, self.myRank, self.device_id, self.hostDevices)

    def getRankFromDevice(self, hostname, device_id):
        if hostname == 'localhost':
            hostname = socket.gethostname()
        # hash
        result = 5381
        for c in hostname:
            result = result * 33 + ord(c)
        rank = 0
        while rank < self.nrank and (result != self.hostHashs[rank] or device_id != self.hostDevices[rank]):
            rank += 1
        assert rank < self.nrank, 'Device %d in host %s not found.' % (
            device_id, hostname)
        return rank

    def getDeviceFromLocalRank(self, local_rank):
        return self.devices[local_rank] if self.devices else local_rank

    def getLocalRankFromDevice(self, device_id):
        return self.devices.index(device_id) if self.devices else device_id

    def ncclInit(self, stream=None):
        return NCCL_Communicator(self, stream=stream)

    def ncclGroupInit(self, devices_context, stream=None):
        return NCCL_Communicator(self, devices_context, stream=stream)

    def __del__(self):
        self.MPI_Finalize()


class NCCL_Communicator():
    def __init__(self, comm, devices_context=None, stream=None):
        '''
            ncclcomm: the NCCL communicator, to use in ncclAllReduce ...
            ncclId: ncclGetUniqueId should be called once when creating a communicator
                    and the Id should be distributed to all ranks in the communicator before calling ncclCommInitRank.
            stream: the stream for NCCL communication
        '''
        self.mpi_communicator = comm
        self.mpicomm = comm.mpicomm
        self.nRanks = comm.nRanks
        self.myRank = comm.myRank
        self.localRank = comm.localRank
        self.device_id = comm.device_id

        if stream == None:
            self.stream = create_stream_handle(
                ndarray.gpu(self.device_id.value))
        else:
            self.stream = stream

        self.ncclId = ncclUniqueId()
        self.ncclcomm = c_int64(0)
        self.ncclSetDevice(self.device_id.value)
        if devices_context is None:
            self.ncclGetUniqueId()
            self.ncclCommInitRank()
        else:
            assert isinstance(
                devices_context, DeviceGroup), "Devices context should be a DeviceGroup."
            group_list = list(devices_context)
            if len(set(group_list)) != len(group_list):
                print("Warning: Repeated ranks are found in the group.")
                group_list = list(set(group_list))

            # the group_list here is as list of ndarray.(Remote)DLContext
            global_rank = self.rank
            global_size = self.nrank
            group_rank = -1
            group_size = len(group_list)
            local_rank = -1

            rank_list = []
            assert group_size <= global_size, "Error: Too many ranks in the group."
            local_rank_cnt = 0
            for i in range(group_size):
                at_local = group_list[i].local
                hostname = 'localhost' if at_local else group_list[i].hostname
                cur_rank = self.mpi_communicator.getRankFromDevice(
                    hostname, group_list[i].device_id)
                if cur_rank == global_rank:
                    group_rank = i
                    local_rank = local_rank_cnt
                    assert self.dev_id == group_list[i].device_id
                elif at_local:
                    local_rank_cnt += 1
                rank_list.append(cur_rank)
                assert cur_rank < global_size, "Error: The range of ranks should be [0, nrank-1]."

            self.nRanks = c_int32(group_size)
            self.myRank = c_int32(group_rank)
            self.localRank = c_int32(local_rank)

            if local_rank >= 0:
                group_id = 1234
                for x in rank_list:
                    group_id += x
                    group_id *= 33
                    group_id %= 10000007
                self.ncclGetGroupUniqueId(
                    (c_int32 * group_size)(*rank_list), c_int32(global_rank), self.nRanks, c_int32(group_id))
                self.ncclCommInitRank()

    @property
    def dev_id(self):
        return self.device_id.value

    @property
    def local_rank(self):
        return self.localRank.value

    @property
    def rank(self):
        return self.myRank.value

    @property
    def nrank(self):
        return self.nRanks.value

    def ncclSetDevice(self, device_id):
        self.device_id.value = device_id
        lib_mpi_nccl.setDevice(self.device_id.value)

    def getRankFromDevice(self, hostname, device_id):
        return self.mpi_communicator.getRankFromDevice(hostname, device_id)

    def ncclGetUniqueId(self, senderRank=0):
        lib_mpi_nccl.getNcclUniqueId(ctypes.byref(
            self.ncclId), self.mpicomm, self.localRank, c_int(senderRank))

    def ncclGetGroupUniqueId(self, group_list, ori_rank, group_size, group_id):
        lib_mpi_nccl.getGroupNcclUniqueId(ctypes.byref(
            self.ncclId), self.mpicomm, ori_rank, group_list, group_size, group_id)

    def ncclCommInitRank(self):
        '''
            Use partial AllReduce to change here.
            self.nRanks is the number of threads to use ncclallreduce
            self.myRank is the rank among these threads. the value must in [0, self.nRank - 1]
        '''
        lib_mpi_nccl.initNcclCommRank(ctypes.byref(self.ncclcomm), self.nRanks, ctypes.byref(
            self.ncclId), self.myRank, self.localRank)

    def dlarrayNcclAllReduce(self, input_arr, output_arr, datatype, reduceop, executor_stream=None):
        lib_mpi_nccl.dlarrayAllReduce(input_arr.handle, output_arr.handle, c_int(datatype.value), c_int(
            reduceop.value), self.ncclcomm, executor_stream.handle if executor_stream else self.stream.handle)

    def dlarrayBroadcast(self, input_arr, output_arr, datatype, root, executor_stream=None):
        lib_mpi_nccl.dlarrayBroadcast(input_arr.handle, output_arr.handle, c_int(datatype.value), c_int(
            root), self.ncclcomm, executor_stream.handle if executor_stream else self.stream.handle)

    def dlarrayAllGather(self, input_arr, output_arr, datatype, executor_stream=None):
        lib_mpi_nccl.dlarrayAllGather(input_arr.handle, output_arr.handle, c_int(
            datatype.value), self.ncclcomm, executor_stream.handle if executor_stream else self.stream.handle)

    def dlarraySend(self, arr, datatype, target, executor_stream=None):
        lib_mpi_nccl.dlarraySend(arr.handle, c_int(datatype.value), c_int(
            target), self.ncclcomm, executor_stream.handle if executor_stream else self.stream.handle)

    def dlarrayRecv(self, arr, datatype, src, executor_stream=None):
        lib_mpi_nccl.dlarrayRecv(arr.handle, c_int(datatype.value), c_int(
            src), self.ncclcomm, executor_stream.handle if executor_stream else self.stream.handle)

    def ncclCommDestroy(self):
        lib_mpi_nccl.commDestroyNccl(ctypes.byref(self.ncclcomm))

    def __del__(self):
        self.ncclCommDestroy()


def mpi_communicator(devices=None):
    return MPI_Communicator(devices=devices)


# NCCL_DEBUG=INFO mpirun --allow-run-as-root -np 4 python mpi_nccl_comm.py
if __name__ == "__main__":
    t = mpi_communicator()
    t = t.ncclInit()

    arr = np.ones(16)*t.localRank.value
    print("before: = ", arr)
    arr = ndarray.array(arr, ctx=ndarray.gpu(t.device_id.value))
    output_arr = np.zeros(16 * t.nRanks.value)

    output_arr = ndarray.array(output_arr, ctx=ndarray.gpu(t.device_id.value))
    t.dlarrayNcclAllReduce(
        arr, arr, ncclDataType_t.ncclFloat32, ncclRedOp_t.ncclSum)

    print("after: = ", arr.asnumpy())