PKU-DAIR
/
Hetu
Not watched
1
0
Fork 0
Code
Releases 0 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Branch: master
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'master'
${ noResults }
Hetu/python/hetu/gpu_ops/DistGCN_15d.py

157 lines
6.5 kB
Raw Permalink Blame History 

  1. from __future__ import absolute_import

  2. from .Node import Op

  3. import math

  4. from .. import ndarray

  5. import numpy as np

  6. from ctypes import *

  7. 

  8. 

  9. def row_num(node_count, rank, size):

  10.     n_per_proc = math.ceil(float(node_count) / size)

  11.     if (node_count % size == 0):

  12.         return int(node_count / size)

  13.     if (rank < size - 1):

  14.         return int(n_per_proc)

  15.     else:

  16.         return int(node_count % n_per_proc)

  17. 

  18. 

  19. def broad_func(node_count, adj_matrix, inputs, rank, size, replication, row_groups, col_groups, ctx, comm=None, stream_handle=None):

  20.     assert size % (replication ** 2) == 0

  21. 

  22.     n_per_proc = math.ceil(float(node_count) / (size // replication))

  23.     proc_node_count = row_num(

  24.         node_count, rank//replication, size // replication)

  25. 

  26.     z_loc = ndarray.empty((proc_node_count, inputs.shape[1]), ctx=ctx)

  27.     tmp = ndarray.empty((proc_node_count, inputs.shape[1]), ctx=ctx)

  28.     inputs_recv = ndarray.empty((int(n_per_proc), inputs.shape[1]), ctx=ctx)

  29. 

  30.     rank_c = rank // replication

  31.     rank_col = rank % replication

  32. 

  33.     stages = size // (replication ** 2)

  34.     node_count_col = stages * n_per_proc

  35.     if rank_col == replication - 1:

  36.         stages = (size // replication) - (replication - 1) * stages

  37.         node_count_col = node_count - (replication - 1) * node_count_col

  38. 

  39.     start_pos = list(range(0, int(node_count_col), int(n_per_proc)))

  40.     end_pos = start_pos[1:]+[int(node_count_col)]

  41. 

  42.     for i in range(stages):

  43.         q = (rank_col * (size // (replication ** 2)) + i) * \

  44.             replication + rank_col

  45.         q_c = q // replication

  46. 

  47.         if q_c == size // replication - 1:

  48.             inputs_recv = ndarray.empty((row_num(

  49.                 node_count, size//replication - 1, size//replication), inputs.shape[1]), ctx=ctx)

  50.         if q == rank:

  51.             inputs.copyto(inputs_recv)

  52. 

  53.         from ..communicator.mpi_nccl_comm import ncclDataType_t, ncclRedOp_t

  54.         if replication > 1:

  55.             col_groups[rank_col].dlarrayBroadcast(

  56.                 inputs_recv, inputs_recv, ncclDataType_t.ncclFloat32, q)

  57.         else:

  58.             comm.dlarrayBroadcast(inputs_recv, inputs_recv,

  59.                                   ncclDataType_t.ncclFloat32, q)

  60. 

  61.         from ..gpu_links import CuSparse_Csrmm, matrix_elementwise_add

  62.         CuSparse_Csrmm(adj_matrix, False, inputs_recv, False, tmp,

  63.                        stream=stream_handle, start_pos=int(start_pos[i]), end_pos=int(end_pos[i]))

  64.         matrix_elementwise_add(z_loc, tmp, z_loc, stream_handle)

  65. 

  66.     if replication > 1:

  67.         row_groups[rank_c].dlarrayNcclAllReduce(

  68.             z_loc, z_loc, ncclDataType_t.ncclFloat32, reduceop=ncclRedOp_t.ncclSum)

  69. 

  70.     return z_loc

  71. 

  72. 

  73. class DistGCN_15dOp(Op):

  74.     def __init__(self, node_A, node_B, node_C, node_Count_Self, node_Count_All, size, replication, device_id, comm, comm_groups=[None, None], need_W=True):

  75.         super().__init__(DistGCN_15dOp, [

  76.             node_A, node_B, node_C], ctx=ndarray.gpu(device_id))

  77.         self.need_W = need_W

  78.         self.node_Count_Self = node_Count_Self

  79.         self.node_Count_All = node_Count_All

  80.         self.replication = replication

  81.         self.size = size

  82.         self.comm = comm

  83.         self.comm_groups = comm_groups

  84.         self.device_id = device_id

  85. 

  86.     def compute(self, input_vals, output_val, stream_handle=None):

  87.         adj_matrix = input_vals[0]

  88.         inputs_H = input_vals[1]

  89.         weight = input_vals[2]

  90.         node_count = self.node_Count_All

  91.         comm = self.comm

  92.         rank = comm.localRank.value

  93.         ctx = ndarray.gpu(self.device_id)

  94. 

  95.         if weight.shape[1] < inputs_H.shape[1]:

  96.             HW = ndarray.empty((inputs_H.shape[0], weight.shape[1]), ctx=ctx)

  97.             if (self.need_W == True):

  98.                 from ..gpu_links import matrix_multiply

  99.                 matrix_multiply(inputs_H, False, weight,

  100.                                 False, HW, stream_handle)

  101.             else:

  102.                 HW = inputs_H

  103.             z = broad_func(node_count, adj_matrix, HW, rank, self.size, self.replication,

  104.                            row_groups=self.comm_groups[0], col_groups=self.comm_groups[1], ctx=ctx, comm=comm, stream_handle=stream_handle)

  105.             z.copyto(output_val)

  106.         else:

  107.             AH = broad_func(node_count, adj_matrix, inputs_H, rank, self.size, self.replication,

  108.                             row_groups=self.comm_groups[0], col_groups=self.comm_groups[1], ctx=ctx, comm=comm, stream_handle=stream_handle)

  109.             z = ndarray.empty((AH.shape[0], weight.shape[1]), ctx=ctx)

  110.             if (self.need_W == True):

  111.                 from ..gpu_links import matrix_multiply

  112.                 matrix_multiply(AH, False, weight, False, z, stream_handle)

  113.             else:

  114.                 z = AH

  115.             z.copyto(output_val)

  116. 

  117.     def gradient(self, output_grad):

  118.         adj_matrix = self.inputs[0]

  119.         inputs_H = self.inputs[1]

  120.         weight = self.inputs[2]

  121.         node_Count_Self = self.node_Count_Self

  122.         node_Count_All = self.node_Count_All

  123.         comm = self.comm

  124.         rank = comm.localRank.value

  125.         ag = distgcn_15d_op(adj_matrix, output_grad, weight, node_Count_Self, node_Count_All,

  126.                             self.size, self.replication, self.device_id, comm, self.comm_groups, need_W=False)

  127. 

  128.         from . import matmul_op

  129.         grad_A = None

  130.         grad_H = matmul_op(ag, weight, trans_B=True)

  131.         grad_weight = matmul_op(inputs_H, ag, trans_A=True)

  132.         from . import groupallreduceCommunicate_op

  133.         if self.replication > 1:

  134.             weight_groups = self.comm_groups[1]

  135.             if len(self.comm_groups) == 3:

  136.                 weight_groups = self.comm_groups[2]

  137.             grad_W = groupallreduceCommunicate_op(

  138.                 grad_weight, weight_groups[rank % self.replication])

  139.         else:

  140.             grad_W = groupallreduceCommunicate_op(grad_weight, comm)

  141.         return [grad_A, grad_H, grad_W]

  142. 

  143.     def infer_shape(self, input_shapes):

  144.         assert len(input_shapes) == 3

  145.         H = input_shapes[1]

  146.         W = input_shapes[2]

  147.         shape_H = H[1]

  148.         shape_W = W[1]

  149.         if (self.need_W == True):

  150.             return (self.node_Count_Self, shape_W)

  151.         else:

  152.             return (self.node_Count_Self, shape_H)

  153. 

  154. 

  155. def distgcn_15d_op(node_A, node_B, node_C, node_Count_Self, node_Count_All, size, replication, device_id, comm, comm_groups=[None, None], need_W=True):

  156.     return DistGCN_15dOp(node_A, node_B, node_C, node_Count_Self, node_Count_All, size, replication, device_id, comm, comm_groups, need_W=need_W)